diff --git a/README.md b/README.md deleted file mode 100644 index 905373c6..00000000 --- a/README.md +++ /dev/null @@ -1,118 +0,0 @@ -This is the commune ai project. - -**This project has three roles: user, owner, and friend like social network** - -**1. Frontend Running** - # Next.js Project - - This is a Next.js project created to demonstrate how to run a Next.js application. - - ## Getting Started - - To get started with this project, follow the steps below: - - 1. Clone the repository: - - ``` - git clone https://github.com/commune-ai/frontend.git - ``` - - 2. Install the dependencies: - - ``` - cd next-js-project - npm install - ``` - - 3. Run the development server: - - ``` - npm run dev - ``` - - Open [http://localhost:3000](http://localhost:3000) in your browser to see the application running. - - ## Available Scripts - - In the project directory, you can run the following scripts: - - ### `npm run dev` - - Runs the app in development mode.
- Open [http://localhost:3000](http://localhost:3000) to view it in the browser. - - The page will reload if you make edits.
- You will also see any lint errors in the console. - - ### `npm run build` - - Builds the app for production to the `.next` folder.
- It correctly bundles React in production mode and optimizes the build for the best performance. - - ### `npm start` - - Starts the app in production mode.
- Open [http://localhost:3000](http://localhost:3000) to view it in the browser. - - ## Learn More - - To learn more about Next.js, check out the [Next.js documentation](https://nextjs.org/docs). - -**2. Backend** - - - go to backend folder and follow the steps - # Django Project - - This is a Django project created to demonstrate how to run a Django application. - - ## Getting Started - - To get started with this project, follow the steps below: - - 1. Clone the repository: - - ``` - git clone https://github.com/your-username/django-project.git - ``` - - 2. Install the dependencies: - - ``` - cd django-project - pip install -r requirements.txt - ``` - - 3. Run the development server: - - ``` - python manage.py runserver - ``` - - Open [http://localhost:8000](http://localhost:8000) in your browser to see the application running. - - ## Available Commands - - In the project directory, you can run the following commands: - - ### `python manage.py runserver` - - Runs the development server.
- Open [http://localhost:8000](http://localhost:8000) to view it in the browser. - - The server will automatically reload if you make changes to the code. - - ### `python manage.py migrate` - - Applies any database migrations needed for the project. - - ### `python manage.py createsuperuser` - - Creates a superuser for the admin interface. - - ## Learn More - - To learn more about Django, check out the [Django documentation](https://docs.djangoproject.com/). - -## License - -This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. diff --git a/backend/.devcontainer/bashrc.override.sh b/backend/.devcontainer/bashrc.override.sh deleted file mode 100644 index bedddf64..00000000 --- a/backend/.devcontainer/bashrc.override.sh +++ /dev/null @@ -1,20 +0,0 @@ - -# -# .bashrc.override.sh -# - -# persistent bash history -HISTFILE=~/.bash_history -PROMPT_COMMAND="history -a; $PROMPT_COMMAND" - -# set some django env vars -source /entrypoint - -# restore default shell options -set +o errexit -set +o pipefail -set +o nounset - -# start ssh-agent -# https://code.visualstudio.com/docs/remote/troubleshooting -eval "$(ssh-agent -s)" diff --git a/backend/.devcontainer/devcontainer.json b/backend/.devcontainer/devcontainer.json deleted file mode 100644 index f3dbd9a5..00000000 --- a/backend/.devcontainer/devcontainer.json +++ /dev/null @@ -1,85 +0,0 @@ -// For format details, see https://containers.dev/implementors/json_reference/ -{ - "name": "backend_dev", - "dockerComposeFile": [ - "../local.yml" - ], - "init": true, - "mounts": [ - { - "source": "./.devcontainer/bash_history", - "target": "/home/dev-user/.bash_history", - "type": "bind" - }, - { - "source": "/tmp", - "target": "/tmp", - "type": "bind" - }, - { - "source": "~/.ssh", - "target": "/home/dev-user/.ssh", - "type": "bind" - } - ], - // Tells devcontainer.json supporting services / tools whether they should run - // /bin/sh -c "while sleep 1000; do :; done" when starting the container instead of the container’s default command - "overrideCommand": false, - "service": "django", - // "remoteEnv": {"PATH": "/home/dev-user/.local/bin:${containerEnv:PATH}"}, - "remoteUser": "dev-user", - "workspaceFolder": "/app", - // Set *default* container specific settings.json values on container create. - "customizations": { - "vscode": { - "settings": { - "editor.formatOnSave": true, - "[python]": { - "analysis.autoImportCompletions": true, - "analysis.typeCheckingMode": "basic", - "defaultInterpreterPath": "/usr/local/bin/python", - "editor.codeActionsOnSave": { - "source.organizeImports": true - }, - // Uncomment when fixed - // https://github.com/microsoft/vscode-remote-release/issues/8474 - // "editor.defaultFormatter": "ms-python.black-formatter", - "formatting.blackPath": "/usr/local/bin/black", - "formatting.provider": "black", - "languageServer": "Pylance", - // "linting.banditPath": "/usr/local/py-utils/bin/bandit", - "linting.enabled": true, - "linting.flake8Enabled": true, - "linting.flake8Path": "/usr/local/bin/flake8", - "linting.mypyEnabled": true, - "linting.mypyPath": "/usr/local/bin/mypy", - "linting.pycodestylePath": "/usr/local/bin/pycodestyle", - // "linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", - "linting.pylintEnabled": true, - "linting.pylintPath": "/usr/local/bin/pylint" - } - }, - // https://code.visualstudio.com/docs/remote/devcontainerjson-reference#_vs-code-specific-properties - // Add the IDs of extensions you want installed when the container is created. - "extensions": [ - "davidanson.vscode-markdownlint", - "mrmlnc.vscode-duplicate", - "visualstudioexptteam.vscodeintellicode", - "visualstudioexptteam.intellicode-api-usage-examples", - // python - "ms-python.python", - "ms-python.vscode-pylance", - "ms-python.isort", - "ms-python.black-formatter", - // django - "batisteo.vscode-django" - ] - } - }, - // Uncomment the next line if you want start specific services in your Docker Compose config. - // "runServices": [], - // Uncomment the next line if you want to keep your containers running after VS Code shuts down. - // "shutdownAction": "none", - // Uncomment the next line to run commands after the container is created. - "postCreateCommand": "cat .devcontainer/bashrc.override.sh >> ~/.bashrc" -} diff --git a/backend/.dockerignore b/backend/.dockerignore deleted file mode 100644 index a602416c..00000000 --- a/backend/.dockerignore +++ /dev/null @@ -1,12 +0,0 @@ -.editorconfig -.gitattributes -.github -.gitignore -.gitlab-ci.yml -.idea -.pre-commit-config.yaml -.readthedocs.yml -.travis.yml -venv -.git -.envs/ diff --git a/backend/.editorconfig b/backend/.editorconfig deleted file mode 100644 index c0ce3426..00000000 --- a/backend/.editorconfig +++ /dev/null @@ -1,27 +0,0 @@ -# http://editorconfig.org - -root = true - -[*] -charset = utf-8 -end_of_line = lf -insert_final_newline = true -trim_trailing_whitespace = true - -[*.{py,rst,ini}] -indent_style = space -indent_size = 4 - -[*.{html,css,scss,json,yml,xml}] -indent_style = space -indent_size = 2 - -[*.md] -trim_trailing_whitespace = false - -[Makefile] -indent_style = tab - -[default.conf] -indent_style = space -indent_size = 2 diff --git a/backend/.envs/.local/.django b/backend/.envs/.local/.django deleted file mode 100644 index 247287be..00000000 --- a/backend/.envs/.local/.django +++ /dev/null @@ -1,14 +0,0 @@ -# General -# ------------------------------------------------------------------------------ -USE_DOCKER=yes -IPYTHONDIR=/app/.ipython -# Redis -# ------------------------------------------------------------------------------ -REDIS_URL=redis://redis:6379/0 - -# Celery -# ------------------------------------------------------------------------------ - -# Flower -CELERY_FLOWER_USER=debug -CELERY_FLOWER_PASSWORD=debug diff --git a/backend/.envs/.local/.postgres b/backend/.envs/.local/.postgres deleted file mode 100644 index 97bc5b27..00000000 --- a/backend/.envs/.local/.postgres +++ /dev/null @@ -1,7 +0,0 @@ -# PostgreSQL -# ------------------------------------------------------------------------------ -POSTGRES_HOST=postgres -POSTGRES_PORT=5432 -POSTGRES_DB=backend -POSTGRES_USER=debug -POSTGRES_PASSWORD=debug diff --git a/backend/.gitattributes b/backend/.gitattributes deleted file mode 100644 index 176a458f..00000000 --- a/backend/.gitattributes +++ /dev/null @@ -1 +0,0 @@ -* text=auto diff --git a/backend/.github/dependabot.yml b/backend/.github/dependabot.yml deleted file mode 100644 index 9e2dcb01..00000000 --- a/backend/.github/dependabot.yml +++ /dev/null @@ -1,91 +0,0 @@ -# Config for Dependabot updates. See Documentation here: -# https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates - -version: 2 -updates: - # Update GitHub actions in workflows - - package-ecosystem: 'github-actions' - directory: '/' - # Every weekday - schedule: - interval: 'daily' - - # Enable version updates for Docker - # We need to specify each Dockerfile in a separate entry because Dependabot doesn't - # support wildcards or recursively checking subdirectories. Check this issue for updates: - # https://github.com/dependabot/dependabot-core/issues/2178 - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/local/django` directory - directory: 'compose/local/django/' - # Every weekday - schedule: - interval: 'daily' - # Ignore minor version updates (3.10 -> 3.11) but update patch versions - ignore: - - dependency-name: '*' - update-types: - - 'version-update:semver-major' - - 'version-update:semver-minor' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/local/docs` directory - directory: 'compose/local/docs/' - # Every weekday - schedule: - interval: 'daily' - # Ignore minor version updates (3.10 -> 3.11) but update patch versions - ignore: - - dependency-name: '*' - update-types: - - 'version-update:semver-major' - - 'version-update:semver-minor' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/local/node` directory - directory: 'compose/local/node/' - # Every weekday - schedule: - interval: 'daily' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/production/aws` directory - directory: 'compose/production/aws/' - # Every weekday - schedule: - interval: 'daily' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/production/django` directory - directory: 'compose/production/django/' - # Every weekday - schedule: - interval: 'daily' - # Ignore minor version updates (3.10 -> 3.11) but update patch versions - ignore: - - dependency-name: '*' - update-types: - - 'version-update:semver-major' - - 'version-update:semver-minor' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/production/postgres` directory - directory: 'compose/production/postgres/' - # Every weekday - schedule: - interval: 'daily' - - - package-ecosystem: 'docker' - # Look for a `Dockerfile` in the `compose/production/traefik` directory - directory: 'compose/production/traefik/' - # Every weekday - schedule: - interval: 'daily' - - # Enable version updates for Python/Pip - Production - - package-ecosystem: 'pip' - # Look for a `requirements.txt` in the `root` directory - # also 'setup.cfg', 'runtime.txt' and 'requirements/*.txt' - directory: '/' - # Every weekday - schedule: - interval: 'daily' diff --git a/backend/.github/workflows/ci.yml b/backend/.github/workflows/ci.yml deleted file mode 100644 index 028c1868..00000000 --- a/backend/.github/workflows/ci.yml +++ /dev/null @@ -1,53 +0,0 @@ -name: CI - -# Enable Buildkit and let compose use it to speed up image building -env: - DOCKER_BUILDKIT: 1 - COMPOSE_DOCKER_CLI_BUILD: 1 - -on: - pull_request: - branches: ['master', 'main'] - paths-ignore: ['docs/**'] - - push: - branches: ['master', 'main'] - paths-ignore: ['docs/**'] - -concurrency: - group: ${{ github.head_ref || github.run_id }} - cancel-in-progress: true - -jobs: - linter: - runs-on: ubuntu-latest - steps: - - name: Checkout Code Repository - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.11' - - name: Run pre-commit - uses: pre-commit/action@v3.0.0 - - # With no caching at all the entire ci process takes 4m 30s to complete! - pytest: - runs-on: ubuntu-latest - - steps: - - name: Checkout Code Repository - uses: actions/checkout@v4 - - - name: Build the Stack - run: docker compose -f local.yml build - - - name: Run DB Migrations - run: docker compose -f local.yml run --rm django python manage.py migrate - - - name: Run Django Tests - run: docker compose -f local.yml run django pytest - - - name: Tear down the Stack - run: docker compose -f local.yml down diff --git a/backend/.gitignore b/backend/.gitignore index b90882dd..6f87abe1 100644 --- a/backend/.gitignore +++ b/backend/.gitignore @@ -1,281 +1,62 @@ -### Python template # Byte-compiled / optimized / DLL files __pycache__/ -*.py[cod] -*$py.class +*.pyc +*.pyo +*.pyd # C extensions *.so # Distribution / packaging -.Python -build/ -develop-eggs/ +.env +.env.* +.idea/ +.vscode/ +.venv/ +venv/ +env/ dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ +build/ *.egg-info/ -.installed.cfg -*.egg - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -.hypothesis/ - -# Translations -*.mo -*.pot - -# Django stuff: -staticfiles/ - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# pyenv -.python-version - -# celery beat schedule file -celerybeat-schedule - -# Environments -.venv -venv/ -ENV/ -myenv -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ - - -### Node template -# Logs -logs +*.egg/ *.log -npm-debug.log* -yarn-debug.log* -yarn-error.log* - -# Runtime data -pids -*.pid -*.seed -*.pid.lock - -# Directory for instrumented libs generated by jscoverage/JSCover -lib-cov - -# Coverage directory used by tools like istanbul -coverage - -# nyc test coverage -.nyc_output - -# Bower dependency directory (https://bower.io/) -bower_components - -# node-waf configuration -.lock-wscript - -# Compiled binary addons (http://nodejs.org/api/addons.html) -build/Release +*.sqlite3 -# Dependency directories -node_modules/ -jspm_packages/ - -# Typescript v1 declaration files -typings/ - -# Optional npm cache directory -.npm - -# Optional eslint cache -.eslintcache - -# Optional REPL history -.node_repl_history - -# Output of 'npm pack' -*.tgz - -# Yarn Integrity file -.yarn-integrity - - -### Linux template +# Development +*.swp +*.bak +*.tmp +*.bak *~ +/.vscode +/.idea +.DS_Store -# temporary files which can be created if a process still has a handle open of a deleted file -.fuse_hidden* - -# KDE directory preferences -.directory - -# Linux trash folder which might appear on any partition or disk -.Trash-* - -# .nfs files are created when an open file is removed but is still being accessed -.nfs* - - -### VisualStudioCode template -.vscode/* -!.vscode/settings.json -!.vscode/tasks.json -!.vscode/launch.json -!.vscode/extensions.json -*.code-workspace - -# Local History for devcontainer -.devcontainer/bash_history - - - - -### Windows template -# Windows thumbnail cache files -Thumbs.db -ehthumbs.db -ehthumbs_vista.db - -# Dump file -*.stackdump - -# Folder config file -Desktop.ini - -# Recycle Bin used on file shares -$RECYCLE.BIN/ - -# Windows Installer files -*.cab -*.msi -*.msm -*.msp - -# Windows shortcuts -*.lnk - - -### macOS template -# General -*.DS_Store -.AppleDouble -.LSOverride - -# Icon must end with two \r -Icon - -# Thumbnails -._* - -# Files that might appear in the root of a volume -.DocumentRevisions-V100 -.fseventsd -.Spotlight-V100 -.TemporaryItems -.Trashes -.VolumeIcon.icns -.com.apple.timemachine.donotpresent - -# Directories potentially created on remote AFP share -.AppleDB -.AppleDesktop -Network Trash Folder -Temporary Items -.apdisk - - -### SublimeText template -# Cache files for Sublime Text -*.tmlanguage.cache -*.tmPreferences.cache -*.stTheme.cache - -# Workspace files are user-specific -*.sublime-workspace - -# Project files should be checked into the repository, unless a significant -# proportion of contributors will probably not be using Sublime Text -# *.sublime-project - -# SFTP configuration file -sftp-config.json - -# Package control specific files -Package Control.last-run -Package Control.ca-list -Package Control.ca-bundle -Package Control.system-ca-bundle -Package Control.cache/ -Package Control.ca-certs/ -Package Control.merged-ca-bundle -Package Control.user-ca-bundle -oscrypto-ca-bundle.crt -bh_unicode_properties.cache - -# Sublime-github package stores a github token in this file -# https://packagecontrol.io/packages/sublime-github -GitHub.sublime-settings - - -### Vim template -# Swap -[._]*.s[a-v][a-z] -[._]*.sw[a-p] -[._]s[a-v][a-z] -[._]sw[a-p] +# Ignore environment configuration files +.env +*.pyc +__pycache__/ -# Session -Session.vim +# Ignore virtual environment +venv +env -# Temporary -.netrwhist +# Ignore cache and compiled files +__pycache__/ +*.pyc +*.pyo +*.pyd -# Auto-generated tag files -tags +# Ignore database files +*.sqlite3 -# Redis dump file -dump.rdb +# Ignore coverage reports +htmlcov/ +.coverage -### Project template -backend/media/ +# Ignore compiled translations +*.mo +*.pot -.pytest_cache/ -.ipython/ -.env -.envs/* -!.envs/.local/ +# Ignore setuptools/distutils metadata directory +*.egg-info/ diff --git a/backend/.pre-commit-config.yaml b/backend/.pre-commit-config.yaml deleted file mode 100644 index d592d961..00000000 --- a/backend/.pre-commit-config.yaml +++ /dev/null @@ -1,57 +0,0 @@ -exclude: '^docs/|/migrations/' -default_stages: [commit] - -repos: - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.5.0 - hooks: - - id: trailing-whitespace - - id: end-of-file-fixer - - id: check-json - - id: check-toml - - id: check-xml - - id: check-yaml - - id: debug-statements - - id: check-builtin-literals - - id: check-case-conflict - - id: check-docstring-first - - id: detect-private-key - - - repo: https://github.com/adamchainz/django-upgrade - rev: '1.15.0' - hooks: - - id: django-upgrade - args: ['--target-version', '4.2'] - - - repo: https://github.com/asottile/pyupgrade - rev: v3.15.0 - hooks: - - id: pyupgrade - args: [--py311-plus] - - - repo: https://github.com/psf/black - rev: 23.9.1 - hooks: - - id: black - - - repo: https://github.com/PyCQA/isort - rev: 5.12.0 - hooks: - - id: isort - - - repo: https://github.com/PyCQA/flake8 - rev: 6.1.0 - hooks: - - id: flake8 - - - repo: https://github.com/Riverside-Healthcare/djLint - rev: v1.34.0 - hooks: - - id: djlint-reformat-django - - id: djlint-django - -# sets up .pre-commit-ci.yaml to ensure pre-commit dependencies stay up to date -ci: - autoupdate_schedule: weekly - skip: [] - submodules: false diff --git a/backend/.readthedocs.yml b/backend/.readthedocs.yml deleted file mode 100644 index d5a8ef66..00000000 --- a/backend/.readthedocs.yml +++ /dev/null @@ -1,20 +0,0 @@ -# Read the Docs configuration file -# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details - -# Required -version: 2 - -# Set the version of Python and other tools you might need -build: - os: ubuntu-22.04 - tools: - python: '3.11' - -# Build documentation in the docs/ directory with Sphinx -sphinx: - configuration: docs/conf.py - -# Python requirements required to build your docs -python: - install: - - requirements: requirements/local.txt diff --git a/backend/README.md b/backend/README.md deleted file mode 100644 index 37d0dd37..00000000 --- a/backend/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# backend - -backend - -[![Built with Cookiecutter Django](https://img.shields.io/badge/built%20with-Cookiecutter%20Django-ff69b4.svg?logo=cookiecutter)](https://github.com/cookiecutter/cookiecutter-django/) -[![Black code style](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/ambv/black) - -## Settings - -Moved to [settings](http://cookiecutter-django.readthedocs.io/en/latest/settings.html). - -## Basic Commands - -### Setting Up Your Users - -- To create a **normal user account**, just go to Sign Up and fill out the form. Once you submit it, you'll see a "Verify Your E-mail Address" page. Go to your console to see a simulated email verification message. Copy the link into your browser. Now the user's email should be verified and ready to go. - -- To create a **superuser account**, use this command: - - $ python manage.py createsuperuser - -For convenience, you can keep your normal user logged in on Chrome and your superuser logged in on Firefox (or similar), so that you can see how the site behaves for both kinds of users. - -### Type checks - -Running type checks with mypy: - - $ mypy backend - -### Test coverage - -To run the tests, check your test coverage, and generate an HTML coverage report: - - $ coverage run -m pytest - $ coverage html - $ open htmlcov/index.html - -#### Running tests with pytest - - $ pytest - -### Live reloading and Sass CSS compilation - -Moved to [Live reloading and SASS compilation](https://cookiecutter-django.readthedocs.io/en/latest/developing-locally.html#sass-compilation-live-reloading). - -### Celery - -This app comes with Celery. - -To run a celery worker: - -```bash -cd backend -celery -A config.celery_app worker -l info -``` - -Please note: For Celery's import magic to work, it is important _where_ the celery commands are run. If you are in the same folder with _manage.py_, you should be right. - -To run [periodic tasks](https://docs.celeryq.dev/en/stable/userguide/periodic-tasks.html), you'll need to start the celery beat scheduler service. You can start it as a standalone process: - -```bash -cd backend -celery -A config.celery_app beat -``` - -or you can embed the beat service inside a worker with the `-B` option (not recommended for production use): - -```bash -cd backend -celery -A config.celery_app worker -B -l info -``` - -## Deployment - -The following details how to deploy this application. - -### Docker - -See detailed [cookiecutter-django Docker documentation](http://cookiecutter-django.readthedocs.io/en/latest/deployment-with-docker.html). diff --git a/backend/alembic.ini b/backend/alembic.ini new file mode 100644 index 00000000..c10d4ca0 --- /dev/null +++ b/backend/alembic.ini @@ -0,0 +1,116 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = alembic + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python>=3.9 or backports.zoneinfo library. +# Any required deps can installed by adding `alembic[tz]` to the pip requirements +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the +# "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to alembic/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the exec runner, execute a binary +# hooks = ruff +# ruff.type = exec +# ruff.executable = %(here)s/.venv/bin/ruff +# ruff.options = --fix REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/backend/alembic/README b/backend/alembic/README new file mode 100644 index 00000000..98e4f9c4 --- /dev/null +++ b/backend/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/backend/alembic/env.py b/backend/alembic/env.py new file mode 100644 index 00000000..fbb98279 --- /dev/null +++ b/backend/alembic/env.py @@ -0,0 +1,79 @@ +from logging.config import fileConfig + +from sqlalchemy import engine_from_config +from sqlalchemy import pool + +from alembic import context + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = None + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure( + connection=connection, target_metadata=target_metadata + ) + + with context.begin_transaction(): + context.run_migrations() + + +# if context.is_offline_mode(): +# run_migrations_offline() +# else: +# run_migrations_online() +run_migrations_offline() \ No newline at end of file diff --git a/backend/alembic/script.py.mako b/backend/alembic/script.py.mako new file mode 100644 index 00000000..fbc4b07d --- /dev/null +++ b/backend/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/backend/backend/__init__.py b/backend/backend/__init__.py deleted file mode 100644 index 9c9b9534..00000000 --- a/backend/backend/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -__version__ = "0.1.0" -__version_info__ = tuple(int(num) if num.isdigit() else num for num in __version__.replace("-", ".", 1).split(".")) diff --git a/backend/backend/conftest.py b/backend/backend/conftest.py deleted file mode 100644 index 627ea6d0..00000000 --- a/backend/backend/conftest.py +++ /dev/null @@ -1,14 +0,0 @@ -import pytest - -from backend.users.models import User -from backend.users.tests.factories import UserFactory - - -@pytest.fixture(autouse=True) -def media_storage(settings, tmpdir): - settings.MEDIA_ROOT = tmpdir.strpath - - -@pytest.fixture -def user(db) -> User: - return UserFactory() diff --git a/backend/backend/contrib/__init__.py b/backend/backend/contrib/__init__.py deleted file mode 100644 index 1c7ecc89..00000000 --- a/backend/backend/contrib/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -""" -To understand why this file is here, please read: - -http://cookiecutter-django.readthedocs.io/en/latest/faq.html#why-is-there-a-django-contrib-sites-directory-in-cookiecutter-django -""" diff --git a/backend/backend/contrib/sites/__init__.py b/backend/backend/contrib/sites/__init__.py deleted file mode 100644 index 1c7ecc89..00000000 --- a/backend/backend/contrib/sites/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -""" -To understand why this file is here, please read: - -http://cookiecutter-django.readthedocs.io/en/latest/faq.html#why-is-there-a-django-contrib-sites-directory-in-cookiecutter-django -""" diff --git a/backend/backend/contrib/sites/migrations/0001_initial.py b/backend/backend/contrib/sites/migrations/0001_initial.py deleted file mode 100644 index 304cd6d7..00000000 --- a/backend/backend/contrib/sites/migrations/0001_initial.py +++ /dev/null @@ -1,42 +0,0 @@ -import django.contrib.sites.models -from django.contrib.sites.models import _simple_domain_name_validator -from django.db import migrations, models - - -class Migration(migrations.Migration): - - dependencies = [] - - operations = [ - migrations.CreateModel( - name="Site", - fields=[ - ( - "id", - models.AutoField( - verbose_name="ID", - serialize=False, - auto_created=True, - primary_key=True, - ), - ), - ( - "domain", - models.CharField( - max_length=100, - verbose_name="domain name", - validators=[_simple_domain_name_validator], - ), - ), - ("name", models.CharField(max_length=50, verbose_name="display name")), - ], - options={ - "ordering": ("domain",), - "db_table": "django_site", - "verbose_name": "site", - "verbose_name_plural": "sites", - }, - bases=(models.Model,), - managers=[("objects", django.contrib.sites.models.SiteManager())], - ) - ] diff --git a/backend/backend/contrib/sites/migrations/0002_alter_domain_unique.py b/backend/backend/contrib/sites/migrations/0002_alter_domain_unique.py deleted file mode 100644 index 2c8d6dac..00000000 --- a/backend/backend/contrib/sites/migrations/0002_alter_domain_unique.py +++ /dev/null @@ -1,20 +0,0 @@ -import django.contrib.sites.models -from django.db import migrations, models - - -class Migration(migrations.Migration): - - dependencies = [("sites", "0001_initial")] - - operations = [ - migrations.AlterField( - model_name="site", - name="domain", - field=models.CharField( - max_length=100, - unique=True, - validators=[django.contrib.sites.models._simple_domain_name_validator], - verbose_name="domain name", - ), - ) - ] diff --git a/backend/backend/contrib/sites/migrations/0003_set_site_domain_and_name.py b/backend/backend/contrib/sites/migrations/0003_set_site_domain_and_name.py deleted file mode 100644 index 0228bd79..00000000 --- a/backend/backend/contrib/sites/migrations/0003_set_site_domain_and_name.py +++ /dev/null @@ -1,63 +0,0 @@ -""" -To understand why this file is here, please read: - -http://cookiecutter-django.readthedocs.io/en/latest/faq.html#why-is-there-a-django-contrib-sites-directory-in-cookiecutter-django -""" -from django.conf import settings -from django.db import migrations - - -def _update_or_create_site_with_sequence(site_model, connection, domain, name): - """Update or create the site with default ID and keep the DB sequence in sync.""" - site, created = site_model.objects.update_or_create( - id=settings.SITE_ID, - defaults={ - "domain": domain, - "name": name, - }, - ) - if created: - # We provided the ID explicitly when creating the Site entry, therefore the DB - # sequence to auto-generate them wasn't used and is now out of sync. If we - # don't do anything, we'll get a unique constraint violation the next time a - # site is created. - # To avoid this, we need to manually update DB sequence and make sure it's - # greater than the maximum value. - max_id = site_model.objects.order_by('-id').first().id - with connection.cursor() as cursor: - cursor.execute("SELECT last_value from django_site_id_seq") - (current_id,) = cursor.fetchone() - if current_id <= max_id: - cursor.execute( - "alter sequence django_site_id_seq restart with %s", - [max_id + 1], - ) - - -def update_site_forward(apps, schema_editor): - """Set site domain and name.""" - Site = apps.get_model("sites", "Site") - _update_or_create_site_with_sequence( - Site, - schema_editor.connection, - "example.com", - "backend", - ) - - -def update_site_backward(apps, schema_editor): - """Revert site domain and name to default.""" - Site = apps.get_model("sites", "Site") - _update_or_create_site_with_sequence( - Site, - schema_editor.connection, - "example.com", - "example.com", - ) - - -class Migration(migrations.Migration): - - dependencies = [("sites", "0002_alter_domain_unique")] - - operations = [migrations.RunPython(update_site_forward, update_site_backward)] diff --git a/backend/backend/contrib/sites/migrations/0004_alter_options_ordering_domain.py b/backend/backend/contrib/sites/migrations/0004_alter_options_ordering_domain.py deleted file mode 100644 index f7118ca8..00000000 --- a/backend/backend/contrib/sites/migrations/0004_alter_options_ordering_domain.py +++ /dev/null @@ -1,21 +0,0 @@ -# Generated by Django 3.1.7 on 2021-02-04 14:49 - -from django.db import migrations - - -class Migration(migrations.Migration): - - dependencies = [ - ("sites", "0003_set_site_domain_and_name"), - ] - - operations = [ - migrations.AlterModelOptions( - name="site", - options={ - "ordering": ["domain"], - "verbose_name": "site", - "verbose_name_plural": "sites", - }, - ), - ] diff --git a/backend/backend/contrib/sites/migrations/__init__.py b/backend/backend/contrib/sites/migrations/__init__.py deleted file mode 100644 index 1c7ecc89..00000000 --- a/backend/backend/contrib/sites/migrations/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -""" -To understand why this file is here, please read: - -http://cookiecutter-django.readthedocs.io/en/latest/faq.html#why-is-there-a-django-contrib-sites-directory-in-cookiecutter-django -""" diff --git a/backend/backend/data_analysis/__init__.py b/backend/backend/data_analysis/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/data_analysis/admin.py b/backend/backend/data_analysis/admin.py deleted file mode 100644 index cc376486..00000000 --- a/backend/backend/data_analysis/admin.py +++ /dev/null @@ -1,5 +0,0 @@ -from django.contrib import admin -from backend.users.models import Company -from .models import TransactionRecordModel -# Register your models here. -admin.site.register(TransactionRecordModel) diff --git a/backend/backend/data_analysis/apps.py b/backend/backend/data_analysis/apps.py deleted file mode 100644 index 01773768..00000000 --- a/backend/backend/data_analysis/apps.py +++ /dev/null @@ -1,6 +0,0 @@ -from django.apps import AppConfig - - -class DataAnalysisConfig(AppConfig): - default_auto_field = "django.db.models.BigAutoField" - name = "backend.data_analysis" diff --git a/backend/backend/data_analysis/data_connectors/__init__.py b/backend/backend/data_analysis/data_connectors/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/data_analysis/data_connectors/urls.py b/backend/backend/data_analysis/data_connectors/urls.py deleted file mode 100644 index 8c960a13..00000000 --- a/backend/backend/data_analysis/data_connectors/urls.py +++ /dev/null @@ -1,2 +0,0 @@ -url_patterns = [ -] diff --git a/backend/backend/data_analysis/data_connectors/views.py b/backend/backend/data_analysis/data_connectors/views.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/data_analysis/migrations/0001_initial.py b/backend/backend/data_analysis/migrations/0001_initial.py deleted file mode 100644 index e2361272..00000000 --- a/backend/backend/data_analysis/migrations/0001_initial.py +++ /dev/null @@ -1,85 +0,0 @@ -# Generated by Django 4.2.6 on 2023-10-09 23:52 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - initial = True - - dependencies = [] - - operations = [ - migrations.CreateModel( - name="Dashboard", - fields=[ - ("dashboard_id", models.AutoField(primary_key=True, serialize=False)), - ("dashboard_name", models.CharField(max_length=255)), - ], - ), - migrations.CreateModel( - name="DataProduct", - fields=[ - ("product_id", models.AutoField(primary_key=True, serialize=False)), - ("product_name", models.CharField(max_length=255)), - ("product_description", models.TextField()), - ], - ), - migrations.CreateModel( - name="Metric", - fields=[ - ("metric_id", models.AutoField(primary_key=True, serialize=False)), - ("metric_name", models.CharField(max_length=255)), - ("metric_configuration", models.TextField()), - ("likes", models.IntegerField(default=0)), - ], - ), - migrations.CreateModel( - name="MetricComment", - fields=[ - ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")), - ("comment", models.TextField()), - ], - ), - migrations.CreateModel( - name="Project", - fields=[ - ("project_id", models.AutoField(primary_key=True, serialize=False)), - ("project_name", models.CharField(max_length=255)), - ], - ), - migrations.CreateModel( - name="Space", - fields=[ - ("space_id", models.AutoField(primary_key=True, serialize=False)), - ("space_name", models.CharField(max_length=255)), - ("space_description", models.TextField()), - ("space_url", models.CharField(max_length=255)), - ("archived", models.BooleanField(default=False)), - ], - ), - migrations.CreateModel( - name="Team", - fields=[ - ("team_id", models.AutoField(primary_key=True, serialize=False)), - ("team_name", models.CharField(max_length=255)), - ("team_description", models.TextField()), - ("team_logo", models.ImageField(blank=True, null=True, upload_to="team_logos/")), - ], - ), - migrations.CreateModel( - name="TeamMember", - fields=[ - ("team_member_id", models.AutoField(primary_key=True, serialize=False)), - ("member_role", models.CharField(choices=[("Admin", "Admin"), ("Member", "Member")], max_length=50)), - ], - ), - migrations.CreateModel( - name="Workflow", - fields=[ - ("workflow_id", models.AutoField(primary_key=True, serialize=False)), - ("workflow_name", models.CharField(max_length=255)), - ("workflow_status", models.CharField(max_length=255)), - ("alert_flag", models.BooleanField(default=False)), - ], - ), - ] diff --git a/backend/backend/data_analysis/migrations/0002_initial.py b/backend/backend/data_analysis/migrations/0002_initial.py deleted file mode 100644 index bd481307..00000000 --- a/backend/backend/data_analysis/migrations/0002_initial.py +++ /dev/null @@ -1,89 +0,0 @@ -# Generated by Django 4.2.6 on 2023-10-09 23:52 - -from django.conf import settings -from django.db import migrations, models -import django.db.models.deletion - - -class Migration(migrations.Migration): - initial = True - - dependencies = [ - migrations.swappable_dependency(settings.AUTH_USER_MODEL), - ("data_analysis", "0001_initial"), - ] - - operations = [ - migrations.AddField( - model_name="workflow", - name="reviewed_by", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="teammember", - name="member_id", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="teammember", - name="team_id", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="data_analysis.team"), - ), - migrations.AddField( - model_name="team", - name="admin_id", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="space", - name="edit_permissions", - field=models.ManyToManyField(related_name="edit_permissions", to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="project", - name="team_id", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="data_analysis.team"), - ), - migrations.AddField( - model_name="project", - name="user_id", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="metriccomment", - name="metric", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="data_analysis.metric"), - ), - migrations.AddField( - model_name="metriccomment", - name="user", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="metric", - name="comments", - field=models.ManyToManyField( - related_name="metric_comments", through="data_analysis.MetricComment", to=settings.AUTH_USER_MODEL - ), - ), - migrations.AddField( - model_name="metric", - name="refined_by", - field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="metric", - name="tags", - field=models.ManyToManyField(related_name="metric_tags", to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="dataproduct", - name="uploaded_data_sources", - field=models.ManyToManyField(related_name="uploaded_data_sources", to=settings.AUTH_USER_MODEL), - ), - migrations.AddField( - model_name="dashboard", - name="metrics_used", - field=models.ManyToManyField(to="data_analysis.metric"), - ), - ] diff --git a/backend/backend/data_analysis/migrations/0003_transactionrecordmodel_and_more.py b/backend/backend/data_analysis/migrations/0003_transactionrecordmodel_and_more.py deleted file mode 100644 index 6414cec8..00000000 --- a/backend/backend/data_analysis/migrations/0003_transactionrecordmodel_and_more.py +++ /dev/null @@ -1,104 +0,0 @@ -# Generated by Django 4.2.6 on 2024-02-28 17:00 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ("data_analysis", "0002_initial"), - ] - - operations = [ - migrations.CreateModel( - name="TransactionRecordModel", - fields=[ - ("transactionId", models.AutoField(primary_key=True, serialize=False)), - ("payType", models.CharField(max_length=255)), - ("amount", models.FloatField()), - ("destinationAddress", models.CharField(max_length=100)), - ("txHash", models.CharField(max_length=100)), - ("created_at", models.DateTimeField(auto_now_add=True)), - ("updated_at", models.DateTimeField(auto_now=True)), - ("deleted_at", models.DateTimeField(blank=True, null=True)), - ], - ), - migrations.RemoveField( - model_name="dataproduct", - name="uploaded_data_sources", - ), - migrations.RemoveField( - model_name="metric", - name="comments", - ), - migrations.RemoveField( - model_name="metric", - name="refined_by", - ), - migrations.RemoveField( - model_name="metric", - name="tags", - ), - migrations.RemoveField( - model_name="metriccomment", - name="metric", - ), - migrations.RemoveField( - model_name="metriccomment", - name="user", - ), - migrations.RemoveField( - model_name="project", - name="team_id", - ), - migrations.RemoveField( - model_name="project", - name="user_id", - ), - migrations.RemoveField( - model_name="space", - name="edit_permissions", - ), - migrations.RemoveField( - model_name="team", - name="admin_id", - ), - migrations.RemoveField( - model_name="teammember", - name="member_id", - ), - migrations.RemoveField( - model_name="teammember", - name="team_id", - ), - migrations.RemoveField( - model_name="workflow", - name="reviewed_by", - ), - migrations.DeleteModel( - name="Dashboard", - ), - migrations.DeleteModel( - name="DataProduct", - ), - migrations.DeleteModel( - name="Metric", - ), - migrations.DeleteModel( - name="MetricComment", - ), - migrations.DeleteModel( - name="Project", - ), - migrations.DeleteModel( - name="Space", - ), - migrations.DeleteModel( - name="Team", - ), - migrations.DeleteModel( - name="TeamMember", - ), - migrations.DeleteModel( - name="Workflow", - ), - ] diff --git a/backend/backend/data_analysis/migrations/__init__.py b/backend/backend/data_analysis/migrations/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/data_analysis/models.py b/backend/backend/data_analysis/models.py deleted file mode 100644 index f651a658..00000000 --- a/backend/backend/data_analysis/models.py +++ /dev/null @@ -1,22 +0,0 @@ -from django.db import models -from django.utils.translation import gettext_lazy as _ - -from backend.users.models import User -from backend.data_analysis import utils -from django.contrib.postgres.fields import ArrayField -from django.db.models import JSONField -from datetime import datetime - -class TransactionRecordModel(models.Model): - transactionId = models.AutoField(primary_key=True) - payType = models.CharField(max_length=255) - amount = models.FloatField() # Change from Decimal to FloatField - destinationAddress = models.CharField(max_length=100) - txHash = models.CharField(max_length=100) - - created_at = models.DateTimeField(auto_now_add=True) - updated_at = models.DateTimeField(auto_now=True) - deleted_at = models.DateTimeField(blank=True, null=True) - - def __str__(self): - return str(self.transactionId) diff --git a/backend/backend/data_analysis/serializers.py b/backend/backend/data_analysis/serializers.py deleted file mode 100644 index dea767a8..00000000 --- a/backend/backend/data_analysis/serializers.py +++ /dev/null @@ -1,10 +0,0 @@ -from rest_framework import serializers -from backend.data_analysis.models import ( - TransactionRecordModel -) - - -class TransactionRecordSerializer(serializers.ModelSerializer): - class Meta: - model = TransactionRecordModel - fields = "__all__" diff --git a/backend/backend/data_analysis/tests.py b/backend/backend/data_analysis/tests.py deleted file mode 100644 index 7ce503c2..00000000 --- a/backend/backend/data_analysis/tests.py +++ /dev/null @@ -1,3 +0,0 @@ -from django.test import TestCase - -# Create your tests here. diff --git a/backend/backend/data_analysis/urls.py b/backend/backend/data_analysis/urls.py deleted file mode 100644 index 4fa25dd5..00000000 --- a/backend/backend/data_analysis/urls.py +++ /dev/null @@ -1,9 +0,0 @@ -from django.urls import path - -from backend.data_analysis.views import ( - SaveTransaction, -) - -urlpatterns = [ - path("saveTransaction/", view=SaveTransaction.as_view(), name="SaveTransaction"), -] diff --git a/backend/backend/data_analysis/utils.py b/backend/backend/data_analysis/utils.py deleted file mode 100644 index 90b789db..00000000 --- a/backend/backend/data_analysis/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -# utils.py - -from cryptography.fernet import Fernet - -def encrypt_data(data, key): - cipher_suite = Fernet(key) - cipher_text = cipher_suite.encrypt(data.encode()) - return cipher_text - -def decrypt_data(cipher_text, key): - cipher_suite = Fernet(key) - plain_text = cipher_suite.decrypt(cipher_text).decode() - return plain_text - -# Generate a new key using cryptography library -from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC -from cryptography.hazmat.backends import default_backend -from cryptography.hazmat.primitives import hashes # Add this import -from base64 import urlsafe_b64encode - -def generate_key(): - password = 'P@ssw0rd1997' # Change this to a strong, secret password - salt = b'salt_123' # Change this to a random, unique salt - kdf = PBKDF2HMAC( - algorithm=hashes.SHA256(), - iterations=100000, - salt=salt, - length=32, - backend=default_backend() - ) - key = urlsafe_b64encode(kdf.derive(password.encode())) - return key - -SECRET_KEY = generate_key() diff --git a/backend/backend/data_analysis/views.py b/backend/backend/data_analysis/views.py deleted file mode 100644 index e27fcd22..00000000 --- a/backend/backend/data_analysis/views.py +++ /dev/null @@ -1,63 +0,0 @@ -import os -import requests -from django.http import JsonResponse - -from rest_framework import status -from rest_framework.response import Response -from rest_framework.parsers import JSONParser -from rest_framework import generics -from rest_framework.exceptions import NotFound - -from rest_framework import permissions, status -from rest_framework.response import Response -from rest_framework.views import APIView - -# import mysql.connector - -from backend.data_analysis.models import ( - TransactionRecordModel, -) -from backend.data_analysis.serializers import ( - TransactionRecordSerializer, -) - - -def send_mail(email, content): - email_params = { - "apikey": os.getenv("MAIL_KEY"), - "from": email, - "to": "tunki1201@gmail.com", - "subject": "This is the content", - "body": f"{content}", - "isTransactional": True, - } - - response = requests.post( - "https://api.elasticemail.com/v2/email/send", - data=email_params, - ) - return response - -class SaveTransaction(APIView): - permission_classes = (permissions.AllowAny,) - authentication_classes = () - - def post(self, request): - payType = request.data.get("payType") - amount = request.data.get("amount") - destinationAddress = request.data.get("destinationAddress") - txHash = request.data.get("txHash") - # Create a new TransactionRecordModel instance and save it - try: - transaction_record = TransactionRecordModel.objects.create( - payType=payType, - amount=amount, - destinationAddress=destinationAddress, - txHash=txHash, - ) - - saved_transactionRecord = TransactionRecordSerializer(transaction_record) - - return Response({"message": "Transaction saved successfully"}, status=status.HTTP_201_CREATED) - except Exception as e: - return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) diff --git a/backend/backend/static/css/project.css b/backend/backend/static/css/project.css deleted file mode 100644 index f1d543da..00000000 --- a/backend/backend/static/css/project.css +++ /dev/null @@ -1,13 +0,0 @@ -/* These styles are generated from project.scss. */ - -.alert-debug { - color: black; - background-color: white; - border-color: #d6e9c6; -} - -.alert-error { - color: #b94a48; - background-color: #f2dede; - border-color: #eed3d7; -} diff --git a/backend/backend/static/fonts/.gitkeep b/backend/backend/static/fonts/.gitkeep deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/static/images/favicons/favicon.ico b/backend/backend/static/images/favicons/favicon.ico deleted file mode 100644 index e1c1dd1a..00000000 Binary files a/backend/backend/static/images/favicons/favicon.ico and /dev/null differ diff --git a/backend/backend/static/js/project.js b/backend/backend/static/js/project.js deleted file mode 100644 index d26d23b9..00000000 --- a/backend/backend/static/js/project.js +++ /dev/null @@ -1 +0,0 @@ -/* Project specific Javascript goes here. */ diff --git a/backend/backend/templates/403.html b/backend/backend/templates/403.html deleted file mode 100644 index 57956032..00000000 --- a/backend/backend/templates/403.html +++ /dev/null @@ -1,13 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Forbidden (403){% endblock title %} -{% block content %} -

Forbidden (403)

-

- {% if exception %} - {{ exception }} - {% else %} - You're not allowed to access this page. - {% endif %} -

-{% endblock content %} diff --git a/backend/backend/templates/403_csrf.html b/backend/backend/templates/403_csrf.html deleted file mode 100644 index 57956032..00000000 --- a/backend/backend/templates/403_csrf.html +++ /dev/null @@ -1,13 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Forbidden (403){% endblock title %} -{% block content %} -

Forbidden (403)

-

- {% if exception %} - {{ exception }} - {% else %} - You're not allowed to access this page. - {% endif %} -

-{% endblock content %} diff --git a/backend/backend/templates/404.html b/backend/backend/templates/404.html deleted file mode 100644 index 5111d3aa..00000000 --- a/backend/backend/templates/404.html +++ /dev/null @@ -1,13 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Page not found{% endblock title %} -{% block content %} -

Page not found

-

- {% if exception %} - {{ exception }} - {% else %} - This is not the page you were looking for. - {% endif %} -

-{% endblock content %} diff --git a/backend/backend/templates/500.html b/backend/backend/templates/500.html deleted file mode 100644 index 1add004d..00000000 --- a/backend/backend/templates/500.html +++ /dev/null @@ -1,10 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Server Error{% endblock title %} -{% block content %} -

Ooops!!! 500

-

Looks like something went wrong!

-

- We track these errors automatically, but if the problem persists feel free to contact us. In the meantime, try refreshing. -

-{% endblock content %} diff --git a/backend/backend/templates/account/account_inactive.html b/backend/backend/templates/account/account_inactive.html deleted file mode 100644 index 036adb40..00000000 --- a/backend/backend/templates/account/account_inactive.html +++ /dev/null @@ -1,11 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Account Inactive" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Account Inactive" %}

-

{% translate "This account is inactive." %}

-{% endblock inner %} diff --git a/backend/backend/templates/account/base.html b/backend/backend/templates/account/base.html deleted file mode 100644 index c336c1cb..00000000 --- a/backend/backend/templates/account/base.html +++ /dev/null @@ -1,13 +0,0 @@ -{% extends "base.html" %} - -{% block title %} - {% block head_title %} - {% endblock head_title %} -{% endblock title %} -{% block content %} -
-
- {% block inner %}{% endblock inner %} -
-
-{% endblock content %} diff --git a/backend/backend/templates/account/email.html b/backend/backend/templates/account/email.html deleted file mode 100644 index 440f80c5..00000000 --- a/backend/backend/templates/account/email.html +++ /dev/null @@ -1,79 +0,0 @@ - -{% extends "account/base.html" %} - -{% load i18n %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Account" %} -{% endblock head_title %} -{% block inner %} -

{% translate "E-mail Addresses" %}

- {% if user.emailaddress_set.all %} -

{% translate "The following e-mail addresses are associated with your account:" %}

-
- {% csrf_token %} -
- {% for emailaddress in user.emailaddress_set.all %} -
- -
- {% endfor %} -
- - - -
-
-
- {% else %} -

- {% translate "Warning:" %} {% translate "You currently do not have any e-mail address set up. You should really add an e-mail address so you can receive notifications, reset your password, etc." %} -

- {% endif %} -

{% translate "Add E-mail Address" %}

-
- {% csrf_token %} - {{ form|crispy }} - -
-{% endblock inner %} -{% block inline_javascript %} - {{ block.super }} - -{% endblock inline_javascript %} diff --git a/backend/backend/templates/account/email_confirm.html b/backend/backend/templates/account/email_confirm.html deleted file mode 100644 index 95ef44f7..00000000 --- a/backend/backend/templates/account/email_confirm.html +++ /dev/null @@ -1,27 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load account %} - -{% block head_title %} - {% translate "Confirm E-mail Address" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Confirm E-mail Address" %}

- {% if confirmation %} - {% user_display confirmation.email_address.user as user_display %} -

- {% blocktranslate with confirmation.email_address.email as email %}Please confirm that {{ email }} is an e-mail address for user {{ user_display }}.{% endblocktranslate %} -

-
- {% csrf_token %} - -
- {% else %} - {% url 'account_email' as email_url %} -

- {% blocktranslate %}This e-mail confirmation link expired or is invalid. Please issue a new e-mail confirmation request.{% endblocktranslate %} -

- {% endif %} -{% endblock inner %} diff --git a/backend/backend/templates/account/login.html b/backend/backend/templates/account/login.html deleted file mode 100644 index ca182952..00000000 --- a/backend/backend/templates/account/login.html +++ /dev/null @@ -1,52 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load account socialaccount %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Sign In" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Sign In" %}

- {% get_providers as socialaccount_providers %} - {% if socialaccount_providers %} -

- {% translate "Please sign in with one of your existing third party accounts:" %} - {% if ACCOUNT_ALLOW_REGISTRATION %} - {% blocktranslate trimmed %} - Or, sign up - for a {{ site_name }} account and sign in below: - {% endblocktranslate %} - {% endif %} -

-
- -
{% translate "or" %}
-
- {% include "socialaccount/snippets/login_extra.html" %} - {% else %} - {% if ACCOUNT_ALLOW_REGISTRATION %} -

- {% blocktranslate trimmed %} - If you have not created an account yet, then please - sign up first. - {% endblocktranslate %} -

- {% endif %} - {% endif %} -
- {% csrf_token %} - {{ form|crispy }} - {% if redirect_field_value %} - - {% endif %} - {% translate "Forgot Password?" %} - -
-{% endblock inner %} diff --git a/backend/backend/templates/account/logout.html b/backend/backend/templates/account/logout.html deleted file mode 100644 index e3140206..00000000 --- a/backend/backend/templates/account/logout.html +++ /dev/null @@ -1,20 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Sign Out" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Sign Out" %}

-

{% translate "Are you sure you want to sign out?" %}

-
- {% csrf_token %} - {% if redirect_field_value %} - - {% endif %} - -
-{% endblock inner %} diff --git a/backend/backend/templates/account/password_change.html b/backend/backend/templates/account/password_change.html deleted file mode 100644 index 029aa549..00000000 --- a/backend/backend/templates/account/password_change.html +++ /dev/null @@ -1,18 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Change Password" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Change Password" %}

-
- {% csrf_token %} - {{ form|crispy }} - -
-{% endblock inner %} diff --git a/backend/backend/templates/account/password_reset.html b/backend/backend/templates/account/password_reset.html deleted file mode 100644 index ee35e75e..00000000 --- a/backend/backend/templates/account/password_reset.html +++ /dev/null @@ -1,28 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load account %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Password Reset" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Password Reset" %}

- {% if user.is_authenticated %} - {% include "account/snippets/already_logged_in.html" %} - {% endif %} -

- {% translate "Forgotten your password? Enter your e-mail address below, and we'll send you an e-mail allowing you to reset it." %} -

-
- {% csrf_token %} - {{ form|crispy }} - -
-

{% blocktranslate %}Please contact us if you have any trouble resetting your password.{% endblocktranslate %}

-{% endblock inner %} diff --git a/backend/backend/templates/account/password_reset_done.html b/backend/backend/templates/account/password_reset_done.html deleted file mode 100644 index 8efc4aa2..00000000 --- a/backend/backend/templates/account/password_reset_done.html +++ /dev/null @@ -1,17 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load account %} - -{% block head_title %} - {% translate "Password Reset" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Password Reset" %}

- {% if user.is_authenticated %} - {% include "account/snippets/already_logged_in.html" %} - {% endif %} -

- {% blocktranslate %}We have sent you an e-mail. Please contact us if you do not receive it within a few minutes.{% endblocktranslate %} -

-{% endblock inner %} diff --git a/backend/backend/templates/account/password_reset_from_key.html b/backend/backend/templates/account/password_reset_from_key.html deleted file mode 100644 index b880bf99..00000000 --- a/backend/backend/templates/account/password_reset_from_key.html +++ /dev/null @@ -1,36 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Change Password" %} -{% endblock head_title %} -{% block inner %} -

- {% if token_fail %} - {% translate "Bad Token" %} - {% else %} - {% translate "Change Password" %} - {% endif %} -

- {% if token_fail %} - {% url 'account_reset_password' as passwd_reset_url %} -

- {% blocktranslate %}The password reset link was invalid, possibly because it has already been used. Please request a new password reset.{% endblocktranslate %} -

- {% else %} - {% if form %} -
- {% csrf_token %} - {{ form|crispy }} - -
- {% else %} -

{% translate "Your password is now changed." %}

- {% endif %} - {% endif %} -{% endblock inner %} diff --git a/backend/backend/templates/account/password_reset_from_key_done.html b/backend/backend/templates/account/password_reset_from_key_done.html deleted file mode 100644 index 28db548c..00000000 --- a/backend/backend/templates/account/password_reset_from_key_done.html +++ /dev/null @@ -1,11 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Change Password" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Change Password" %}

-

{% translate "Your password is now changed." %}

-{% endblock inner %} diff --git a/backend/backend/templates/account/password_set.html b/backend/backend/templates/account/password_set.html deleted file mode 100644 index 56fb6188..00000000 --- a/backend/backend/templates/account/password_set.html +++ /dev/null @@ -1,21 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Set Password" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Set Password" %}

-
- {% csrf_token %} - {{ form|crispy }} - -
-{% endblock inner %} diff --git a/backend/backend/templates/account/signup.html b/backend/backend/templates/account/signup.html deleted file mode 100644 index 04bcbf9a..00000000 --- a/backend/backend/templates/account/signup.html +++ /dev/null @@ -1,27 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} -{% load crispy_forms_tags %} - -{% block head_title %} - {% translate "Signup" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Sign Up" %}

-

- {% blocktranslate %}Already have an account? Then please sign in.{% endblocktranslate %} -

-
- {% csrf_token %} - {{ form|crispy }} - {% if redirect_field_value %} - - {% endif %} - -
-{% endblock inner %} diff --git a/backend/backend/templates/account/signup_closed.html b/backend/backend/templates/account/signup_closed.html deleted file mode 100644 index 94c4327e..00000000 --- a/backend/backend/templates/account/signup_closed.html +++ /dev/null @@ -1,11 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Sign Up Closed" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Sign Up Closed" %}

-

{% translate "We are sorry, but the sign up is currently closed." %}

-{% endblock inner %} diff --git a/backend/backend/templates/account/verification_sent.html b/backend/backend/templates/account/verification_sent.html deleted file mode 100644 index de2a33f3..00000000 --- a/backend/backend/templates/account/verification_sent.html +++ /dev/null @@ -1,13 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Verify Your E-mail Address" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Verify Your E-mail Address" %}

-

- {% blocktranslate %}We have sent an e-mail to you for verification. Follow the link provided to finalize the signup process. Please contact us if you do not receive it within a few minutes.{% endblocktranslate %} -

-{% endblock inner %} diff --git a/backend/backend/templates/account/verified_email_required.html b/backend/backend/templates/account/verified_email_required.html deleted file mode 100644 index aab0877f..00000000 --- a/backend/backend/templates/account/verified_email_required.html +++ /dev/null @@ -1,24 +0,0 @@ -{% extends "account/base.html" %} - -{% load i18n %} - -{% block head_title %} - {% translate "Verify Your E-mail Address" %} -{% endblock head_title %} -{% block inner %} -

{% translate "Verify Your E-mail Address" %}

- {% url 'account_email' as email_url %} -

- {% blocktranslate %}This part of the site requires us to verify that -you are who you claim to be. For this purpose, we require that you -verify ownership of your e-mail address. {% endblocktranslate %} -

-

- {% blocktranslate %}We have sent an e-mail to you for -verification. Please click on the link inside this e-mail. Please -contact us if you do not receive it within a few minutes.{% endblocktranslate %} -

-

- {% blocktranslate %}Note: you can still change your e-mail address.{% endblocktranslate %} -

-{% endblock inner %} diff --git a/backend/backend/templates/base.html b/backend/backend/templates/base.html deleted file mode 100644 index 3ab2cb9e..00000000 --- a/backend/backend/templates/base.html +++ /dev/null @@ -1,136 +0,0 @@ - -{% load static i18n %} -{% get_current_language as LANGUAGE_CODE %} - - - - - - {% block title %} - backend - {% endblock title %} - - - - - - {% block css %} - - - - - - - - - -{% endblock css %} - -{# Placed at the top of the document so pages load faster with defer #} -{% block javascript %} - - - - - - - - - - - - - -{% endblock javascript %} - - -
- -
-
- {% if messages %} - {% for message in messages %} -
- {{ message }} - -
- {% endfor %} - {% endif %} - {% block content %} -

Use this document as a way to quick start any new project.

- {% endblock content %} -
- - {% block modal %} - {% endblock modal %} - {% block inline_javascript %} - {% comment %} - Script tags with only code, no src (defer by default). To run - with a "defer" so that you run inline code: - - {% endcomment %} - {% endblock inline_javascript %} - - diff --git a/backend/backend/templates/pages/about.html b/backend/backend/templates/pages/about.html deleted file mode 100644 index 30703483..00000000 --- a/backend/backend/templates/pages/about.html +++ /dev/null @@ -1,3 +0,0 @@ -{% extends "base.html" %} - - diff --git a/backend/backend/templates/pages/home.html b/backend/backend/templates/pages/home.html deleted file mode 100644 index 30703483..00000000 --- a/backend/backend/templates/pages/home.html +++ /dev/null @@ -1,3 +0,0 @@ -{% extends "base.html" %} - - diff --git a/backend/backend/templates/users/user_detail.html b/backend/backend/templates/users/user_detail.html deleted file mode 100644 index 3edbaa69..00000000 --- a/backend/backend/templates/users/user_detail.html +++ /dev/null @@ -1,38 +0,0 @@ -{% extends "base.html" %} - -{% load static %} - -{% block title %} - User: - - {{ object.name }} - - -{% endblock title %} -{% block content %} -
-
-
-

- - - {{ object.name }} - -

-
-
- {% if object == request.user %} - -
-
- My Info - E-Mail - -
-
- - {% endif %} -
-{% endblock content %} diff --git a/backend/backend/templates/users/user_form.html b/backend/backend/templates/users/user_form.html deleted file mode 100644 index 8e233270..00000000 --- a/backend/backend/templates/users/user_form.html +++ /dev/null @@ -1,31 +0,0 @@ -{% extends "base.html" %} - -{% load crispy_forms_tags %} - -{% block title %} - - - {{ user.name }} - - -{% endblock title %} -{% block content %} -

- - - {{ user.name }} - - -

-
- {% csrf_token %} - {{ form|crispy }} -
-
- -
-
-
-{% endblock content %} diff --git a/backend/backend/users/__init__.py b/backend/backend/users/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/users/adapters.py b/backend/backend/users/adapters.py deleted file mode 100644 index 0f82f1e5..00000000 --- a/backend/backend/users/adapters.py +++ /dev/null @@ -1,37 +0,0 @@ -from __future__ import annotations - -import typing - -from allauth.account.adapter import DefaultAccountAdapter -from allauth.socialaccount.adapter import DefaultSocialAccountAdapter -from django.conf import settings -from django.http import HttpRequest - -if typing.TYPE_CHECKING: - from allauth.socialaccount.models import SocialLogin - from backend.users.models import User - - -class AccountAdapter(DefaultAccountAdapter): - def is_open_for_signup(self, request: HttpRequest) -> bool: - return getattr(settings, "ACCOUNT_ALLOW_REGISTRATION", True) - - -class SocialAccountAdapter(DefaultSocialAccountAdapter): - def is_open_for_signup(self, request: HttpRequest, sociallogin: SocialLogin) -> bool: - return getattr(settings, "ACCOUNT_ALLOW_REGISTRATION", True) - - def populate_user(self, request: HttpRequest, sociallogin: SocialLogin, data: dict[str, typing.Any]) -> User: - """ - Populates user information from social provider info. - - See: https://django-allauth.readthedocs.io/en/latest/advanced.html?#creating-and-populating-user-instances - """ - user = sociallogin.user - if name := data.get("name"): - user.name = name - elif first_name := data.get("first_name"): - user.name = first_name - if last_name := data.get("last_name"): - user.name += f" {last_name}" - return super().populate_user(request, sociallogin, data) diff --git a/backend/backend/users/admin.py b/backend/backend/users/admin.py deleted file mode 100644 index 00f683b6..00000000 --- a/backend/backend/users/admin.py +++ /dev/null @@ -1,49 +0,0 @@ -from django.conf import settings -from django.contrib import admin -from django.contrib.auth import admin as auth_admin -from django.contrib.auth import get_user_model, decorators -from django.utils.translation import gettext_lazy as _ - -from backend.users.forms import UserAdminChangeForm, UserAdminCreationForm - -User = get_user_model() - -if settings.DJANGO_ADMIN_FORCE_ALLAUTH: - # Force the `admin` sign in process to go through the `django-allauth` workflow: - # https://django-allauth.readthedocs.io/en/stable/advanced.html#admin - admin.site.login = decorators.login_required(admin.site.login) # type: ignore[method-assign] - - -@admin.register(User) -class UserAdmin(auth_admin.UserAdmin): - form = UserAdminChangeForm - add_form = UserAdminCreationForm - fieldsets = ( - (None, {"fields": ("email", "password")}), - (_("Personal info"), {"fields": ("name",)}), - ( - _("Permissions"), - { - "fields": ( - "is_active", - "is_staff", - "is_superuser", - "groups", - "user_permissions", - ), - }, - ), - (_("Important dates"), {"fields": ("last_login", "date_joined")}), - ) - list_display = ["email", "name", "is_superuser"] - search_fields = ["name"] - ordering = ["id"] - add_fieldsets = ( - ( - None, - { - "classes": ("wide",), - "fields": ("email", "password1", "password2"), - }, - ), - ) diff --git a/backend/backend/users/api/__init__.py b/backend/backend/users/api/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/users/api/serializers.py b/backend/backend/users/api/serializers.py deleted file mode 100644 index 81bfea72..00000000 --- a/backend/backend/users/api/serializers.py +++ /dev/null @@ -1,22 +0,0 @@ -from django.contrib.auth import get_user_model -from rest_framework import serializers - -from backend.users.models import User as UserType - - -User = get_user_model() - - -class UserSerializer(serializers.ModelSerializer[UserType]): - class Meta: - model = User - fields = ["name", "url"] - - extra_kwargs = { - "url": {"view_name": "api:user-detail", "lookup_field": "pk"}, - } - - -class LoginSerializer(serializers.Serializer): - email = serializers.EmailField() - password = serializers.CharField() diff --git a/backend/backend/users/api/urls.py b/backend/backend/users/api/urls.py deleted file mode 100644 index 98dbd64d..00000000 --- a/backend/backend/users/api/urls.py +++ /dev/null @@ -1,27 +0,0 @@ -from django.urls import path -from .views import ( - LoginView, - RegisterView, - GetUserView, - MailVerifyView, - ForgetPasswordView, - ResetPasswordView, - loginWithGoogle, - GetUserInfo, - UpdateUserProfile, -) - -# from .views import * - -urlpatterns = [ - path("login/", view=LoginView.as_view(), name="login"), - path("register/", view=RegisterView.as_view(), name="register"), - path("mail-verify/", view=MailVerifyView.as_view(), name="mailverify"), - path("forget-password/", view=ForgetPasswordView.as_view(), name="forgetpassword"), - path("reset-password/", view=ResetPasswordView.as_view(), name="resetpassword"), - path("getUser/", view=GetUserView.as_view(), name="getUserInfo"), - path("getUserInfo/", view=GetUserInfo.as_view(), name="getUserInfoData"), - path("getUserName/", view=GetUserInfo.as_view(), name="getUserInfoData"), - path("loginWithGoogle/", view=loginWithGoogle.as_view(), name="loginWithGoogle"), - path("profile/", view=UpdateUserProfile.as_view(), name="UpdateUserProfile"), -] diff --git a/backend/backend/users/api/views.py b/backend/backend/users/api/views.py deleted file mode 100644 index 42095ec5..00000000 --- a/backend/backend/users/api/views.py +++ /dev/null @@ -1,294 +0,0 @@ -import os -import requests -from django.contrib.auth import get_user_model, authenticate, login -from rest_framework import status -from rest_framework.decorators import action -from rest_framework.mixins import ListModelMixin, RetrieveModelMixin, UpdateModelMixin -from rest_framework.response import Response -from rest_framework.viewsets import GenericViewSet - -from rest_framework import permissions, status -from rest_framework.response import Response -from rest_framework.views import APIView -from rest_framework.permissions import AllowAny -from rest_framework_simplejwt.tokens import RefreshToken - -from .serializers import UserSerializer, LoginSerializer - -User = get_user_model() - - -def send_mail(email, content): - email_params = { - "apikey": os.getenv("MAIL_KEY"), - "from": "tunki1201@gmail.com", - "to": email, - "subject": "Email Verify - Metricloop Account", - "body": f"{content}", - "isTransactional": True, - } - - response = requests.post( - "https://api.elasticemail.com/v2/email/send", - data=email_params, - ) - return response - - -class UserViewSet(RetrieveModelMixin, ListModelMixin, UpdateModelMixin, GenericViewSet): - serializer_class = UserSerializer - queryset = User.objects.all() - lookup_field = "pk" - - def get_queryset(self, *args, **kwargs): - assert isinstance(self.request.user.id, int) - return self.queryset.filter(id=self.request.user.id) - - @action(detail=False) - def me(self, request): - serializer = UserSerializer(request.user, context={"request": request}) - return Response(status=status.HTTP_200_OK, data=serializer.data) - - -class RegisterView(APIView): - permission_classes = (permissions.AllowAny,) - authentication_classes = () - - def post(self, request): - try: - data = request.data - username = data["fullName"] - email = data["email"] - password = data["password"] - confirmPassword = data["confirmPassword"] - - if password == confirmPassword: - if len(password) >= 6: - if not User.objects.filter(email=email).exists(): - user = User.objects.create_user( - email=email, - fullname=username, - password=password, - ) - if User.objects.filter(email=email).exists(): - # mail verify - refresh = RefreshToken.for_user(user) - response = send_mail( - email, f"{os.getenv('FRONT_URL')}/mail-verify/?token={str(refresh.access_token)}" - ) - if response.status_code == 200: - return Response( - {"success": "User created successfully and sent verification link."}, - status=status.HTTP_201_CREATED, - ) - else: - return Response( - {"success": "Resend verification link."}, - status=status.HTTP_201_CREATED, - ) - else: - return Response( - {"error": "Something went wrong creating user"}, - status=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - else: - return Response( - {"error": "Username already exists"}, - status=status.HTTP_400_BAD_REQUEST, - ) - else: - return Response( - {"error": "Password must be at least 6 characters long"}, - status=status.HTTP_400_BAD_REQUEST, - ) - else: - return Response( - {"error": "Passwords do not match"}, - status=status.HTTP_400_BAD_REQUEST, - ) - except Exception as e: - print(e) - return Response( - {"error": "Something went wrong"}, - status=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - -class MailVerifyView(APIView): - def post(self, request): - user = request.user - user.mail_verify = True - user.save() - return Response({"success": "Email verified"}, status=status.HTTP_200_OK) - - -class ForgetPasswordView(APIView): - permission_classes = (AllowAny,) - authentication_classes = () - - def post(self, request): - # email = request.data.get("email") - email = request.data["email"] - try: - user = User.objects.get(email=email) - refresh = RefreshToken.for_user(user) - response = send_mail(email, f"{os.getenv('FRONT_URL')}/reset-password/?token={str(refresh.access_token)}") - if response.status_code == 200: - return Response( - {"success": "sent verification link."}, - status=status.HTTP_201_CREATED, - ) - else: - return Response( - {"success": "Resend verification link."}, - status=status.HTTP_201_CREATED, - ) - except Exception as e: - print(e) - return Response({"error": "user does not exist"}, status=status.HTTP_400_BAD_REQUEST) - - -class ResetPasswordView(APIView): - def post(self, request): - user = request.user - password = request.data["password"] - confirm_password = request.data["confirmPassword"] - if password == confirm_password: - user.password = password - user.save() - return Response({"success": "reset password"}, status=status.HTTP_200_OK) - else: - return Response({"error": "password not matched"}) - - -class LoginView(APIView): - permission_classes = (AllowAny,) - authentication_classes = () - - def post(self, request): - serializer = LoginSerializer(data=request.data) - if serializer.is_valid(): - print(serializer.validated_data) - email = serializer.validated_data["email"] - password = serializer.validated_data["password"] - user = authenticate(request, email=email, password=password) - print(user) - if user is not None: - if user.mail_verify: - login(request, user) - refresh = RefreshToken.for_user(user) - - return Response( - { - "status": { - "type": "success", - "message": "Welcome back! You have successfully logged in.", - }, - "result": { - "token": str(refresh.access_token), - "user": { - "username": user.fullname, - "email": user.email, - # Include any other user fields as needed - }, - }, - "navigate": "/home", - }, - ) - else: - return Response( - { - "error": "Email not verified", - }, - status=status.HTTP_401_UNAUTHORIZED, - ) - else: - return Response( - { - "error": "Invalid email or password", - }, - status=status.HTTP_400_BAD_REQUEST, - ) - else: - return Response(serializer.error_messages) - - -class GetUserView(APIView): - def post(self, request): - user = request.user - print(user) - find_user = User.objects.get(email=user) - - if find_user: - refresh = RefreshToken.for_user(find_user) - return Response( - { - "result": { - "token": str(refresh.access_token), - "user": { - "username": user.fullname, - "email": user.email, - # Include any other user fields as needed - }, - } - }, - status=status.HTTP_200_OK, - ) - else: - return Response({"error": "User does not exist. Please login"}, status=status.HTTP_400_BAD_REQUEST) - - -class GetUserInfo(APIView): - def post(self, request): - # Retrieve all users from the database - all_users = User.objects.all() - - # Extract relevant information (username and email) from each user object - users_info = [{"username": user.fullname, "email": user.email} for user in all_users] - - return Response({"user": users_info}, status=status.HTTP_200_OK) - - -class loginWithGoogle(APIView): - permission_classes = (AllowAny,) - authentication_classes = () - - def post(self, request): - name = request.data.get("name") - email = request.data.get("email") - - # Check if the user already exists - user, created = User.objects.get_or_create(fullname=name, email=email) - if created: - user.fullname = name - user.mail_verify = True - user.save() - - return Response("User saved successfully", status=status.HTTP_200_OK) - - -class UpdateUserProfile(APIView): - parser_classes = (AllowAny,) - authentication_classes = () - - def post(self, request): - # Assuming you have the required fields in your user model - full_name = request.data.get("full_name") - username = request.data.get("username") - avatar = request.data.get("avatar") - - # Assuming you are using the default User model - - if full_name: - User.full_name = full_name - if username: - User.username = username - - if avatar: - # Handle avatar upload or save the file path to the user's profile - # You may want to use a serializer to handle file upload and save logic - User.data_url = avatar - # Save the user instance - User.save() - - return Response({"message": "Profile updated successfully"}) diff --git a/backend/backend/users/apps.py b/backend/backend/users/apps.py deleted file mode 100644 index e7fc1887..00000000 --- a/backend/backend/users/apps.py +++ /dev/null @@ -1,13 +0,0 @@ -from django.apps import AppConfig -from django.utils.translation import gettext_lazy as _ - - -class UsersConfig(AppConfig): - name = "backend.users" - verbose_name = _("Users") - - def ready(self): - try: - import backend.users.signals # noqa: F401 - except ImportError: - pass diff --git a/backend/backend/users/context_processors.py b/backend/backend/users/context_processors.py deleted file mode 100644 index e2633aec..00000000 --- a/backend/backend/users/context_processors.py +++ /dev/null @@ -1,8 +0,0 @@ -from django.conf import settings - - -def allauth_settings(request): - """Expose some settings from django-allauth in templates.""" - return { - "ACCOUNT_ALLOW_REGISTRATION": settings.ACCOUNT_ALLOW_REGISTRATION, - } diff --git a/backend/backend/users/forms.py b/backend/backend/users/forms.py deleted file mode 100644 index 604b32b3..00000000 --- a/backend/backend/users/forms.py +++ /dev/null @@ -1,45 +0,0 @@ -from allauth.account.forms import SignupForm -from allauth.socialaccount.forms import SignupForm as SocialSignupForm -from django.contrib.auth import forms as admin_forms -from django.contrib.auth import get_user_model -from django.forms import EmailField -from django.utils.translation import gettext_lazy as _ - -User = get_user_model() - - -class UserAdminChangeForm(admin_forms.UserChangeForm): - class Meta(admin_forms.UserChangeForm.Meta): - model = User - field_classes = {"email": EmailField} - - -class UserAdminCreationForm(admin_forms.UserCreationForm): - """ - Form for User Creation in the Admin Area. - To change user signup, see UserSignupForm and UserSocialSignupForm. - """ - - class Meta(admin_forms.UserCreationForm.Meta): - model = User - fields = ("email",) - field_classes = {"email": EmailField} - error_messages = { - "email": {"unique": _("This email has already been taken.")}, - } - - -class UserSignupForm(SignupForm): - """ - Form that will be rendered on a user sign up section/screen. - Default fields will be added automatically. - Check UserSocialSignupForm for accounts created from social. - """ - - -class UserSocialSignupForm(SocialSignupForm): - """ - Renders the form when user has signed up using social accounts. - Default fields will be added automatically. - See UserSignupForm otherwise. - """ diff --git a/backend/backend/users/managers.py b/backend/backend/users/managers.py deleted file mode 100644 index 017ab14e..00000000 --- a/backend/backend/users/managers.py +++ /dev/null @@ -1,34 +0,0 @@ -from django.contrib.auth.hashers import make_password -from django.contrib.auth.models import UserManager as DjangoUserManager - - -class UserManager(DjangoUserManager): - """Custom manager for the User model.""" - - def _create_user(self, email: str, password: str | None, **extra_fields): - """ - Create and save a user with the given email and password. - """ - if not email: - raise ValueError("The given email must be set") - email = self.normalize_email(email) - user = self.model(email=email, **extra_fields) - user.password = make_password(password) - user.save(using=self._db) - return user - - def create_user(self, email: str, password: str | None = None, **extra_fields): - extra_fields.setdefault("is_staff", False) - extra_fields.setdefault("is_superuser", False) - return self._create_user(email, password, **extra_fields) - - def create_superuser(self, email: str, password: str | None = None, **extra_fields): - extra_fields.setdefault("is_staff", True) - extra_fields.setdefault("is_superuser", True) - - if extra_fields.get("is_staff") is not True: - raise ValueError("Superuser must have is_staff=True.") - if extra_fields.get("is_superuser") is not True: - raise ValueError("Superuser must have is_superuser=True.") - - return self._create_user(email, password, **extra_fields) diff --git a/backend/backend/users/migrations/0001_initial.py b/backend/backend/users/migrations/0001_initial.py deleted file mode 100644 index 11106fde..00000000 --- a/backend/backend/users/migrations/0001_initial.py +++ /dev/null @@ -1,112 +0,0 @@ -# Generated by Django 4.2.6 on 2023-10-09 23:52 - -import backend.users.managers -from django.db import migrations, models -import django.db.models.deletion -import django.utils.timezone - - -class Migration(migrations.Migration): - initial = True - - dependencies = [ - ("auth", "0012_alter_user_first_name_max_length"), - ] - - operations = [ - migrations.CreateModel( - name="Company", - fields=[ - ("company_id", models.AutoField(primary_key=True, serialize=False)), - ("company_email", models.EmailField(max_length=254, unique=True)), - ("company_name", models.CharField(max_length=255)), - ("company_logo", models.ImageField(blank=True, null=True, upload_to="company_logos/")), - ], - ), - migrations.CreateModel( - name="User", - fields=[ - ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")), - ("password", models.CharField(max_length=128, verbose_name="password")), - ("last_login", models.DateTimeField(blank=True, null=True, verbose_name="last login")), - ( - "is_superuser", - models.BooleanField( - default=False, - help_text="Designates that this user has all permissions without explicitly assigning them.", - verbose_name="superuser status", - ), - ), - ( - "is_staff", - models.BooleanField( - default=False, - help_text="Designates whether the user can log into this admin site.", - verbose_name="staff status", - ), - ), - ( - "is_active", - models.BooleanField( - default=True, - help_text="Designates whether this user should be treated as active. Unselect this instead of deleting accounts.", - verbose_name="active", - ), - ), - ("date_joined", models.DateTimeField(default=django.utils.timezone.now, verbose_name="date joined")), - ("fullname", models.CharField(default="John Doe", max_length=255)), - ("name", models.CharField(blank=True, max_length=255, verbose_name="Name of User")), - ("email", models.EmailField(max_length=254, unique=True, verbose_name="email address")), - ( - "user_type", - models.CharField( - choices=[ - ("Data Engineer", "Data Engineer"), - ("Data Analyst", "Data Analyst"), - ("Data Scientist", "Data Scientist"), - ("Subject Matter Expert", "Subject Matter Expert"), - ], - max_length=50, - ), - ), - ("profile_info", models.TextField(blank=True)), - ("user_logo", models.ImageField(blank=True, null=True, upload_to="user_logos/")), - ( - "company_id", - models.ForeignKey( - blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to="users.company" - ), - ), - ( - "groups", - models.ManyToManyField( - blank=True, - help_text="The groups this user belongs to. A user will get all permissions granted to each of their groups.", - related_name="user_set", - related_query_name="user", - to="auth.group", - verbose_name="groups", - ), - ), - ( - "user_permissions", - models.ManyToManyField( - blank=True, - help_text="Specific permissions for this user.", - related_name="user_set", - related_query_name="user", - to="auth.permission", - verbose_name="user permissions", - ), - ), - ], - options={ - "verbose_name": "user", - "verbose_name_plural": "users", - "abstract": False, - }, - managers=[ - ("objects", backend.users.managers.UserManager()), - ], - ), - ] diff --git a/backend/backend/users/migrations/0002_user_mail_verify.py b/backend/backend/users/migrations/0002_user_mail_verify.py deleted file mode 100644 index 4cf248b2..00000000 --- a/backend/backend/users/migrations/0002_user_mail_verify.py +++ /dev/null @@ -1,17 +0,0 @@ -# Generated by Django 4.2.6 on 2023-10-11 12:08 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ("users", "0001_initial"), - ] - - operations = [ - migrations.AddField( - model_name="user", - name="mail_verify", - field=models.BooleanField(default=False), - ), - ] diff --git a/backend/backend/users/migrations/0003_user_data_url_user_path.py b/backend/backend/users/migrations/0003_user_data_url_user_path.py deleted file mode 100644 index 9089b51f..00000000 --- a/backend/backend/users/migrations/0003_user_data_url_user_path.py +++ /dev/null @@ -1,22 +0,0 @@ -# Generated by Django 4.2.6 on 2024-02-28 17:00 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ("users", "0002_user_mail_verify"), - ] - - operations = [ - migrations.AddField( - model_name="user", - name="data_url", - field=models.CharField(max_length=255, null=True), - ), - migrations.AddField( - model_name="user", - name="path", - field=models.FileField(blank=True, null=True, upload_to="dataurl", verbose_name="upload Field"), - ), - ] diff --git a/backend/backend/users/migrations/__init__.py b/backend/backend/users/migrations/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/users/models.py b/backend/backend/users/models.py deleted file mode 100644 index 35b9fcdb..00000000 --- a/backend/backend/users/models.py +++ /dev/null @@ -1,77 +0,0 @@ -from django.contrib.auth.models import AbstractUser -from django.db.models import CharField, EmailField -from django.db import models -from django.urls import reverse -from django.utils.translation import gettext_lazy as _ - -from backend.users.managers import UserManager -from django.contrib.auth.models import BaseUserManager - -class EncryptedUserManager(BaseUserManager): - def create_user(self, email, password=None, **extra_fields): - if not email: - raise ValueError("The Email field must be set") - email = self.normalize_email(email) - user = self.model(email=email, **extra_fields) - user.set_password(password) - user.save(using=self._db) - return user - - def create_superuser(self, email, password=None, **extra_fields): - extra_fields.setdefault("is_staff", True) - extra_fields.setdefault("is_superuser", True) - - return self.create_user(email, password, **extra_fields) - -class Company(models.Model): - company_id = models.AutoField(primary_key=True) - company_email = models.EmailField(unique=True) - company_name = models.CharField(max_length=255) - company_logo = models.ImageField(upload_to="company_logos/", null=True, blank=True) - - def __str__(self): - return self.company_name - - -class User(AbstractUser): - """ - Default custom user model for backend. - If adding fields that need to be filled at user signup, - check forms.SignupForm and forms.SocialSignupForms accordingly. - """ - - USER_TYPE_CHOICES = ( - ("Data Engineer", "Data Engineer"), - ("Data Analyst", "Data Analyst"), - ("Data Scientist", "Data Scientist"), - ("Subject Matter Expert", "Subject Matter Expert"), - ) - - # First and last name do not cover name patterns around the globe - fullname = CharField(max_length=255, default="John Doe") - name = CharField(_("Name of User"), blank=True, max_length=255) - first_name = None # type: ignore - last_name = None # type: ignore - email = EmailField(_("email address"), unique=True) - username = None # type: ignore - user_type = CharField(max_length=50, choices=USER_TYPE_CHOICES) - profile_info = models.TextField(blank=True) - company_id = models.ForeignKey(Company, on_delete=models.SET_NULL, null=True, blank=True) - user_logo = models.ImageField(upload_to="user_logos/", null=True, blank=True) - mail_verify = models.BooleanField(default=False) - path = models.FileField(_("upload Field"), upload_to="dataurl", null=True, blank=True) - data_url = models.CharField(max_length=255, null=True) - - USERNAME_FIELD = "email" - REQUIRED_FIELDS = [] - - objects = UserManager() - - def get_absolute_url(self) -> str: - """Get URL for user's detail view. - - Returns: - str: URL for user detail. - - """ - return reverse("users:detail", kwargs={"pk": self.id}) diff --git a/backend/backend/users/tasks.py b/backend/backend/users/tasks.py deleted file mode 100644 index c99341c5..00000000 --- a/backend/backend/users/tasks.py +++ /dev/null @@ -1,11 +0,0 @@ -from django.contrib.auth import get_user_model - -from config import celery_app - -User = get_user_model() - - -@celery_app.task() -def get_users_count(): - """A pointless Celery task to demonstrate usage.""" - return User.objects.count() diff --git a/backend/backend/users/tests/__init__.py b/backend/backend/users/tests/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/users/tests/factories.py b/backend/backend/users/tests/factories.py deleted file mode 100644 index 106d342e..00000000 --- a/backend/backend/users/tests/factories.py +++ /dev/null @@ -1,38 +0,0 @@ -from collections.abc import Sequence -from typing import Any - -from django.contrib.auth import get_user_model -from factory import Faker, post_generation -from factory.django import DjangoModelFactory - - -class UserFactory(DjangoModelFactory): - email = Faker("email") - name = Faker("name") - - @post_generation - def password(self, create: bool, extracted: Sequence[Any], **kwargs): - password = ( - extracted - if extracted - else Faker( - "password", - length=42, - special_chars=True, - digits=True, - upper_case=True, - lower_case=True, - ).evaluate(None, None, extra={"locale": None}) - ) - self.set_password(password) - - @classmethod - def _after_postgeneration(cls, instance, create, results=None): - """Save again the instance if creating and at least one hook ran.""" - if create and results and not cls._meta.skip_postgeneration_save: - # Some post-generation hooks ran, and may have modified us. - instance.save() - - class Meta: - model = get_user_model() - django_get_or_create = ["email"] diff --git a/backend/backend/users/tests/test_admin.py b/backend/backend/users/tests/test_admin.py deleted file mode 100644 index 3f261084..00000000 --- a/backend/backend/users/tests/test_admin.py +++ /dev/null @@ -1,37 +0,0 @@ -from django.urls import reverse - -from backend.users.models import User - - -class TestUserAdmin: - def test_changelist(self, admin_client): - url = reverse("admin:users_user_changelist") - response = admin_client.get(url) - assert response.status_code == 200 - - def test_search(self, admin_client): - url = reverse("admin:users_user_changelist") - response = admin_client.get(url, data={"q": "test"}) - assert response.status_code == 200 - - def test_add(self, admin_client): - url = reverse("admin:users_user_add") - response = admin_client.get(url) - assert response.status_code == 200 - - response = admin_client.post( - url, - data={ - "email": "new-admin@example.com", - "password1": "My_R@ndom-P@ssw0rd", - "password2": "My_R@ndom-P@ssw0rd", - }, - ) - assert response.status_code == 302 - assert User.objects.filter(email="new-admin@example.com").exists() - - def test_view_user(self, admin_client): - user = User.objects.get(email="admin@example.com") - url = reverse("admin:users_user_change", kwargs={"object_id": user.pk}) - response = admin_client.get(url) - assert response.status_code == 200 diff --git a/backend/backend/users/tests/test_drf_urls.py b/backend/backend/users/tests/test_drf_urls.py deleted file mode 100644 index 7c9e329f..00000000 --- a/backend/backend/users/tests/test_drf_urls.py +++ /dev/null @@ -1,18 +0,0 @@ -from django.urls import resolve, reverse - -from backend.users.models import User - - -def test_user_detail(user: User): - assert reverse("api:user-detail", kwargs={"pk": user.pk}) == f"/api/users/{user.pk}/" - assert resolve(f"/api/users/{user.pk}/").view_name == "api:user-detail" - - -def test_user_list(): - assert reverse("api:user-list") == "/api/users/" - assert resolve("/api/users/").view_name == "api:user-list" - - -def test_user_me(): - assert reverse("api:user-me") == "/api/users/me/" - assert resolve("/api/users/me/").view_name == "api:user-me" diff --git a/backend/backend/users/tests/test_drf_views.py b/backend/backend/users/tests/test_drf_views.py deleted file mode 100644 index d63f3d2c..00000000 --- a/backend/backend/users/tests/test_drf_views.py +++ /dev/null @@ -1,34 +0,0 @@ -import pytest -from rest_framework.test import APIRequestFactory - -from backend.users.api.views import UserViewSet -from backend.users.models import User - - -class TestUserViewSet: - @pytest.fixture - def api_rf(self) -> APIRequestFactory: - return APIRequestFactory() - - def test_get_queryset(self, user: User, api_rf: APIRequestFactory): - view = UserViewSet() - request = api_rf.get("/fake-url/") - request.user = user - - view.request = request - - assert user in view.get_queryset() - - def test_me(self, user: User, api_rf: APIRequestFactory): - view = UserViewSet() - request = api_rf.get("/fake-url/") - request.user = user - - view.request = request - - response = view.me(request) # type: ignore - - assert response.data == { - "url": f"http://testserver/api/users/{user.pk}/", - "name": user.name, - } diff --git a/backend/backend/users/tests/test_forms.py b/backend/backend/users/tests/test_forms.py deleted file mode 100644 index ace38d12..00000000 --- a/backend/backend/users/tests/test_forms.py +++ /dev/null @@ -1,36 +0,0 @@ -""" -Module for all Form Tests. -""" -from django.utils.translation import gettext_lazy as _ - -from backend.users.forms import UserAdminCreationForm -from backend.users.models import User - - -class TestUserAdminCreationForm: - """ - Test class for all tests related to the UserAdminCreationForm - """ - - def test_username_validation_error_msg(self, user: User): - """ - Tests UserAdminCreation Form's unique validator functions correctly by testing: - 1) A new user with an existing username cannot be added. - 2) Only 1 error is raised by the UserCreation Form - 3) The desired error message is raised - """ - - # The user already exists, - # hence cannot be created. - form = UserAdminCreationForm( - { - "email": user.email, - "password1": user.password, - "password2": user.password, - } - ) - - assert not form.is_valid() - assert len(form.errors) == 1 - assert "email" in form.errors - assert form.errors["email"][0] == _("This email has already been taken.") diff --git a/backend/backend/users/tests/test_managers.py b/backend/backend/users/tests/test_managers.py deleted file mode 100644 index fcc555b9..00000000 --- a/backend/backend/users/tests/test_managers.py +++ /dev/null @@ -1,55 +0,0 @@ -from io import StringIO - -import pytest -from django.core.management import call_command - -from backend.users.models import User - - -@pytest.mark.django_db -class TestUserManager: - def test_create_user(self): - user = User.objects.create_user( - email="john@example.com", - password="something-r@nd0m!", - ) - assert user.email == "john@example.com" - assert not user.is_staff - assert not user.is_superuser - assert user.check_password("something-r@nd0m!") - assert user.username is None - - def test_create_superuser(self): - user = User.objects.create_superuser( - email="admin@example.com", - password="something-r@nd0m!", - ) - assert user.email == "admin@example.com" - assert user.is_staff - assert user.is_superuser - assert user.username is None - - def test_create_superuser_username_ignored(self): - user = User.objects.create_superuser( - email="test@example.com", - password="something-r@nd0m!", - ) - assert user.username is None - - -@pytest.mark.django_db -def test_createsuperuser_command(): - """Ensure createsuperuser command works with our custom manager.""" - out = StringIO() - command_result = call_command( - "createsuperuser", - "--email", - "henry@example.com", - interactive=False, - stdout=out, - ) - - assert command_result is None - assert out.getvalue() == "Superuser created successfully.\n" - user = User.objects.get(email="henry@example.com") - assert not user.has_usable_password() diff --git a/backend/backend/users/tests/test_models.py b/backend/backend/users/tests/test_models.py deleted file mode 100644 index eb0327d5..00000000 --- a/backend/backend/users/tests/test_models.py +++ /dev/null @@ -1,5 +0,0 @@ -from backend.users.models import User - - -def test_user_get_absolute_url(user: User): - assert user.get_absolute_url() == f"/users/{user.pk}/" diff --git a/backend/backend/users/tests/test_swagger.py b/backend/backend/users/tests/test_swagger.py deleted file mode 100644 index f97658b5..00000000 --- a/backend/backend/users/tests/test_swagger.py +++ /dev/null @@ -1,21 +0,0 @@ -import pytest -from django.urls import reverse - - -def test_swagger_accessible_by_admin(admin_client): - url = reverse("api-docs") - response = admin_client.get(url) - assert response.status_code == 200 - - -@pytest.mark.django_db -def test_swagger_ui_not_accessible_by_normal_user(client): - url = reverse("api-docs") - response = client.get(url) - assert response.status_code == 403 - - -def test_api_schema_generated_successfully(admin_client): - url = reverse("api-schema") - response = admin_client.get(url) - assert response.status_code == 200 diff --git a/backend/backend/users/tests/test_tasks.py b/backend/backend/users/tests/test_tasks.py deleted file mode 100644 index df748098..00000000 --- a/backend/backend/users/tests/test_tasks.py +++ /dev/null @@ -1,16 +0,0 @@ -import pytest -from celery.result import EagerResult - -from backend.users.tasks import get_users_count -from backend.users.tests.factories import UserFactory - -pytestmark = pytest.mark.django_db - - -def test_user_count(settings): - """A basic test to execute the get_users_count Celery task.""" - UserFactory.create_batch(3) - settings.CELERY_TASK_ALWAYS_EAGER = True - task_result = get_users_count.delay() - assert isinstance(task_result, EagerResult) - assert task_result.result == 3 diff --git a/backend/backend/users/tests/test_urls.py b/backend/backend/users/tests/test_urls.py deleted file mode 100644 index a7146e42..00000000 --- a/backend/backend/users/tests/test_urls.py +++ /dev/null @@ -1,18 +0,0 @@ -from django.urls import resolve, reverse - -from backend.users.models import User - - -def test_detail(user: User): - assert reverse("users:detail", kwargs={"pk": user.pk}) == f"/users/{user.pk}/" - assert resolve(f"/users/{user.pk}/").view_name == "users:detail" - - -def test_update(): - assert reverse("users:update") == "/users/~update/" - assert resolve("/users/~update/").view_name == "users:update" - - -def test_redirect(): - assert reverse("users:redirect") == "/users/~redirect/" - assert resolve("/users/~redirect/").view_name == "users:redirect" diff --git a/backend/backend/users/tests/test_views.py b/backend/backend/users/tests/test_views.py deleted file mode 100644 index 77dc4ae2..00000000 --- a/backend/backend/users/tests/test_views.py +++ /dev/null @@ -1,100 +0,0 @@ -import pytest -from django.conf import settings -from django.contrib import messages -from django.contrib.auth.models import AnonymousUser -from django.contrib.messages.middleware import MessageMiddleware -from django.contrib.sessions.middleware import SessionMiddleware -from django.http import HttpRequest, HttpResponseRedirect -from django.test import RequestFactory -from django.urls import reverse -from django.utils.translation import gettext_lazy as _ - -from backend.users.forms import UserAdminChangeForm -from backend.users.models import User -from backend.users.tests.factories import UserFactory -from backend.users.views import ( - UserRedirectView, - UserUpdateView, - user_detail_view, -) - -pytestmark = pytest.mark.django_db - - -class TestUserUpdateView: - """ - TODO: - extracting view initialization code as class-scoped fixture - would be great if only pytest-django supported non-function-scoped - fixture db access -- this is a work-in-progress for now: - https://github.com/pytest-dev/pytest-django/pull/258 - """ - - def dummy_get_response(self, request: HttpRequest): - return None - - def test_get_success_url(self, user: User, rf: RequestFactory): - view = UserUpdateView() - request = rf.get("/fake-url/") - request.user = user - - view.request = request - assert view.get_success_url() == f"/users/{user.pk}/" - - def test_get_object(self, user: User, rf: RequestFactory): - view = UserUpdateView() - request = rf.get("/fake-url/") - request.user = user - - view.request = request - - assert view.get_object() == user - - def test_form_valid(self, user: User, rf: RequestFactory): - view = UserUpdateView() - request = rf.get("/fake-url/") - - # Add the session/message middleware to the request - SessionMiddleware(self.dummy_get_response).process_request(request) - MessageMiddleware(self.dummy_get_response).process_request(request) - request.user = user - - view.request = request - - # Initialize the form - form = UserAdminChangeForm() - form.cleaned_data = {} - form.instance = user - view.form_valid(form) - - messages_sent = [m.message for m in messages.get_messages(request)] - assert messages_sent == [_("Information successfully updated")] - - -class TestUserRedirectView: - def test_get_redirect_url(self, user: User, rf: RequestFactory): - view = UserRedirectView() - request = rf.get("/fake-url") - request.user = user - - view.request = request - assert view.get_redirect_url() == f"/users/{user.pk}/" - - -class TestUserDetailView: - def test_authenticated(self, user: User, rf: RequestFactory): - request = rf.get("/fake-url/") - request.user = UserFactory() - response = user_detail_view(request, pk=user.pk) - - assert response.status_code == 200 - - def test_not_authenticated(self, user: User, rf: RequestFactory): - request = rf.get("/fake-url/") - request.user = AnonymousUser() - response = user_detail_view(request, pk=user.pk) - login_url = reverse(settings.LOGIN_URL) - - assert isinstance(response, HttpResponseRedirect) - assert response.status_code == 302 - assert response.url == f"{login_url}?next=/fake-url/" diff --git a/backend/backend/users/urls.py b/backend/backend/users/urls.py deleted file mode 100644 index 715f83a1..00000000 --- a/backend/backend/users/urls.py +++ /dev/null @@ -1,14 +0,0 @@ -from django.urls import path - -from backend.users.views import ( - user_detail_view, - user_redirect_view, - user_update_view, -) - -app_name = "users" -urlpatterns = [ - path("~redirect/", view=user_redirect_view, name="redirect"), - path("~update/", view=user_update_view, name="update"), - path("/", view=user_detail_view, name="detail"), -] diff --git a/backend/backend/users/utils.py b/backend/backend/users/utils.py deleted file mode 100644 index 1cfc99d5..00000000 --- a/backend/backend/users/utils.py +++ /dev/null @@ -1,33 +0,0 @@ -# utils.py - -from cryptography.fernet import Fernet - - -def encrypt_data(data, key): - cipher_suite = Fernet(key) - cipher_text = cipher_suite.encrypt(data.encode()) - return cipher_text - - -def decrypt_data(cipher_text, key): - cipher_suite = Fernet(key) - plain_text = cipher_suite.decrypt(cipher_text).decode() - return plain_text - - -# Generate a new key using cryptography library -from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC -from cryptography.hazmat.backends import default_backend -from cryptography.hazmat.primitives import hashes # Add this import -from base64 import urlsafe_b64encode - - -def generate_key(): - password = "P@ssw0rd1997" # Change this to a strong, secret password - salt = b"salt_123" # Change this to a random, unique salt - kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), iterations=100000, salt=salt, length=32, backend=default_backend()) - key = urlsafe_b64encode(kdf.derive(password.encode())) - return key - - -SECRET_KEY = generate_key() diff --git a/backend/backend/users/views.py b/backend/backend/users/views.py deleted file mode 100644 index a6d8475c..00000000 --- a/backend/backend/users/views.py +++ /dev/null @@ -1,43 +0,0 @@ -from django.contrib.auth import get_user_model -from django.contrib.auth.mixins import LoginRequiredMixin -from django.contrib.messages.views import SuccessMessageMixin -from django.urls import reverse -from django.utils.translation import gettext_lazy as _ -from django.views.generic import DetailView, RedirectView, UpdateView - -User = get_user_model() - - -class UserDetailView(LoginRequiredMixin, DetailView): - model = User - slug_field = "id" - slug_url_kwarg = "id" - - -user_detail_view = UserDetailView.as_view() - - -class UserUpdateView(LoginRequiredMixin, SuccessMessageMixin, UpdateView): - model = User - fields = ["name"] - success_message = _("Information successfully updated") - - def get_success_url(self): - assert self.request.user.is_authenticated # for mypy to know that the user is authenticated - return self.request.user.get_absolute_url() - - def get_object(self): - return self.request.user - - -user_update_view = UserUpdateView.as_view() - - -class UserRedirectView(LoginRequiredMixin, RedirectView): - permanent = False - - def get_redirect_url(self): - return reverse("users:detail", kwargs={"pk": self.request.user.pk}) - - -user_redirect_view = UserRedirectView.as_view() diff --git a/backend/backend/utils/__init__.py b/backend/backend/utils/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/backend/utils/storages.py b/backend/backend/utils/storages.py deleted file mode 100644 index d9a321c5..00000000 --- a/backend/backend/utils/storages.py +++ /dev/null @@ -1,11 +0,0 @@ -from storages.backends.s3 import S3Storage - - -class StaticS3Storage(S3Storage): - location = "static" - default_acl = "public-read" - - -class MediaS3Storage(S3Storage): - location = "media" - file_overwrite = False diff --git a/backend/comfy_nodes/chainner_models/__init__.py b/backend/comfy_nodes/chainner_models/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/comfy_nodes/chainner_models/architecture/DAT.py b/backend/comfy_nodes/chainner_models/architecture/DAT.py deleted file mode 100644 index 0bcc26ef..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/DAT.py +++ /dev/null @@ -1,1182 +0,0 @@ -# pylint: skip-file -import math -import re - -import numpy as np -import torch -import torch.nn as nn -import torch.utils.checkpoint as checkpoint -from einops import rearrange -from einops.layers.torch import Rearrange -from torch import Tensor -from torch.nn import functional as F - -from .timm.drop import DropPath -from .timm.weight_init import trunc_normal_ - - -def img2windows(img, H_sp, W_sp): - """ - Input: Image (B, C, H, W) - Output: Window Partition (B', N, C) - """ - B, C, H, W = img.shape - img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp) - img_perm = ( - img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C) - ) - return img_perm - - -def windows2img(img_splits_hw, H_sp, W_sp, H, W): - """ - Input: Window Partition (B', N, C) - Output: Image (B, H, W, C) - """ - B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp)) - - img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1) - img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) - return img - - -class SpatialGate(nn.Module): - """Spatial-Gate. - Args: - dim (int): Half of input channels. - """ - - def __init__(self, dim): - super().__init__() - self.norm = nn.LayerNorm(dim) - self.conv = nn.Conv2d( - dim, dim, kernel_size=3, stride=1, padding=1, groups=dim - ) # DW Conv - - def forward(self, x, H, W): - # Split - x1, x2 = x.chunk(2, dim=-1) - B, N, C = x.shape - x2 = ( - self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C // 2, H, W)) - .flatten(2) - .transpose(-1, -2) - .contiguous() - ) - - return x1 * x2 - - -class SGFN(nn.Module): - """Spatial-Gate Feed-Forward Network. - Args: - in_features (int): Number of input channels. - hidden_features (int | None): Number of hidden channels. Default: None - out_features (int | None): Number of output channels. Default: None - act_layer (nn.Module): Activation layer. Default: nn.GELU - drop (float): Dropout rate. Default: 0.0 - """ - - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.sg = SpatialGate(hidden_features // 2) - self.fc2 = nn.Linear(hidden_features // 2, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x, H, W): - """ - Input: x: (B, H*W, C), H, W - Output: x: (B, H*W, C) - """ - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - - x = self.sg(x, H, W) - x = self.drop(x) - - x = self.fc2(x) - x = self.drop(x) - return x - - -class DynamicPosBias(nn.Module): - # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py - """Dynamic Relative Position Bias. - Args: - dim (int): Number of input channels. - num_heads (int): Number of attention heads. - residual (bool): If True, use residual strage to connect conv. - """ - - def __init__(self, dim, num_heads, residual): - super().__init__() - self.residual = residual - self.num_heads = num_heads - self.pos_dim = dim // 4 - self.pos_proj = nn.Linear(2, self.pos_dim) - self.pos1 = nn.Sequential( - nn.LayerNorm(self.pos_dim), - nn.ReLU(inplace=True), - nn.Linear(self.pos_dim, self.pos_dim), - ) - self.pos2 = nn.Sequential( - nn.LayerNorm(self.pos_dim), - nn.ReLU(inplace=True), - nn.Linear(self.pos_dim, self.pos_dim), - ) - self.pos3 = nn.Sequential( - nn.LayerNorm(self.pos_dim), - nn.ReLU(inplace=True), - nn.Linear(self.pos_dim, self.num_heads), - ) - - def forward(self, biases): - if self.residual: - pos = self.pos_proj(biases) # 2Gh-1 * 2Gw-1, heads - pos = pos + self.pos1(pos) - pos = pos + self.pos2(pos) - pos = self.pos3(pos) - else: - pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases)))) - return pos - - -class Spatial_Attention(nn.Module): - """Spatial Window Self-Attention. - It supports rectangle window (containing square window). - Args: - dim (int): Number of input channels. - idx (int): The indentix of window. (0/1) - split_size (tuple(int)): Height and Width of spatial window. - dim_out (int | None): The dimension of the attention output. Default: None - num_heads (int): Number of attention heads. Default: 6 - attn_drop (float): Dropout ratio of attention weight. Default: 0.0 - proj_drop (float): Dropout ratio of output. Default: 0.0 - qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set - position_bias (bool): The dynamic relative position bias. Default: True - """ - - def __init__( - self, - dim, - idx, - split_size=[8, 8], - dim_out=None, - num_heads=6, - attn_drop=0.0, - proj_drop=0.0, - qk_scale=None, - position_bias=True, - ): - super().__init__() - self.dim = dim - self.dim_out = dim_out or dim - self.split_size = split_size - self.num_heads = num_heads - self.idx = idx - self.position_bias = position_bias - - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - - if idx == 0: - H_sp, W_sp = self.split_size[0], self.split_size[1] - elif idx == 1: - W_sp, H_sp = self.split_size[0], self.split_size[1] - else: - print("ERROR MODE", idx) - exit(0) - self.H_sp = H_sp - self.W_sp = W_sp - - if self.position_bias: - self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False) - # generate mother-set - position_bias_h = torch.arange(1 - self.H_sp, self.H_sp) - position_bias_w = torch.arange(1 - self.W_sp, self.W_sp) - biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w])) - biases = biases.flatten(1).transpose(0, 1).contiguous().float() - self.register_buffer("rpe_biases", biases) - - # get pair-wise relative position index for each token inside the window - coords_h = torch.arange(self.H_sp) - coords_w = torch.arange(self.W_sp) - coords = torch.stack(torch.meshgrid([coords_h, coords_w])) - coords_flatten = torch.flatten(coords, 1) - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] - relative_coords = relative_coords.permute(1, 2, 0).contiguous() - relative_coords[:, :, 0] += self.H_sp - 1 - relative_coords[:, :, 1] += self.W_sp - 1 - relative_coords[:, :, 0] *= 2 * self.W_sp - 1 - relative_position_index = relative_coords.sum(-1) - self.register_buffer("relative_position_index", relative_position_index) - - self.attn_drop = nn.Dropout(attn_drop) - - def im2win(self, x, H, W): - B, N, C = x.shape - x = x.transpose(-2, -1).contiguous().view(B, C, H, W) - x = img2windows(x, self.H_sp, self.W_sp) - x = ( - x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads) - .permute(0, 2, 1, 3) - .contiguous() - ) - return x - - def forward(self, qkv, H, W, mask=None): - """ - Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size - Output: x (B, H, W, C) - """ - q, k, v = qkv[0], qkv[1], qkv[2] - - B, L, C = q.shape - assert L == H * W, "flatten img_tokens has wrong size" - - # partition the q,k,v, image to window - q = self.im2win(q, H, W) - k = self.im2win(k, H, W) - v = self.im2win(v, H, W) - - q = q * self.scale - attn = q @ k.transpose(-2, -1) # B head N C @ B head C N --> B head N N - - # calculate drpe - if self.position_bias: - pos = self.pos(self.rpe_biases) - # select position bias - relative_position_bias = pos[self.relative_position_index.view(-1)].view( - self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1 - ) - relative_position_bias = relative_position_bias.permute( - 2, 0, 1 - ).contiguous() - attn = attn + relative_position_bias.unsqueeze(0) - - N = attn.shape[3] - - # use mask for shift window - if mask is not None: - nW = mask.shape[0] - attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze( - 0 - ) - attn = attn.view(-1, self.num_heads, N, N) - - attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype) - attn = self.attn_drop(attn) - - x = attn @ v - x = x.transpose(1, 2).reshape( - -1, self.H_sp * self.W_sp, C - ) # B head N N @ B head N C - - # merge the window, window to image - x = windows2img(x, self.H_sp, self.W_sp, H, W) # B H' W' C - - return x - - -class Adaptive_Spatial_Attention(nn.Module): - # The implementation builds on CAT code https://github.com/Zhengchen1999/CAT - """Adaptive Spatial Self-Attention - Args: - dim (int): Number of input channels. - num_heads (int): Number of attention heads. Default: 6 - split_size (tuple(int)): Height and Width of spatial window. - shift_size (tuple(int)): Shift size for spatial window. - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. - drop (float): Dropout rate. Default: 0.0 - attn_drop (float): Attention dropout rate. Default: 0.0 - rg_idx (int): The indentix of Residual Group (RG) - b_idx (int): The indentix of Block in each RG - """ - - def __init__( - self, - dim, - num_heads, - reso=64, - split_size=[8, 8], - shift_size=[1, 2], - qkv_bias=False, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - rg_idx=0, - b_idx=0, - ): - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.split_size = split_size - self.shift_size = shift_size - self.b_idx = b_idx - self.rg_idx = rg_idx - self.patches_resolution = reso - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - - assert ( - 0 <= self.shift_size[0] < self.split_size[0] - ), "shift_size must in 0-split_size0" - assert ( - 0 <= self.shift_size[1] < self.split_size[1] - ), "shift_size must in 0-split_size1" - - self.branch_num = 2 - - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(drop) - - self.attns = nn.ModuleList( - [ - Spatial_Attention( - dim // 2, - idx=i, - split_size=split_size, - num_heads=num_heads // 2, - dim_out=dim // 2, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=drop, - position_bias=True, - ) - for i in range(self.branch_num) - ] - ) - - if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or ( - self.rg_idx % 2 != 0 and self.b_idx % 4 == 0 - ): - attn_mask = self.calculate_mask( - self.patches_resolution, self.patches_resolution - ) - self.register_buffer("attn_mask_0", attn_mask[0]) - self.register_buffer("attn_mask_1", attn_mask[1]) - else: - attn_mask = None - self.register_buffer("attn_mask_0", None) - self.register_buffer("attn_mask_1", None) - - self.dwconv = nn.Sequential( - nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim), - nn.BatchNorm2d(dim), - nn.GELU(), - ) - self.channel_interaction = nn.Sequential( - nn.AdaptiveAvgPool2d(1), - nn.Conv2d(dim, dim // 8, kernel_size=1), - nn.BatchNorm2d(dim // 8), - nn.GELU(), - nn.Conv2d(dim // 8, dim, kernel_size=1), - ) - self.spatial_interaction = nn.Sequential( - nn.Conv2d(dim, dim // 16, kernel_size=1), - nn.BatchNorm2d(dim // 16), - nn.GELU(), - nn.Conv2d(dim // 16, 1, kernel_size=1), - ) - - def calculate_mask(self, H, W): - # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py - # calculate attention mask for shift window - img_mask_0 = torch.zeros((1, H, W, 1)) # 1 H W 1 idx=0 - img_mask_1 = torch.zeros((1, H, W, 1)) # 1 H W 1 idx=1 - h_slices_0 = ( - slice(0, -self.split_size[0]), - slice(-self.split_size[0], -self.shift_size[0]), - slice(-self.shift_size[0], None), - ) - w_slices_0 = ( - slice(0, -self.split_size[1]), - slice(-self.split_size[1], -self.shift_size[1]), - slice(-self.shift_size[1], None), - ) - - h_slices_1 = ( - slice(0, -self.split_size[1]), - slice(-self.split_size[1], -self.shift_size[1]), - slice(-self.shift_size[1], None), - ) - w_slices_1 = ( - slice(0, -self.split_size[0]), - slice(-self.split_size[0], -self.shift_size[0]), - slice(-self.shift_size[0], None), - ) - cnt = 0 - for h in h_slices_0: - for w in w_slices_0: - img_mask_0[:, h, w, :] = cnt - cnt += 1 - cnt = 0 - for h in h_slices_1: - for w in w_slices_1: - img_mask_1[:, h, w, :] = cnt - cnt += 1 - - # calculate mask for window-0 - img_mask_0 = img_mask_0.view( - 1, - H // self.split_size[0], - self.split_size[0], - W // self.split_size[1], - self.split_size[1], - 1, - ) - img_mask_0 = ( - img_mask_0.permute(0, 1, 3, 2, 4, 5) - .contiguous() - .view(-1, self.split_size[0], self.split_size[1], 1) - ) # nW, sw[0], sw[1], 1 - mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1]) - attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2) - attn_mask_0 = attn_mask_0.masked_fill( - attn_mask_0 != 0, float(-100.0) - ).masked_fill(attn_mask_0 == 0, float(0.0)) - - # calculate mask for window-1 - img_mask_1 = img_mask_1.view( - 1, - H // self.split_size[1], - self.split_size[1], - W // self.split_size[0], - self.split_size[0], - 1, - ) - img_mask_1 = ( - img_mask_1.permute(0, 1, 3, 2, 4, 5) - .contiguous() - .view(-1, self.split_size[1], self.split_size[0], 1) - ) # nW, sw[1], sw[0], 1 - mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0]) - attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2) - attn_mask_1 = attn_mask_1.masked_fill( - attn_mask_1 != 0, float(-100.0) - ).masked_fill(attn_mask_1 == 0, float(0.0)) - - return attn_mask_0, attn_mask_1 - - def forward(self, x, H, W): - """ - Input: x: (B, H*W, C), H, W - Output: x: (B, H*W, C) - """ - B, L, C = x.shape - assert L == H * W, "flatten img_tokens has wrong size" - - qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3) # 3, B, HW, C - # V without partition - v = qkv[2].transpose(-2, -1).contiguous().view(B, C, H, W) - - # image padding - max_split_size = max(self.split_size[0], self.split_size[1]) - pad_l = pad_t = 0 - pad_r = (max_split_size - W % max_split_size) % max_split_size - pad_b = (max_split_size - H % max_split_size) % max_split_size - - qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2) # 3B C H W - qkv = ( - F.pad(qkv, (pad_l, pad_r, pad_t, pad_b)) - .reshape(3, B, C, -1) - .transpose(-2, -1) - ) # l r t b - _H = pad_b + H - _W = pad_r + W - _L = _H * _W - - # window-0 and window-1 on split channels [C/2, C/2]; for square windows (e.g., 8x8), window-0 and window-1 can be merged - # shift in block: (0, 4, 8, ...), (2, 6, 10, ...), (0, 4, 8, ...), (2, 6, 10, ...), ... - if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or ( - self.rg_idx % 2 != 0 and self.b_idx % 4 == 0 - ): - qkv = qkv.view(3, B, _H, _W, C) - qkv_0 = torch.roll( - qkv[:, :, :, :, : C // 2], - shifts=(-self.shift_size[0], -self.shift_size[1]), - dims=(2, 3), - ) - qkv_0 = qkv_0.view(3, B, _L, C // 2) - qkv_1 = torch.roll( - qkv[:, :, :, :, C // 2 :], - shifts=(-self.shift_size[1], -self.shift_size[0]), - dims=(2, 3), - ) - qkv_1 = qkv_1.view(3, B, _L, C // 2) - - if self.patches_resolution != _H or self.patches_resolution != _W: - mask_tmp = self.calculate_mask(_H, _W) - x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device)) - x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device)) - else: - x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0) - x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1) - - x1 = torch.roll( - x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2) - ) - x2 = torch.roll( - x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2) - ) - x1 = x1[:, :H, :W, :].reshape(B, L, C // 2) - x2 = x2[:, :H, :W, :].reshape(B, L, C // 2) - # attention output - attened_x = torch.cat([x1, x2], dim=2) - - else: - x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape( - B, L, C // 2 - ) - x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape( - B, L, C // 2 - ) - # attention output - attened_x = torch.cat([x1, x2], dim=2) - - # convolution output - conv_x = self.dwconv(v) - - # Adaptive Interaction Module (AIM) - # C-Map (before sigmoid) - channel_map = ( - self.channel_interaction(conv_x) - .permute(0, 2, 3, 1) - .contiguous() - .view(B, 1, C) - ) - # S-Map (before sigmoid) - attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W) - spatial_map = self.spatial_interaction(attention_reshape) - - # C-I - attened_x = attened_x * torch.sigmoid(channel_map) - # S-I - conv_x = torch.sigmoid(spatial_map) * conv_x - conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, L, C) - - x = attened_x + conv_x - - x = self.proj(x) - x = self.proj_drop(x) - - return x - - -class Adaptive_Channel_Attention(nn.Module): - # The implementation builds on XCiT code https://github.com/facebookresearch/xcit - """Adaptive Channel Self-Attention - Args: - dim (int): Number of input channels. - num_heads (int): Number of attention heads. Default: 6 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. - attn_drop (float): Attention dropout rate. Default: 0.0 - drop_path (float): Stochastic depth rate. Default: 0.0 - """ - - def __init__( - self, - dim, - num_heads=8, - qkv_bias=False, - qk_scale=None, - attn_drop=0.0, - proj_drop=0.0, - ): - super().__init__() - self.num_heads = num_heads - self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - - self.dwconv = nn.Sequential( - nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim), - nn.BatchNorm2d(dim), - nn.GELU(), - ) - self.channel_interaction = nn.Sequential( - nn.AdaptiveAvgPool2d(1), - nn.Conv2d(dim, dim // 8, kernel_size=1), - nn.BatchNorm2d(dim // 8), - nn.GELU(), - nn.Conv2d(dim // 8, dim, kernel_size=1), - ) - self.spatial_interaction = nn.Sequential( - nn.Conv2d(dim, dim // 16, kernel_size=1), - nn.BatchNorm2d(dim // 16), - nn.GELU(), - nn.Conv2d(dim // 16, 1, kernel_size=1), - ) - - def forward(self, x, H, W): - """ - Input: x: (B, H*W, C), H, W - Output: x: (B, H*W, C) - """ - B, N, C = x.shape - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) - qkv = qkv.permute(2, 0, 3, 1, 4) - q, k, v = qkv[0], qkv[1], qkv[2] - - q = q.transpose(-2, -1) - k = k.transpose(-2, -1) - v = v.transpose(-2, -1) - - v_ = v.reshape(B, C, N).contiguous().view(B, C, H, W) - - q = torch.nn.functional.normalize(q, dim=-1) - k = torch.nn.functional.normalize(k, dim=-1) - - attn = (q @ k.transpose(-2, -1)) * self.temperature - attn = attn.softmax(dim=-1) - attn = self.attn_drop(attn) - - # attention output - attened_x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C) - - # convolution output - conv_x = self.dwconv(v_) - - # Adaptive Interaction Module (AIM) - # C-Map (before sigmoid) - attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W) - channel_map = self.channel_interaction(attention_reshape) - # S-Map (before sigmoid) - spatial_map = ( - self.spatial_interaction(conv_x) - .permute(0, 2, 3, 1) - .contiguous() - .view(B, N, 1) - ) - - # S-I - attened_x = attened_x * torch.sigmoid(spatial_map) - # C-I - conv_x = conv_x * torch.sigmoid(channel_map) - conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, N, C) - - x = attened_x + conv_x - - x = self.proj(x) - x = self.proj_drop(x) - - return x - - -class DATB(nn.Module): - def __init__( - self, - dim, - num_heads, - reso=64, - split_size=[2, 4], - shift_size=[1, 2], - expansion_factor=4.0, - qkv_bias=False, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - act_layer=nn.GELU, - norm_layer=nn.LayerNorm, - rg_idx=0, - b_idx=0, - ): - super().__init__() - - self.norm1 = norm_layer(dim) - - if b_idx % 2 == 0: - # DSTB - self.attn = Adaptive_Spatial_Attention( - dim, - num_heads=num_heads, - reso=reso, - split_size=split_size, - shift_size=shift_size, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - rg_idx=rg_idx, - b_idx=b_idx, - ) - else: - # DCTB - self.attn = Adaptive_Channel_Attention( - dim, - num_heads=num_heads, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=drop, - ) - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - - ffn_hidden_dim = int(dim * expansion_factor) - self.ffn = SGFN( - in_features=dim, - hidden_features=ffn_hidden_dim, - out_features=dim, - act_layer=act_layer, - ) - self.norm2 = norm_layer(dim) - - def forward(self, x, x_size): - """ - Input: x: (B, H*W, C), x_size: (H, W) - Output: x: (B, H*W, C) - """ - H, W = x_size - x = x + self.drop_path(self.attn(self.norm1(x), H, W)) - x = x + self.drop_path(self.ffn(self.norm2(x), H, W)) - - return x - - -class ResidualGroup(nn.Module): - """ResidualGroup - Args: - dim (int): Number of input channels. - reso (int): Input resolution. - num_heads (int): Number of attention heads. - split_size (tuple(int)): Height and Width of spatial window. - expansion_factor (float): Ratio of ffn hidden dim to embedding dim. - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None - drop (float): Dropout rate. Default: 0 - attn_drop(float): Attention dropout rate. Default: 0 - drop_paths (float | None): Stochastic depth rate. - act_layer (nn.Module): Activation layer. Default: nn.GELU - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm - depth (int): Number of dual aggregation Transformer blocks in residual group. - use_chk (bool): Whether to use checkpointing to save memory. - resi_connection: The convolutional block before residual connection. '1conv'/'3conv' - """ - - def __init__( - self, - dim, - reso, - num_heads, - split_size=[2, 4], - expansion_factor=4.0, - qkv_bias=False, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_paths=None, - act_layer=nn.GELU, - norm_layer=nn.LayerNorm, - depth=2, - use_chk=False, - resi_connection="1conv", - rg_idx=0, - ): - super().__init__() - self.use_chk = use_chk - self.reso = reso - - self.blocks = nn.ModuleList( - [ - DATB( - dim=dim, - num_heads=num_heads, - reso=reso, - split_size=split_size, - shift_size=[split_size[0] // 2, split_size[1] // 2], - expansion_factor=expansion_factor, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_paths[i], - act_layer=act_layer, - norm_layer=norm_layer, - rg_idx=rg_idx, - b_idx=i, - ) - for i in range(depth) - ] - ) - - if resi_connection == "1conv": - self.conv = nn.Conv2d(dim, dim, 3, 1, 1) - elif resi_connection == "3conv": - self.conv = nn.Sequential( - nn.Conv2d(dim, dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim, 3, 1, 1), - ) - - def forward(self, x, x_size): - """ - Input: x: (B, H*W, C), x_size: (H, W) - Output: x: (B, H*W, C) - """ - H, W = x_size - res = x - for blk in self.blocks: - if self.use_chk: - x = checkpoint.checkpoint(blk, x, x_size) - else: - x = blk(x, x_size) - x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W) - x = self.conv(x) - x = rearrange(x, "b c h w -> b (h w) c") - x = res + x - - return x - - -class Upsample(nn.Sequential): - """Upsample module. - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - """ - - def __init__(self, scale, num_feat): - m = [] - if (scale & (scale - 1)) == 0: # scale = 2^n - for _ in range(int(math.log(scale, 2))): - m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(2)) - elif scale == 3: - m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(3)) - else: - raise ValueError( - f"scale {scale} is not supported. " "Supported scales: 2^n and 3." - ) - super(Upsample, self).__init__(*m) - - -class UpsampleOneStep(nn.Sequential): - """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) - Used in lightweight SR to save parameters. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - - """ - - def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): - self.num_feat = num_feat - self.input_resolution = input_resolution - m = [] - m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1)) - m.append(nn.PixelShuffle(scale)) - super(UpsampleOneStep, self).__init__(*m) - - def flops(self): - h, w = self.input_resolution - flops = h * w * self.num_feat * 3 * 9 - return flops - - -class DAT(nn.Module): - """Dual Aggregation Transformer - Args: - img_size (int): Input image size. Default: 64 - in_chans (int): Number of input image channels. Default: 3 - embed_dim (int): Patch embedding dimension. Default: 180 - depths (tuple(int)): Depth of each residual group (number of DATB in each RG). - split_size (tuple(int)): Height and Width of spatial window. - num_heads (tuple(int)): Number of attention heads in different residual groups. - expansion_factor (float): Ratio of ffn hidden dim to embedding dim. Default: 4 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None - drop_rate (float): Dropout rate. Default: 0 - attn_drop_rate (float): Attention dropout rate. Default: 0 - drop_path_rate (float): Stochastic depth rate. Default: 0.1 - act_layer (nn.Module): Activation layer. Default: nn.GELU - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm - use_chk (bool): Whether to use checkpointing to save memory. - upscale: Upscale factor. 2/3/4 for image SR - img_range: Image range. 1. or 255. - resi_connection: The convolutional block before residual connection. '1conv'/'3conv' - """ - - def __init__(self, state_dict): - super().__init__() - - # defaults - img_size = 64 - in_chans = 3 - embed_dim = 180 - split_size = [2, 4] - depth = [2, 2, 2, 2] - num_heads = [2, 2, 2, 2] - expansion_factor = 4.0 - qkv_bias = True - qk_scale = None - drop_rate = 0.0 - attn_drop_rate = 0.0 - drop_path_rate = 0.1 - act_layer = nn.GELU - norm_layer = nn.LayerNorm - use_chk = False - upscale = 2 - img_range = 1.0 - resi_connection = "1conv" - upsampler = "pixelshuffle" - - self.model_arch = "DAT" - self.sub_type = "SR" - self.state = state_dict - - state_keys = state_dict.keys() - if "conv_before_upsample.0.weight" in state_keys: - if "conv_up1.weight" in state_keys: - upsampler = "nearest+conv" - else: - upsampler = "pixelshuffle" - supports_fp16 = False - elif "upsample.0.weight" in state_keys: - upsampler = "pixelshuffledirect" - else: - upsampler = "" - - num_feat = ( - state_dict.get("conv_before_upsample.0.weight", None).shape[1] - if state_dict.get("conv_before_upsample.weight", None) - else 64 - ) - - num_in_ch = state_dict["conv_first.weight"].shape[1] - in_chans = num_in_ch - if "conv_last.weight" in state_keys: - num_out_ch = state_dict["conv_last.weight"].shape[0] - else: - num_out_ch = num_in_ch - - upscale = 1 - if upsampler == "nearest+conv": - upsample_keys = [ - x for x in state_keys if "conv_up" in x and "bias" not in x - ] - - for upsample_key in upsample_keys: - upscale *= 2 - elif upsampler == "pixelshuffle": - upsample_keys = [ - x - for x in state_keys - if "upsample" in x and "conv" not in x and "bias" not in x - ] - for upsample_key in upsample_keys: - shape = state_dict[upsample_key].shape[0] - upscale *= math.sqrt(shape // num_feat) - upscale = int(upscale) - elif upsampler == "pixelshuffledirect": - upscale = int( - math.sqrt(state_dict["upsample.0.bias"].shape[0] // num_out_ch) - ) - - max_layer_num = 0 - max_block_num = 0 - for key in state_keys: - result = re.match(r"layers.(\d*).blocks.(\d*).norm1.weight", key) - if result: - layer_num, block_num = result.groups() - max_layer_num = max(max_layer_num, int(layer_num)) - max_block_num = max(max_block_num, int(block_num)) - - depth = [max_block_num + 1 for _ in range(max_layer_num + 1)] - - if "layers.0.blocks.1.attn.temperature" in state_keys: - num_heads_num = state_dict["layers.0.blocks.1.attn.temperature"].shape[0] - num_heads = [num_heads_num for _ in range(max_layer_num + 1)] - else: - num_heads = depth - - embed_dim = state_dict["conv_first.weight"].shape[0] - expansion_factor = float( - state_dict["layers.0.blocks.0.ffn.fc1.weight"].shape[0] / embed_dim - ) - - # TODO: could actually count the layers, but this should do - if "layers.0.conv.4.weight" in state_keys: - resi_connection = "3conv" - else: - resi_connection = "1conv" - - if "layers.0.blocks.2.attn.attn_mask_0" in state_keys: - attn_mask_0_x, attn_mask_0_y, attn_mask_0_z = state_dict[ - "layers.0.blocks.2.attn.attn_mask_0" - ].shape - - img_size = int(math.sqrt(attn_mask_0_x * attn_mask_0_y)) - - if "layers.0.blocks.0.attn.attns.0.rpe_biases" in state_keys: - split_sizes = ( - state_dict["layers.0.blocks.0.attn.attns.0.rpe_biases"][-1] + 1 - ) - split_size = [int(x) for x in split_sizes] - - self.in_nc = num_in_ch - self.out_nc = num_out_ch - self.num_feat = num_feat - self.embed_dim = embed_dim - self.num_heads = num_heads - self.depth = depth - self.scale = upscale - self.upsampler = upsampler - self.img_size = img_size - self.img_range = img_range - self.expansion_factor = expansion_factor - self.resi_connection = resi_connection - self.split_size = split_size - - self.supports_fp16 = False # Too much weirdness to support this at the moment - self.supports_bfp16 = True - self.min_size_restriction = 16 - - num_in_ch = in_chans - num_out_ch = in_chans - num_feat = 64 - self.img_range = img_range - if in_chans == 3: - rgb_mean = (0.4488, 0.4371, 0.4040) - self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) - else: - self.mean = torch.zeros(1, 1, 1, 1) - self.upscale = upscale - self.upsampler = upsampler - - # ------------------------- 1, Shallow Feature Extraction ------------------------- # - self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) - - # ------------------------- 2, Deep Feature Extraction ------------------------- # - self.num_layers = len(depth) - self.use_chk = use_chk - self.num_features = ( - self.embed_dim - ) = embed_dim # num_features for consistency with other models - heads = num_heads - - self.before_RG = nn.Sequential( - Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim) - ) - - curr_dim = embed_dim - dpr = [ - x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth)) - ] # stochastic depth decay rule - - self.layers = nn.ModuleList() - for i in range(self.num_layers): - layer = ResidualGroup( - dim=embed_dim, - num_heads=heads[i], - reso=img_size, - split_size=split_size, - expansion_factor=expansion_factor, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop_rate, - attn_drop=attn_drop_rate, - drop_paths=dpr[sum(depth[:i]) : sum(depth[: i + 1])], - act_layer=act_layer, - norm_layer=norm_layer, - depth=depth[i], - use_chk=use_chk, - resi_connection=resi_connection, - rg_idx=i, - ) - self.layers.append(layer) - - self.norm = norm_layer(curr_dim) - # build the last conv layer in deep feature extraction - if resi_connection == "1conv": - self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) - elif resi_connection == "3conv": - # to save parameters and memory - self.conv_after_body = nn.Sequential( - nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1), - ) - - # ------------------------- 3, Reconstruction ------------------------- # - if self.upsampler == "pixelshuffle": - # for classical SR - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR (to save parameters) - self.upsample = UpsampleOneStep( - upscale, embed_dim, num_out_ch, (img_size, img_size) - ) - - self.apply(self._init_weights) - self.load_state_dict(state_dict, strict=True) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance( - m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d) - ): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - - def forward_features(self, x): - _, _, H, W = x.shape - x_size = [H, W] - x = self.before_RG(x) - for layer in self.layers: - x = layer(x, x_size) - x = self.norm(x) - x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W) - - return x - - def forward(self, x): - """ - Input: x: (B, C, H, W) - """ - self.mean = self.mean.type_as(x) - x = (x - self.mean) * self.img_range - - if self.upsampler == "pixelshuffle": - # for image SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.upsample(x) - - x = x / self.img_range + self.mean - return x diff --git a/backend/comfy_nodes/chainner_models/architecture/HAT.py b/backend/comfy_nodes/chainner_models/architecture/HAT.py deleted file mode 100644 index 66947421..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/HAT.py +++ /dev/null @@ -1,1277 +0,0 @@ -# pylint: skip-file -# HAT from https://github.com/XPixelGroup/HAT/blob/main/hat/archs/hat_arch.py -import math -import re - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange - -from .timm.helpers import to_2tuple -from .timm.weight_init import trunc_normal_ - - -def drop_path(x, drop_prob: float = 0.0, training: bool = False): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py - """ - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * ( - x.ndim - 1 - ) # work with diff dim tensors, not just 2D ConvNets - random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) - random_tensor.floor_() # binarize - output = x.div(keep_prob) * random_tensor - return output - - -class DropPath(nn.Module): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py - """ - - def __init__(self, drop_prob=None): - super(DropPath, self).__init__() - self.drop_prob = drop_prob - - def forward(self, x): - return drop_path(x, self.drop_prob, self.training) # type: ignore - - -class ChannelAttention(nn.Module): - """Channel attention used in RCAN. - Args: - num_feat (int): Channel number of intermediate features. - squeeze_factor (int): Channel squeeze factor. Default: 16. - """ - - def __init__(self, num_feat, squeeze_factor=16): - super(ChannelAttention, self).__init__() - self.attention = nn.Sequential( - nn.AdaptiveAvgPool2d(1), - nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0), - nn.ReLU(inplace=True), - nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0), - nn.Sigmoid(), - ) - - def forward(self, x): - y = self.attention(x) - return x * y - - -class CAB(nn.Module): - def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30): - super(CAB, self).__init__() - - self.cab = nn.Sequential( - nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1), - nn.GELU(), - nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1), - ChannelAttention(num_feat, squeeze_factor), - ) - - def forward(self, x): - return self.cab(x) - - -class Mlp(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -def window_partition(x, window_size): - """ - Args: - x: (b, h, w, c) - window_size (int): window size - Returns: - windows: (num_windows*b, window_size, window_size, c) - """ - b, h, w, c = x.shape - x = x.view(b, h // window_size, window_size, w // window_size, window_size, c) - windows = ( - x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c) - ) - return windows - - -def window_reverse(windows, window_size, h, w): - """ - Args: - windows: (num_windows*b, window_size, window_size, c) - window_size (int): Window size - h (int): Height of image - w (int): Width of image - Returns: - x: (b, h, w, c) - """ - b = int(windows.shape[0] / (h * w / window_size / window_size)) - x = windows.view( - b, h // window_size, w // window_size, window_size, window_size, -1 - ) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1) - return x - - -class WindowAttention(nn.Module): - r"""Window based multi-head self attention (W-MSA) module with relative position bias. - It supports both of shifted and non-shifted window. - Args: - dim (int): Number of input channels. - window_size (tuple[int]): The height and width of the window. - num_heads (int): Number of attention heads. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set - attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 - proj_drop (float, optional): Dropout ratio of output. Default: 0.0 - """ - - def __init__( - self, - dim, - window_size, - num_heads, - qkv_bias=True, - qk_scale=None, - attn_drop=0.0, - proj_drop=0.0, - ): - super().__init__() - self.dim = dim - self.window_size = window_size # Wh, Ww - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - - # define a parameter table of relative position bias - self.relative_position_bias_table = nn.Parameter( # type: ignore - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) - ) # 2*Wh-1 * 2*Ww-1, nH - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - - self.proj_drop = nn.Dropout(proj_drop) - - trunc_normal_(self.relative_position_bias_table, std=0.02) - self.softmax = nn.Softmax(dim=-1) - - def forward(self, x, rpi, mask=None): - """ - Args: - x: input features with shape of (num_windows*b, n, c) - mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None - """ - b_, n, c = x.shape - qkv = ( - self.qkv(x) - .reshape(b_, n, 3, self.num_heads, c // self.num_heads) - .permute(2, 0, 3, 1, 4) - ) - q, k, v = ( - qkv[0], - qkv[1], - qkv[2], - ) # make torchscript happy (cannot use tensor as tuple) - - q = q * self.scale - attn = q @ k.transpose(-2, -1) - - relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view( - self.window_size[0] * self.window_size[1], - self.window_size[0] * self.window_size[1], - -1, - ) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute( - 2, 0, 1 - ).contiguous() # nH, Wh*Ww, Wh*Ww - attn = attn + relative_position_bias.unsqueeze(0) - - if mask is not None: - nw = mask.shape[0] - attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze( - 1 - ).unsqueeze(0) - attn = attn.view(-1, self.num_heads, n, n) - attn = self.softmax(attn) - else: - attn = self.softmax(attn) - - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(b_, n, c) - x = self.proj(x) - x = self.proj_drop(x) - return x - - -class HAB(nn.Module): - r"""Hybrid Attention Block. - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - num_heads (int): Number of attention heads. - window_size (int): Window size. - shift_size (int): Shift size for SW-MSA. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float, optional): Stochastic depth rate. Default: 0.0 - act_layer (nn.Module, optional): Activation layer. Default: nn.GELU - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__( - self, - dim, - input_resolution, - num_heads, - window_size=7, - shift_size=0, - compress_ratio=3, - squeeze_factor=30, - conv_scale=0.01, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - act_layer=nn.GELU, - norm_layer=nn.LayerNorm, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.num_heads = num_heads - self.window_size = window_size - self.shift_size = shift_size - self.mlp_ratio = mlp_ratio - if min(self.input_resolution) <= self.window_size: - # if window size is larger than input resolution, we don't partition windows - self.shift_size = 0 - self.window_size = min(self.input_resolution) - assert ( - 0 <= self.shift_size < self.window_size - ), "shift_size must in 0-window_size" - - self.norm1 = norm_layer(dim) - self.attn = WindowAttention( - dim, - window_size=to_2tuple(self.window_size), - num_heads=num_heads, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=drop, - ) - - self.conv_scale = conv_scale - self.conv_block = CAB( - num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor - ) - - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp( - in_features=dim, - hidden_features=mlp_hidden_dim, - act_layer=act_layer, - drop=drop, - ) - - def forward(self, x, x_size, rpi_sa, attn_mask): - h, w = x_size - b, _, c = x.shape - # assert seq_len == h * w, "input feature has wrong size" - - shortcut = x - x = self.norm1(x) - x = x.view(b, h, w, c) - - # Conv_X - conv_x = self.conv_block(x.permute(0, 3, 1, 2)) - conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c) - - # cyclic shift - if self.shift_size > 0: - shifted_x = torch.roll( - x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) - ) - attn_mask = attn_mask - else: - shifted_x = x - attn_mask = None - - # partition windows - x_windows = window_partition( - shifted_x, self.window_size - ) # nw*b, window_size, window_size, c - x_windows = x_windows.view( - -1, self.window_size * self.window_size, c - ) # nw*b, window_size*window_size, c - - # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size - attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask) - - # merge windows - attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c) - shifted_x = window_reverse(attn_windows, self.window_size, h, w) # b h' w' c - - # reverse cyclic shift - if self.shift_size > 0: - attn_x = torch.roll( - shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) - ) - else: - attn_x = shifted_x - attn_x = attn_x.view(b, h * w, c) - - # FFN - x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale - x = x + self.drop_path(self.mlp(self.norm2(x))) - - return x - - -class PatchMerging(nn.Module): - r"""Patch Merging Layer. - Args: - input_resolution (tuple[int]): Resolution of input feature. - dim (int): Number of input channels. - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) - self.norm = norm_layer(4 * dim) - - def forward(self, x): - """ - x: b, h*w, c - """ - h, w = self.input_resolution - b, seq_len, c = x.shape - assert seq_len == h * w, "input feature has wrong size" - assert h % 2 == 0 and w % 2 == 0, f"x size ({h}*{w}) are not even." - - x = x.view(b, h, w, c) - - x0 = x[:, 0::2, 0::2, :] # b h/2 w/2 c - x1 = x[:, 1::2, 0::2, :] # b h/2 w/2 c - x2 = x[:, 0::2, 1::2, :] # b h/2 w/2 c - x3 = x[:, 1::2, 1::2, :] # b h/2 w/2 c - x = torch.cat([x0, x1, x2, x3], -1) # b h/2 w/2 4*c - x = x.view(b, -1, 4 * c) # b h/2*w/2 4*c - - x = self.norm(x) - x = self.reduction(x) - - return x - - -class OCAB(nn.Module): - # overlapping cross-attention block - - def __init__( - self, - dim, - input_resolution, - window_size, - overlap_ratio, - num_heads, - qkv_bias=True, - qk_scale=None, - mlp_ratio=2, - norm_layer=nn.LayerNorm, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.window_size = window_size - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - self.overlap_win_size = int(window_size * overlap_ratio) + window_size - - self.norm1 = norm_layer(dim) - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.unfold = nn.Unfold( - kernel_size=(self.overlap_win_size, self.overlap_win_size), - stride=window_size, - padding=(self.overlap_win_size - window_size) // 2, - ) - - # define a parameter table of relative position bias - self.relative_position_bias_table = nn.Parameter( # type: ignore - torch.zeros( - (window_size + self.overlap_win_size - 1) - * (window_size + self.overlap_win_size - 1), - num_heads, - ) - ) # 2*Wh-1 * 2*Ww-1, nH - - trunc_normal_(self.relative_position_bias_table, std=0.02) - self.softmax = nn.Softmax(dim=-1) - - self.proj = nn.Linear(dim, dim) - - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp( - in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU - ) - - def forward(self, x, x_size, rpi): - h, w = x_size - b, _, c = x.shape - - shortcut = x - x = self.norm1(x) - x = x.view(b, h, w, c) - - qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2) # 3, b, c, h, w - q = qkv[0].permute(0, 2, 3, 1) # b, h, w, c - kv = torch.cat((qkv[1], qkv[2]), dim=1) # b, 2*c, h, w - - # partition windows - q_windows = window_partition( - q, self.window_size - ) # nw*b, window_size, window_size, c - q_windows = q_windows.view( - -1, self.window_size * self.window_size, c - ) # nw*b, window_size*window_size, c - - kv_windows = self.unfold(kv) # b, c*w*w, nw - kv_windows = rearrange( - kv_windows, - "b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch", - nc=2, - ch=c, - owh=self.overlap_win_size, - oww=self.overlap_win_size, - ).contiguous() # 2, nw*b, ow*ow, c - # Do the above rearrangement without the rearrange function - # kv_windows = kv_windows.view( - # 2, b, self.overlap_win_size, self.overlap_win_size, c, -1 - # ) - # kv_windows = kv_windows.permute(0, 5, 1, 2, 3, 4).contiguous() - # kv_windows = kv_windows.view( - # 2, -1, self.overlap_win_size * self.overlap_win_size, c - # ) - - k_windows, v_windows = kv_windows[0], kv_windows[1] # nw*b, ow*ow, c - - b_, nq, _ = q_windows.shape - _, n, _ = k_windows.shape - d = self.dim // self.num_heads - q = q_windows.reshape(b_, nq, self.num_heads, d).permute( - 0, 2, 1, 3 - ) # nw*b, nH, nq, d - k = k_windows.reshape(b_, n, self.num_heads, d).permute( - 0, 2, 1, 3 - ) # nw*b, nH, n, d - v = v_windows.reshape(b_, n, self.num_heads, d).permute( - 0, 2, 1, 3 - ) # nw*b, nH, n, d - - q = q * self.scale - attn = q @ k.transpose(-2, -1) - - relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view( - self.window_size * self.window_size, - self.overlap_win_size * self.overlap_win_size, - -1, - ) # ws*ws, wse*wse, nH - relative_position_bias = relative_position_bias.permute( - 2, 0, 1 - ).contiguous() # nH, ws*ws, wse*wse - attn = attn + relative_position_bias.unsqueeze(0) - - attn = self.softmax(attn) - attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim) - - # merge windows - attn_windows = attn_windows.view( - -1, self.window_size, self.window_size, self.dim - ) - x = window_reverse(attn_windows, self.window_size, h, w) # b h w c - x = x.view(b, h * w, self.dim) - - x = self.proj(x) + shortcut - - x = x + self.mlp(self.norm2(x)) - return x - - -class AttenBlocks(nn.Module): - """A series of attention blocks for one RHAG. - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - compress_ratio, - squeeze_factor, - conv_scale, - overlap_ratio, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.depth = depth - self.use_checkpoint = use_checkpoint - - # build blocks - self.blocks = nn.ModuleList( - [ - HAB( - dim=dim, - input_resolution=input_resolution, - num_heads=num_heads, - window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - compress_ratio=compress_ratio, - squeeze_factor=squeeze_factor, - conv_scale=conv_scale, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path[i] - if isinstance(drop_path, list) - else drop_path, - norm_layer=norm_layer, - ) - for i in range(depth) - ] - ) - - # OCAB - self.overlap_attn = OCAB( - dim=dim, - input_resolution=input_resolution, - window_size=window_size, - overlap_ratio=overlap_ratio, - num_heads=num_heads, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - mlp_ratio=mlp_ratio, # type: ignore - norm_layer=norm_layer, - ) - - # patch merging layer - if downsample is not None: - self.downsample = downsample( - input_resolution, dim=dim, norm_layer=norm_layer - ) - else: - self.downsample = None - - def forward(self, x, x_size, params): - for blk in self.blocks: - x = blk(x, x_size, params["rpi_sa"], params["attn_mask"]) - - x = self.overlap_attn(x, x_size, params["rpi_oca"]) - - if self.downsample is not None: - x = self.downsample(x) - return x - - -class RHAG(nn.Module): - """Residual Hybrid Attention Group (RHAG). - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - img_size: Input image size. - patch_size: Patch size. - resi_connection: The convolutional block before residual connection. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - compress_ratio, - squeeze_factor, - conv_scale, - overlap_ratio, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - img_size=224, - patch_size=4, - resi_connection="1conv", - ): - super(RHAG, self).__init__() - - self.dim = dim - self.input_resolution = input_resolution - - self.residual_group = AttenBlocks( - dim=dim, - input_resolution=input_resolution, - depth=depth, - num_heads=num_heads, - window_size=window_size, - compress_ratio=compress_ratio, - squeeze_factor=squeeze_factor, - conv_scale=conv_scale, - overlap_ratio=overlap_ratio, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path, - norm_layer=norm_layer, - downsample=downsample, - use_checkpoint=use_checkpoint, - ) - - if resi_connection == "1conv": - self.conv = nn.Conv2d(dim, dim, 3, 1, 1) - elif resi_connection == "identity": - self.conv = nn.Identity() - - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=0, - embed_dim=dim, - norm_layer=None, - ) - - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=0, - embed_dim=dim, - norm_layer=None, - ) - - def forward(self, x, x_size, params): - return ( - self.patch_embed( - self.conv( - self.patch_unembed(self.residual_group(x, x_size, params), x_size) - ) - ) - + x - ) - - -class PatchEmbed(nn.Module): - r"""Image to Patch Embedding - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [ - img_size[0] // patch_size[0], # type: ignore - img_size[1] // patch_size[1], # type: ignore - ] - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - - def forward(self, x): - x = x.flatten(2).transpose(1, 2) # b Ph*Pw c - if self.norm is not None: - x = self.norm(x) - return x - - -class PatchUnEmbed(nn.Module): - r"""Image to Patch Unembedding - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [ - img_size[0] // patch_size[0], # type: ignore - img_size[1] // patch_size[1], # type: ignore - ] - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - def forward(self, x, x_size): - x = ( - x.transpose(1, 2) - .contiguous() - .view(x.shape[0], self.embed_dim, x_size[0], x_size[1]) - ) # b Ph*Pw c - return x - - -class Upsample(nn.Sequential): - """Upsample module. - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - """ - - def __init__(self, scale, num_feat): - m = [] - if (scale & (scale - 1)) == 0: # scale = 2^n - for _ in range(int(math.log(scale, 2))): - m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(2)) - elif scale == 3: - m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(3)) - else: - raise ValueError( - f"scale {scale} is not supported. " "Supported scales: 2^n and 3." - ) - super(Upsample, self).__init__(*m) - - -class HAT(nn.Module): - r"""Hybrid Attention Transformer - A PyTorch implementation of : `Activating More Pixels in Image Super-Resolution Transformer`. - Some codes are based on SwinIR. - Args: - img_size (int | tuple(int)): Input image size. Default 64 - patch_size (int | tuple(int)): Patch size. Default: 1 - in_chans (int): Number of input image channels. Default: 3 - embed_dim (int): Patch embedding dimension. Default: 96 - depths (tuple(int)): Depth of each Swin Transformer layer. - num_heads (tuple(int)): Number of attention heads in different layers. - window_size (int): Window size. Default: 7 - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None - drop_rate (float): Dropout rate. Default: 0 - attn_drop_rate (float): Attention dropout rate. Default: 0 - drop_path_rate (float): Stochastic depth rate. Default: 0.1 - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. - ape (bool): If True, add absolute position embedding to the patch embedding. Default: False - patch_norm (bool): If True, add normalization after patch embedding. Default: True - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False - upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction - img_range: Image range. 1. or 255. - upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None - resi_connection: The convolutional block before residual connection. '1conv'/'3conv' - """ - - def __init__( - self, - state_dict, - **kwargs, - ): - super(HAT, self).__init__() - - # Defaults - img_size = 64 - patch_size = 1 - in_chans = 3 - embed_dim = 96 - depths = (6, 6, 6, 6) - num_heads = (6, 6, 6, 6) - window_size = 7 - compress_ratio = 3 - squeeze_factor = 30 - conv_scale = 0.01 - overlap_ratio = 0.5 - mlp_ratio = 4.0 - qkv_bias = True - qk_scale = None - drop_rate = 0.0 - attn_drop_rate = 0.0 - drop_path_rate = 0.1 - norm_layer = nn.LayerNorm - ape = False - patch_norm = True - use_checkpoint = False - upscale = 2 - img_range = 1.0 - upsampler = "" - resi_connection = "1conv" - - self.state = state_dict - self.model_arch = "HAT" - self.sub_type = "SR" - self.supports_fp16 = False - self.support_bf16 = True - self.min_size_restriction = 16 - - state_keys = list(state_dict.keys()) - - num_feat = state_dict["conv_last.weight"].shape[1] - in_chans = state_dict["conv_first.weight"].shape[1] - num_out_ch = state_dict["conv_last.weight"].shape[0] - embed_dim = state_dict["conv_first.weight"].shape[0] - - if "conv_before_upsample.0.weight" in state_keys: - if "conv_up1.weight" in state_keys: - upsampler = "nearest+conv" - else: - upsampler = "pixelshuffle" - supports_fp16 = False - elif "upsample.0.weight" in state_keys: - upsampler = "pixelshuffledirect" - else: - upsampler = "" - upscale = 1 - if upsampler == "nearest+conv": - upsample_keys = [ - x for x in state_keys if "conv_up" in x and "bias" not in x - ] - - for upsample_key in upsample_keys: - upscale *= 2 - elif upsampler == "pixelshuffle": - upsample_keys = [ - x - for x in state_keys - if "upsample" in x and "conv" not in x and "bias" not in x - ] - for upsample_key in upsample_keys: - shape = self.state[upsample_key].shape[0] - upscale *= math.sqrt(shape // num_feat) - upscale = int(upscale) - elif upsampler == "pixelshuffledirect": - upscale = int( - math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch) - ) - - max_layer_num = 0 - max_block_num = 0 - for key in state_keys: - result = re.match( - r"layers.(\d*).residual_group.blocks.(\d*).conv_block.cab.0.weight", key - ) - if result: - layer_num, block_num = result.groups() - max_layer_num = max(max_layer_num, int(layer_num)) - max_block_num = max(max_block_num, int(block_num)) - - depths = [max_block_num + 1 for _ in range(max_layer_num + 1)] - - if ( - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - in state_keys - ): - num_heads_num = self.state[ - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - ].shape[-1] - num_heads = [num_heads_num for _ in range(max_layer_num + 1)] - else: - num_heads = depths - - mlp_ratio = float( - self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0] - / embed_dim - ) - - # TODO: could actually count the layers, but this should do - if "layers.0.conv.4.weight" in state_keys: - resi_connection = "3conv" - else: - resi_connection = "1conv" - - window_size = int(math.sqrt(self.state["relative_position_index_SA"].shape[0])) - - # Not sure if this is needed or used at all anywhere in HAT's config - if "layers.0.residual_group.blocks.1.attn_mask" in state_keys: - img_size = int( - math.sqrt( - self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0] - ) - * window_size - ) - - self.window_size = window_size - self.shift_size = window_size // 2 - self.overlap_ratio = overlap_ratio - - self.in_nc = in_chans - self.out_nc = num_out_ch - self.num_feat = num_feat - self.embed_dim = embed_dim - self.num_heads = num_heads - self.depths = depths - self.window_size = window_size - self.mlp_ratio = mlp_ratio - self.scale = upscale - self.upsampler = upsampler - self.img_size = img_size - self.img_range = img_range - self.resi_connection = resi_connection - - num_in_ch = in_chans - # num_out_ch = in_chans - # num_feat = 64 - self.img_range = img_range - if in_chans == 3: - rgb_mean = (0.4488, 0.4371, 0.4040) - self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) - else: - self.mean = torch.zeros(1, 1, 1, 1) - self.upscale = upscale - self.upsampler = upsampler - - # relative position index - relative_position_index_SA = self.calculate_rpi_sa() - relative_position_index_OCA = self.calculate_rpi_oca() - self.register_buffer("relative_position_index_SA", relative_position_index_SA) - self.register_buffer("relative_position_index_OCA", relative_position_index_OCA) - - # ------------------------- 1, shallow feature extraction ------------------------- # - self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) - - # ------------------------- 2, deep feature extraction ------------------------- # - self.num_layers = len(depths) - self.embed_dim = embed_dim - self.ape = ape - self.patch_norm = patch_norm - self.num_features = embed_dim - self.mlp_ratio = mlp_ratio - - # split image into non-overlapping patches - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - num_patches = self.patch_embed.num_patches - patches_resolution = self.patch_embed.patches_resolution - self.patches_resolution = patches_resolution - - # merge non-overlapping patches into image - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - - # absolute position embedding - if self.ape: - self.absolute_pos_embed = nn.Parameter( # type: ignore[arg-type] - torch.zeros(1, num_patches, embed_dim) - ) - trunc_normal_(self.absolute_pos_embed, std=0.02) - - self.pos_drop = nn.Dropout(p=drop_rate) - - # stochastic depth - dpr = [ - x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) - ] # stochastic depth decay rule - - # build Residual Hybrid Attention Groups (RHAG) - self.layers = nn.ModuleList() - for i_layer in range(self.num_layers): - layer = RHAG( - dim=embed_dim, - input_resolution=(patches_resolution[0], patches_resolution[1]), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - compress_ratio=compress_ratio, - squeeze_factor=squeeze_factor, - conv_scale=conv_scale, - overlap_ratio=overlap_ratio, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop_rate, - attn_drop=attn_drop_rate, - drop_path=dpr[ - sum(depths[:i_layer]) : sum(depths[: i_layer + 1]) # type: ignore - ], # no impact on SR results - norm_layer=norm_layer, - downsample=None, - use_checkpoint=use_checkpoint, - img_size=img_size, - patch_size=patch_size, - resi_connection=resi_connection, - ) - self.layers.append(layer) - self.norm = norm_layer(self.num_features) - - # build the last conv layer in deep feature extraction - if resi_connection == "1conv": - self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) - elif resi_connection == "identity": - self.conv_after_body = nn.Identity() - - # ------------------------- 3, high quality image reconstruction ------------------------- # - if self.upsampler == "pixelshuffle": - # for classical SR - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - self.apply(self._init_weights) - self.load_state_dict(self.state, strict=False) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - - def calculate_rpi_sa(self): - # calculate relative position index for SA - coords_h = torch.arange(self.window_size) - coords_w = torch.arange(self.window_size) - coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww - coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = ( - coords_flatten[:, :, None] - coords_flatten[:, None, :] - ) # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute( - 1, 2, 0 - ).contiguous() # Wh*Ww, Wh*Ww, 2 - relative_coords[:, :, 0] += self.window_size - 1 # shift to start from 0 - relative_coords[:, :, 1] += self.window_size - 1 - relative_coords[:, :, 0] *= 2 * self.window_size - 1 - relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww - return relative_position_index - - def calculate_rpi_oca(self): - # calculate relative position index for OCA - window_size_ori = self.window_size - window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size) - - coords_h = torch.arange(window_size_ori) - coords_w = torch.arange(window_size_ori) - coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, ws, ws - coords_ori_flatten = torch.flatten(coords_ori, 1) # 2, ws*ws - - coords_h = torch.arange(window_size_ext) - coords_w = torch.arange(window_size_ext) - coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, wse, wse - coords_ext_flatten = torch.flatten(coords_ext, 1) # 2, wse*wse - - relative_coords = ( - coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None] - ) # 2, ws*ws, wse*wse - - relative_coords = relative_coords.permute( - 1, 2, 0 - ).contiguous() # ws*ws, wse*wse, 2 - relative_coords[:, :, 0] += ( - window_size_ori - window_size_ext + 1 - ) # shift to start from 0 - relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1 - - relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1 - relative_position_index = relative_coords.sum(-1) - return relative_position_index - - def calculate_mask(self, x_size): - # calculate attention mask for SW-MSA - h, w = x_size - img_mask = torch.zeros((1, h, w, 1)) # 1 h w 1 - h_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - w_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - mask_windows = window_partition( - img_mask, self.window_size - ) # nw, window_size, window_size, 1 - mask_windows = mask_windows.view(-1, self.window_size * self.window_size) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( - attn_mask == 0, float(0.0) - ) - - return attn_mask - - @torch.jit.ignore # type: ignore - def no_weight_decay(self): - return {"absolute_pos_embed"} - - @torch.jit.ignore # type: ignore - def no_weight_decay_keywords(self): - return {"relative_position_bias_table"} - - def check_image_size(self, x): - _, _, h, w = x.size() - mod_pad_h = (self.window_size - h % self.window_size) % self.window_size - mod_pad_w = (self.window_size - w % self.window_size) % self.window_size - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect") - return x - - def forward_features(self, x): - x_size = (x.shape[2], x.shape[3]) - - # Calculate attention mask and relative position index in advance to speed up inference. - # The original code is very time-cosuming for large window size. - attn_mask = self.calculate_mask(x_size).to(x.device) - params = { - "attn_mask": attn_mask, - "rpi_sa": self.relative_position_index_SA, - "rpi_oca": self.relative_position_index_OCA, - } - - x = self.patch_embed(x) - if self.ape: - x = x + self.absolute_pos_embed - x = self.pos_drop(x) - - for layer in self.layers: - x = layer(x, x_size, params) - - x = self.norm(x) # b seq_len c - x = self.patch_unembed(x, x_size) - - return x - - def forward(self, x): - H, W = x.shape[2:] - self.mean = self.mean.type_as(x) - x = (x - self.mean) * self.img_range - x = self.check_image_size(x) - - if self.upsampler == "pixelshuffle": - # for classical SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - - x = x / self.img_range + self.mean - - return x[:, :, : H * self.upscale, : W * self.upscale] diff --git a/backend/comfy_nodes/chainner_models/architecture/LICENSE-DAT b/backend/comfy_nodes/chainner_models/architecture/LICENSE-DAT deleted file mode 100644 index 261eeb9e..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-DAT +++ /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/backend/comfy_nodes/chainner_models/architecture/LICENSE-ESRGAN b/backend/comfy_nodes/chainner_models/architecture/LICENSE-ESRGAN deleted file mode 100644 index 261eeb9e..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-ESRGAN +++ /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/backend/comfy_nodes/chainner_models/architecture/LICENSE-HAT b/backend/comfy_nodes/chainner_models/architecture/LICENSE-HAT deleted file mode 100644 index 003e97e9..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-HAT +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2022 Xiangyu Chen - -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/backend/comfy_nodes/chainner_models/architecture/LICENSE-RealESRGAN b/backend/comfy_nodes/chainner_models/architecture/LICENSE-RealESRGAN deleted file mode 100644 index 552a1eea..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-RealESRGAN +++ /dev/null @@ -1,29 +0,0 @@ -BSD 3-Clause License - -Copyright (c) 2021, Xintao Wang -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 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/backend/comfy_nodes/chainner_models/architecture/LICENSE-SCUNet b/backend/comfy_nodes/chainner_models/architecture/LICENSE-SCUNet deleted file mode 100644 index ff75c988..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-SCUNet +++ /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 2022 Kai Zhang (cskaizhang@gmail.com, https://cszn.github.io/). All rights reserved. - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT 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/backend/comfy_nodes/chainner_models/architecture/LICENSE-SPSR b/backend/comfy_nodes/chainner_models/architecture/LICENSE-SPSR deleted file mode 100644 index 3245f3f9..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-SPSR +++ /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 2018-2022 BasicSR Authors - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT 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/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwiftSRGAN b/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwiftSRGAN deleted file mode 100644 index 0e259d42..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwiftSRGAN +++ /dev/null @@ -1,121 +0,0 @@ -Creative Commons Legal Code - -CC0 1.0 Universal - - CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE - LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN - ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS - INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES - REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS - PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM - THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED - HEREUNDER. - -Statement of Purpose - -The laws of most jurisdictions throughout the world automatically confer -exclusive Copyright and Related Rights (defined below) upon the creator -and subsequent owner(s) (each and all, an "owner") of an original work of -authorship and/or a database (each, a "Work"). - -Certain owners wish to permanently relinquish those rights to a Work for -the purpose of contributing to a commons of creative, cultural and -scientific works ("Commons") that the public can reliably and without fear -of later claims of infringement build upon, modify, incorporate in other -works, reuse and redistribute as freely as possible in any form whatsoever -and for any purposes, including without limitation commercial purposes. -These owners may contribute to the Commons to promote the ideal of a free -culture and the further production of creative, cultural and scientific -works, or to gain reputation or greater distribution for their Work in -part through the use and efforts of others. - -For these and/or other purposes and motivations, and without any -expectation of additional consideration or compensation, the person -associating CC0 with a Work (the "Affirmer"), to the extent that he or she -is an owner of Copyright and Related Rights in the Work, voluntarily -elects to apply CC0 to the Work and publicly distribute the Work under its -terms, with knowledge of his or her Copyright and Related Rights in the -Work and the meaning and intended legal effect of CC0 on those rights. - -1. Copyright and Related Rights. A Work made available under CC0 may be -protected by copyright and related or neighboring rights ("Copyright and -Related Rights"). Copyright and Related Rights include, but are not -limited to, the following: - - i. the right to reproduce, adapt, distribute, perform, display, - communicate, and translate a Work; - ii. moral rights retained by the original author(s) and/or performer(s); -iii. publicity and privacy rights pertaining to a person's image or - likeness depicted in a Work; - iv. rights protecting against unfair competition in regards to a Work, - subject to the limitations in paragraph 4(a), below; - v. rights protecting the extraction, dissemination, use and reuse of data - in a Work; - vi. database rights (such as those arising under Directive 96/9/EC of the - European Parliament and of the Council of 11 March 1996 on the legal - protection of databases, and under any national implementation - thereof, including any amended or successor version of such - directive); and -vii. other similar, equivalent or corresponding rights throughout the - world based on applicable law or treaty, and any national - implementations thereof. - -2. Waiver. To the greatest extent permitted by, but not in contravention -of, applicable law, Affirmer hereby overtly, fully, permanently, -irrevocably and unconditionally waives, abandons, and surrenders all of -Affirmer's Copyright and Related Rights and associated claims and causes -of action, whether now known or unknown (including existing as well as -future claims and causes of action), in the Work (i) in all territories -worldwide, (ii) for the maximum duration provided by applicable law or -treaty (including future time extensions), (iii) in any current or future -medium and for any number of copies, and (iv) for any purpose whatsoever, -including without limitation commercial, advertising or promotional -purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each -member of the public at large and to the detriment of Affirmer's heirs and -successors, fully intending that such Waiver shall not be subject to -revocation, rescission, cancellation, termination, or any other legal or -equitable action to disrupt the quiet enjoyment of the Work by the public -as contemplated by Affirmer's express Statement of Purpose. - -3. Public License Fallback. Should any part of the Waiver for any reason -be judged legally invalid or ineffective under applicable law, then the -Waiver shall be preserved to the maximum extent permitted taking into -account Affirmer's express Statement of Purpose. In addition, to the -extent the Waiver is so judged Affirmer hereby grants to each affected -person a royalty-free, non transferable, non sublicensable, non exclusive, -irrevocable and unconditional license to exercise Affirmer's Copyright and -Related Rights in the Work (i) in all territories worldwide, (ii) for the -maximum duration provided by applicable law or treaty (including future -time extensions), (iii) in any current or future medium and for any number -of copies, and (iv) for any purpose whatsoever, including without -limitation commercial, advertising or promotional purposes (the -"License"). The License shall be deemed effective as of the date CC0 was -applied by Affirmer to the Work. Should any part of the License for any -reason be judged legally invalid or ineffective under applicable law, such -partial invalidity or ineffectiveness shall not invalidate the remainder -of the License, and in such case Affirmer hereby affirms that he or she -will not (i) exercise any of his or her remaining Copyright and Related -Rights in the Work or (ii) assert any associated claims and causes of -action with respect to the Work, in either case contrary to Affirmer's -express Statement of Purpose. - -4. Limitations and Disclaimers. - - a. No trademark or patent rights held by Affirmer are waived, abandoned, - surrendered, licensed or otherwise affected by this document. - b. Affirmer offers the Work as-is and makes no representations or - warranties of any kind concerning the Work, express, implied, - statutory or otherwise, including without limitation warranties of - title, merchantability, fitness for a particular purpose, non - infringement, or the absence of latent or other defects, accuracy, or - the present or absence of errors, whether or not discoverable, all to - the greatest extent permissible under applicable law. - c. Affirmer disclaims responsibility for clearing rights of other persons - that may apply to the Work or any use thereof, including without - limitation any person's Copyright and Related Rights in the Work. - Further, Affirmer disclaims responsibility for obtaining any necessary - consents, permissions or other rights required for any use of the - Work. - d. Affirmer understands and acknowledges that Creative Commons is not a - party to this document and has no duty or obligation with respect to - this CC0 or use of the Work. diff --git a/backend/comfy_nodes/chainner_models/architecture/LICENSE-Swin2SR b/backend/comfy_nodes/chainner_models/architecture/LICENSE-Swin2SR deleted file mode 100644 index e5e4ee06..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-Swin2SR +++ /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 [2021] [SwinIR Authors] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT 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/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwinIR b/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwinIR deleted file mode 100644 index e5e4ee06..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-SwinIR +++ /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 [2021] [SwinIR Authors] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT 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/backend/comfy_nodes/chainner_models/architecture/LICENSE-lama b/backend/comfy_nodes/chainner_models/architecture/LICENSE-lama deleted file mode 100644 index ca822bb5..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LICENSE-lama +++ /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 [2021] Samsung Research - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT 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/backend/comfy_nodes/chainner_models/architecture/LaMa.py b/backend/comfy_nodes/chainner_models/architecture/LaMa.py deleted file mode 100644 index a781f3e4..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/LaMa.py +++ /dev/null @@ -1,694 +0,0 @@ -# pylint: skip-file -""" -Model adapted from advimman's lama project: https://github.com/advimman/lama -""" - -# Fast Fourier Convolution NeurIPS 2020 -# original implementation https://github.com/pkumivision/FFC/blob/main/model_zoo/ffc.py -# paper https://proceedings.neurips.cc/paper/2020/file/2fd5d41ec6cfab47e32164d5624269b1-Paper.pdf - -from typing import List - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torchvision.transforms.functional import InterpolationMode, rotate - - -class LearnableSpatialTransformWrapper(nn.Module): - def __init__(self, impl, pad_coef=0.5, angle_init_range=80, train_angle=True): - super().__init__() - self.impl = impl - self.angle = torch.rand(1) * angle_init_range - if train_angle: - self.angle = nn.Parameter(self.angle, requires_grad=True) - self.pad_coef = pad_coef - - def forward(self, x): - if torch.is_tensor(x): - return self.inverse_transform(self.impl(self.transform(x)), x) - elif isinstance(x, tuple): - x_trans = tuple(self.transform(elem) for elem in x) - y_trans = self.impl(x_trans) - return tuple( - self.inverse_transform(elem, orig_x) for elem, orig_x in zip(y_trans, x) - ) - else: - raise ValueError(f"Unexpected input type {type(x)}") - - def transform(self, x): - height, width = x.shape[2:] - pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef) - x_padded = F.pad(x, [pad_w, pad_w, pad_h, pad_h], mode="reflect") - x_padded_rotated = rotate( - x_padded, self.angle.to(x_padded), InterpolationMode.BILINEAR, fill=0 - ) - - return x_padded_rotated - - def inverse_transform(self, y_padded_rotated, orig_x): - height, width = orig_x.shape[2:] - pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef) - - y_padded = rotate( - y_padded_rotated, - -self.angle.to(y_padded_rotated), - InterpolationMode.BILINEAR, - fill=0, - ) - y_height, y_width = y_padded.shape[2:] - y = y_padded[:, :, pad_h : y_height - pad_h, pad_w : y_width - pad_w] - return y - - -class SELayer(nn.Module): - def __init__(self, channel, reduction=16): - super(SELayer, self).__init__() - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Sequential( - nn.Linear(channel, channel // reduction, bias=False), - nn.ReLU(inplace=True), - nn.Linear(channel // reduction, channel, bias=False), - nn.Sigmoid(), - ) - - def forward(self, x): - b, c, _, _ = x.size() - y = self.avg_pool(x).view(b, c) - y = self.fc(y).view(b, c, 1, 1) - res = x * y.expand_as(x) - return res - - -class FourierUnit(nn.Module): - def __init__( - self, - in_channels, - out_channels, - groups=1, - spatial_scale_factor=None, - spatial_scale_mode="bilinear", - spectral_pos_encoding=False, - use_se=False, - se_kwargs=None, - ffc3d=False, - fft_norm="ortho", - ): - # bn_layer not used - super(FourierUnit, self).__init__() - self.groups = groups - - self.conv_layer = torch.nn.Conv2d( - in_channels=in_channels * 2 + (2 if spectral_pos_encoding else 0), - out_channels=out_channels * 2, - kernel_size=1, - stride=1, - padding=0, - groups=self.groups, - bias=False, - ) - self.bn = torch.nn.BatchNorm2d(out_channels * 2) - self.relu = torch.nn.ReLU(inplace=True) - - # squeeze and excitation block - self.use_se = use_se - if use_se: - if se_kwargs is None: - se_kwargs = {} - self.se = SELayer(self.conv_layer.in_channels, **se_kwargs) - - self.spatial_scale_factor = spatial_scale_factor - self.spatial_scale_mode = spatial_scale_mode - self.spectral_pos_encoding = spectral_pos_encoding - self.ffc3d = ffc3d - self.fft_norm = fft_norm - - def forward(self, x): - half_check = False - if x.type() == "torch.cuda.HalfTensor": - # half only works on gpu anyway - half_check = True - - batch = x.shape[0] - - if self.spatial_scale_factor is not None: - orig_size = x.shape[-2:] - x = F.interpolate( - x, - scale_factor=self.spatial_scale_factor, - mode=self.spatial_scale_mode, - align_corners=False, - ) - - # (batch, c, h, w/2+1, 2) - fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1) - if half_check == True: - ffted = torch.fft.rfftn( - x.float(), dim=fft_dim, norm=self.fft_norm - ) # .type(torch.cuda.HalfTensor) - else: - ffted = torch.fft.rfftn(x, dim=fft_dim, norm=self.fft_norm) - - ffted = torch.stack((ffted.real, ffted.imag), dim=-1) - ffted = ffted.permute(0, 1, 4, 2, 3).contiguous() # (batch, c, 2, h, w/2+1) - ffted = ffted.view( - ( - batch, - -1, - ) - + ffted.size()[3:] - ) - - if self.spectral_pos_encoding: - height, width = ffted.shape[-2:] - coords_vert = ( - torch.linspace(0, 1, height)[None, None, :, None] - .expand(batch, 1, height, width) - .to(ffted) - ) - coords_hor = ( - torch.linspace(0, 1, width)[None, None, None, :] - .expand(batch, 1, height, width) - .to(ffted) - ) - ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1) - - if self.use_se: - ffted = self.se(ffted) - - if half_check == True: - ffted = self.conv_layer(ffted.half()) # (batch, c*2, h, w/2+1) - else: - ffted = self.conv_layer( - ffted - ) # .type(torch.cuda.FloatTensor) # (batch, c*2, h, w/2+1) - - ffted = self.relu(self.bn(ffted)) - # forcing to be always float - ffted = ffted.float() - - ffted = ( - ffted.view( - ( - batch, - -1, - 2, - ) - + ffted.size()[2:] - ) - .permute(0, 1, 3, 4, 2) - .contiguous() - ) # (batch,c, t, h, w/2+1, 2) - - ffted = torch.complex(ffted[..., 0], ffted[..., 1]) - - ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:] - output = torch.fft.irfftn( - ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm - ) - - if half_check == True: - output = output.half() - - if self.spatial_scale_factor is not None: - output = F.interpolate( - output, - size=orig_size, - mode=self.spatial_scale_mode, - align_corners=False, - ) - - return output - - -class SpectralTransform(nn.Module): - def __init__( - self, - in_channels, - out_channels, - stride=1, - groups=1, - enable_lfu=True, - separable_fu=False, - **fu_kwargs, - ): - # bn_layer not used - super(SpectralTransform, self).__init__() - self.enable_lfu = enable_lfu - if stride == 2: - self.downsample = nn.AvgPool2d(kernel_size=(2, 2), stride=2) - else: - self.downsample = nn.Identity() - - self.stride = stride - self.conv1 = nn.Sequential( - nn.Conv2d( - in_channels, out_channels // 2, kernel_size=1, groups=groups, bias=False - ), - nn.BatchNorm2d(out_channels // 2), - nn.ReLU(inplace=True), - ) - fu_class = FourierUnit - self.fu = fu_class(out_channels // 2, out_channels // 2, groups, **fu_kwargs) - if self.enable_lfu: - self.lfu = fu_class(out_channels // 2, out_channels // 2, groups) - self.conv2 = torch.nn.Conv2d( - out_channels // 2, out_channels, kernel_size=1, groups=groups, bias=False - ) - - def forward(self, x): - x = self.downsample(x) - x = self.conv1(x) - output = self.fu(x) - - if self.enable_lfu: - _, c, h, _ = x.shape - split_no = 2 - split_s = h // split_no - xs = torch.cat( - torch.split(x[:, : c // 4], split_s, dim=-2), dim=1 - ).contiguous() - xs = torch.cat(torch.split(xs, split_s, dim=-1), dim=1).contiguous() - xs = self.lfu(xs) - xs = xs.repeat(1, 1, split_no, split_no).contiguous() - else: - xs = 0 - - output = self.conv2(x + output + xs) - - return output - - -class FFC(nn.Module): - def __init__( - self, - in_channels, - out_channels, - kernel_size, - ratio_gin, - ratio_gout, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=False, - enable_lfu=True, - padding_type="reflect", - gated=False, - **spectral_kwargs, - ): - super(FFC, self).__init__() - - assert stride == 1 or stride == 2, "Stride should be 1 or 2." - self.stride = stride - - in_cg = int(in_channels * ratio_gin) - in_cl = in_channels - in_cg - out_cg = int(out_channels * ratio_gout) - out_cl = out_channels - out_cg - # groups_g = 1 if groups == 1 else int(groups * ratio_gout) - # groups_l = 1 if groups == 1 else groups - groups_g - - self.ratio_gin = ratio_gin - self.ratio_gout = ratio_gout - self.global_in_num = in_cg - - module = nn.Identity if in_cl == 0 or out_cl == 0 else nn.Conv2d - self.convl2l = module( - in_cl, - out_cl, - kernel_size, - stride, - padding, - dilation, - groups, - bias, - padding_mode=padding_type, - ) - module = nn.Identity if in_cl == 0 or out_cg == 0 else nn.Conv2d - self.convl2g = module( - in_cl, - out_cg, - kernel_size, - stride, - padding, - dilation, - groups, - bias, - padding_mode=padding_type, - ) - module = nn.Identity if in_cg == 0 or out_cl == 0 else nn.Conv2d - self.convg2l = module( - in_cg, - out_cl, - kernel_size, - stride, - padding, - dilation, - groups, - bias, - padding_mode=padding_type, - ) - module = nn.Identity if in_cg == 0 or out_cg == 0 else SpectralTransform - self.convg2g = module( - in_cg, - out_cg, - stride, - 1 if groups == 1 else groups // 2, - enable_lfu, - **spectral_kwargs, - ) - - self.gated = gated - module = ( - nn.Identity if in_cg == 0 or out_cl == 0 or not self.gated else nn.Conv2d - ) - self.gate = module(in_channels, 2, 1) - - def forward(self, x): - x_l, x_g = x if type(x) is tuple else (x, 0) - out_xl, out_xg = 0, 0 - - if self.gated: - total_input_parts = [x_l] - if torch.is_tensor(x_g): - total_input_parts.append(x_g) - total_input = torch.cat(total_input_parts, dim=1) - - gates = torch.sigmoid(self.gate(total_input)) - g2l_gate, l2g_gate = gates.chunk(2, dim=1) - else: - g2l_gate, l2g_gate = 1, 1 - - if self.ratio_gout != 1: - out_xl = self.convl2l(x_l) + self.convg2l(x_g) * g2l_gate - if self.ratio_gout != 0: - out_xg = self.convl2g(x_l) * l2g_gate + self.convg2g(x_g) - - return out_xl, out_xg - - -class FFC_BN_ACT(nn.Module): - def __init__( - self, - in_channels, - out_channels, - kernel_size, - ratio_gin, - ratio_gout, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=False, - norm_layer=nn.BatchNorm2d, - activation_layer=nn.Identity, - padding_type="reflect", - enable_lfu=True, - **kwargs, - ): - super(FFC_BN_ACT, self).__init__() - self.ffc = FFC( - in_channels, - out_channels, - kernel_size, - ratio_gin, - ratio_gout, - stride, - padding, - dilation, - groups, - bias, - enable_lfu, - padding_type=padding_type, - **kwargs, - ) - lnorm = nn.Identity if ratio_gout == 1 else norm_layer - gnorm = nn.Identity if ratio_gout == 0 else norm_layer - global_channels = int(out_channels * ratio_gout) - self.bn_l = lnorm(out_channels - global_channels) - self.bn_g = gnorm(global_channels) - - lact = nn.Identity if ratio_gout == 1 else activation_layer - gact = nn.Identity if ratio_gout == 0 else activation_layer - self.act_l = lact(inplace=True) - self.act_g = gact(inplace=True) - - def forward(self, x): - x_l, x_g = self.ffc(x) - x_l = self.act_l(self.bn_l(x_l)) - x_g = self.act_g(self.bn_g(x_g)) - return x_l, x_g - - -class FFCResnetBlock(nn.Module): - def __init__( - self, - dim, - padding_type, - norm_layer, - activation_layer=nn.ReLU, - dilation=1, - spatial_transform_kwargs=None, - inline=False, - **conv_kwargs, - ): - super().__init__() - self.conv1 = FFC_BN_ACT( - dim, - dim, - kernel_size=3, - padding=dilation, - dilation=dilation, - norm_layer=norm_layer, - activation_layer=activation_layer, - padding_type=padding_type, - **conv_kwargs, - ) - self.conv2 = FFC_BN_ACT( - dim, - dim, - kernel_size=3, - padding=dilation, - dilation=dilation, - norm_layer=norm_layer, - activation_layer=activation_layer, - padding_type=padding_type, - **conv_kwargs, - ) - if spatial_transform_kwargs is not None: - self.conv1 = LearnableSpatialTransformWrapper( - self.conv1, **spatial_transform_kwargs - ) - self.conv2 = LearnableSpatialTransformWrapper( - self.conv2, **spatial_transform_kwargs - ) - self.inline = inline - - def forward(self, x): - if self.inline: - x_l, x_g = ( - x[:, : -self.conv1.ffc.global_in_num], - x[:, -self.conv1.ffc.global_in_num :], - ) - else: - x_l, x_g = x if type(x) is tuple else (x, 0) - - id_l, id_g = x_l, x_g - - x_l, x_g = self.conv1((x_l, x_g)) - x_l, x_g = self.conv2((x_l, x_g)) - - x_l, x_g = id_l + x_l, id_g + x_g - out = x_l, x_g - if self.inline: - out = torch.cat(out, dim=1) - return out - - -class ConcatTupleLayer(nn.Module): - def forward(self, x): - assert isinstance(x, tuple) - x_l, x_g = x - assert torch.is_tensor(x_l) or torch.is_tensor(x_g) - if not torch.is_tensor(x_g): - return x_l - return torch.cat(x, dim=1) - - -class FFCResNetGenerator(nn.Module): - def __init__( - self, - input_nc, - output_nc, - ngf=64, - n_downsampling=3, - n_blocks=18, - norm_layer=nn.BatchNorm2d, - padding_type="reflect", - activation_layer=nn.ReLU, - up_norm_layer=nn.BatchNorm2d, - up_activation=nn.ReLU(True), - init_conv_kwargs={}, - downsample_conv_kwargs={}, - resnet_conv_kwargs={}, - spatial_transform_layers=None, - spatial_transform_kwargs={}, - max_features=1024, - out_ffc=False, - out_ffc_kwargs={}, - ): - assert n_blocks >= 0 - super().__init__() - """ - init_conv_kwargs = {'ratio_gin': 0, 'ratio_gout': 0, 'enable_lfu': False} - downsample_conv_kwargs = {'ratio_gin': '${generator.init_conv_kwargs.ratio_gout}', 'ratio_gout': '${generator.downsample_conv_kwargs.ratio_gin}', 'enable_lfu': False} - resnet_conv_kwargs = {'ratio_gin': 0.75, 'ratio_gout': '${generator.resnet_conv_kwargs.ratio_gin}', 'enable_lfu': False} - spatial_transform_kwargs = {} - out_ffc_kwargs = {} - """ - """ - print(input_nc, output_nc, ngf, n_downsampling, n_blocks, norm_layer, - padding_type, activation_layer, - up_norm_layer, up_activation, - spatial_transform_layers, - add_out_act, max_features, out_ffc, file=sys.stderr) - - 4 3 64 3 18 - reflect - - ReLU(inplace=True) - None sigmoid 1024 False - """ - init_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False} - downsample_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False} - resnet_conv_kwargs = { - "ratio_gin": 0.75, - "ratio_gout": 0.75, - "enable_lfu": False, - } - spatial_transform_kwargs = {} - out_ffc_kwargs = {} - - model = [ - nn.ReflectionPad2d(3), - FFC_BN_ACT( - input_nc, - ngf, - kernel_size=7, - padding=0, - norm_layer=norm_layer, - activation_layer=activation_layer, - **init_conv_kwargs, - ), - ] - - ### downsample - for i in range(n_downsampling): - mult = 2**i - if i == n_downsampling - 1: - cur_conv_kwargs = dict(downsample_conv_kwargs) - cur_conv_kwargs["ratio_gout"] = resnet_conv_kwargs.get("ratio_gin", 0) - else: - cur_conv_kwargs = downsample_conv_kwargs - model += [ - FFC_BN_ACT( - min(max_features, ngf * mult), - min(max_features, ngf * mult * 2), - kernel_size=3, - stride=2, - padding=1, - norm_layer=norm_layer, - activation_layer=activation_layer, - **cur_conv_kwargs, - ) - ] - - mult = 2**n_downsampling - feats_num_bottleneck = min(max_features, ngf * mult) - - ### resnet blocks - for i in range(n_blocks): - cur_resblock = FFCResnetBlock( - feats_num_bottleneck, - padding_type=padding_type, - activation_layer=activation_layer, - norm_layer=norm_layer, - **resnet_conv_kwargs, - ) - if spatial_transform_layers is not None and i in spatial_transform_layers: - cur_resblock = LearnableSpatialTransformWrapper( - cur_resblock, **spatial_transform_kwargs - ) - model += [cur_resblock] - - model += [ConcatTupleLayer()] - - ### upsample - for i in range(n_downsampling): - mult = 2 ** (n_downsampling - i) - model += [ - nn.ConvTranspose2d( - min(max_features, ngf * mult), - min(max_features, int(ngf * mult / 2)), - kernel_size=3, - stride=2, - padding=1, - output_padding=1, - ), - up_norm_layer(min(max_features, int(ngf * mult / 2))), - up_activation, - ] - - if out_ffc: - model += [ - FFCResnetBlock( - ngf, - padding_type=padding_type, - activation_layer=activation_layer, - norm_layer=norm_layer, - inline=True, - **out_ffc_kwargs, - ) - ] - - model += [ - nn.ReflectionPad2d(3), - nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), - ] - model.append(nn.Sigmoid()) - self.model = nn.Sequential(*model) - - def forward(self, image, mask): - return self.model(torch.cat([image, mask], dim=1)) - - -class LaMa(nn.Module): - def __init__(self, state_dict) -> None: - super(LaMa, self).__init__() - self.model_arch = "LaMa" - self.sub_type = "Inpaint" - self.in_nc = 4 - self.out_nc = 3 - self.scale = 1 - - self.min_size = None - self.pad_mod = 8 - self.pad_to_square = False - - self.model = FFCResNetGenerator(self.in_nc, self.out_nc) - self.state = { - k.replace("generator.model", "model.model"): v - for k, v in state_dict.items() - } - - self.supports_fp16 = False - self.support_bf16 = True - - self.load_state_dict(self.state, strict=False) - - def forward(self, img, mask): - masked_img = img * (1 - mask) - inpainted_mask = mask * self.model.forward(masked_img, mask) - result = inpainted_mask + (1 - mask) * img - return result diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/ChannelAttention.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/ChannelAttention.py deleted file mode 100644 index f4d52aa1..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/ChannelAttention.py +++ /dev/null @@ -1,110 +0,0 @@ -import math - -import torch.nn as nn - - -class CA_layer(nn.Module): - def __init__(self, channel, reduction=16): - super(CA_layer, self).__init__() - # global average pooling - self.gap = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Sequential( - nn.Conv2d(channel, channel // reduction, kernel_size=(1, 1), bias=False), - nn.GELU(), - nn.Conv2d(channel // reduction, channel, kernel_size=(1, 1), bias=False), - # nn.Sigmoid() - ) - - def forward(self, x): - y = self.fc(self.gap(x)) - return x * y.expand_as(x) - - -class Simple_CA_layer(nn.Module): - def __init__(self, channel): - super(Simple_CA_layer, self).__init__() - self.gap = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Conv2d( - in_channels=channel, - out_channels=channel, - kernel_size=1, - padding=0, - stride=1, - groups=1, - bias=True, - ) - - def forward(self, x): - return x * self.fc(self.gap(x)) - - -class ECA_layer(nn.Module): - """Constructs a ECA module. - Args: - channel: Number of channels of the input feature map - k_size: Adaptive selection of kernel size - """ - - def __init__(self, channel): - super(ECA_layer, self).__init__() - - b = 1 - gamma = 2 - k_size = int(abs(math.log(channel, 2) + b) / gamma) - k_size = k_size if k_size % 2 else k_size + 1 - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.conv = nn.Conv1d( - 1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False - ) - # self.sigmoid = nn.Sigmoid() - - def forward(self, x): - # x: input features with shape [b, c, h, w] - # b, c, h, w = x.size() - - # feature descriptor on the global spatial information - y = self.avg_pool(x) - - # Two different branches of ECA module - y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1) - - # Multi-scale information fusion - # y = self.sigmoid(y) - - return x * y.expand_as(x) - - -class ECA_MaxPool_layer(nn.Module): - """Constructs a ECA module. - Args: - channel: Number of channels of the input feature map - k_size: Adaptive selection of kernel size - """ - - def __init__(self, channel): - super(ECA_MaxPool_layer, self).__init__() - - b = 1 - gamma = 2 - k_size = int(abs(math.log(channel, 2) + b) / gamma) - k_size = k_size if k_size % 2 else k_size + 1 - self.max_pool = nn.AdaptiveMaxPool2d(1) - self.conv = nn.Conv1d( - 1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False - ) - # self.sigmoid = nn.Sigmoid() - - def forward(self, x): - # x: input features with shape [b, c, h, w] - # b, c, h, w = x.size() - - # feature descriptor on the global spatial information - y = self.max_pool(x) - - # Two different branches of ECA module - y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1) - - # Multi-scale information fusion - # y = self.sigmoid(y) - - return x * y.expand_as(x) diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/LICENSE b/backend/comfy_nodes/chainner_models/architecture/OmniSR/LICENSE deleted file mode 100644 index 261eeb9e..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/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/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSA.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSA.py deleted file mode 100644 index d7a12969..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSA.py +++ /dev/null @@ -1,577 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: OSA.py -# Created Date: Tuesday April 28th 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Sunday, 23rd April 2023 3:07:42 pm -# Modified By: Chen Xuanhong -# Copyright (c) 2020 Shanghai Jiao Tong University -############################################################# - -import torch -import torch.nn.functional as F -from einops import rearrange, repeat -from einops.layers.torch import Rearrange, Reduce -from torch import einsum, nn - -from .layernorm import LayerNorm2d - -# helpers - - -def exists(val): - return val is not None - - -def default(val, d): - return val if exists(val) else d - - -def cast_tuple(val, length=1): - return val if isinstance(val, tuple) else ((val,) * length) - - -# helper classes - - -class PreNormResidual(nn.Module): - def __init__(self, dim, fn): - super().__init__() - self.norm = nn.LayerNorm(dim) - self.fn = fn - - def forward(self, x): - return self.fn(self.norm(x)) + x - - -class Conv_PreNormResidual(nn.Module): - def __init__(self, dim, fn): - super().__init__() - self.norm = LayerNorm2d(dim) - self.fn = fn - - def forward(self, x): - return self.fn(self.norm(x)) + x - - -class FeedForward(nn.Module): - def __init__(self, dim, mult=2, dropout=0.0): - super().__init__() - inner_dim = int(dim * mult) - self.net = nn.Sequential( - nn.Linear(dim, inner_dim), - nn.GELU(), - nn.Dropout(dropout), - nn.Linear(inner_dim, dim), - nn.Dropout(dropout), - ) - - def forward(self, x): - return self.net(x) - - -class Conv_FeedForward(nn.Module): - def __init__(self, dim, mult=2, dropout=0.0): - super().__init__() - inner_dim = int(dim * mult) - self.net = nn.Sequential( - nn.Conv2d(dim, inner_dim, 1, 1, 0), - nn.GELU(), - nn.Dropout(dropout), - nn.Conv2d(inner_dim, dim, 1, 1, 0), - nn.Dropout(dropout), - ) - - def forward(self, x): - return self.net(x) - - -class Gated_Conv_FeedForward(nn.Module): - def __init__(self, dim, mult=1, bias=False, dropout=0.0): - super().__init__() - - hidden_features = int(dim * mult) - - self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias) - - self.dwconv = nn.Conv2d( - hidden_features * 2, - hidden_features * 2, - kernel_size=3, - stride=1, - padding=1, - groups=hidden_features * 2, - bias=bias, - ) - - self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias) - - def forward(self, x): - x = self.project_in(x) - x1, x2 = self.dwconv(x).chunk(2, dim=1) - x = F.gelu(x1) * x2 - x = self.project_out(x) - return x - - -# MBConv - - -class SqueezeExcitation(nn.Module): - def __init__(self, dim, shrinkage_rate=0.25): - super().__init__() - hidden_dim = int(dim * shrinkage_rate) - - self.gate = nn.Sequential( - Reduce("b c h w -> b c", "mean"), - nn.Linear(dim, hidden_dim, bias=False), - nn.SiLU(), - nn.Linear(hidden_dim, dim, bias=False), - nn.Sigmoid(), - Rearrange("b c -> b c 1 1"), - ) - - def forward(self, x): - return x * self.gate(x) - - -class MBConvResidual(nn.Module): - def __init__(self, fn, dropout=0.0): - super().__init__() - self.fn = fn - self.dropsample = Dropsample(dropout) - - def forward(self, x): - out = self.fn(x) - out = self.dropsample(out) - return out + x - - -class Dropsample(nn.Module): - def __init__(self, prob=0): - super().__init__() - self.prob = prob - - def forward(self, x): - device = x.device - - if self.prob == 0.0 or (not self.training): - return x - - keep_mask = ( - torch.FloatTensor((x.shape[0], 1, 1, 1), device=device).uniform_() - > self.prob - ) - return x * keep_mask / (1 - self.prob) - - -def MBConv( - dim_in, dim_out, *, downsample, expansion_rate=4, shrinkage_rate=0.25, dropout=0.0 -): - hidden_dim = int(expansion_rate * dim_out) - stride = 2 if downsample else 1 - - net = nn.Sequential( - nn.Conv2d(dim_in, hidden_dim, 1), - # nn.BatchNorm2d(hidden_dim), - nn.GELU(), - nn.Conv2d( - hidden_dim, hidden_dim, 3, stride=stride, padding=1, groups=hidden_dim - ), - # nn.BatchNorm2d(hidden_dim), - nn.GELU(), - SqueezeExcitation(hidden_dim, shrinkage_rate=shrinkage_rate), - nn.Conv2d(hidden_dim, dim_out, 1), - # nn.BatchNorm2d(dim_out) - ) - - if dim_in == dim_out and not downsample: - net = MBConvResidual(net, dropout=dropout) - - return net - - -# attention related classes -class Attention(nn.Module): - def __init__( - self, - dim, - dim_head=32, - dropout=0.0, - window_size=7, - with_pe=True, - ): - super().__init__() - assert ( - dim % dim_head - ) == 0, "dimension should be divisible by dimension per head" - - self.heads = dim // dim_head - self.scale = dim_head**-0.5 - self.with_pe = with_pe - - self.to_qkv = nn.Linear(dim, dim * 3, bias=False) - - self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout)) - - self.to_out = nn.Sequential( - nn.Linear(dim, dim, bias=False), nn.Dropout(dropout) - ) - - # relative positional bias - if self.with_pe: - self.rel_pos_bias = nn.Embedding((2 * window_size - 1) ** 2, self.heads) - - pos = torch.arange(window_size) - grid = torch.stack(torch.meshgrid(pos, pos)) - grid = rearrange(grid, "c i j -> (i j) c") - rel_pos = rearrange(grid, "i ... -> i 1 ...") - rearrange( - grid, "j ... -> 1 j ..." - ) - rel_pos += window_size - 1 - rel_pos_indices = (rel_pos * torch.tensor([2 * window_size - 1, 1])).sum( - dim=-1 - ) - - self.register_buffer("rel_pos_indices", rel_pos_indices, persistent=False) - - def forward(self, x): - batch, height, width, window_height, window_width, _, device, h = ( - *x.shape, - x.device, - self.heads, - ) - - # flatten - - x = rearrange(x, "b x y w1 w2 d -> (b x y) (w1 w2) d") - - # project for queries, keys, values - - q, k, v = self.to_qkv(x).chunk(3, dim=-1) - - # split heads - - q, k, v = map(lambda t: rearrange(t, "b n (h d ) -> b h n d", h=h), (q, k, v)) - - # scale - - q = q * self.scale - - # sim - - sim = einsum("b h i d, b h j d -> b h i j", q, k) - - # add positional bias - if self.with_pe: - bias = self.rel_pos_bias(self.rel_pos_indices) - sim = sim + rearrange(bias, "i j h -> h i j") - - # attention - - attn = self.attend(sim) - - # aggregate - - out = einsum("b h i j, b h j d -> b h i d", attn, v) - - # merge heads - - out = rearrange( - out, "b h (w1 w2) d -> b w1 w2 (h d)", w1=window_height, w2=window_width - ) - - # combine heads out - - out = self.to_out(out) - return rearrange(out, "(b x y) ... -> b x y ...", x=height, y=width) - - -class Block_Attention(nn.Module): - def __init__( - self, - dim, - dim_head=32, - bias=False, - dropout=0.0, - window_size=7, - with_pe=True, - ): - super().__init__() - assert ( - dim % dim_head - ) == 0, "dimension should be divisible by dimension per head" - - self.heads = dim // dim_head - self.ps = window_size - self.scale = dim_head**-0.5 - self.with_pe = with_pe - - self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias) - self.qkv_dwconv = nn.Conv2d( - dim * 3, - dim * 3, - kernel_size=3, - stride=1, - padding=1, - groups=dim * 3, - bias=bias, - ) - - self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout)) - - self.to_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) - - def forward(self, x): - # project for queries, keys, values - b, c, h, w = x.shape - - qkv = self.qkv_dwconv(self.qkv(x)) - q, k, v = qkv.chunk(3, dim=1) - - # split heads - - q, k, v = map( - lambda t: rearrange( - t, - "b (h d) (x w1) (y w2) -> (b x y) h (w1 w2) d", - h=self.heads, - w1=self.ps, - w2=self.ps, - ), - (q, k, v), - ) - - # scale - - q = q * self.scale - - # sim - - sim = einsum("b h i d, b h j d -> b h i j", q, k) - - # attention - attn = self.attend(sim) - - # aggregate - - out = einsum("b h i j, b h j d -> b h i d", attn, v) - - # merge heads - out = rearrange( - out, - "(b x y) head (w1 w2) d -> b (head d) (x w1) (y w2)", - x=h // self.ps, - y=w // self.ps, - head=self.heads, - w1=self.ps, - w2=self.ps, - ) - - out = self.to_out(out) - return out - - -class Channel_Attention(nn.Module): - def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7): - super(Channel_Attention, self).__init__() - self.heads = heads - - self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) - - self.ps = window_size - - self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias) - self.qkv_dwconv = nn.Conv2d( - dim * 3, - dim * 3, - kernel_size=3, - stride=1, - padding=1, - groups=dim * 3, - bias=bias, - ) - self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) - - def forward(self, x): - b, c, h, w = x.shape - - qkv = self.qkv_dwconv(self.qkv(x)) - qkv = qkv.chunk(3, dim=1) - - q, k, v = map( - lambda t: rearrange( - t, - "b (head d) (h ph) (w pw) -> b (h w) head d (ph pw)", - ph=self.ps, - pw=self.ps, - head=self.heads, - ), - qkv, - ) - - q = F.normalize(q, dim=-1) - k = F.normalize(k, dim=-1) - - attn = (q @ k.transpose(-2, -1)) * self.temperature - attn = attn.softmax(dim=-1) - out = attn @ v - - out = rearrange( - out, - "b (h w) head d (ph pw) -> b (head d) (h ph) (w pw)", - h=h // self.ps, - w=w // self.ps, - ph=self.ps, - pw=self.ps, - head=self.heads, - ) - - out = self.project_out(out) - - return out - - -class Channel_Attention_grid(nn.Module): - def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7): - super(Channel_Attention_grid, self).__init__() - self.heads = heads - - self.temperature = nn.Parameter(torch.ones(heads, 1, 1)) - - self.ps = window_size - - self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias) - self.qkv_dwconv = nn.Conv2d( - dim * 3, - dim * 3, - kernel_size=3, - stride=1, - padding=1, - groups=dim * 3, - bias=bias, - ) - self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) - - def forward(self, x): - b, c, h, w = x.shape - - qkv = self.qkv_dwconv(self.qkv(x)) - qkv = qkv.chunk(3, dim=1) - - q, k, v = map( - lambda t: rearrange( - t, - "b (head d) (h ph) (w pw) -> b (ph pw) head d (h w)", - ph=self.ps, - pw=self.ps, - head=self.heads, - ), - qkv, - ) - - q = F.normalize(q, dim=-1) - k = F.normalize(k, dim=-1) - - attn = (q @ k.transpose(-2, -1)) * self.temperature - attn = attn.softmax(dim=-1) - out = attn @ v - - out = rearrange( - out, - "b (ph pw) head d (h w) -> b (head d) (h ph) (w pw)", - h=h // self.ps, - w=w // self.ps, - ph=self.ps, - pw=self.ps, - head=self.heads, - ) - - out = self.project_out(out) - - return out - - -class OSA_Block(nn.Module): - def __init__( - self, - channel_num=64, - bias=True, - ffn_bias=True, - window_size=8, - with_pe=False, - dropout=0.0, - ): - super(OSA_Block, self).__init__() - - w = window_size - - self.layer = nn.Sequential( - MBConv( - channel_num, - channel_num, - downsample=False, - expansion_rate=1, - shrinkage_rate=0.25, - ), - Rearrange( - "b d (x w1) (y w2) -> b x y w1 w2 d", w1=w, w2=w - ), # block-like attention - PreNormResidual( - channel_num, - Attention( - dim=channel_num, - dim_head=channel_num // 4, - dropout=dropout, - window_size=window_size, - with_pe=with_pe, - ), - ), - Rearrange("b x y w1 w2 d -> b d (x w1) (y w2)"), - Conv_PreNormResidual( - channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout) - ), - # channel-like attention - Conv_PreNormResidual( - channel_num, - Channel_Attention( - dim=channel_num, heads=4, dropout=dropout, window_size=window_size - ), - ), - Conv_PreNormResidual( - channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout) - ), - Rearrange( - "b d (w1 x) (w2 y) -> b x y w1 w2 d", w1=w, w2=w - ), # grid-like attention - PreNormResidual( - channel_num, - Attention( - dim=channel_num, - dim_head=channel_num // 4, - dropout=dropout, - window_size=window_size, - with_pe=with_pe, - ), - ), - Rearrange("b x y w1 w2 d -> b d (w1 x) (w2 y)"), - Conv_PreNormResidual( - channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout) - ), - # channel-like attention - Conv_PreNormResidual( - channel_num, - Channel_Attention_grid( - dim=channel_num, heads=4, dropout=dropout, window_size=window_size - ), - ), - Conv_PreNormResidual( - channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout) - ), - ) - - def forward(self, x): - out = self.layer(x) - return out diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSAG.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSAG.py deleted file mode 100644 index 477e81f9..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/OSAG.py +++ /dev/null @@ -1,60 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: OSAG.py -# Created Date: Tuesday April 28th 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Sunday, 23rd April 2023 3:08:49 pm -# Modified By: Chen Xuanhong -# Copyright (c) 2020 Shanghai Jiao Tong University -############################################################# - - -import torch.nn as nn - -from .esa import ESA -from .OSA import OSA_Block - - -class OSAG(nn.Module): - def __init__( - self, - channel_num=64, - bias=True, - block_num=4, - ffn_bias=False, - window_size=0, - pe=False, - ): - super(OSAG, self).__init__() - - # print("window_size: %d" % (window_size)) - # print("with_pe", pe) - # print("ffn_bias: %d" % (ffn_bias)) - - # block_script_name = kwargs.get("block_script_name", "OSA") - # block_class_name = kwargs.get("block_class_name", "OSA_Block") - - # script_name = "." + block_script_name - # package = __import__(script_name, fromlist=True) - block_class = OSA_Block # getattr(package, block_class_name) - group_list = [] - for _ in range(block_num): - temp_res = block_class( - channel_num, - bias, - ffn_bias=ffn_bias, - window_size=window_size, - with_pe=pe, - ) - group_list.append(temp_res) - group_list.append(nn.Conv2d(channel_num, channel_num, 1, 1, 0, bias=bias)) - self.residual_layer = nn.Sequential(*group_list) - esa_channel = max(channel_num // 4, 16) - self.esa = ESA(esa_channel, channel_num) - - def forward(self, x): - out = self.residual_layer(x) - out = out + x - return self.esa(out) diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/OmniSR.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/OmniSR.py deleted file mode 100644 index 1e1c3f35..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/OmniSR.py +++ /dev/null @@ -1,143 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: OmniSR.py -# Created Date: Tuesday April 28th 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Sunday, 23rd April 2023 3:06:36 pm -# Modified By: Chen Xuanhong -# Copyright (c) 2020 Shanghai Jiao Tong University -############################################################# - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from .OSAG import OSAG -from .pixelshuffle import pixelshuffle_block - - -class OmniSR(nn.Module): - def __init__( - self, - state_dict, - **kwargs, - ): - super(OmniSR, self).__init__() - self.state = state_dict - - bias = True # Fine to assume this for now - block_num = 1 # Fine to assume this for now - ffn_bias = True - pe = True - - num_feat = state_dict["input.weight"].shape[0] or 64 - num_in_ch = state_dict["input.weight"].shape[1] or 3 - num_out_ch = num_in_ch # we can just assume this for now. pixelshuffle smh - - pixelshuffle_shape = state_dict["up.0.weight"].shape[0] - up_scale = math.sqrt(pixelshuffle_shape / num_out_ch) - if up_scale - int(up_scale) > 0: - print( - "out_nc is probably different than in_nc, scale calculation might be wrong" - ) - up_scale = int(up_scale) - res_num = 0 - for key in state_dict.keys(): - if "residual_layer" in key: - temp_res_num = int(key.split(".")[1]) - if temp_res_num > res_num: - res_num = temp_res_num - res_num = res_num + 1 # zero-indexed - - residual_layer = [] - self.res_num = res_num - - if ( - "residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight" - in state_dict.keys() - ): - rel_pos_bias_weight = state_dict[ - "residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight" - ].shape[0] - self.window_size = int((math.sqrt(rel_pos_bias_weight) + 1) / 2) - else: - self.window_size = 8 - - self.up_scale = up_scale - - for _ in range(res_num): - temp_res = OSAG( - channel_num=num_feat, - bias=bias, - block_num=block_num, - ffn_bias=ffn_bias, - window_size=self.window_size, - pe=pe, - ) - residual_layer.append(temp_res) - self.residual_layer = nn.Sequential(*residual_layer) - self.input = nn.Conv2d( - in_channels=num_in_ch, - out_channels=num_feat, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ) - self.output = nn.Conv2d( - in_channels=num_feat, - out_channels=num_feat, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ) - self.up = pixelshuffle_block(num_feat, num_out_ch, up_scale, bias=bias) - - # self.tail = pixelshuffle_block(num_feat,num_out_ch,up_scale,bias=bias) - - # for m in self.modules(): - # if isinstance(m, nn.Conv2d): - # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - # m.weight.data.normal_(0, sqrt(2. / n)) - - # chaiNNer specific stuff - self.model_arch = "OmniSR" - self.sub_type = "SR" - self.in_nc = num_in_ch - self.out_nc = num_out_ch - self.num_feat = num_feat - self.scale = up_scale - - self.supports_fp16 = True # TODO: Test this - self.supports_bfp16 = True - self.min_size_restriction = 16 - - self.load_state_dict(state_dict, strict=False) - - def check_image_size(self, x): - _, _, h, w = x.size() - # import pdb; pdb.set_trace() - mod_pad_h = (self.window_size - h % self.window_size) % self.window_size - mod_pad_w = (self.window_size - w % self.window_size) % self.window_size - # x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect') - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "constant", 0) - return x - - def forward(self, x): - H, W = x.shape[2:] - x = self.check_image_size(x) - - residual = self.input(x) - out = self.residual_layer(residual) - - # origin - out = torch.add(self.output(out), residual) - out = self.up(out) - - out = out[:, :, : H * self.up_scale, : W * self.up_scale] - return out diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/esa.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/esa.py deleted file mode 100644 index f9ce7f7a..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/esa.py +++ /dev/null @@ -1,294 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: esa.py -# Created Date: Tuesday April 28th 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Thursday, 20th April 2023 9:28:06 am -# Modified By: Chen Xuanhong -# Copyright (c) 2020 Shanghai Jiao Tong University -############################################################# - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from .layernorm import LayerNorm2d - - -def moment(x, dim=(2, 3), k=2): - assert len(x.size()) == 4 - mean = torch.mean(x, dim=dim).unsqueeze(-1).unsqueeze(-1) - mk = (1 / (x.size(2) * x.size(3))) * torch.sum(torch.pow(x - mean, k), dim=dim) - return mk - - -class ESA(nn.Module): - """ - Modification of Enhanced Spatial Attention (ESA), which is proposed by - `Residual Feature Aggregation Network for Image Super-Resolution` - Note: `conv_max` and `conv3_` are NOT used here, so the corresponding codes - are deleted. - """ - - def __init__(self, esa_channels, n_feats, conv=nn.Conv2d): - super(ESA, self).__init__() - f = esa_channels - self.conv1 = conv(n_feats, f, kernel_size=1) - self.conv_f = conv(f, f, kernel_size=1) - self.conv2 = conv(f, f, kernel_size=3, stride=2, padding=0) - self.conv3 = conv(f, f, kernel_size=3, padding=1) - self.conv4 = conv(f, n_feats, kernel_size=1) - self.sigmoid = nn.Sigmoid() - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - c1_ = self.conv1(x) - c1 = self.conv2(c1_) - v_max = F.max_pool2d(c1, kernel_size=7, stride=3) - c3 = self.conv3(v_max) - c3 = F.interpolate( - c3, (x.size(2), x.size(3)), mode="bilinear", align_corners=False - ) - cf = self.conv_f(c1_) - c4 = self.conv4(c3 + cf) - m = self.sigmoid(c4) - return x * m - - -class LK_ESA(nn.Module): - def __init__( - self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True - ): - super(LK_ESA, self).__init__() - f = esa_channels - self.conv1 = conv(n_feats, f, kernel_size=1) - self.conv_f = conv(f, f, kernel_size=1) - - kernel_size = 17 - kernel_expand = kernel_expand - padding = kernel_size // 2 - - self.vec_conv = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(1, kernel_size), - padding=(0, padding), - groups=2, - bias=bias, - ) - self.vec_conv3x1 = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(1, 3), - padding=(0, 1), - groups=2, - bias=bias, - ) - - self.hor_conv = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(kernel_size, 1), - padding=(padding, 0), - groups=2, - bias=bias, - ) - self.hor_conv1x3 = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(3, 1), - padding=(1, 0), - groups=2, - bias=bias, - ) - - self.conv4 = conv(f, n_feats, kernel_size=1) - self.sigmoid = nn.Sigmoid() - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - c1_ = self.conv1(x) - - res = self.vec_conv(c1_) + self.vec_conv3x1(c1_) - res = self.hor_conv(res) + self.hor_conv1x3(res) - - cf = self.conv_f(c1_) - c4 = self.conv4(res + cf) - m = self.sigmoid(c4) - return x * m - - -class LK_ESA_LN(nn.Module): - def __init__( - self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True - ): - super(LK_ESA_LN, self).__init__() - f = esa_channels - self.conv1 = conv(n_feats, f, kernel_size=1) - self.conv_f = conv(f, f, kernel_size=1) - - kernel_size = 17 - kernel_expand = kernel_expand - padding = kernel_size // 2 - - self.norm = LayerNorm2d(n_feats) - - self.vec_conv = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(1, kernel_size), - padding=(0, padding), - groups=2, - bias=bias, - ) - self.vec_conv3x1 = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(1, 3), - padding=(0, 1), - groups=2, - bias=bias, - ) - - self.hor_conv = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(kernel_size, 1), - padding=(padding, 0), - groups=2, - bias=bias, - ) - self.hor_conv1x3 = nn.Conv2d( - in_channels=f * kernel_expand, - out_channels=f * kernel_expand, - kernel_size=(3, 1), - padding=(1, 0), - groups=2, - bias=bias, - ) - - self.conv4 = conv(f, n_feats, kernel_size=1) - self.sigmoid = nn.Sigmoid() - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - c1_ = self.norm(x) - c1_ = self.conv1(c1_) - - res = self.vec_conv(c1_) + self.vec_conv3x1(c1_) - res = self.hor_conv(res) + self.hor_conv1x3(res) - - cf = self.conv_f(c1_) - c4 = self.conv4(res + cf) - m = self.sigmoid(c4) - return x * m - - -class AdaGuidedFilter(nn.Module): - def __init__( - self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True - ): - super(AdaGuidedFilter, self).__init__() - - self.gap = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Conv2d( - in_channels=n_feats, - out_channels=1, - kernel_size=1, - padding=0, - stride=1, - groups=1, - bias=True, - ) - - self.r = 5 - - def box_filter(self, x, r): - channel = x.shape[1] - kernel_size = 2 * r + 1 - weight = 1.0 / (kernel_size**2) - box_kernel = weight * torch.ones( - (channel, 1, kernel_size, kernel_size), dtype=torch.float32, device=x.device - ) - output = F.conv2d(x, weight=box_kernel, stride=1, padding=r, groups=channel) - return output - - def forward(self, x): - _, _, H, W = x.shape - N = self.box_filter( - torch.ones((1, 1, H, W), dtype=x.dtype, device=x.device), self.r - ) - - # epsilon = self.fc(self.gap(x)) - # epsilon = torch.pow(epsilon, 2) - epsilon = 1e-2 - - mean_x = self.box_filter(x, self.r) / N - var_x = self.box_filter(x * x, self.r) / N - mean_x * mean_x - - A = var_x / (var_x + epsilon) - b = (1 - A) * mean_x - m = A * x + b - - # mean_A = self.box_filter(A, self.r) / N - # mean_b = self.box_filter(b, self.r) / N - # m = mean_A * x + mean_b - return x * m - - -class AdaConvGuidedFilter(nn.Module): - def __init__( - self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True - ): - super(AdaConvGuidedFilter, self).__init__() - f = esa_channels - - self.conv_f = conv(f, f, kernel_size=1) - - kernel_size = 17 - kernel_expand = kernel_expand - padding = kernel_size // 2 - - self.vec_conv = nn.Conv2d( - in_channels=f, - out_channels=f, - kernel_size=(1, kernel_size), - padding=(0, padding), - groups=f, - bias=bias, - ) - - self.hor_conv = nn.Conv2d( - in_channels=f, - out_channels=f, - kernel_size=(kernel_size, 1), - padding=(padding, 0), - groups=f, - bias=bias, - ) - - self.gap = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Conv2d( - in_channels=f, - out_channels=f, - kernel_size=1, - padding=0, - stride=1, - groups=1, - bias=True, - ) - - def forward(self, x): - y = self.vec_conv(x) - y = self.hor_conv(y) - - sigma = torch.pow(y, 2) - epsilon = self.fc(self.gap(y)) - - weight = sigma / (sigma + epsilon) - - m = weight * x + (1 - weight) - - return x * m diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/layernorm.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/layernorm.py deleted file mode 100644 index 731a25f7..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/layernorm.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: layernorm.py -# Created Date: Tuesday April 28th 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Thursday, 20th April 2023 9:28:20 am -# Modified By: Chen Xuanhong -# Copyright (c) 2020 Shanghai Jiao Tong University -############################################################# - -import torch -import torch.nn as nn - - -class LayerNormFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x, weight, bias, eps): - ctx.eps = eps - N, C, H, W = x.size() - mu = x.mean(1, keepdim=True) - var = (x - mu).pow(2).mean(1, keepdim=True) - y = (x - mu) / (var + eps).sqrt() - ctx.save_for_backward(y, var, weight) - y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1) - return y - - @staticmethod - def backward(ctx, grad_output): - eps = ctx.eps - - N, C, H, W = grad_output.size() - y, var, weight = ctx.saved_variables - g = grad_output * weight.view(1, C, 1, 1) - mean_g = g.mean(dim=1, keepdim=True) - - mean_gy = (g * y).mean(dim=1, keepdim=True) - gx = 1.0 / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g) - return ( - gx, - (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), - grad_output.sum(dim=3).sum(dim=2).sum(dim=0), - None, - ) - - -class LayerNorm2d(nn.Module): - def __init__(self, channels, eps=1e-6): - super(LayerNorm2d, self).__init__() - self.register_parameter("weight", nn.Parameter(torch.ones(channels))) - self.register_parameter("bias", nn.Parameter(torch.zeros(channels))) - self.eps = eps - - def forward(self, x): - return LayerNormFunction.apply(x, self.weight, self.bias, self.eps) - - -class GRN(nn.Module): - """GRN (Global Response Normalization) layer""" - - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1)) - self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1)) - - def forward(self, x): - Gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True) - Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6) - return self.gamma * (x * Nx) + self.beta + x diff --git a/backend/comfy_nodes/chainner_models/architecture/OmniSR/pixelshuffle.py b/backend/comfy_nodes/chainner_models/architecture/OmniSR/pixelshuffle.py deleted file mode 100644 index 4260fb7c..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/OmniSR/pixelshuffle.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding:utf-8 -*- -############################################################# -# File: pixelshuffle.py -# Created Date: Friday July 1st 2022 -# Author: Chen Xuanhong -# Email: chenxuanhongzju@outlook.com -# Last Modified: Friday, 1st July 2022 10:18:39 am -# Modified By: Chen Xuanhong -# Copyright (c) 2022 Shanghai Jiao Tong University -############################################################# - -import torch.nn as nn - - -def pixelshuffle_block( - in_channels, out_channels, upscale_factor=2, kernel_size=3, bias=False -): - """ - Upsample features according to `upscale_factor`. - """ - padding = kernel_size // 2 - conv = nn.Conv2d( - in_channels, - out_channels * (upscale_factor**2), - kernel_size, - padding=1, - bias=bias, - ) - pixel_shuffle = nn.PixelShuffle(upscale_factor) - return nn.Sequential(*[conv, pixel_shuffle]) diff --git a/backend/comfy_nodes/chainner_models/architecture/RRDB.py b/backend/comfy_nodes/chainner_models/architecture/RRDB.py deleted file mode 100644 index b50db7c2..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/RRDB.py +++ /dev/null @@ -1,296 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -import functools -import math -import re -from collections import OrderedDict - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from . import block as B - - -# Borrowed from https://github.com/rlaphoenix/VSGAN/blob/master/vsgan/archs/ESRGAN.py -# Which enhanced stuff that was already here -class RRDBNet(nn.Module): - def __init__( - self, - state_dict, - norm=None, - act: str = "leakyrelu", - upsampler: str = "upconv", - mode: B.ConvMode = "CNA", - ) -> None: - """ - ESRGAN - Enhanced Super-Resolution Generative Adversarial Networks. - By Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Yu Qiao, - and Chen Change Loy. - This is old-arch Residual in Residual Dense Block Network and is not - the newest revision that's available at github.com/xinntao/ESRGAN. - This is on purpose, the newest Network has severely limited the - potential use of the Network with no benefits. - This network supports model files from both new and old-arch. - Args: - norm: Normalization layer - act: Activation layer - upsampler: Upsample layer. upconv, pixel_shuffle - mode: Convolution mode - """ - super(RRDBNet, self).__init__() - self.model_arch = "ESRGAN" - self.sub_type = "SR" - - self.state = state_dict - self.norm = norm - self.act = act - self.upsampler = upsampler - self.mode = mode - - self.state_map = { - # currently supports old, new, and newer RRDBNet arch models - # ESRGAN, BSRGAN/RealSR, Real-ESRGAN - "model.0.weight": ("conv_first.weight",), - "model.0.bias": ("conv_first.bias",), - "model.1.sub./NB/.weight": ("trunk_conv.weight", "conv_body.weight"), - "model.1.sub./NB/.bias": ("trunk_conv.bias", "conv_body.bias"), - r"model.1.sub.\1.RDB\2.conv\3.0.\4": ( - r"RRDB_trunk\.(\d+)\.RDB(\d)\.conv(\d+)\.(weight|bias)", - r"body\.(\d+)\.rdb(\d)\.conv(\d+)\.(weight|bias)", - ), - } - if "params_ema" in self.state: - self.state = self.state["params_ema"] - # self.model_arch = "RealESRGAN" - self.num_blocks = self.get_num_blocks() - self.plus = any("conv1x1" in k for k in self.state.keys()) - if self.plus: - self.model_arch = "ESRGAN+" - - self.state = self.new_to_old_arch(self.state) - - self.key_arr = list(self.state.keys()) - - self.in_nc: int = self.state[self.key_arr[0]].shape[1] - self.out_nc: int = self.state[self.key_arr[-1]].shape[0] - - self.scale: int = self.get_scale() - self.num_filters: int = self.state[self.key_arr[0]].shape[0] - - c2x2 = False - if self.state["model.0.weight"].shape[-2] == 2: - c2x2 = True - self.scale = round(math.sqrt(self.scale / 4)) - self.model_arch = "ESRGAN-2c2" - - self.supports_fp16 = True - self.supports_bfp16 = True - self.min_size_restriction = None - - # Detect if pixelunshuffle was used (Real-ESRGAN) - if self.in_nc in (self.out_nc * 4, self.out_nc * 16) and self.out_nc in ( - self.in_nc / 4, - self.in_nc / 16, - ): - self.shuffle_factor = int(math.sqrt(self.in_nc / self.out_nc)) - else: - self.shuffle_factor = None - - upsample_block = { - "upconv": B.upconv_block, - "pixel_shuffle": B.pixelshuffle_block, - }.get(self.upsampler) - if upsample_block is None: - raise NotImplementedError(f"Upsample mode [{self.upsampler}] is not found") - - if self.scale == 3: - upsample_blocks = upsample_block( - in_nc=self.num_filters, - out_nc=self.num_filters, - upscale_factor=3, - act_type=self.act, - c2x2=c2x2, - ) - else: - upsample_blocks = [ - upsample_block( - in_nc=self.num_filters, - out_nc=self.num_filters, - act_type=self.act, - c2x2=c2x2, - ) - for _ in range(int(math.log(self.scale, 2))) - ] - - self.model = B.sequential( - # fea conv - B.conv_block( - in_nc=self.in_nc, - out_nc=self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - c2x2=c2x2, - ), - B.ShortcutBlock( - B.sequential( - # rrdb blocks - *[ - B.RRDB( - nf=self.num_filters, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=self.norm, - act_type=self.act, - mode="CNA", - plus=self.plus, - c2x2=c2x2, - ) - for _ in range(self.num_blocks) - ], - # lr conv - B.conv_block( - in_nc=self.num_filters, - out_nc=self.num_filters, - kernel_size=3, - norm_type=self.norm, - act_type=None, - mode=self.mode, - c2x2=c2x2, - ), - ) - ), - *upsample_blocks, - # hr_conv0 - B.conv_block( - in_nc=self.num_filters, - out_nc=self.num_filters, - kernel_size=3, - norm_type=None, - act_type=self.act, - c2x2=c2x2, - ), - # hr_conv1 - B.conv_block( - in_nc=self.num_filters, - out_nc=self.out_nc, - kernel_size=3, - norm_type=None, - act_type=None, - c2x2=c2x2, - ), - ) - - # Adjust these properties for calculations outside of the model - if self.shuffle_factor: - self.in_nc //= self.shuffle_factor**2 - self.scale //= self.shuffle_factor - - self.load_state_dict(self.state, strict=False) - - def new_to_old_arch(self, state): - """Convert a new-arch model state dictionary to an old-arch dictionary.""" - if "params_ema" in state: - state = state["params_ema"] - - if "conv_first.weight" not in state: - # model is already old arch, this is a loose check, but should be sufficient - return state - - # add nb to state keys - for kind in ("weight", "bias"): - self.state_map[f"model.1.sub.{self.num_blocks}.{kind}"] = self.state_map[ - f"model.1.sub./NB/.{kind}" - ] - del self.state_map[f"model.1.sub./NB/.{kind}"] - - old_state = OrderedDict() - for old_key, new_keys in self.state_map.items(): - for new_key in new_keys: - if r"\1" in old_key: - for k, v in state.items(): - sub = re.sub(new_key, old_key, k) - if sub != k: - old_state[sub] = v - else: - if new_key in state: - old_state[old_key] = state[new_key] - - # upconv layers - max_upconv = 0 - for key in state.keys(): - match = re.match(r"(upconv|conv_up)(\d)\.(weight|bias)", key) - if match is not None: - _, key_num, key_type = match.groups() - old_state[f"model.{int(key_num) * 3}.{key_type}"] = state[key] - max_upconv = max(max_upconv, int(key_num) * 3) - - # final layers - for key in state.keys(): - if key in ("HRconv.weight", "conv_hr.weight"): - old_state[f"model.{max_upconv + 2}.weight"] = state[key] - elif key in ("HRconv.bias", "conv_hr.bias"): - old_state[f"model.{max_upconv + 2}.bias"] = state[key] - elif key in ("conv_last.weight",): - old_state[f"model.{max_upconv + 4}.weight"] = state[key] - elif key in ("conv_last.bias",): - old_state[f"model.{max_upconv + 4}.bias"] = state[key] - - # Sort by first numeric value of each layer - def compare(item1, item2): - parts1 = item1.split(".") - parts2 = item2.split(".") - int1 = int(parts1[1]) - int2 = int(parts2[1]) - return int1 - int2 - - sorted_keys = sorted(old_state.keys(), key=functools.cmp_to_key(compare)) - - # Rebuild the output dict in the right order - out_dict = OrderedDict((k, old_state[k]) for k in sorted_keys) - - return out_dict - - def get_scale(self, min_part: int = 6) -> int: - n = 0 - for part in list(self.state): - parts = part.split(".")[1:] - if len(parts) == 2: - part_num = int(parts[0]) - if part_num > min_part and parts[1] == "weight": - n += 1 - return 2**n - - def get_num_blocks(self) -> int: - nbs = [] - state_keys = self.state_map[r"model.1.sub.\1.RDB\2.conv\3.0.\4"] + ( - r"model\.\d+\.sub\.(\d+)\.RDB(\d+)\.conv(\d+)\.0\.(weight|bias)", - ) - for state_key in state_keys: - for k in self.state: - m = re.search(state_key, k) - if m: - nbs.append(int(m.group(1))) - if nbs: - break - return max(*nbs) + 1 - - def forward(self, x): - if self.shuffle_factor: - _, _, h, w = x.size() - mod_pad_h = ( - self.shuffle_factor - h % self.shuffle_factor - ) % self.shuffle_factor - mod_pad_w = ( - self.shuffle_factor - w % self.shuffle_factor - ) % self.shuffle_factor - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect") - x = torch.pixel_unshuffle(x, downscale_factor=self.shuffle_factor) - x = self.model(x) - return x[:, :, : h * self.scale, : w * self.scale] - return self.model(x) diff --git a/backend/comfy_nodes/chainner_models/architecture/SCUNet.py b/backend/comfy_nodes/chainner_models/architecture/SCUNet.py deleted file mode 100644 index b8354a87..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/SCUNet.py +++ /dev/null @@ -1,455 +0,0 @@ -# pylint: skip-file -# ----------------------------------------------------------------------------------- -# SCUNet: Practical Blind Denoising via Swin-Conv-UNet and Data Synthesis, https://arxiv.org/abs/2203.13278 -# Zhang, Kai and Li, Yawei and Liang, Jingyun and Cao, Jiezhang and Zhang, Yulun and Tang, Hao and Timofte, Radu and Van Gool, Luc -# ----------------------------------------------------------------------------------- - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange -from einops.layers.torch import Rearrange - -from .timm.drop import DropPath -from .timm.weight_init import trunc_normal_ - - -# Borrowed from https://github.com/cszn/SCUNet/blob/main/models/network_scunet.py -class WMSA(nn.Module): - """Self-attention module in Swin Transformer""" - - def __init__(self, input_dim, output_dim, head_dim, window_size, type): - super(WMSA, self).__init__() - self.input_dim = input_dim - self.output_dim = output_dim - self.head_dim = head_dim - self.scale = self.head_dim**-0.5 - self.n_heads = input_dim // head_dim - self.window_size = window_size - self.type = type - self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True) - - self.relative_position_params = nn.Parameter( - torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads) - ) - # TODO recover - # self.relative_position_params = nn.Parameter(torch.zeros(self.n_heads, 2 * window_size - 1, 2 * window_size -1)) - self.relative_position_params = nn.Parameter( - torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads) - ) - - self.linear = nn.Linear(self.input_dim, self.output_dim) - - trunc_normal_(self.relative_position_params, std=0.02) - self.relative_position_params = torch.nn.Parameter( - self.relative_position_params.view( - 2 * window_size - 1, 2 * window_size - 1, self.n_heads - ) - .transpose(1, 2) - .transpose(0, 1) - ) - - def generate_mask(self, h, w, p, shift): - """generating the mask of SW-MSA - Args: - shift: shift parameters in CyclicShift. - Returns: - attn_mask: should be (1 1 w p p), - """ - # supporting square. - attn_mask = torch.zeros( - h, - w, - p, - p, - p, - p, - dtype=torch.bool, - device=self.relative_position_params.device, - ) - if self.type == "W": - return attn_mask - - s = p - shift - attn_mask[-1, :, :s, :, s:, :] = True - attn_mask[-1, :, s:, :, :s, :] = True - attn_mask[:, -1, :, :s, :, s:] = True - attn_mask[:, -1, :, s:, :, :s] = True - attn_mask = rearrange( - attn_mask, "w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)" - ) - return attn_mask - - def forward(self, x): - """Forward pass of Window Multi-head Self-attention module. - Args: - x: input tensor with shape of [b h w c]; - attn_mask: attention mask, fill -inf where the value is True; - Returns: - output: tensor shape [b h w c] - """ - if self.type != "W": - x = torch.roll( - x, - shifts=(-(self.window_size // 2), -(self.window_size // 2)), - dims=(1, 2), - ) - - x = rearrange( - x, - "b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c", - p1=self.window_size, - p2=self.window_size, - ) - h_windows = x.size(1) - w_windows = x.size(2) - # square validation - # assert h_windows == w_windows - - x = rearrange( - x, - "b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c", - p1=self.window_size, - p2=self.window_size, - ) - qkv = self.embedding_layer(x) - q, k, v = rearrange( - qkv, "b nw np (threeh c) -> threeh b nw np c", c=self.head_dim - ).chunk(3, dim=0) - sim = torch.einsum("hbwpc,hbwqc->hbwpq", q, k) * self.scale - # Adding learnable relative embedding - sim = sim + rearrange(self.relative_embedding(), "h p q -> h 1 1 p q") - # Using Attn Mask to distinguish different subwindows. - if self.type != "W": - attn_mask = self.generate_mask( - h_windows, w_windows, self.window_size, shift=self.window_size // 2 - ) - sim = sim.masked_fill_(attn_mask, float("-inf")) - - probs = nn.functional.softmax(sim, dim=-1) - output = torch.einsum("hbwij,hbwjc->hbwic", probs, v) - output = rearrange(output, "h b w p c -> b w p (h c)") - output = self.linear(output) - output = rearrange( - output, - "b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c", - w1=h_windows, - p1=self.window_size, - ) - - if self.type != "W": - output = torch.roll( - output, - shifts=(self.window_size // 2, self.window_size // 2), - dims=(1, 2), - ) - - return output - - def relative_embedding(self): - cord = torch.tensor( - np.array( - [ - [i, j] - for i in range(self.window_size) - for j in range(self.window_size) - ] - ) - ) - relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1 - # negative is allowed - return self.relative_position_params[ - :, relation[:, :, 0].long(), relation[:, :, 1].long() - ] - - -class Block(nn.Module): - def __init__( - self, - input_dim, - output_dim, - head_dim, - window_size, - drop_path, - type="W", - input_resolution=None, - ): - """SwinTransformer Block""" - super(Block, self).__init__() - self.input_dim = input_dim - self.output_dim = output_dim - assert type in ["W", "SW"] - self.type = type - if input_resolution <= window_size: - self.type = "W" - - self.ln1 = nn.LayerNorm(input_dim) - self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type) - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - self.ln2 = nn.LayerNorm(input_dim) - self.mlp = nn.Sequential( - nn.Linear(input_dim, 4 * input_dim), - nn.GELU(), - nn.Linear(4 * input_dim, output_dim), - ) - - def forward(self, x): - x = x + self.drop_path(self.msa(self.ln1(x))) - x = x + self.drop_path(self.mlp(self.ln2(x))) - return x - - -class ConvTransBlock(nn.Module): - def __init__( - self, - conv_dim, - trans_dim, - head_dim, - window_size, - drop_path, - type="W", - input_resolution=None, - ): - """SwinTransformer and Conv Block""" - super(ConvTransBlock, self).__init__() - self.conv_dim = conv_dim - self.trans_dim = trans_dim - self.head_dim = head_dim - self.window_size = window_size - self.drop_path = drop_path - self.type = type - self.input_resolution = input_resolution - - assert self.type in ["W", "SW"] - if self.input_resolution <= self.window_size: - self.type = "W" - - self.trans_block = Block( - self.trans_dim, - self.trans_dim, - self.head_dim, - self.window_size, - self.drop_path, - self.type, - self.input_resolution, - ) - self.conv1_1 = nn.Conv2d( - self.conv_dim + self.trans_dim, - self.conv_dim + self.trans_dim, - 1, - 1, - 0, - bias=True, - ) - self.conv1_2 = nn.Conv2d( - self.conv_dim + self.trans_dim, - self.conv_dim + self.trans_dim, - 1, - 1, - 0, - bias=True, - ) - - self.conv_block = nn.Sequential( - nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False), - nn.ReLU(True), - nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False), - ) - - def forward(self, x): - conv_x, trans_x = torch.split( - self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1 - ) - conv_x = self.conv_block(conv_x) + conv_x - trans_x = Rearrange("b c h w -> b h w c")(trans_x) - trans_x = self.trans_block(trans_x) - trans_x = Rearrange("b h w c -> b c h w")(trans_x) - res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1)) - x = x + res - - return x - - -class SCUNet(nn.Module): - def __init__( - self, - state_dict, - in_nc=3, - config=[4, 4, 4, 4, 4, 4, 4], - dim=64, - drop_path_rate=0.0, - input_resolution=256, - ): - super(SCUNet, self).__init__() - self.model_arch = "SCUNet" - self.sub_type = "SR" - - self.num_filters: int = 0 - - self.state = state_dict - self.config = config - self.dim = dim - self.head_dim = 32 - self.window_size = 8 - - self.in_nc = in_nc - self.out_nc = self.in_nc - self.scale = 1 - self.supports_fp16 = True - - # drop path rate for each layer - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))] - - self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)] - - begin = 0 - self.m_down1 = [ - ConvTransBlock( - dim // 2, - dim // 2, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution, - ) - for i in range(config[0]) - ] + [nn.Conv2d(dim, 2 * dim, 2, 2, 0, bias=False)] - - begin += config[0] - self.m_down2 = [ - ConvTransBlock( - dim, - dim, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution // 2, - ) - for i in range(config[1]) - ] + [nn.Conv2d(2 * dim, 4 * dim, 2, 2, 0, bias=False)] - - begin += config[1] - self.m_down3 = [ - ConvTransBlock( - 2 * dim, - 2 * dim, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution // 4, - ) - for i in range(config[2]) - ] + [nn.Conv2d(4 * dim, 8 * dim, 2, 2, 0, bias=False)] - - begin += config[2] - self.m_body = [ - ConvTransBlock( - 4 * dim, - 4 * dim, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution // 8, - ) - for i in range(config[3]) - ] - - begin += config[3] - self.m_up3 = [ - nn.ConvTranspose2d(8 * dim, 4 * dim, 2, 2, 0, bias=False), - ] + [ - ConvTransBlock( - 2 * dim, - 2 * dim, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution // 4, - ) - for i in range(config[4]) - ] - - begin += config[4] - self.m_up2 = [ - nn.ConvTranspose2d(4 * dim, 2 * dim, 2, 2, 0, bias=False), - ] + [ - ConvTransBlock( - dim, - dim, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution // 2, - ) - for i in range(config[5]) - ] - - begin += config[5] - self.m_up1 = [ - nn.ConvTranspose2d(2 * dim, dim, 2, 2, 0, bias=False), - ] + [ - ConvTransBlock( - dim // 2, - dim // 2, - self.head_dim, - self.window_size, - dpr[i + begin], - "W" if not i % 2 else "SW", - input_resolution, - ) - for i in range(config[6]) - ] - - self.m_tail = [nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False)] - - self.m_head = nn.Sequential(*self.m_head) - self.m_down1 = nn.Sequential(*self.m_down1) - self.m_down2 = nn.Sequential(*self.m_down2) - self.m_down3 = nn.Sequential(*self.m_down3) - self.m_body = nn.Sequential(*self.m_body) - self.m_up3 = nn.Sequential(*self.m_up3) - self.m_up2 = nn.Sequential(*self.m_up2) - self.m_up1 = nn.Sequential(*self.m_up1) - self.m_tail = nn.Sequential(*self.m_tail) - # self.apply(self._init_weights) - self.load_state_dict(state_dict, strict=True) - - def check_image_size(self, x): - _, _, h, w = x.size() - mod_pad_h = (64 - h % 64) % 64 - mod_pad_w = (64 - w % 64) % 64 - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect") - return x - - def forward(self, x0): - h, w = x0.size()[-2:] - x0 = self.check_image_size(x0) - - x1 = self.m_head(x0) - x2 = self.m_down1(x1) - x3 = self.m_down2(x2) - x4 = self.m_down3(x3) - x = self.m_body(x4) - x = self.m_up3(x + x4) - x = self.m_up2(x + x3) - x = self.m_up1(x + x2) - x = self.m_tail(x + x1) - - x = x[:, :, :h, :w] - return x - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) diff --git a/backend/comfy_nodes/chainner_models/architecture/SPSR.py b/backend/comfy_nodes/chainner_models/architecture/SPSR.py deleted file mode 100644 index c3cefff1..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/SPSR.py +++ /dev/null @@ -1,383 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from . import block as B - - -class Get_gradient_nopadding(nn.Module): - def __init__(self): - super(Get_gradient_nopadding, self).__init__() - kernel_v = [[0, -1, 0], [0, 0, 0], [0, 1, 0]] - kernel_h = [[0, 0, 0], [-1, 0, 1], [0, 0, 0]] - kernel_h = torch.FloatTensor(kernel_h).unsqueeze(0).unsqueeze(0) - kernel_v = torch.FloatTensor(kernel_v).unsqueeze(0).unsqueeze(0) - self.weight_h = nn.Parameter(data=kernel_h, requires_grad=False) # type: ignore - - self.weight_v = nn.Parameter(data=kernel_v, requires_grad=False) # type: ignore - - def forward(self, x): - x_list = [] - for i in range(x.shape[1]): - x_i = x[:, i] - x_i_v = F.conv2d(x_i.unsqueeze(1), self.weight_v, padding=1) - x_i_h = F.conv2d(x_i.unsqueeze(1), self.weight_h, padding=1) - x_i = torch.sqrt(torch.pow(x_i_v, 2) + torch.pow(x_i_h, 2) + 1e-6) - x_list.append(x_i) - - x = torch.cat(x_list, dim=1) - - return x - - -class SPSRNet(nn.Module): - def __init__( - self, - state_dict, - norm=None, - act: str = "leakyrelu", - upsampler: str = "upconv", - mode: B.ConvMode = "CNA", - ): - super(SPSRNet, self).__init__() - self.model_arch = "SPSR" - self.sub_type = "SR" - - self.state = state_dict - self.norm = norm - self.act = act - self.upsampler = upsampler - self.mode = mode - - self.num_blocks = self.get_num_blocks() - - self.in_nc: int = self.state["model.0.weight"].shape[1] - self.out_nc: int = self.state["f_HR_conv1.0.bias"].shape[0] - - self.scale = self.get_scale(4) - self.num_filters: int = self.state["model.0.weight"].shape[0] - - self.supports_fp16 = True - self.supports_bfp16 = True - self.min_size_restriction = None - - n_upscale = int(math.log(self.scale, 2)) - if self.scale == 3: - n_upscale = 1 - - fea_conv = B.conv_block( - self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None - ) - rb_blocks = [ - B.RRDB( - self.num_filters, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - for _ in range(self.num_blocks) - ] - LR_conv = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=norm, - act_type=None, - mode=mode, - ) - - if upsampler == "upconv": - upsample_block = B.upconv_block - elif upsampler == "pixelshuffle": - upsample_block = B.pixelshuffle_block - else: - raise NotImplementedError(f"upsample mode [{upsampler}] is not found") - if self.scale == 3: - a_upsampler = upsample_block( - self.num_filters, self.num_filters, 3, act_type=act - ) - else: - a_upsampler = [ - upsample_block(self.num_filters, self.num_filters, act_type=act) - for _ in range(n_upscale) - ] - self.HR_conv0_new = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=act, - ) - self.HR_conv1_new = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - - self.model = B.sequential( - fea_conv, - B.ShortcutBlockSPSR(B.sequential(*rb_blocks, LR_conv)), - *a_upsampler, - self.HR_conv0_new, - ) - - self.get_g_nopadding = Get_gradient_nopadding() - - self.b_fea_conv = B.conv_block( - self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None - ) - - self.b_concat_1 = B.conv_block( - 2 * self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - self.b_block_1 = B.RRDB( - self.num_filters * 2, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - - self.b_concat_2 = B.conv_block( - 2 * self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - self.b_block_2 = B.RRDB( - self.num_filters * 2, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - - self.b_concat_3 = B.conv_block( - 2 * self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - self.b_block_3 = B.RRDB( - self.num_filters * 2, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - - self.b_concat_4 = B.conv_block( - 2 * self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - self.b_block_4 = B.RRDB( - self.num_filters * 2, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - - self.b_LR_conv = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=norm, - act_type=None, - mode=mode, - ) - - if upsampler == "upconv": - upsample_block = B.upconv_block - elif upsampler == "pixelshuffle": - upsample_block = B.pixelshuffle_block - else: - raise NotImplementedError(f"upsample mode [{upsampler}] is not found") - if self.scale == 3: - b_upsampler = upsample_block( - self.num_filters, self.num_filters, 3, act_type=act - ) - else: - b_upsampler = [ - upsample_block(self.num_filters, self.num_filters, act_type=act) - for _ in range(n_upscale) - ] - - b_HR_conv0 = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=act, - ) - b_HR_conv1 = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - - self.b_module = B.sequential(*b_upsampler, b_HR_conv0, b_HR_conv1) - - self.conv_w = B.conv_block( - self.num_filters, self.out_nc, kernel_size=1, norm_type=None, act_type=None - ) - - self.f_concat = B.conv_block( - self.num_filters * 2, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=None, - ) - - self.f_block = B.RRDB( - self.num_filters * 2, - kernel_size=3, - gc=32, - stride=1, - bias=True, - pad_type="zero", - norm_type=norm, - act_type=act, - mode="CNA", - ) - - self.f_HR_conv0 = B.conv_block( - self.num_filters, - self.num_filters, - kernel_size=3, - norm_type=None, - act_type=act, - ) - self.f_HR_conv1 = B.conv_block( - self.num_filters, self.out_nc, kernel_size=3, norm_type=None, act_type=None - ) - - self.load_state_dict(self.state, strict=False) - - def get_scale(self, min_part: int = 4) -> int: - n = 0 - for part in list(self.state): - parts = part.split(".") - if len(parts) == 3: - part_num = int(parts[1]) - if part_num > min_part and parts[0] == "model" and parts[2] == "weight": - n += 1 - return 2**n - - def get_num_blocks(self) -> int: - nb = 0 - for part in list(self.state): - parts = part.split(".") - n_parts = len(parts) - if n_parts == 5 and parts[2] == "sub": - nb = int(parts[3]) - return nb - - def forward(self, x): - x_grad = self.get_g_nopadding(x) - x = self.model[0](x) - - x, block_list = self.model[1](x) - - x_ori = x - for i in range(5): - x = block_list[i](x) - x_fea1 = x - - for i in range(5): - x = block_list[i + 5](x) - x_fea2 = x - - for i in range(5): - x = block_list[i + 10](x) - x_fea3 = x - - for i in range(5): - x = block_list[i + 15](x) - x_fea4 = x - - x = block_list[20:](x) - # short cut - x = x_ori + x - x = self.model[2:](x) - x = self.HR_conv1_new(x) - - x_b_fea = self.b_fea_conv(x_grad) - x_cat_1 = torch.cat([x_b_fea, x_fea1], dim=1) - - x_cat_1 = self.b_block_1(x_cat_1) - x_cat_1 = self.b_concat_1(x_cat_1) - - x_cat_2 = torch.cat([x_cat_1, x_fea2], dim=1) - - x_cat_2 = self.b_block_2(x_cat_2) - x_cat_2 = self.b_concat_2(x_cat_2) - - x_cat_3 = torch.cat([x_cat_2, x_fea3], dim=1) - - x_cat_3 = self.b_block_3(x_cat_3) - x_cat_3 = self.b_concat_3(x_cat_3) - - x_cat_4 = torch.cat([x_cat_3, x_fea4], dim=1) - - x_cat_4 = self.b_block_4(x_cat_4) - x_cat_4 = self.b_concat_4(x_cat_4) - - x_cat_4 = self.b_LR_conv(x_cat_4) - - # short cut - x_cat_4 = x_cat_4 + x_b_fea - x_branch = self.b_module(x_cat_4) - - # x_out_branch = self.conv_w(x_branch) - ######## - x_branch_d = x_branch - x_f_cat = torch.cat([x_branch_d, x], dim=1) - x_f_cat = self.f_block(x_f_cat) - x_out = self.f_concat(x_f_cat) - x_out = self.f_HR_conv0(x_out) - x_out = self.f_HR_conv1(x_out) - - ######### - # return x_out_branch, x_out, x_grad - return x_out diff --git a/backend/comfy_nodes/chainner_models/architecture/SRVGG.py b/backend/comfy_nodes/chainner_models/architecture/SRVGG.py deleted file mode 100644 index 7a8ec37a..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/SRVGG.py +++ /dev/null @@ -1,114 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -import math - -import torch.nn as nn -import torch.nn.functional as F - - -class SRVGGNetCompact(nn.Module): - """A compact VGG-style network structure for super-resolution. - It is a compact network structure, which performs upsampling in the last layer and no convolution is - conducted on the HR feature space. - Args: - num_in_ch (int): Channel number of inputs. Default: 3. - num_out_ch (int): Channel number of outputs. Default: 3. - num_feat (int): Channel number of intermediate features. Default: 64. - num_conv (int): Number of convolution layers in the body network. Default: 16. - upscale (int): Upsampling factor. Default: 4. - act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu. - """ - - def __init__( - self, - state_dict, - act_type: str = "prelu", - ): - super(SRVGGNetCompact, self).__init__() - self.model_arch = "SRVGG (RealESRGAN)" - self.sub_type = "SR" - - self.act_type = act_type - - self.state = state_dict - - if "params" in self.state: - self.state = self.state["params"] - - self.key_arr = list(self.state.keys()) - - self.in_nc = self.get_in_nc() - self.num_feat = self.get_num_feats() - self.num_conv = self.get_num_conv() - self.out_nc = self.in_nc # :( - self.pixelshuffle_shape = None # Defined in get_scale() - self.scale = self.get_scale() - - self.supports_fp16 = True - self.supports_bfp16 = True - self.min_size_restriction = None - - self.body = nn.ModuleList() - # the first conv - self.body.append(nn.Conv2d(self.in_nc, self.num_feat, 3, 1, 1)) - # the first activation - if act_type == "relu": - activation = nn.ReLU(inplace=True) - elif act_type == "prelu": - activation = nn.PReLU(num_parameters=self.num_feat) - elif act_type == "leakyrelu": - activation = nn.LeakyReLU(negative_slope=0.1, inplace=True) - self.body.append(activation) # type: ignore - - # the body structure - for _ in range(self.num_conv): - self.body.append(nn.Conv2d(self.num_feat, self.num_feat, 3, 1, 1)) - # activation - if act_type == "relu": - activation = nn.ReLU(inplace=True) - elif act_type == "prelu": - activation = nn.PReLU(num_parameters=self.num_feat) - elif act_type == "leakyrelu": - activation = nn.LeakyReLU(negative_slope=0.1, inplace=True) - self.body.append(activation) # type: ignore - - # the last conv - self.body.append(nn.Conv2d(self.num_feat, self.pixelshuffle_shape, 3, 1, 1)) # type: ignore - # upsample - self.upsampler = nn.PixelShuffle(self.scale) - - self.load_state_dict(self.state, strict=False) - - def get_num_conv(self) -> int: - return (int(self.key_arr[-1].split(".")[1]) - 2) // 2 - - def get_num_feats(self) -> int: - return self.state[self.key_arr[0]].shape[0] - - def get_in_nc(self) -> int: - return self.state[self.key_arr[0]].shape[1] - - def get_scale(self) -> int: - self.pixelshuffle_shape = self.state[self.key_arr[-1]].shape[0] - # Assume out_nc is the same as in_nc - # I cant think of a better way to do that - self.out_nc = self.in_nc - scale = math.sqrt(self.pixelshuffle_shape / self.out_nc) - if scale - int(scale) > 0: - print( - "out_nc is probably different than in_nc, scale calculation might be wrong" - ) - scale = int(scale) - return scale - - def forward(self, x): - out = x - for i in range(0, len(self.body)): - out = self.body[i](out) - - out = self.upsampler(out) - # add the nearest upsampled image, so that the network learns the residual - base = F.interpolate(x, scale_factor=self.scale, mode="nearest") - out += base - return out diff --git a/backend/comfy_nodes/chainner_models/architecture/SwiftSRGAN.py b/backend/comfy_nodes/chainner_models/architecture/SwiftSRGAN.py deleted file mode 100644 index dbb7725b..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/SwiftSRGAN.py +++ /dev/null @@ -1,161 +0,0 @@ -# From https://github.com/Koushik0901/Swift-SRGAN/blob/master/swift-srgan/models.py - -import torch -from torch import nn - - -class SeperableConv2d(nn.Module): - def __init__( - self, in_channels, out_channels, kernel_size, stride=1, padding=1, bias=True - ): - super(SeperableConv2d, self).__init__() - self.depthwise = nn.Conv2d( - in_channels, - in_channels, - kernel_size=kernel_size, - stride=stride, - groups=in_channels, - bias=bias, - padding=padding, - ) - self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias) - - def forward(self, x): - return self.pointwise(self.depthwise(x)) - - -class ConvBlock(nn.Module): - def __init__( - self, - in_channels, - out_channels, - use_act=True, - use_bn=True, - discriminator=False, - **kwargs, - ): - super(ConvBlock, self).__init__() - - self.use_act = use_act - self.cnn = SeperableConv2d(in_channels, out_channels, **kwargs, bias=not use_bn) - self.bn = nn.BatchNorm2d(out_channels) if use_bn else nn.Identity() - self.act = ( - nn.LeakyReLU(0.2, inplace=True) - if discriminator - else nn.PReLU(num_parameters=out_channels) - ) - - def forward(self, x): - return self.act(self.bn(self.cnn(x))) if self.use_act else self.bn(self.cnn(x)) - - -class UpsampleBlock(nn.Module): - def __init__(self, in_channels, scale_factor): - super(UpsampleBlock, self).__init__() - - self.conv = SeperableConv2d( - in_channels, - in_channels * scale_factor**2, - kernel_size=3, - stride=1, - padding=1, - ) - self.ps = nn.PixelShuffle( - scale_factor - ) # (in_channels * 4, H, W) -> (in_channels, H*2, W*2) - self.act = nn.PReLU(num_parameters=in_channels) - - def forward(self, x): - return self.act(self.ps(self.conv(x))) - - -class ResidualBlock(nn.Module): - def __init__(self, in_channels): - super(ResidualBlock, self).__init__() - - self.block1 = ConvBlock( - in_channels, in_channels, kernel_size=3, stride=1, padding=1 - ) - self.block2 = ConvBlock( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, use_act=False - ) - - def forward(self, x): - out = self.block1(x) - out = self.block2(out) - return out + x - - -class Generator(nn.Module): - """Swift-SRGAN Generator - Args: - in_channels (int): number of input image channels. - num_channels (int): number of hidden channels. - num_blocks (int): number of residual blocks. - upscale_factor (int): factor to upscale the image [2x, 4x, 8x]. - Returns: - torch.Tensor: super resolution image - """ - - def __init__( - self, - state_dict, - ): - super(Generator, self).__init__() - self.model_arch = "Swift-SRGAN" - self.sub_type = "SR" - self.state = state_dict - if "model" in self.state: - self.state = self.state["model"] - - self.in_nc: int = self.state["initial.cnn.depthwise.weight"].shape[0] - self.out_nc: int = self.state["final_conv.pointwise.weight"].shape[0] - self.num_filters: int = self.state["initial.cnn.pointwise.weight"].shape[0] - self.num_blocks = len( - set([x.split(".")[1] for x in self.state.keys() if "residual" in x]) - ) - self.scale: int = 2 ** len( - set([x.split(".")[1] for x in self.state.keys() if "upsampler" in x]) - ) - - in_channels = self.in_nc - num_channels = self.num_filters - num_blocks = self.num_blocks - upscale_factor = self.scale - - self.supports_fp16 = True - self.supports_bfp16 = True - self.min_size_restriction = None - - self.initial = ConvBlock( - in_channels, num_channels, kernel_size=9, stride=1, padding=4, use_bn=False - ) - self.residual = nn.Sequential( - *[ResidualBlock(num_channels) for _ in range(num_blocks)] - ) - self.convblock = ConvBlock( - num_channels, - num_channels, - kernel_size=3, - stride=1, - padding=1, - use_act=False, - ) - self.upsampler = nn.Sequential( - *[ - UpsampleBlock(num_channels, scale_factor=2) - for _ in range(upscale_factor // 2) - ] - ) - self.final_conv = SeperableConv2d( - num_channels, in_channels, kernel_size=9, stride=1, padding=4 - ) - - self.load_state_dict(self.state, strict=False) - - def forward(self, x): - initial = self.initial(x) - x = self.residual(initial) - x = self.convblock(x) + initial - x = self.upsampler(x) - return (torch.tanh(self.final_conv(x)) + 1) / 2 diff --git a/backend/comfy_nodes/chainner_models/architecture/Swin2SR.py b/backend/comfy_nodes/chainner_models/architecture/Swin2SR.py deleted file mode 100644 index cb57ecfc..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/Swin2SR.py +++ /dev/null @@ -1,1377 +0,0 @@ -# pylint: skip-file -# ----------------------------------------------------------------------------------- -# Swin2SR: Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration, https://arxiv.org/abs/2209.11345 -# Written by Conde and Choi et al. -# From: https://raw.githubusercontent.com/mv-lab/swin2sr/main/models/network_swin2sr.py -# ----------------------------------------------------------------------------------- - -import math -import re - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint as checkpoint - -# Originally from the timm package -from .timm.drop import DropPath -from .timm.helpers import to_2tuple -from .timm.weight_init import trunc_normal_ - - -class Mlp(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -def window_partition(x, window_size): - """ - Args: - x: (B, H, W, C) - window_size (int): window size - Returns: - windows: (num_windows*B, window_size, window_size, C) - """ - B, H, W, C = x.shape - x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = ( - x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) - ) - return windows - - -def window_reverse(windows, window_size, H, W): - """ - Args: - windows: (num_windows*B, window_size, window_size, C) - window_size (int): Window size - H (int): Height of image - W (int): Width of image - Returns: - x: (B, H, W, C) - """ - B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view( - B, H // window_size, W // window_size, window_size, window_size, -1 - ) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) - return x - - -class WindowAttention(nn.Module): - r"""Window based multi-head self attention (W-MSA) module with relative position bias. - It supports both of shifted and non-shifted window. - Args: - dim (int): Number of input channels. - window_size (tuple[int]): The height and width of the window. - num_heads (int): Number of attention heads. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 - proj_drop (float, optional): Dropout ratio of output. Default: 0.0 - pretrained_window_size (tuple[int]): The height and width of the window in pre-training. - """ - - def __init__( - self, - dim, - window_size, - num_heads, - qkv_bias=True, - attn_drop=0.0, - proj_drop=0.0, - pretrained_window_size=[0, 0], - ): - super().__init__() - self.dim = dim - self.window_size = window_size # Wh, Ww - self.pretrained_window_size = pretrained_window_size - self.num_heads = num_heads - - self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True) # type: ignore - - # mlp to generate continuous relative position bias - self.cpb_mlp = nn.Sequential( - nn.Linear(2, 512, bias=True), - nn.ReLU(inplace=True), - nn.Linear(512, num_heads, bias=False), - ) - - # get relative_coords_table - relative_coords_h = torch.arange( - -(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32 - ) - relative_coords_w = torch.arange( - -(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32 - ) - relative_coords_table = ( - torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w])) - .permute(1, 2, 0) - .contiguous() - .unsqueeze(0) - ) # 1, 2*Wh-1, 2*Ww-1, 2 - if pretrained_window_size[0] > 0: - relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1 - relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1 - else: - relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1 - relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1 - relative_coords_table *= 8 # normalize to -8, 8 - relative_coords_table = ( - torch.sign(relative_coords_table) - * torch.log2(torch.abs(relative_coords_table) + 1.0) - / np.log2(8) - ) - - self.register_buffer("relative_coords_table", relative_coords_table) - - # get pair-wise relative position index for each token inside the window - coords_h = torch.arange(self.window_size[0]) - coords_w = torch.arange(self.window_size[1]) - coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww - coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = ( - coords_flatten[:, :, None] - coords_flatten[:, None, :] - ) # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute( - 1, 2, 0 - ).contiguous() # Wh*Ww, Wh*Ww, 2 - relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 - relative_coords[:, :, 1] += self.window_size[1] - 1 - relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 - relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww - self.register_buffer("relative_position_index", relative_position_index) - - self.qkv = nn.Linear(dim, dim * 3, bias=False) - if qkv_bias: - self.q_bias = nn.Parameter(torch.zeros(dim)) # type: ignore - self.v_bias = nn.Parameter(torch.zeros(dim)) # type: ignore - else: - self.q_bias = None - self.v_bias = None - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - self.softmax = nn.Softmax(dim=-1) - - def forward(self, x, mask=None): - """ - Args: - x: input features with shape of (num_windows*B, N, C) - mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None - """ - B_, N, C = x.shape - qkv_bias = None - if self.q_bias is not None: - qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias)) # type: ignore - qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) - qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) - q, k, v = ( - qkv[0], - qkv[1], - qkv[2], - ) # make torchscript happy (cannot use tensor as tuple) - - # cosine attention - attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) - logit_scale = torch.clamp( - self.logit_scale, - max=torch.log(torch.tensor(1.0 / 0.01)).to(self.logit_scale.device), - ).exp() - attn = attn * logit_scale - - relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view( - -1, self.num_heads - ) - relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view( # type: ignore - self.window_size[0] * self.window_size[1], - self.window_size[0] * self.window_size[1], - -1, - ) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute( - 2, 0, 1 - ).contiguous() # nH, Wh*Ww, Wh*Ww - relative_position_bias = 16 * torch.sigmoid(relative_position_bias) - attn = attn + relative_position_bias.unsqueeze(0) - - if mask is not None: - nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( - 1 - ).unsqueeze(0) - attn = attn.view(-1, self.num_heads, N, N) - attn = self.softmax(attn) - else: - attn = self.softmax(attn) - - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B_, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - def extra_repr(self) -> str: - return ( - f"dim={self.dim}, window_size={self.window_size}, " - f"pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}" - ) - - def flops(self, N): - # calculate flops for 1 window with token length of N - flops = 0 - # qkv = self.qkv(x) - flops += N * self.dim * 3 * self.dim - # attn = (q @ k.transpose(-2, -1)) - flops += self.num_heads * N * (self.dim // self.num_heads) * N - # x = (attn @ v) - flops += self.num_heads * N * N * (self.dim // self.num_heads) - # x = self.proj(x) - flops += N * self.dim * self.dim - return flops - - -class SwinTransformerBlock(nn.Module): - r"""Swin Transformer Block. - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resulotion. - num_heads (int): Number of attention heads. - window_size (int): Window size. - shift_size (int): Shift size for SW-MSA. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float, optional): Stochastic depth rate. Default: 0.0 - act_layer (nn.Module, optional): Activation layer. Default: nn.GELU - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - pretrained_window_size (int): Window size in pre-training. - """ - - def __init__( - self, - dim, - input_resolution, - num_heads, - window_size=7, - shift_size=0, - mlp_ratio=4.0, - qkv_bias=True, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - act_layer=nn.GELU, - norm_layer=nn.LayerNorm, - pretrained_window_size=0, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.num_heads = num_heads - self.window_size = window_size - self.shift_size = shift_size - self.mlp_ratio = mlp_ratio - if min(self.input_resolution) <= self.window_size: - # if window size is larger than input resolution, we don't partition windows - self.shift_size = 0 - self.window_size = min(self.input_resolution) - assert ( - 0 <= self.shift_size < self.window_size - ), "shift_size must in 0-window_size" - - self.norm1 = norm_layer(dim) - self.attn = WindowAttention( - dim, - window_size=to_2tuple(self.window_size), - num_heads=num_heads, - qkv_bias=qkv_bias, - attn_drop=attn_drop, - proj_drop=drop, - pretrained_window_size=to_2tuple(pretrained_window_size), - ) - - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp( - in_features=dim, - hidden_features=mlp_hidden_dim, - act_layer=act_layer, - drop=drop, - ) - - if self.shift_size > 0: - attn_mask = self.calculate_mask(self.input_resolution) - else: - attn_mask = None - - self.register_buffer("attn_mask", attn_mask) - - def calculate_mask(self, x_size): - # calculate attention mask for SW-MSA - H, W = x_size - img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - w_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - mask_windows = window_partition( - img_mask, self.window_size - ) # nW, window_size, window_size, 1 - mask_windows = mask_windows.view(-1, self.window_size * self.window_size) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( - attn_mask == 0, float(0.0) - ) - - return attn_mask - - def forward(self, x, x_size): - H, W = x_size - B, L, C = x.shape - # assert L == H * W, "input feature has wrong size" - - shortcut = x - x = x.view(B, H, W, C) - - # cyclic shift - if self.shift_size > 0: - shifted_x = torch.roll( - x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) - ) - else: - shifted_x = x - - # partition windows - x_windows = window_partition( - shifted_x, self.window_size - ) # nW*B, window_size, window_size, C - x_windows = x_windows.view( - -1, self.window_size * self.window_size, C - ) # nW*B, window_size*window_size, C - - # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size - if self.input_resolution == x_size: - attn_windows = self.attn( - x_windows, mask=self.attn_mask - ) # nW*B, window_size*window_size, C - else: - attn_windows = self.attn( - x_windows, mask=self.calculate_mask(x_size).to(x.device) - ) - - # merge windows - attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) - shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C - - # reverse cyclic shift - if self.shift_size > 0: - x = torch.roll( - shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) - ) - else: - x = shifted_x - x = x.view(B, H * W, C) - x = shortcut + self.drop_path(self.norm1(x)) - - # FFN - x = x + self.drop_path(self.norm2(self.mlp(x))) - - return x - - def extra_repr(self) -> str: - return ( - f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" - ) - - def flops(self): - flops = 0 - H, W = self.input_resolution - # norm1 - flops += self.dim * H * W - # W-MSA/SW-MSA - nW = H * W / self.window_size / self.window_size - flops += nW * self.attn.flops(self.window_size * self.window_size) - # mlp - flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio - # norm2 - flops += self.dim * H * W - return flops - - -class PatchMerging(nn.Module): - r"""Patch Merging Layer. - Args: - input_resolution (tuple[int]): Resolution of input feature. - dim (int): Number of input channels. - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) - self.norm = norm_layer(2 * dim) - - def forward(self, x): - """ - x: B, H*W, C - """ - H, W = self.input_resolution - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even." - - x = x.view(B, H, W, C) - - x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C - x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C - x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C - x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C - x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C - x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C - - x = self.reduction(x) - x = self.norm(x) - - return x - - def extra_repr(self) -> str: - return f"input_resolution={self.input_resolution}, dim={self.dim}" - - def flops(self): - H, W = self.input_resolution - flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim - flops += H * W * self.dim // 2 - return flops - - -class BasicLayer(nn.Module): - """A basic Swin Transformer layer for one stage. - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - pretrained_window_size (int): Local window size in pre-training. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - mlp_ratio=4.0, - qkv_bias=True, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - pretrained_window_size=0, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.depth = depth - self.use_checkpoint = use_checkpoint - - # build blocks - self.blocks = nn.ModuleList( - [ - SwinTransformerBlock( - dim=dim, - input_resolution=input_resolution, - num_heads=num_heads, - window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path[i] - if isinstance(drop_path, list) - else drop_path, - norm_layer=norm_layer, - pretrained_window_size=pretrained_window_size, - ) - for i in range(depth) - ] - ) - - # patch merging layer - if downsample is not None: - self.downsample = downsample( - input_resolution, dim=dim, norm_layer=norm_layer - ) - else: - self.downsample = None - - def forward(self, x, x_size): - for blk in self.blocks: - if self.use_checkpoint: - x = checkpoint.checkpoint(blk, x, x_size) - else: - x = blk(x, x_size) - if self.downsample is not None: - x = self.downsample(x) - return x - - def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" - - def flops(self): - flops = 0 - for blk in self.blocks: - flops += blk.flops() # type: ignore - if self.downsample is not None: - flops += self.downsample.flops() - return flops - - def _init_respostnorm(self): - for blk in self.blocks: - nn.init.constant_(blk.norm1.bias, 0) # type: ignore - nn.init.constant_(blk.norm1.weight, 0) # type: ignore - nn.init.constant_(blk.norm2.bias, 0) # type: ignore - nn.init.constant_(blk.norm2.weight, 0) # type: ignore - - -class PatchEmbed(nn.Module): - r"""Image to Patch Embedding - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] # type: ignore - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - self.proj = nn.Conv2d( - in_chans, embed_dim, kernel_size=patch_size, stride=patch_size # type: ignore - ) - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - - def forward(self, x): - B, C, H, W = x.shape - # FIXME look at relaxing size constraints - # assert H == self.img_size[0] and W == self.img_size[1], - # f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." - x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C - if self.norm is not None: - x = self.norm(x) - return x - - def flops(self): - Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) # type: ignore - if self.norm is not None: - flops += Ho * Wo * self.embed_dim - return flops - - -class RSTB(nn.Module): - """Residual Swin Transformer Block (RSTB). - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - img_size: Input image size. - patch_size: Patch size. - resi_connection: The convolutional block before residual connection. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - mlp_ratio=4.0, - qkv_bias=True, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - img_size=224, - patch_size=4, - resi_connection="1conv", - ): - super(RSTB, self).__init__() - - self.dim = dim - self.input_resolution = input_resolution - - self.residual_group = BasicLayer( - dim=dim, - input_resolution=input_resolution, - depth=depth, - num_heads=num_heads, - window_size=window_size, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path, - norm_layer=norm_layer, - downsample=downsample, - use_checkpoint=use_checkpoint, - ) - - if resi_connection == "1conv": - self.conv = nn.Conv2d(dim, dim, 3, 1, 1) - elif resi_connection == "3conv": - # to save parameters and memory - self.conv = nn.Sequential( - nn.Conv2d(dim, dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim, 3, 1, 1), - ) - - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=dim, - embed_dim=dim, - norm_layer=None, - ) - - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=dim, - embed_dim=dim, - norm_layer=None, - ) - - def forward(self, x, x_size): - return ( - self.patch_embed( - self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size)) - ) - + x - ) - - def flops(self): - flops = 0 - flops += self.residual_group.flops() - H, W = self.input_resolution - flops += H * W * self.dim * self.dim * 9 - flops += self.patch_embed.flops() - flops += self.patch_unembed.flops() - - return flops - - -class PatchUnEmbed(nn.Module): - r"""Image to Patch Unembedding - - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] # type: ignore - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - def forward(self, x, x_size): - B, HW, C = x.shape - x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C - return x - - def flops(self): - flops = 0 - return flops - - -class Upsample(nn.Sequential): - """Upsample module. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - """ - - def __init__(self, scale, num_feat): - m = [] - if (scale & (scale - 1)) == 0: # scale = 2^n - for _ in range(int(math.log(scale, 2))): - m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(2)) - elif scale == 3: - m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(3)) - else: - raise ValueError( - f"scale {scale} is not supported. " "Supported scales: 2^n and 3." - ) - super(Upsample, self).__init__(*m) - - -class Upsample_hf(nn.Sequential): - """Upsample module. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - """ - - def __init__(self, scale, num_feat): - m = [] - if (scale & (scale - 1)) == 0: # scale = 2^n - for _ in range(int(math.log(scale, 2))): - m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(2)) - elif scale == 3: - m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(3)) - else: - raise ValueError( - f"scale {scale} is not supported. " "Supported scales: 2^n and 3." - ) - super(Upsample_hf, self).__init__(*m) - - -class UpsampleOneStep(nn.Sequential): - """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) - Used in lightweight SR to save parameters. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - - """ - - def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): - self.num_feat = num_feat - self.input_resolution = input_resolution - m = [] - m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1)) - m.append(nn.PixelShuffle(scale)) - super(UpsampleOneStep, self).__init__(*m) - - def flops(self): - H, W = self.input_resolution # type: ignore - flops = H * W * self.num_feat * 3 * 9 - return flops - - -class Swin2SR(nn.Module): - r"""Swin2SR - A PyTorch impl of : `Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration`. - - Args: - img_size (int | tuple(int)): Input image size. Default 64 - patch_size (int | tuple(int)): Patch size. Default: 1 - in_chans (int): Number of input image channels. Default: 3 - embed_dim (int): Patch embedding dimension. Default: 96 - depths (tuple(int)): Depth of each Swin Transformer layer. - num_heads (tuple(int)): Number of attention heads in different layers. - window_size (int): Window size. Default: 7 - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - drop_rate (float): Dropout rate. Default: 0 - attn_drop_rate (float): Attention dropout rate. Default: 0 - drop_path_rate (float): Stochastic depth rate. Default: 0.1 - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. - ape (bool): If True, add absolute position embedding to the patch embedding. Default: False - patch_norm (bool): If True, add normalization after patch embedding. Default: True - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False - upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction - img_range: Image range. 1. or 255. - upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None - resi_connection: The convolutional block before residual connection. '1conv'/'3conv' - """ - - def __init__( - self, - state_dict, - **kwargs, - ): - super(Swin2SR, self).__init__() - - # Defaults - img_size = 128 - patch_size = 1 - in_chans = 3 - embed_dim = 96 - depths = [6, 6, 6, 6] - num_heads = [6, 6, 6, 6] - window_size = 7 - mlp_ratio = 4.0 - qkv_bias = True - drop_rate = 0.0 - attn_drop_rate = 0.0 - drop_path_rate = 0.1 - norm_layer = nn.LayerNorm - ape = False - patch_norm = True - use_checkpoint = False - upscale = 2 - img_range = 1.0 - upsampler = "" - resi_connection = "1conv" - num_in_ch = in_chans - num_out_ch = in_chans - num_feat = 64 - - self.model_arch = "Swin2SR" - self.sub_type = "SR" - self.state = state_dict - if "params_ema" in self.state: - self.state = self.state["params_ema"] - elif "params" in self.state: - self.state = self.state["params"] - - state_keys = self.state.keys() - - if "conv_before_upsample.0.weight" in state_keys: - if "conv_aux.weight" in state_keys: - upsampler = "pixelshuffle_aux" - elif "conv_up1.weight" in state_keys: - upsampler = "nearest+conv" - else: - upsampler = "pixelshuffle" - supports_fp16 = False - elif "upsample.0.weight" in state_keys: - upsampler = "pixelshuffledirect" - else: - upsampler = "" - - num_feat = ( - self.state.get("conv_before_upsample.0.weight", None).shape[1] - if self.state.get("conv_before_upsample.weight", None) - else 64 - ) - - num_in_ch = self.state["conv_first.weight"].shape[1] - in_chans = num_in_ch - if "conv_last.weight" in state_keys: - num_out_ch = self.state["conv_last.weight"].shape[0] - else: - num_out_ch = num_in_ch - - upscale = 1 - if upsampler == "nearest+conv": - upsample_keys = [ - x for x in state_keys if "conv_up" in x and "bias" not in x - ] - - for upsample_key in upsample_keys: - upscale *= 2 - elif upsampler == "pixelshuffle" or upsampler == "pixelshuffle_aux": - upsample_keys = [ - x - for x in state_keys - if "upsample" in x and "conv" not in x and "bias" not in x - ] - for upsample_key in upsample_keys: - shape = self.state[upsample_key].shape[0] - upscale *= math.sqrt(shape // num_feat) - upscale = int(upscale) - elif upsampler == "pixelshuffledirect": - upscale = int( - math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch) - ) - - max_layer_num = 0 - max_block_num = 0 - for key in state_keys: - result = re.match( - r"layers.(\d*).residual_group.blocks.(\d*).norm1.weight", key - ) - if result: - layer_num, block_num = result.groups() - max_layer_num = max(max_layer_num, int(layer_num)) - max_block_num = max(max_block_num, int(block_num)) - - depths = [max_block_num + 1 for _ in range(max_layer_num + 1)] - - if ( - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - in state_keys - ): - num_heads_num = self.state[ - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - ].shape[-1] - num_heads = [num_heads_num for _ in range(max_layer_num + 1)] - else: - num_heads = depths - - embed_dim = self.state["conv_first.weight"].shape[0] - - mlp_ratio = float( - self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0] - / embed_dim - ) - - # TODO: could actually count the layers, but this should do - if "layers.0.conv.4.weight" in state_keys: - resi_connection = "3conv" - else: - resi_connection = "1conv" - - window_size = int( - math.sqrt( - self.state[ - "layers.0.residual_group.blocks.0.attn.relative_position_index" - ].shape[0] - ) - ) - - if "layers.0.residual_group.blocks.1.attn_mask" in state_keys: - img_size = int( - math.sqrt( - self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0] - ) - * window_size - ) - - # The JPEG models are the only ones with window-size 7, and they also use this range - img_range = 255.0 if window_size == 7 else 1.0 - - self.in_nc = num_in_ch - self.out_nc = num_out_ch - self.num_feat = num_feat - self.embed_dim = embed_dim - self.num_heads = num_heads - self.depths = depths - self.window_size = window_size - self.mlp_ratio = mlp_ratio - self.scale = upscale - self.upsampler = upsampler - self.img_size = img_size - self.img_range = img_range - self.resi_connection = resi_connection - - self.supports_fp16 = False # Too much weirdness to support this at the moment - self.supports_bfp16 = True - self.min_size_restriction = 16 - - ## END AUTO DETECTION - - if in_chans == 3: - rgb_mean = (0.4488, 0.4371, 0.4040) - self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) - else: - self.mean = torch.zeros(1, 1, 1, 1) - self.upscale = upscale - self.upsampler = upsampler - self.window_size = window_size - - ##################################################################################################### - ################################### 1, shallow feature extraction ################################### - self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) - - ##################################################################################################### - ################################### 2, deep feature extraction ###################################### - self.num_layers = len(depths) - self.embed_dim = embed_dim - self.ape = ape - self.patch_norm = patch_norm - self.num_features = embed_dim - self.mlp_ratio = mlp_ratio - - # split image into non-overlapping patches - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - num_patches = self.patch_embed.num_patches - patches_resolution = self.patch_embed.patches_resolution - self.patches_resolution = patches_resolution - - # merge non-overlapping patches into image - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - - # absolute position embedding - if self.ape: - self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) # type: ignore - trunc_normal_(self.absolute_pos_embed, std=0.02) - - self.pos_drop = nn.Dropout(p=drop_rate) - - # stochastic depth - dpr = [ - x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) - ] # stochastic depth decay rule - - # build Residual Swin Transformer blocks (RSTB) - self.layers = nn.ModuleList() - for i_layer in range(self.num_layers): - layer = RSTB( - dim=embed_dim, - input_resolution=(patches_resolution[0], patches_resolution[1]), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, - drop=drop_rate, - attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], # type: ignore # no impact on SR results - norm_layer=norm_layer, - downsample=None, - use_checkpoint=use_checkpoint, - img_size=img_size, - patch_size=patch_size, - resi_connection=resi_connection, - ) - self.layers.append(layer) - - if self.upsampler == "pixelshuffle_hf": - self.layers_hf = nn.ModuleList() - for i_layer in range(self.num_layers): - layer = RSTB( - dim=embed_dim, - input_resolution=(patches_resolution[0], patches_resolution[1]), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, - drop=drop_rate, - attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], # type: ignore # no impact on SR results # type: ignore - norm_layer=norm_layer, - downsample=None, - use_checkpoint=use_checkpoint, - img_size=img_size, - patch_size=patch_size, - resi_connection=resi_connection, - ) - self.layers_hf.append(layer) - - self.norm = norm_layer(self.num_features) - - # build the last conv layer in deep feature extraction - if resi_connection == "1conv": - self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) - elif resi_connection == "3conv": - # to save parameters and memory - self.conv_after_body = nn.Sequential( - nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1), - ) - - ##################################################################################################### - ################################ 3, high quality image reconstruction ################################ - if self.upsampler == "pixelshuffle": - # for classical SR - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - elif self.upsampler == "pixelshuffle_aux": - self.conv_bicubic = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.conv_aux = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - self.conv_after_aux = nn.Sequential( - nn.Conv2d(3, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - elif self.upsampler == "pixelshuffle_hf": - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.upsample_hf = Upsample_hf(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - self.conv_first_hf = nn.Sequential( - nn.Conv2d(num_feat, embed_dim, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.conv_after_body_hf = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) - self.conv_before_upsample_hf = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.conv_last_hf = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR (to save parameters) - self.upsample = UpsampleOneStep( - upscale, - embed_dim, - num_out_ch, - (patches_resolution[0], patches_resolution[1]), - ) - elif self.upsampler == "nearest+conv": - # for real-world SR (less artifacts) - assert self.upscale == 4, "only support x4 now." - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - else: - # for image denoising and JPEG compression artifact reduction - self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1) - - self.apply(self._init_weights) - - self.load_state_dict(state_dict) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - - @torch.jit.ignore # type: ignore - def no_weight_decay(self): - return {"absolute_pos_embed"} - - @torch.jit.ignore # type: ignore - def no_weight_decay_keywords(self): - return {"relative_position_bias_table"} - - def check_image_size(self, x): - _, _, h, w = x.size() - mod_pad_h = (self.window_size - h % self.window_size) % self.window_size - mod_pad_w = (self.window_size - w % self.window_size) % self.window_size - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect") - return x - - def forward_features(self, x): - x_size = (x.shape[2], x.shape[3]) - x = self.patch_embed(x) - if self.ape: - x = x + self.absolute_pos_embed - x = self.pos_drop(x) - - for layer in self.layers: - x = layer(x, x_size) - - x = self.norm(x) # B L C - x = self.patch_unembed(x, x_size) - - return x - - def forward_features_hf(self, x): - x_size = (x.shape[2], x.shape[3]) - x = self.patch_embed(x) - if self.ape: - x = x + self.absolute_pos_embed - x = self.pos_drop(x) - - for layer in self.layers_hf: - x = layer(x, x_size) - - x = self.norm(x) # B L C - x = self.patch_unembed(x, x_size) - - return x - - def forward(self, x): - H, W = x.shape[2:] - x = self.check_image_size(x) - - self.mean = self.mean.type_as(x) - x = (x - self.mean) * self.img_range - - if self.upsampler == "pixelshuffle": - # for classical SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - elif self.upsampler == "pixelshuffle_aux": - bicubic = F.interpolate( - x, - size=(H * self.upscale, W * self.upscale), - mode="bicubic", - align_corners=False, - ) - bicubic = self.conv_bicubic(bicubic) - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - aux = self.conv_aux(x) # b, 3, LR_H, LR_W - x = self.conv_after_aux(aux) - x = ( - self.upsample(x)[:, :, : H * self.upscale, : W * self.upscale] - + bicubic[:, :, : H * self.upscale, : W * self.upscale] - ) - x = self.conv_last(x) - aux = aux / self.img_range + self.mean - elif self.upsampler == "pixelshuffle_hf": - # for classical SR with HF - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x_before = self.conv_before_upsample(x) - x_out = self.conv_last(self.upsample(x_before)) - - x_hf = self.conv_first_hf(x_before) - x_hf = self.conv_after_body_hf(self.forward_features_hf(x_hf)) + x_hf - x_hf = self.conv_before_upsample_hf(x_hf) - x_hf = self.conv_last_hf(self.upsample_hf(x_hf)) - x = x_out + x_hf - x_hf = x_hf / self.img_range + self.mean - - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.upsample(x) - elif self.upsampler == "nearest+conv": - # for real-world SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.lrelu( - self.conv_up1( - torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest") - ) - ) - x = self.lrelu( - self.conv_up2( - torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest") - ) - ) - x = self.conv_last(self.lrelu(self.conv_hr(x))) - else: - # for image denoising and JPEG compression artifact reduction - x_first = self.conv_first(x) - res = self.conv_after_body(self.forward_features(x_first)) + x_first - x = x + self.conv_last(res) - - x = x / self.img_range + self.mean - if self.upsampler == "pixelshuffle_aux": - # NOTE: I removed an "aux" output here. not sure what that was for - return x[:, :, : H * self.upscale, : W * self.upscale] # type: ignore - - elif self.upsampler == "pixelshuffle_hf": - x_out = x_out / self.img_range + self.mean # type: ignore - return x_out[:, :, : H * self.upscale, : W * self.upscale], x[:, :, : H * self.upscale, : W * self.upscale], x_hf[:, :, : H * self.upscale, : W * self.upscale] # type: ignore - - else: - return x[:, :, : H * self.upscale, : W * self.upscale] - - def flops(self): - flops = 0 - H, W = self.patches_resolution - flops += H * W * 3 * self.embed_dim * 9 - flops += self.patch_embed.flops() - for i, layer in enumerate(self.layers): - flops += layer.flops() # type: ignore - flops += H * W * 3 * self.embed_dim * self.embed_dim - flops += self.upsample.flops() # type: ignore - return flops diff --git a/backend/comfy_nodes/chainner_models/architecture/SwinIR.py b/backend/comfy_nodes/chainner_models/architecture/SwinIR.py deleted file mode 100644 index 439dcbcb..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/SwinIR.py +++ /dev/null @@ -1,1224 +0,0 @@ -# pylint: skip-file -# ----------------------------------------------------------------------------------- -# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257 -# Originally Written by Ze Liu, Modified by Jingyun Liang. -# ----------------------------------------------------------------------------------- - -import math -import re - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint as checkpoint - -# Originally from the timm package -from .timm.drop import DropPath -from .timm.helpers import to_2tuple -from .timm.weight_init import trunc_normal_ - - -class Mlp(nn.Module): - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.0, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -def window_partition(x, window_size): - """ - Args: - x: (B, H, W, C) - window_size (int): window size - - Returns: - windows: (num_windows*B, window_size, window_size, C) - """ - B, H, W, C = x.shape - x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = ( - x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) - ) - return windows - - -def window_reverse(windows, window_size, H, W): - """ - Args: - windows: (num_windows*B, window_size, window_size, C) - window_size (int): Window size - H (int): Height of image - W (int): Width of image - - Returns: - x: (B, H, W, C) - """ - B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view( - B, H // window_size, W // window_size, window_size, window_size, -1 - ) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) - return x - - -class WindowAttention(nn.Module): - r"""Window based multi-head self attention (W-MSA) module with relative position bias. - It supports both of shifted and non-shifted window. - - Args: - dim (int): Number of input channels. - window_size (tuple[int]): The height and width of the window. - num_heads (int): Number of attention heads. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set - attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 - proj_drop (float, optional): Dropout ratio of output. Default: 0.0 - """ - - def __init__( - self, - dim, - window_size, - num_heads, - qkv_bias=True, - qk_scale=None, - attn_drop=0.0, - proj_drop=0.0, - ): - super().__init__() - self.dim = dim - self.window_size = window_size # Wh, Ww - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = qk_scale or head_dim**-0.5 - - # define a parameter table of relative position bias - self.relative_position_bias_table = nn.Parameter( # type: ignore - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) - ) # 2*Wh-1 * 2*Ww-1, nH - - # get pair-wise relative position index for each token inside the window - coords_h = torch.arange(self.window_size[0]) - coords_w = torch.arange(self.window_size[1]) - coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww - coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = ( - coords_flatten[:, :, None] - coords_flatten[:, None, :] - ) # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute( - 1, 2, 0 - ).contiguous() # Wh*Ww, Wh*Ww, 2 - relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 - relative_coords[:, :, 1] += self.window_size[1] - 1 - relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 - relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww - self.register_buffer("relative_position_index", relative_position_index) - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - - self.proj_drop = nn.Dropout(proj_drop) - - trunc_normal_(self.relative_position_bias_table, std=0.02) - self.softmax = nn.Softmax(dim=-1) - - def forward(self, x, mask=None): - """ - Args: - x: input features with shape of (num_windows*B, N, C) - mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None - """ - B_, N, C = x.shape - qkv = ( - self.qkv(x) - .reshape(B_, N, 3, self.num_heads, C // self.num_heads) - .permute(2, 0, 3, 1, 4) - ) - q, k, v = ( - qkv[0], - qkv[1], - qkv[2], - ) # make torchscript happy (cannot use tensor as tuple) - - q = q * self.scale - attn = q @ k.transpose(-2, -1) - - relative_position_bias = self.relative_position_bias_table[ - self.relative_position_index.view(-1) # type: ignore - ].view( - self.window_size[0] * self.window_size[1], - self.window_size[0] * self.window_size[1], - -1, - ) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute( - 2, 0, 1 - ).contiguous() # nH, Wh*Ww, Wh*Ww - attn = attn + relative_position_bias.unsqueeze(0) - - if mask is not None: - nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze( - 1 - ).unsqueeze(0) - attn = attn.view(-1, self.num_heads, N, N) - attn = self.softmax(attn) - else: - attn = self.softmax(attn) - - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B_, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - def extra_repr(self) -> str: - return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}" - - def flops(self, N): - # calculate flops for 1 window with token length of N - flops = 0 - # qkv = self.qkv(x) - flops += N * self.dim * 3 * self.dim - # attn = (q @ k.transpose(-2, -1)) - flops += self.num_heads * N * (self.dim // self.num_heads) * N - # x = (attn @ v) - flops += self.num_heads * N * N * (self.dim // self.num_heads) - # x = self.proj(x) - flops += N * self.dim * self.dim - return flops - - -class SwinTransformerBlock(nn.Module): - r"""Swin Transformer Block. - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resulotion. - num_heads (int): Number of attention heads. - window_size (int): Window size. - shift_size (int): Shift size for SW-MSA. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float, optional): Stochastic depth rate. Default: 0.0 - act_layer (nn.Module, optional): Activation layer. Default: nn.GELU - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__( - self, - dim, - input_resolution, - num_heads, - window_size=7, - shift_size=0, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - act_layer=nn.GELU, - norm_layer=nn.LayerNorm, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.num_heads = num_heads - self.window_size = window_size - self.shift_size = shift_size - self.mlp_ratio = mlp_ratio - if min(self.input_resolution) <= self.window_size: - # if window size is larger than input resolution, we don't partition windows - self.shift_size = 0 - self.window_size = min(self.input_resolution) - assert ( - 0 <= self.shift_size < self.window_size - ), "shift_size must in 0-window_size" - - self.norm1 = norm_layer(dim) - self.attn = WindowAttention( - dim, - window_size=to_2tuple(self.window_size), - num_heads=num_heads, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - attn_drop=attn_drop, - proj_drop=drop, - ) - - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp( - in_features=dim, - hidden_features=mlp_hidden_dim, - act_layer=act_layer, - drop=drop, - ) - - if self.shift_size > 0: - attn_mask = self.calculate_mask(self.input_resolution) - else: - attn_mask = None - - self.register_buffer("attn_mask", attn_mask) - - def calculate_mask(self, x_size): - # calculate attention mask for SW-MSA - H, W = x_size - img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - w_slices = ( - slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None), - ) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - mask_windows = window_partition( - img_mask, self.window_size - ) # nW, window_size, window_size, 1 - mask_windows = mask_windows.view(-1, self.window_size * self.window_size) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( - attn_mask == 0, float(0.0) - ) - - return attn_mask - - def forward(self, x, x_size): - H, W = x_size - B, L, C = x.shape - # assert L == H * W, "input feature has wrong size" - - shortcut = x - x = self.norm1(x) - x = x.view(B, H, W, C) - - # cyclic shift - if self.shift_size > 0: - shifted_x = torch.roll( - x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2) - ) - else: - shifted_x = x - - # partition windows - x_windows = window_partition( - shifted_x, self.window_size - ) # nW*B, window_size, window_size, C - x_windows = x_windows.view( - -1, self.window_size * self.window_size, C - ) # nW*B, window_size*window_size, C - - # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size - if self.input_resolution == x_size: - attn_windows = self.attn( - x_windows, mask=self.attn_mask - ) # nW*B, window_size*window_size, C - else: - attn_windows = self.attn( - x_windows, mask=self.calculate_mask(x_size).to(x.device) - ) - - # merge windows - attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) - shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C - - # reverse cyclic shift - if self.shift_size > 0: - x = torch.roll( - shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2) - ) - else: - x = shifted_x - x = x.view(B, H * W, C) - - # FFN - x = shortcut + self.drop_path(x) - x = x + self.drop_path(self.mlp(self.norm2(x))) - - return x - - def extra_repr(self) -> str: - return ( - f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" - ) - - def flops(self): - flops = 0 - H, W = self.input_resolution - # norm1 - flops += self.dim * H * W - # W-MSA/SW-MSA - nW = H * W / self.window_size / self.window_size - flops += nW * self.attn.flops(self.window_size * self.window_size) - # mlp - flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio - # norm2 - flops += self.dim * H * W - return flops - - -class PatchMerging(nn.Module): - r"""Patch Merging Layer. - - Args: - input_resolution (tuple[int]): Resolution of input feature. - dim (int): Number of input channels. - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) - self.norm = norm_layer(4 * dim) - - def forward(self, x): - """ - x: B, H*W, C - """ - H, W = self.input_resolution - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even." - - x = x.view(B, H, W, C) - - x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C - x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C - x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C - x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C - x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C - x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C - - x = self.norm(x) - x = self.reduction(x) - - return x - - def extra_repr(self) -> str: - return f"input_resolution={self.input_resolution}, dim={self.dim}" - - def flops(self): - H, W = self.input_resolution - flops = H * W * self.dim - flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim - return flops - - -class BasicLayer(nn.Module): - """A basic Swin Transformer layer for one stage. - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - ): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.depth = depth - self.use_checkpoint = use_checkpoint - - # build blocks - self.blocks = nn.ModuleList( - [ - SwinTransformerBlock( - dim=dim, - input_resolution=input_resolution, - num_heads=num_heads, - window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path[i] - if isinstance(drop_path, list) - else drop_path, - norm_layer=norm_layer, - ) - for i in range(depth) - ] - ) - - # patch merging layer - if downsample is not None: - self.downsample = downsample( - input_resolution, dim=dim, norm_layer=norm_layer - ) - else: - self.downsample = None - - def forward(self, x, x_size): - for blk in self.blocks: - if self.use_checkpoint: - x = checkpoint.checkpoint(blk, x, x_size) - else: - x = blk(x, x_size) - if self.downsample is not None: - x = self.downsample(x) - return x - - def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" - - def flops(self): - flops = 0 - for blk in self.blocks: - flops += blk.flops() # type: ignore - if self.downsample is not None: - flops += self.downsample.flops() - return flops - - -class RSTB(nn.Module): - """Residual Swin Transformer Block (RSTB). - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - img_size: Input image size. - patch_size: Patch size. - resi_connection: The convolutional block before residual connection. - """ - - def __init__( - self, - dim, - input_resolution, - depth, - num_heads, - window_size, - mlp_ratio=4.0, - qkv_bias=True, - qk_scale=None, - drop=0.0, - attn_drop=0.0, - drop_path=0.0, - norm_layer=nn.LayerNorm, - downsample=None, - use_checkpoint=False, - img_size=224, - patch_size=4, - resi_connection="1conv", - ): - super(RSTB, self).__init__() - - self.dim = dim - self.input_resolution = input_resolution - - self.residual_group = BasicLayer( - dim=dim, - input_resolution=input_resolution, - depth=depth, - num_heads=num_heads, - window_size=window_size, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path, - norm_layer=norm_layer, - downsample=downsample, - use_checkpoint=use_checkpoint, - ) - - if resi_connection == "1conv": - self.conv = nn.Conv2d(dim, dim, 3, 1, 1) - elif resi_connection == "3conv": - # to save parameters and memory - self.conv = nn.Sequential( - nn.Conv2d(dim, dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(dim // 4, dim, 3, 1, 1), - ) - - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=0, - embed_dim=dim, - norm_layer=None, - ) - - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=0, - embed_dim=dim, - norm_layer=None, - ) - - def forward(self, x, x_size): - return ( - self.patch_embed( - self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size)) - ) - + x - ) - - def flops(self): - flops = 0 - flops += self.residual_group.flops() - H, W = self.input_resolution - flops += H * W * self.dim * self.dim * 9 - flops += self.patch_embed.flops() - flops += self.patch_unembed.flops() - - return flops - - -class PatchEmbed(nn.Module): - r"""Image to Patch Embedding - - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [ - img_size[0] // patch_size[0], # type: ignore - img_size[1] // patch_size[1], # type: ignore - ] - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - - def forward(self, x): - x = x.flatten(2).transpose(1, 2) # B Ph*Pw C - if self.norm is not None: - x = self.norm(x) - return x - - def flops(self): - flops = 0 - H, W = self.img_size - if self.norm is not None: - flops += H * W * self.embed_dim # type: ignore - return flops - - -class PatchUnEmbed(nn.Module): - r"""Image to Patch Unembedding - - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__( - self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None - ): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [ - img_size[0] // patch_size[0], # type: ignore - img_size[1] // patch_size[1], # type: ignore - ] - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - def forward(self, x, x_size): - B, HW, C = x.shape - x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C - return x - - def flops(self): - flops = 0 - return flops - - -class Upsample(nn.Sequential): - """Upsample module. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - """ - - def __init__(self, scale, num_feat): - m = [] - if (scale & (scale - 1)) == 0: # scale = 2^n - for _ in range(int(math.log(scale, 2))): - m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(2)) - elif scale == 3: - m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) - m.append(nn.PixelShuffle(3)) - else: - raise ValueError( - f"scale {scale} is not supported. " "Supported scales: 2^n and 3." - ) - super(Upsample, self).__init__(*m) - - -class UpsampleOneStep(nn.Sequential): - """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) - Used in lightweight SR to save parameters. - - Args: - scale (int): Scale factor. Supported scales: 2^n and 3. - num_feat (int): Channel number of intermediate features. - - """ - - def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): - self.num_feat = num_feat - self.input_resolution = input_resolution - m = [] - m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1)) - m.append(nn.PixelShuffle(scale)) - super(UpsampleOneStep, self).__init__(*m) - - def flops(self): - H, W = self.input_resolution # type: ignore - flops = H * W * self.num_feat * 3 * 9 - return flops - - -class SwinIR(nn.Module): - r"""SwinIR - A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer. - - Args: - img_size (int | tuple(int)): Input image size. Default 64 - patch_size (int | tuple(int)): Patch size. Default: 1 - in_chans (int): Number of input image channels. Default: 3 - embed_dim (int): Patch embedding dimension. Default: 96 - depths (tuple(int)): Depth of each Swin Transformer layer. - num_heads (tuple(int)): Number of attention heads in different layers. - window_size (int): Window size. Default: 7 - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None - drop_rate (float): Dropout rate. Default: 0 - attn_drop_rate (float): Attention dropout rate. Default: 0 - drop_path_rate (float): Stochastic depth rate. Default: 0.1 - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. - ape (bool): If True, add absolute position embedding to the patch embedding. Default: False - patch_norm (bool): If True, add normalization after patch embedding. Default: True - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False - upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction - img_range: Image range. 1. or 255. - upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None - resi_connection: The convolutional block before residual connection. '1conv'/'3conv' - """ - - def __init__( - self, - state_dict, - **kwargs, - ): - super(SwinIR, self).__init__() - - # Defaults - img_size = 64 - patch_size = 1 - in_chans = 3 - embed_dim = 96 - depths = [6, 6, 6, 6] - num_heads = [6, 6, 6, 6] - window_size = 7 - mlp_ratio = 4.0 - qkv_bias = True - qk_scale = None - drop_rate = 0.0 - attn_drop_rate = 0.0 - drop_path_rate = 0.1 - norm_layer = nn.LayerNorm - ape = False - patch_norm = True - use_checkpoint = False - upscale = 2 - img_range = 1.0 - upsampler = "" - resi_connection = "1conv" - num_feat = 64 - num_in_ch = in_chans - num_out_ch = in_chans - supports_fp16 = True - self.start_unshuffle = 1 - - self.model_arch = "SwinIR" - self.sub_type = "SR" - self.state = state_dict - if "params_ema" in self.state: - self.state = self.state["params_ema"] - elif "params" in self.state: - self.state = self.state["params"] - - state_keys = self.state.keys() - - if "conv_before_upsample.0.weight" in state_keys: - if "conv_up1.weight" in state_keys: - upsampler = "nearest+conv" - else: - upsampler = "pixelshuffle" - supports_fp16 = False - elif "upsample.0.weight" in state_keys: - upsampler = "pixelshuffledirect" - else: - upsampler = "" - - num_feat = ( - self.state.get("conv_before_upsample.0.weight", None).shape[1] - if self.state.get("conv_before_upsample.weight", None) - else 64 - ) - - if "conv_first.1.weight" in self.state: - self.state["conv_first.weight"] = self.state.pop("conv_first.1.weight") - self.state["conv_first.bias"] = self.state.pop("conv_first.1.bias") - self.start_unshuffle = round(math.sqrt(self.state["conv_first.weight"].shape[1] // 3)) - - num_in_ch = self.state["conv_first.weight"].shape[1] - in_chans = num_in_ch - if "conv_last.weight" in state_keys: - num_out_ch = self.state["conv_last.weight"].shape[0] - else: - num_out_ch = num_in_ch - - upscale = 1 - if upsampler == "nearest+conv": - upsample_keys = [ - x for x in state_keys if "conv_up" in x and "bias" not in x - ] - - for upsample_key in upsample_keys: - upscale *= 2 - elif upsampler == "pixelshuffle": - upsample_keys = [ - x - for x in state_keys - if "upsample" in x and "conv" not in x and "bias" not in x - ] - for upsample_key in upsample_keys: - shape = self.state[upsample_key].shape[0] - upscale *= math.sqrt(shape // num_feat) - upscale = int(upscale) - elif upsampler == "pixelshuffledirect": - upscale = int( - math.sqrt(self.state["upsample.0.bias"].shape[0] // num_out_ch) - ) - - max_layer_num = 0 - max_block_num = 0 - for key in state_keys: - result = re.match( - r"layers.(\d*).residual_group.blocks.(\d*).norm1.weight", key - ) - if result: - layer_num, block_num = result.groups() - max_layer_num = max(max_layer_num, int(layer_num)) - max_block_num = max(max_block_num, int(block_num)) - - depths = [max_block_num + 1 for _ in range(max_layer_num + 1)] - - if ( - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - in state_keys - ): - num_heads_num = self.state[ - "layers.0.residual_group.blocks.0.attn.relative_position_bias_table" - ].shape[-1] - num_heads = [num_heads_num for _ in range(max_layer_num + 1)] - else: - num_heads = depths - - embed_dim = self.state["conv_first.weight"].shape[0] - - mlp_ratio = float( - self.state["layers.0.residual_group.blocks.0.mlp.fc1.bias"].shape[0] - / embed_dim - ) - - # TODO: could actually count the layers, but this should do - if "layers.0.conv.4.weight" in state_keys: - resi_connection = "3conv" - else: - resi_connection = "1conv" - - window_size = int( - math.sqrt( - self.state[ - "layers.0.residual_group.blocks.0.attn.relative_position_index" - ].shape[0] - ) - ) - - if "layers.0.residual_group.blocks.1.attn_mask" in state_keys: - img_size = int( - math.sqrt( - self.state["layers.0.residual_group.blocks.1.attn_mask"].shape[0] - ) - * window_size - ) - - # The JPEG models are the only ones with window-size 7, and they also use this range - img_range = 255.0 if window_size == 7 else 1.0 - - self.in_nc = num_in_ch - self.out_nc = num_out_ch - self.num_feat = num_feat - self.embed_dim = embed_dim - self.num_heads = num_heads - self.depths = depths - self.window_size = window_size - self.mlp_ratio = mlp_ratio - self.scale = upscale / self.start_unshuffle - self.upsampler = upsampler - self.img_size = img_size - self.img_range = img_range - self.resi_connection = resi_connection - - self.supports_fp16 = False # Too much weirdness to support this at the moment - self.supports_bfp16 = True - self.min_size_restriction = 16 - - self.img_range = img_range - if in_chans == 3: - rgb_mean = (0.4488, 0.4371, 0.4040) - self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) - else: - self.mean = torch.zeros(1, 1, 1, 1) - self.upscale = upscale - self.upsampler = upsampler - self.window_size = window_size - - ##################################################################################################### - ################################### 1, shallow feature extraction ################################### - self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) - - ##################################################################################################### - ################################### 2, deep feature extraction ###################################### - self.num_layers = len(depths) - self.embed_dim = embed_dim - self.ape = ape - self.patch_norm = patch_norm - self.num_features = embed_dim - self.mlp_ratio = mlp_ratio - - # split image into non-overlapping patches - self.patch_embed = PatchEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - num_patches = self.patch_embed.num_patches - patches_resolution = self.patch_embed.patches_resolution - self.patches_resolution = patches_resolution - - # merge non-overlapping patches into image - self.patch_unembed = PatchUnEmbed( - img_size=img_size, - patch_size=patch_size, - in_chans=embed_dim, - embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None, - ) - - # absolute position embedding - if self.ape: - self.absolute_pos_embed = nn.Parameter( # type: ignore - torch.zeros(1, num_patches, embed_dim) - ) - trunc_normal_(self.absolute_pos_embed, std=0.02) - - self.pos_drop = nn.Dropout(p=drop_rate) - - # stochastic depth - dpr = [ - x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) - ] # stochastic depth decay rule - - # build Residual Swin Transformer blocks (RSTB) - self.layers = nn.ModuleList() - for i_layer in range(self.num_layers): - layer = RSTB( - dim=embed_dim, - input_resolution=(patches_resolution[0], patches_resolution[1]), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, - qk_scale=qk_scale, - drop=drop_rate, - attn_drop=attn_drop_rate, - drop_path=dpr[ - sum(depths[:i_layer]) : sum(depths[: i_layer + 1]) # type: ignore - ], # no impact on SR results - norm_layer=norm_layer, - downsample=None, - use_checkpoint=use_checkpoint, - img_size=img_size, - patch_size=patch_size, - resi_connection=resi_connection, - ) - self.layers.append(layer) - self.norm = norm_layer(self.num_features) - - # build the last conv layer in deep feature extraction - if resi_connection == "1conv": - self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) - elif resi_connection == "3conv": - # to save parameters and memory - self.conv_after_body = nn.Sequential( - nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1), - ) - - ##################################################################################################### - ################################ 3, high quality image reconstruction ################################ - if self.upsampler == "pixelshuffle": - # for classical SR - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.upsample = Upsample(upscale, num_feat) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR (to save parameters) - self.upsample = UpsampleOneStep( - upscale, - embed_dim, - num_out_ch, - (patches_resolution[0], patches_resolution[1]), - ) - elif self.upsampler == "nearest+conv": - # for real-world SR (less artifacts) - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) - ) - self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - if self.upscale == 4: - self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - elif self.upscale == 8: - self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - else: - # for image denoising and JPEG compression artifact reduction - self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1) - - self.apply(self._init_weights) - self.load_state_dict(self.state, strict=False) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - - @torch.jit.ignore # type: ignore - def no_weight_decay(self): - return {"absolute_pos_embed"} - - @torch.jit.ignore # type: ignore - def no_weight_decay_keywords(self): - return {"relative_position_bias_table"} - - def check_image_size(self, x): - _, _, h, w = x.size() - mod_pad_h = (self.window_size - h % self.window_size) % self.window_size - mod_pad_w = (self.window_size - w % self.window_size) % self.window_size - x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "reflect") - return x - - def forward_features(self, x): - x_size = (x.shape[2], x.shape[3]) - x = self.patch_embed(x) - if self.ape: - x = x + self.absolute_pos_embed - x = self.pos_drop(x) - - for layer in self.layers: - x = layer(x, x_size) - - x = self.norm(x) # B L C - x = self.patch_unembed(x, x_size) - - return x - - def forward(self, x): - H, W = x.shape[2:] - x = self.check_image_size(x) - - self.mean = self.mean.type_as(x) - x = (x - self.mean) * self.img_range - - if self.start_unshuffle > 1: - x = torch.nn.functional.pixel_unshuffle(x, self.start_unshuffle) - - if self.upsampler == "pixelshuffle": - # for classical SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - elif self.upsampler == "pixelshuffledirect": - # for lightweight SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.upsample(x) - elif self.upsampler == "nearest+conv": - # for real-world SR - x = self.conv_first(x) - x = self.conv_after_body(self.forward_features(x)) + x - x = self.conv_before_upsample(x) - x = self.lrelu( - self.conv_up1( - torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest") # type: ignore - ) - ) - if self.upscale == 4: - x = self.lrelu( - self.conv_up2( - torch.nn.functional.interpolate( # type: ignore - x, scale_factor=2, mode="nearest" - ) - ) - ) - elif self.upscale == 8: - x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) - x = self.lrelu(self.conv_up3(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) - x = self.conv_last(self.lrelu(self.conv_hr(x))) - else: - # for image denoising and JPEG compression artifact reduction - x_first = self.conv_first(x) - res = self.conv_after_body(self.forward_features(x_first)) + x_first - x = x + self.conv_last(res) - - x = x / self.img_range + self.mean - - return x[:, :, : H * self.upscale, : W * self.upscale] - - def flops(self): - flops = 0 - H, W = self.patches_resolution - flops += H * W * 3 * self.embed_dim * 9 - flops += self.patch_embed.flops() - for i, layer in enumerate(self.layers): - flops += layer.flops() # type: ignore - flops += H * W * 3 * self.embed_dim * self.embed_dim - flops += self.upsample.flops() # type: ignore - return flops diff --git a/backend/comfy_nodes/chainner_models/architecture/__init__.py b/backend/comfy_nodes/chainner_models/architecture/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/comfy_nodes/chainner_models/architecture/block.py b/backend/comfy_nodes/chainner_models/architecture/block.py deleted file mode 100644 index d7bc5d22..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/block.py +++ /dev/null @@ -1,546 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -from __future__ import annotations - -from collections import OrderedDict -try: - from typing import Literal -except ImportError: - from typing_extensions import Literal - -import torch -import torch.nn as nn - -#################### -# Basic blocks -#################### - - -def act(act_type: str, inplace=True, neg_slope=0.2, n_prelu=1): - # helper selecting activation - # neg_slope: for leakyrelu and init of prelu - # n_prelu: for p_relu num_parameters - act_type = act_type.lower() - if act_type == "relu": - layer = nn.ReLU(inplace) - elif act_type == "leakyrelu": - layer = nn.LeakyReLU(neg_slope, inplace) - elif act_type == "prelu": - layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope) - else: - raise NotImplementedError( - "activation layer [{:s}] is not found".format(act_type) - ) - return layer - - -def norm(norm_type: str, nc: int): - # helper selecting normalization layer - norm_type = norm_type.lower() - if norm_type == "batch": - layer = nn.BatchNorm2d(nc, affine=True) - elif norm_type == "instance": - layer = nn.InstanceNorm2d(nc, affine=False) - else: - raise NotImplementedError( - "normalization layer [{:s}] is not found".format(norm_type) - ) - return layer - - -def pad(pad_type: str, padding): - # helper selecting padding layer - # if padding is 'zero', do by conv layers - pad_type = pad_type.lower() - if padding == 0: - return None - if pad_type == "reflect": - layer = nn.ReflectionPad2d(padding) - elif pad_type == "replicate": - layer = nn.ReplicationPad2d(padding) - else: - raise NotImplementedError( - "padding layer [{:s}] is not implemented".format(pad_type) - ) - return layer - - -def get_valid_padding(kernel_size, dilation): - kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) - padding = (kernel_size - 1) // 2 - return padding - - -class ConcatBlock(nn.Module): - # Concat the output of a submodule to its input - def __init__(self, submodule): - super(ConcatBlock, self).__init__() - self.sub = submodule - - def forward(self, x): - output = torch.cat((x, self.sub(x)), dim=1) - return output - - def __repr__(self): - tmpstr = "Identity .. \n|" - modstr = self.sub.__repr__().replace("\n", "\n|") - tmpstr = tmpstr + modstr - return tmpstr - - -class ShortcutBlock(nn.Module): - # Elementwise sum the output of a submodule to its input - def __init__(self, submodule): - super(ShortcutBlock, self).__init__() - self.sub = submodule - - def forward(self, x): - output = x + self.sub(x) - return output - - def __repr__(self): - tmpstr = "Identity + \n|" - modstr = self.sub.__repr__().replace("\n", "\n|") - tmpstr = tmpstr + modstr - return tmpstr - - -class ShortcutBlockSPSR(nn.Module): - # Elementwise sum the output of a submodule to its input - def __init__(self, submodule): - super(ShortcutBlockSPSR, self).__init__() - self.sub = submodule - - def forward(self, x): - return x, self.sub - - def __repr__(self): - tmpstr = "Identity + \n|" - modstr = self.sub.__repr__().replace("\n", "\n|") - tmpstr = tmpstr + modstr - return tmpstr - - -def sequential(*args): - # Flatten Sequential. It unwraps nn.Sequential. - if len(args) == 1: - if isinstance(args[0], OrderedDict): - raise NotImplementedError("sequential does not support OrderedDict input.") - return args[0] # No sequential is needed. - modules = [] - for module in args: - if isinstance(module, nn.Sequential): - for submodule in module.children(): - modules.append(submodule) - elif isinstance(module, nn.Module): - modules.append(module) - return nn.Sequential(*modules) - - -ConvMode = Literal["CNA", "NAC", "CNAC"] - - -# 2x2x2 Conv Block -def conv_block_2c2( - in_nc, - out_nc, - act_type="relu", -): - return sequential( - nn.Conv2d(in_nc, out_nc, kernel_size=2, padding=1), - nn.Conv2d(out_nc, out_nc, kernel_size=2, padding=0), - act(act_type) if act_type else None, - ) - - -def conv_block( - in_nc: int, - out_nc: int, - kernel_size, - stride=1, - dilation=1, - groups=1, - bias=True, - pad_type="zero", - norm_type: str | None = None, - act_type: str | None = "relu", - mode: ConvMode = "CNA", - c2x2=False, -): - """ - Conv layer with padding, normalization, activation - mode: CNA --> Conv -> Norm -> Act - NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16) - """ - - if c2x2: - return conv_block_2c2(in_nc, out_nc, act_type=act_type) - - assert mode in ("CNA", "NAC", "CNAC"), "Wrong conv mode [{:s}]".format(mode) - padding = get_valid_padding(kernel_size, dilation) - p = pad(pad_type, padding) if pad_type and pad_type != "zero" else None - padding = padding if pad_type == "zero" else 0 - - c = nn.Conv2d( - in_nc, - out_nc, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias, - groups=groups, - ) - a = act(act_type) if act_type else None - if mode in ("CNA", "CNAC"): - n = norm(norm_type, out_nc) if norm_type else None - return sequential(p, c, n, a) - elif mode == "NAC": - if norm_type is None and act_type is not None: - a = act(act_type, inplace=False) - # Important! - # input----ReLU(inplace)----Conv--+----output - # |________________________| - # inplace ReLU will modify the input, therefore wrong output - n = norm(norm_type, in_nc) if norm_type else None - return sequential(n, a, p, c) - else: - assert False, f"Invalid conv mode {mode}" - - -#################### -# Useful blocks -#################### - - -class ResNetBlock(nn.Module): - """ - ResNet Block, 3-3 style - with extra residual scaling used in EDSR - (Enhanced Deep Residual Networks for Single Image Super-Resolution, CVPRW 17) - """ - - def __init__( - self, - in_nc, - mid_nc, - out_nc, - kernel_size=3, - stride=1, - dilation=1, - groups=1, - bias=True, - pad_type="zero", - norm_type=None, - act_type="relu", - mode: ConvMode = "CNA", - res_scale=1, - ): - super(ResNetBlock, self).__init__() - conv0 = conv_block( - in_nc, - mid_nc, - kernel_size, - stride, - dilation, - groups, - bias, - pad_type, - norm_type, - act_type, - mode, - ) - if mode == "CNA": - act_type = None - if mode == "CNAC": # Residual path: |-CNAC-| - act_type = None - norm_type = None - conv1 = conv_block( - mid_nc, - out_nc, - kernel_size, - stride, - dilation, - groups, - bias, - pad_type, - norm_type, - act_type, - mode, - ) - # if in_nc != out_nc: - # self.project = conv_block(in_nc, out_nc, 1, stride, dilation, 1, bias, pad_type, \ - # None, None) - # print('Need a projecter in ResNetBlock.') - # else: - # self.project = lambda x:x - self.res = sequential(conv0, conv1) - self.res_scale = res_scale - - def forward(self, x): - res = self.res(x).mul(self.res_scale) - return x + res - - -class RRDB(nn.Module): - """ - Residual in Residual Dense Block - (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks) - """ - - def __init__( - self, - nf, - kernel_size=3, - gc=32, - stride=1, - bias: bool = True, - pad_type="zero", - norm_type=None, - act_type="leakyrelu", - mode: ConvMode = "CNA", - _convtype="Conv2D", - _spectral_norm=False, - plus=False, - c2x2=False, - ): - super(RRDB, self).__init__() - self.RDB1 = ResidualDenseBlock_5C( - nf, - kernel_size, - gc, - stride, - bias, - pad_type, - norm_type, - act_type, - mode, - plus=plus, - c2x2=c2x2, - ) - self.RDB2 = ResidualDenseBlock_5C( - nf, - kernel_size, - gc, - stride, - bias, - pad_type, - norm_type, - act_type, - mode, - plus=plus, - c2x2=c2x2, - ) - self.RDB3 = ResidualDenseBlock_5C( - nf, - kernel_size, - gc, - stride, - bias, - pad_type, - norm_type, - act_type, - mode, - plus=plus, - c2x2=c2x2, - ) - - def forward(self, x): - out = self.RDB1(x) - out = self.RDB2(out) - out = self.RDB3(out) - return out * 0.2 + x - - -class ResidualDenseBlock_5C(nn.Module): - """ - Residual Dense Block - style: 5 convs - The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18) - Modified options that can be used: - - "Partial Convolution based Padding" arXiv:1811.11718 - - "Spectral normalization" arXiv:1802.05957 - - "ICASSP 2020 - ESRGAN+ : Further Improving ESRGAN" N. C. - {Rakotonirina} and A. {Rasoanaivo} - - Args: - nf (int): Channel number of intermediate features (num_feat). - gc (int): Channels for each growth (num_grow_ch: growth channel, - i.e. intermediate channels). - convtype (str): the type of convolution to use. Default: 'Conv2D' - gaussian_noise (bool): enable the ESRGAN+ gaussian noise (no new - trainable parameters) - plus (bool): enable the additional residual paths from ESRGAN+ - (adds trainable parameters) - """ - - def __init__( - self, - nf=64, - kernel_size=3, - gc=32, - stride=1, - bias: bool = True, - pad_type="zero", - norm_type=None, - act_type="leakyrelu", - mode: ConvMode = "CNA", - plus=False, - c2x2=False, - ): - super(ResidualDenseBlock_5C, self).__init__() - - ## + - self.conv1x1 = conv1x1(nf, gc) if plus else None - ## + - - self.conv1 = conv_block( - nf, - gc, - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=act_type, - mode=mode, - c2x2=c2x2, - ) - self.conv2 = conv_block( - nf + gc, - gc, - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=act_type, - mode=mode, - c2x2=c2x2, - ) - self.conv3 = conv_block( - nf + 2 * gc, - gc, - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=act_type, - mode=mode, - c2x2=c2x2, - ) - self.conv4 = conv_block( - nf + 3 * gc, - gc, - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=act_type, - mode=mode, - c2x2=c2x2, - ) - if mode == "CNA": - last_act = None - else: - last_act = act_type - self.conv5 = conv_block( - nf + 4 * gc, - nf, - 3, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=last_act, - mode=mode, - c2x2=c2x2, - ) - - def forward(self, x): - x1 = self.conv1(x) - x2 = self.conv2(torch.cat((x, x1), 1)) - if self.conv1x1: - # pylint: disable=not-callable - x2 = x2 + self.conv1x1(x) # + - x3 = self.conv3(torch.cat((x, x1, x2), 1)) - x4 = self.conv4(torch.cat((x, x1, x2, x3), 1)) - if self.conv1x1: - x4 = x4 + x2 # + - x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) - return x5 * 0.2 + x - - -def conv1x1(in_planes, out_planes, stride=1): - return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) - - -#################### -# Upsampler -#################### - - -def pixelshuffle_block( - in_nc: int, - out_nc: int, - upscale_factor=2, - kernel_size=3, - stride=1, - bias=True, - pad_type="zero", - norm_type: str | None = None, - act_type="relu", -): - """ - Pixel shuffle layer - (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional - Neural Network, CVPR17) - """ - conv = conv_block( - in_nc, - out_nc * (upscale_factor**2), - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=None, - act_type=None, - ) - pixel_shuffle = nn.PixelShuffle(upscale_factor) - - n = norm(norm_type, out_nc) if norm_type else None - a = act(act_type) if act_type else None - return sequential(conv, pixel_shuffle, n, a) - - -def upconv_block( - in_nc: int, - out_nc: int, - upscale_factor=2, - kernel_size=3, - stride=1, - bias=True, - pad_type="zero", - norm_type: str | None = None, - act_type="relu", - mode="nearest", - c2x2=False, -): - # Up conv - # described in https://distill.pub/2016/deconv-checkerboard/ - upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode) - conv = conv_block( - in_nc, - out_nc, - kernel_size, - stride, - bias=bias, - pad_type=pad_type, - norm_type=norm_type, - act_type=act_type, - c2x2=c2x2, - ) - return sequential(upsample, conv) diff --git a/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-GFPGAN b/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-GFPGAN deleted file mode 100644 index 5ac273fd..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-GFPGAN +++ /dev/null @@ -1,351 +0,0 @@ -Tencent is pleased to support the open source community by making GFPGAN available. - -Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved. - -GFPGAN is licensed under the Apache License Version 2.0 except for the third-party components listed below. - - -Terms of the Apache License Version 2.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: - -You must give any other recipients of the Work or Derivative Works a copy of this License; and - -You must cause any modified files to carry prominent notices stating that You changed the files; and - -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 - -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 - - - -Other dependencies and licenses: - - -Open Source Software licensed under the Apache 2.0 license and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. basicsr -Copyright 2018-2020 BasicSR Authors - - -This BasicSR project is released under the Apache 2.0 license. - -A copy of Apache 2.0 is included in this file. - -StyleGAN2 -The codes are modified from the repository stylegan2-pytorch. Many thanks to the author - Kim Seonghyeon 😊 for translating from the official TensorFlow codes to PyTorch ones. Here is the license of stylegan2-pytorch. -The official repository is https://github.com/NVlabs/stylegan2, and here is the NVIDIA license. -DFDNet -The codes are largely modified from the repository DFDNet. Their license is Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License. - -Terms of the Nvidia License: ---------------------------------------------- - -1. Definitions - -"Licensor" means any person or entity that distributes its Work. - -"Software" means the original work of authorship made available under -this License. - -"Work" means the Software and any additions to or derivative works of -the Software that are made available under this License. - -"Nvidia Processors" means any central processing unit (CPU), graphics -processing unit (GPU), field-programmable gate array (FPGA), -application-specific integrated circuit (ASIC) or any combination -thereof designed, made, sold, or provided by Nvidia or its affiliates. - -The terms "reproduce," "reproduction," "derivative works," and -"distribution" have the meaning as provided under U.S. copyright law; -provided, however, that 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. - -Works, including the Software, are "made available" under this License -by including in or with the Work either (a) a copyright notice -referencing the applicability of this License to the Work, or (b) a -copy of this License. - -2. License Grants - - 2.1 Copyright Grant. Subject to the terms and conditions of this - License, each Licensor grants to you a perpetual, worldwide, - non-exclusive, royalty-free, copyright license to reproduce, - prepare derivative works of, publicly display, publicly perform, - sublicense and distribute its Work and any resulting derivative - works in any form. - -3. Limitations - - 3.1 Redistribution. You may reproduce or distribute the Work only - if (a) you do so under this License, (b) you include a complete - copy of this License with your distribution, and (c) you retain - without modification any copyright, patent, trademark, or - attribution notices that are present in the Work. - - 3.2 Derivative Works. You may specify that additional or different - terms apply to the use, reproduction, and distribution of your - derivative works of the Work ("Your Terms") only if (a) Your Terms - provide that the use limitation in Section 3.3 applies to your - derivative works, and (b) you identify the specific derivative - works that are subject to Your Terms. Notwithstanding Your Terms, - this License (including the redistribution requirements in Section - 3.1) will continue to apply to the Work itself. - - 3.3 Use Limitation. The Work and any derivative works thereof only - may be used or intended for use non-commercially. The Work or - derivative works thereof may be used or intended for use by Nvidia - or its affiliates commercially or non-commercially. As used herein, - "non-commercially" means for research or evaluation purposes only. - - 3.4 Patent Claims. If you bring or threaten to bring a patent claim - against any Licensor (including any claim, cross-claim or - counterclaim in a lawsuit) to enforce any patents that you allege - are infringed by any Work, then your rights under this License from - such Licensor (including the grants in Sections 2.1 and 2.2) will - terminate immediately. - - 3.5 Trademarks. This License does not grant any rights to use any - Licensor's or its affiliates' names, logos, or trademarks, except - as necessary to reproduce the notices described in this License. - - 3.6 Termination. If you violate any term of this License, then your - rights under this License (including the grants in Sections 2.1 and - 2.2) will terminate immediately. - -4. Disclaimer of Warranty. - -THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF -MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR -NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER -THIS LICENSE. - -5. Limitation of Liability. - -EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL -THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE -SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, -INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF -OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK -(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, -LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER -COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF -THE POSSIBILITY OF SUCH DAMAGES. - -MIT License - -Copyright (c) 2019 Kim Seonghyeon - -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. - - - -Open Source Software licensed under the BSD 3-Clause license: ---------------------------------------------- -1. torchvision -Copyright (c) Soumith Chintala 2016, -All rights reserved. - -2. torch -Copyright (c) 2016- Facebook, Inc (Adam Paszke) -Copyright (c) 2014- Facebook, Inc (Soumith Chintala) -Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) -Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) -Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) -Copyright (c) 2011-2013 NYU (Clement Farabet) -Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) -Copyright (c) 2006 Idiap Research Institute (Samy Bengio) -Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) - - -Terms of the BSD 3-Clause License: ---------------------------------------------- -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. - - - -Open Source Software licensed under the BSD 3-Clause License and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. numpy -Copyright (c) 2005-2020, NumPy Developers. -All rights reserved. - -A copy of BSD 3-Clause License is included in this file. - -The NumPy repository and source distributions bundle several libraries that are -compatibly licensed. We list these here. - -Name: Numpydoc -Files: doc/sphinxext/numpydoc/* -License: BSD-2-Clause - For details, see doc/sphinxext/LICENSE.txt - -Name: scipy-sphinx-theme -Files: doc/scipy-sphinx-theme/* -License: BSD-3-Clause AND PSF-2.0 AND Apache-2.0 - For details, see doc/scipy-sphinx-theme/LICENSE.txt - -Name: lapack-lite -Files: numpy/linalg/lapack_lite/* -License: BSD-3-Clause - For details, see numpy/linalg/lapack_lite/LICENSE.txt - -Name: tempita -Files: tools/npy_tempita/* -License: MIT - For details, see tools/npy_tempita/license.txt - -Name: dragon4 -Files: numpy/core/src/multiarray/dragon4.c -License: MIT - For license text, see numpy/core/src/multiarray/dragon4.c - - - -Open Source Software licensed under the MIT license: ---------------------------------------------- -1. facexlib -Copyright (c) 2020 Xintao Wang - -2. opencv-python -Copyright (c) Olli-Pekka Heinisuo -Please note that only files in cv2 package are used. - - -Terms of the MIT License: ---------------------------------------------- -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. - - - -Open Source Software licensed under the MIT license and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. tqdm -Copyright (c) 2013 noamraph - -`tqdm` is a product of collaborative work. -Unless otherwise stated, all authors (see commit logs) retain copyright -for their respective work, and release the work under the MIT licence -(text below). - -Exceptions or notable authors are listed below -in reverse chronological order: - -* files: * - MPLv2.0 2015-2020 (c) Casper da Costa-Luis - [casperdcl](https://github.com/casperdcl). -* files: tqdm/_tqdm.py - MIT 2016 (c) [PR #96] on behalf of Google Inc. -* files: tqdm/_tqdm.py setup.py README.rst MANIFEST.in .gitignore - MIT 2013 (c) Noam Yorav-Raphael, original author. - -[PR #96]: https://github.com/tqdm/tqdm/pull/96 - - -Mozilla Public Licence (MPL) v. 2.0 - Exhibit A ------------------------------------------------ - -This Source Code Form is subject to the terms of the -Mozilla Public License, v. 2.0. -If a copy of the MPL was not distributed with this file, -You can obtain one at https://mozilla.org/MPL/2.0/. - - -MIT License (MIT) ------------------ - -Copyright (c) 2013 noamraph - -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/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-RestoreFormer b/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-RestoreFormer deleted file mode 100644 index 5ac273fd..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-RestoreFormer +++ /dev/null @@ -1,351 +0,0 @@ -Tencent is pleased to support the open source community by making GFPGAN available. - -Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved. - -GFPGAN is licensed under the Apache License Version 2.0 except for the third-party components listed below. - - -Terms of the Apache License Version 2.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: - -You must give any other recipients of the Work or Derivative Works a copy of this License; and - -You must cause any modified files to carry prominent notices stating that You changed the files; and - -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 - -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 - - - -Other dependencies and licenses: - - -Open Source Software licensed under the Apache 2.0 license and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. basicsr -Copyright 2018-2020 BasicSR Authors - - -This BasicSR project is released under the Apache 2.0 license. - -A copy of Apache 2.0 is included in this file. - -StyleGAN2 -The codes are modified from the repository stylegan2-pytorch. Many thanks to the author - Kim Seonghyeon 😊 for translating from the official TensorFlow codes to PyTorch ones. Here is the license of stylegan2-pytorch. -The official repository is https://github.com/NVlabs/stylegan2, and here is the NVIDIA license. -DFDNet -The codes are largely modified from the repository DFDNet. Their license is Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License. - -Terms of the Nvidia License: ---------------------------------------------- - -1. Definitions - -"Licensor" means any person or entity that distributes its Work. - -"Software" means the original work of authorship made available under -this License. - -"Work" means the Software and any additions to or derivative works of -the Software that are made available under this License. - -"Nvidia Processors" means any central processing unit (CPU), graphics -processing unit (GPU), field-programmable gate array (FPGA), -application-specific integrated circuit (ASIC) or any combination -thereof designed, made, sold, or provided by Nvidia or its affiliates. - -The terms "reproduce," "reproduction," "derivative works," and -"distribution" have the meaning as provided under U.S. copyright law; -provided, however, that 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. - -Works, including the Software, are "made available" under this License -by including in or with the Work either (a) a copyright notice -referencing the applicability of this License to the Work, or (b) a -copy of this License. - -2. License Grants - - 2.1 Copyright Grant. Subject to the terms and conditions of this - License, each Licensor grants to you a perpetual, worldwide, - non-exclusive, royalty-free, copyright license to reproduce, - prepare derivative works of, publicly display, publicly perform, - sublicense and distribute its Work and any resulting derivative - works in any form. - -3. Limitations - - 3.1 Redistribution. You may reproduce or distribute the Work only - if (a) you do so under this License, (b) you include a complete - copy of this License with your distribution, and (c) you retain - without modification any copyright, patent, trademark, or - attribution notices that are present in the Work. - - 3.2 Derivative Works. You may specify that additional or different - terms apply to the use, reproduction, and distribution of your - derivative works of the Work ("Your Terms") only if (a) Your Terms - provide that the use limitation in Section 3.3 applies to your - derivative works, and (b) you identify the specific derivative - works that are subject to Your Terms. Notwithstanding Your Terms, - this License (including the redistribution requirements in Section - 3.1) will continue to apply to the Work itself. - - 3.3 Use Limitation. The Work and any derivative works thereof only - may be used or intended for use non-commercially. The Work or - derivative works thereof may be used or intended for use by Nvidia - or its affiliates commercially or non-commercially. As used herein, - "non-commercially" means for research or evaluation purposes only. - - 3.4 Patent Claims. If you bring or threaten to bring a patent claim - against any Licensor (including any claim, cross-claim or - counterclaim in a lawsuit) to enforce any patents that you allege - are infringed by any Work, then your rights under this License from - such Licensor (including the grants in Sections 2.1 and 2.2) will - terminate immediately. - - 3.5 Trademarks. This License does not grant any rights to use any - Licensor's or its affiliates' names, logos, or trademarks, except - as necessary to reproduce the notices described in this License. - - 3.6 Termination. If you violate any term of this License, then your - rights under this License (including the grants in Sections 2.1 and - 2.2) will terminate immediately. - -4. Disclaimer of Warranty. - -THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF -MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR -NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER -THIS LICENSE. - -5. Limitation of Liability. - -EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL -THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE -SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, -INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF -OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK -(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, -LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER -COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF -THE POSSIBILITY OF SUCH DAMAGES. - -MIT License - -Copyright (c) 2019 Kim Seonghyeon - -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. - - - -Open Source Software licensed under the BSD 3-Clause license: ---------------------------------------------- -1. torchvision -Copyright (c) Soumith Chintala 2016, -All rights reserved. - -2. torch -Copyright (c) 2016- Facebook, Inc (Adam Paszke) -Copyright (c) 2014- Facebook, Inc (Soumith Chintala) -Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) -Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) -Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) -Copyright (c) 2011-2013 NYU (Clement Farabet) -Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) -Copyright (c) 2006 Idiap Research Institute (Samy Bengio) -Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) - - -Terms of the BSD 3-Clause License: ---------------------------------------------- -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. - - - -Open Source Software licensed under the BSD 3-Clause License and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. numpy -Copyright (c) 2005-2020, NumPy Developers. -All rights reserved. - -A copy of BSD 3-Clause License is included in this file. - -The NumPy repository and source distributions bundle several libraries that are -compatibly licensed. We list these here. - -Name: Numpydoc -Files: doc/sphinxext/numpydoc/* -License: BSD-2-Clause - For details, see doc/sphinxext/LICENSE.txt - -Name: scipy-sphinx-theme -Files: doc/scipy-sphinx-theme/* -License: BSD-3-Clause AND PSF-2.0 AND Apache-2.0 - For details, see doc/scipy-sphinx-theme/LICENSE.txt - -Name: lapack-lite -Files: numpy/linalg/lapack_lite/* -License: BSD-3-Clause - For details, see numpy/linalg/lapack_lite/LICENSE.txt - -Name: tempita -Files: tools/npy_tempita/* -License: MIT - For details, see tools/npy_tempita/license.txt - -Name: dragon4 -Files: numpy/core/src/multiarray/dragon4.c -License: MIT - For license text, see numpy/core/src/multiarray/dragon4.c - - - -Open Source Software licensed under the MIT license: ---------------------------------------------- -1. facexlib -Copyright (c) 2020 Xintao Wang - -2. opencv-python -Copyright (c) Olli-Pekka Heinisuo -Please note that only files in cv2 package are used. - - -Terms of the MIT License: ---------------------------------------------- -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. - - - -Open Source Software licensed under the MIT license and Other Licenses of the Third-Party Components therein: ---------------------------------------------- -1. tqdm -Copyright (c) 2013 noamraph - -`tqdm` is a product of collaborative work. -Unless otherwise stated, all authors (see commit logs) retain copyright -for their respective work, and release the work under the MIT licence -(text below). - -Exceptions or notable authors are listed below -in reverse chronological order: - -* files: * - MPLv2.0 2015-2020 (c) Casper da Costa-Luis - [casperdcl](https://github.com/casperdcl). -* files: tqdm/_tqdm.py - MIT 2016 (c) [PR #96] on behalf of Google Inc. -* files: tqdm/_tqdm.py setup.py README.rst MANIFEST.in .gitignore - MIT 2013 (c) Noam Yorav-Raphael, original author. - -[PR #96]: https://github.com/tqdm/tqdm/pull/96 - - -Mozilla Public Licence (MPL) v. 2.0 - Exhibit A ------------------------------------------------ - -This Source Code Form is subject to the terms of the -Mozilla Public License, v. 2.0. -If a copy of the MPL was not distributed with this file, -You can obtain one at https://mozilla.org/MPL/2.0/. - - -MIT License (MIT) ------------------ - -Copyright (c) 2013 noamraph - -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/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-codeformer b/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-codeformer deleted file mode 100644 index be6c4ed8..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/LICENSE-codeformer +++ /dev/null @@ -1,35 +0,0 @@ -S-Lab License 1.0 - -Copyright 2022 S-Lab - -Redistribution and use for non-commercial purpose 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. - -In the event that redistribution and/or use for commercial purpose in -source or binary forms, with or without modification is required, -please contact the contributor(s) of the work. diff --git a/backend/comfy_nodes/chainner_models/architecture/face/arcface_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/arcface_arch.py deleted file mode 100644 index b548af05..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/arcface_arch.py +++ /dev/null @@ -1,265 +0,0 @@ -import torch.nn as nn - - -def conv3x3(inplanes, outplanes, stride=1): - """A simple wrapper for 3x3 convolution with padding. - - Args: - inplanes (int): Channel number of inputs. - outplanes (int): Channel number of outputs. - stride (int): Stride in convolution. Default: 1. - """ - return nn.Conv2d( - inplanes, outplanes, kernel_size=3, stride=stride, padding=1, bias=False - ) - - -class BasicBlock(nn.Module): - """Basic residual block used in the ResNetArcFace architecture. - - Args: - inplanes (int): Channel number of inputs. - planes (int): Channel number of outputs. - stride (int): Stride in convolution. Default: 1. - downsample (nn.Module): The downsample module. Default: None. - """ - - expansion = 1 # output channel expansion ratio - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(BasicBlock, self).__init__() - self.conv1 = conv3x3(inplanes, planes, stride) - self.bn1 = nn.BatchNorm2d(planes) - self.relu = nn.ReLU(inplace=True) - self.conv2 = conv3x3(planes, planes) - self.bn2 = nn.BatchNorm2d(planes) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class IRBlock(nn.Module): - """Improved residual block (IR Block) used in the ResNetArcFace architecture. - - Args: - inplanes (int): Channel number of inputs. - planes (int): Channel number of outputs. - stride (int): Stride in convolution. Default: 1. - downsample (nn.Module): The downsample module. Default: None. - use_se (bool): Whether use the SEBlock (squeeze and excitation block). Default: True. - """ - - expansion = 1 # output channel expansion ratio - - def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True): - super(IRBlock, self).__init__() - self.bn0 = nn.BatchNorm2d(inplanes) - self.conv1 = conv3x3(inplanes, inplanes) - self.bn1 = nn.BatchNorm2d(inplanes) - self.prelu = nn.PReLU() - self.conv2 = conv3x3(inplanes, planes, stride) - self.bn2 = nn.BatchNorm2d(planes) - self.downsample = downsample - self.stride = stride - self.use_se = use_se - if self.use_se: - self.se = SEBlock(planes) - - def forward(self, x): - residual = x - out = self.bn0(x) - out = self.conv1(out) - out = self.bn1(out) - out = self.prelu(out) - - out = self.conv2(out) - out = self.bn2(out) - if self.use_se: - out = self.se(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.prelu(out) - - return out - - -class Bottleneck(nn.Module): - """Bottleneck block used in the ResNetArcFace architecture. - - Args: - inplanes (int): Channel number of inputs. - planes (int): Channel number of outputs. - stride (int): Stride in convolution. Default: 1. - downsample (nn.Module): The downsample module. Default: None. - """ - - expansion = 4 # output channel expansion ratio - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(Bottleneck, self).__init__() - self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d( - planes, planes, kernel_size=3, stride=stride, padding=1, bias=False - ) - self.bn2 = nn.BatchNorm2d(planes) - self.conv3 = nn.Conv2d( - planes, planes * self.expansion, kernel_size=1, bias=False - ) - self.bn3 = nn.BatchNorm2d(planes * self.expansion) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class SEBlock(nn.Module): - """The squeeze-and-excitation block (SEBlock) used in the IRBlock. - - Args: - channel (int): Channel number of inputs. - reduction (int): Channel reduction ration. Default: 16. - """ - - def __init__(self, channel, reduction=16): - super(SEBlock, self).__init__() - self.avg_pool = nn.AdaptiveAvgPool2d( - 1 - ) # pool to 1x1 without spatial information - self.fc = nn.Sequential( - nn.Linear(channel, channel // reduction), - nn.PReLU(), - nn.Linear(channel // reduction, channel), - nn.Sigmoid(), - ) - - def forward(self, x): - b, c, _, _ = x.size() - y = self.avg_pool(x).view(b, c) - y = self.fc(y).view(b, c, 1, 1) - return x * y - - -class ResNetArcFace(nn.Module): - """ArcFace with ResNet architectures. - - Ref: ArcFace: Additive Angular Margin Loss for Deep Face Recognition. - - Args: - block (str): Block used in the ArcFace architecture. - layers (tuple(int)): Block numbers in each layer. - use_se (bool): Whether use the SEBlock (squeeze and excitation block). Default: True. - """ - - def __init__(self, block, layers, use_se=True): - if block == "IRBlock": - block = IRBlock - self.inplanes = 64 - self.use_se = use_se - super(ResNetArcFace, self).__init__() - - self.conv1 = nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False) - self.bn1 = nn.BatchNorm2d(64) - self.prelu = nn.PReLU() - self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) - self.layer1 = self._make_layer(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=2) - self.bn4 = nn.BatchNorm2d(512) - self.dropout = nn.Dropout() - self.fc5 = nn.Linear(512 * 8 * 8, 512) - self.bn5 = nn.BatchNorm1d(512) - - # initialization - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.xavier_normal_(m.weight) - elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.xavier_normal_(m.weight) - nn.init.constant_(m.bias, 0) - - def _make_layer(self, block, planes, num_blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - bias=False, - ), - nn.BatchNorm2d(planes * block.expansion), - ) - layers = [] - layers.append( - block(self.inplanes, planes, stride, downsample, use_se=self.use_se) - ) - self.inplanes = planes - for _ in range(1, num_blocks): - layers.append(block(self.inplanes, planes, use_se=self.use_se)) - - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - x = self.prelu(x) - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - x = self.bn4(x) - x = self.dropout(x) - x = x.view(x.size(0), -1) - x = self.fc5(x) - x = self.bn5(x) - - return x diff --git a/backend/comfy_nodes/chainner_models/architecture/face/codeformer.py b/backend/comfy_nodes/chainner_models/architecture/face/codeformer.py deleted file mode 100644 index 06614007..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/codeformer.py +++ /dev/null @@ -1,790 +0,0 @@ -""" -Modified from https://github.com/sczhou/CodeFormer -VQGAN code, adapted from the original created by the Unleashing Transformers authors: -https://github.com/samb-t/unleashing-transformers/blob/master/models/vqgan.py -This verison of the arch specifically was gathered from an old version of GFPGAN. If this is a problem, please contact me. -""" -import math -from typing import Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F -import logging as logger -from torch import Tensor - - -class VectorQuantizer(nn.Module): - def __init__(self, codebook_size, emb_dim, beta): - super(VectorQuantizer, self).__init__() - self.codebook_size = codebook_size # number of embeddings - self.emb_dim = emb_dim # dimension of embedding - self.beta = beta # commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2 - self.embedding = nn.Embedding(self.codebook_size, self.emb_dim) - self.embedding.weight.data.uniform_( - -1.0 / self.codebook_size, 1.0 / self.codebook_size - ) - - def forward(self, z): - # reshape z -> (batch, height, width, channel) and flatten - z = z.permute(0, 2, 3, 1).contiguous() - z_flattened = z.view(-1, self.emb_dim) - - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - d = ( - (z_flattened**2).sum(dim=1, keepdim=True) - + (self.embedding.weight**2).sum(1) - - 2 * torch.matmul(z_flattened, self.embedding.weight.t()) - ) - - mean_distance = torch.mean(d) - # find closest encodings - # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1) - min_encoding_scores, min_encoding_indices = torch.topk( - d, 1, dim=1, largest=False - ) - # [0-1], higher score, higher confidence - min_encoding_scores = torch.exp(-min_encoding_scores / 10) - - min_encodings = torch.zeros( - min_encoding_indices.shape[0], self.codebook_size - ).to(z) - min_encodings.scatter_(1, min_encoding_indices, 1) - - # get quantized latent vectors - z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape) - # compute loss for embedding - loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean( - (z_q - z.detach()) ** 2 - ) - # preserve gradients - z_q = z + (z_q - z).detach() - - # perplexity - e_mean = torch.mean(min_encodings, dim=0) - perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10))) - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return ( - z_q, - loss, - { - "perplexity": perplexity, - "min_encodings": min_encodings, - "min_encoding_indices": min_encoding_indices, - "min_encoding_scores": min_encoding_scores, - "mean_distance": mean_distance, - }, - ) - - def get_codebook_feat(self, indices, shape): - # input indices: batch*token_num -> (batch*token_num)*1 - # shape: batch, height, width, channel - indices = indices.view(-1, 1) - min_encodings = torch.zeros(indices.shape[0], self.codebook_size).to(indices) - min_encodings.scatter_(1, indices, 1) - # get quantized latent vectors - z_q = torch.matmul(min_encodings.float(), self.embedding.weight) - - if shape is not None: # reshape back to match original input shape - z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous() - - return z_q - - -class GumbelQuantizer(nn.Module): - def __init__( - self, - codebook_size, - emb_dim, - num_hiddens, - straight_through=False, - kl_weight=5e-4, - temp_init=1.0, - ): - super().__init__() - self.codebook_size = codebook_size # number of embeddings - self.emb_dim = emb_dim # dimension of embedding - self.straight_through = straight_through - self.temperature = temp_init - self.kl_weight = kl_weight - self.proj = nn.Conv2d( - num_hiddens, codebook_size, 1 - ) # projects last encoder layer to quantized logits - self.embed = nn.Embedding(codebook_size, emb_dim) - - def forward(self, z): - hard = self.straight_through if self.training else True - - logits = self.proj(z) - - soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=hard) - - z_q = torch.einsum("b n h w, n d -> b d h w", soft_one_hot, self.embed.weight) - - # + kl divergence to the prior loss - qy = F.softmax(logits, dim=1) - diff = ( - self.kl_weight - * torch.sum(qy * torch.log(qy * self.codebook_size + 1e-10), dim=1).mean() - ) - min_encoding_indices = soft_one_hot.argmax(dim=1) - - return z_q, diff, {"min_encoding_indices": min_encoding_indices} - - -class Downsample(nn.Module): - def __init__(self, in_channels): - super().__init__() - self.conv = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=2, padding=0 - ) - - def forward(self, x): - pad = (0, 1, 0, 1) - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) - return x - - -class Upsample(nn.Module): - def __init__(self, in_channels): - super().__init__() - self.conv = nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1 - ) - - def forward(self, x): - x = F.interpolate(x, scale_factor=2.0, mode="nearest") - x = self.conv(x) - - return x - - -class AttnBlock(nn.Module): - def __init__(self, in_channels): - super().__init__() - self.in_channels = in_channels - - self.norm = normalize(in_channels) - self.q = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.k = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.v = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.proj_out = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - - def forward(self, x): - h_ = x - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - # compute attention - b, c, h, w = q.shape - q = q.reshape(b, c, h * w) - q = q.permute(0, 2, 1) - k = k.reshape(b, c, h * w) - w_ = torch.bmm(q, k) - w_ = w_ * (int(c) ** (-0.5)) - w_ = F.softmax(w_, dim=2) - - # attend to values - v = v.reshape(b, c, h * w) - w_ = w_.permute(0, 2, 1) - h_ = torch.bmm(v, w_) - h_ = h_.reshape(b, c, h, w) - - h_ = self.proj_out(h_) - - return x + h_ - - -class Encoder(nn.Module): - def __init__( - self, - in_channels, - nf, - out_channels, - ch_mult, - num_res_blocks, - resolution, - attn_resolutions, - ): - super().__init__() - self.nf = nf - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.attn_resolutions = attn_resolutions - - curr_res = self.resolution - in_ch_mult = (1,) + tuple(ch_mult) - - blocks = [] - # initial convultion - blocks.append(nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1)) - - # residual and downsampling blocks, with attention on smaller res (16x16) - for i in range(self.num_resolutions): - block_in_ch = nf * in_ch_mult[i] - block_out_ch = nf * ch_mult[i] - for _ in range(self.num_res_blocks): - blocks.append(ResBlock(block_in_ch, block_out_ch)) - block_in_ch = block_out_ch - if curr_res in attn_resolutions: - blocks.append(AttnBlock(block_in_ch)) - - if i != self.num_resolutions - 1: - blocks.append(Downsample(block_in_ch)) - curr_res = curr_res // 2 - - # non-local attention block - blocks.append(ResBlock(block_in_ch, block_in_ch)) # type: ignore - blocks.append(AttnBlock(block_in_ch)) # type: ignore - blocks.append(ResBlock(block_in_ch, block_in_ch)) # type: ignore - - # normalise and convert to latent size - blocks.append(normalize(block_in_ch)) # type: ignore - blocks.append( - nn.Conv2d(block_in_ch, out_channels, kernel_size=3, stride=1, padding=1) # type: ignore - ) - self.blocks = nn.ModuleList(blocks) - - def forward(self, x): - for block in self.blocks: - x = block(x) - - return x - - -class Generator(nn.Module): - def __init__(self, nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim): - super().__init__() - self.nf = nf - self.ch_mult = ch_mult - self.num_resolutions = len(self.ch_mult) - self.num_res_blocks = res_blocks - self.resolution = img_size - self.attn_resolutions = attn_resolutions - self.in_channels = emb_dim - self.out_channels = 3 - block_in_ch = self.nf * self.ch_mult[-1] - curr_res = self.resolution // 2 ** (self.num_resolutions - 1) - - blocks = [] - # initial conv - blocks.append( - nn.Conv2d(self.in_channels, block_in_ch, kernel_size=3, stride=1, padding=1) - ) - - # non-local attention block - blocks.append(ResBlock(block_in_ch, block_in_ch)) - blocks.append(AttnBlock(block_in_ch)) - blocks.append(ResBlock(block_in_ch, block_in_ch)) - - for i in reversed(range(self.num_resolutions)): - block_out_ch = self.nf * self.ch_mult[i] - - for _ in range(self.num_res_blocks): - blocks.append(ResBlock(block_in_ch, block_out_ch)) - block_in_ch = block_out_ch - - if curr_res in self.attn_resolutions: - blocks.append(AttnBlock(block_in_ch)) - - if i != 0: - blocks.append(Upsample(block_in_ch)) - curr_res = curr_res * 2 - - blocks.append(normalize(block_in_ch)) - blocks.append( - nn.Conv2d( - block_in_ch, self.out_channels, kernel_size=3, stride=1, padding=1 - ) - ) - - self.blocks = nn.ModuleList(blocks) - - def forward(self, x): - for block in self.blocks: - x = block(x) - - return x - - -class VQAutoEncoder(nn.Module): - def __init__( - self, - img_size, - nf, - ch_mult, - quantizer="nearest", - res_blocks=2, - attn_resolutions=[16], - codebook_size=1024, - emb_dim=256, - beta=0.25, - gumbel_straight_through=False, - gumbel_kl_weight=1e-8, - model_path=None, - ): - super().__init__() - self.in_channels = 3 - self.nf = nf - self.n_blocks = res_blocks - self.codebook_size = codebook_size - self.embed_dim = emb_dim - self.ch_mult = ch_mult - self.resolution = img_size - self.attn_resolutions = attn_resolutions - self.quantizer_type = quantizer - self.encoder = Encoder( - self.in_channels, - self.nf, - self.embed_dim, - self.ch_mult, - self.n_blocks, - self.resolution, - self.attn_resolutions, - ) - if self.quantizer_type == "nearest": - self.beta = beta # 0.25 - self.quantize = VectorQuantizer( - self.codebook_size, self.embed_dim, self.beta - ) - elif self.quantizer_type == "gumbel": - self.gumbel_num_hiddens = emb_dim - self.straight_through = gumbel_straight_through - self.kl_weight = gumbel_kl_weight - self.quantize = GumbelQuantizer( - self.codebook_size, - self.embed_dim, - self.gumbel_num_hiddens, - self.straight_through, - self.kl_weight, - ) - self.generator = Generator( - nf, ch_mult, res_blocks, img_size, attn_resolutions, emb_dim - ) - - if model_path is not None: - chkpt = torch.load(model_path, map_location="cpu") - if "params_ema" in chkpt: - self.load_state_dict( - torch.load(model_path, map_location="cpu")["params_ema"] - ) - logger.info(f"vqgan is loaded from: {model_path} [params_ema]") - elif "params" in chkpt: - self.load_state_dict( - torch.load(model_path, map_location="cpu")["params"] - ) - logger.info(f"vqgan is loaded from: {model_path} [params]") - else: - raise ValueError("Wrong params!") - - def forward(self, x): - x = self.encoder(x) - quant, codebook_loss, quant_stats = self.quantize(x) - x = self.generator(quant) - return x, codebook_loss, quant_stats - - -def calc_mean_std(feat, eps=1e-5): - """Calculate mean and std for adaptive_instance_normalization. - Args: - feat (Tensor): 4D tensor. - eps (float): A small value added to the variance to avoid - divide-by-zero. Default: 1e-5. - """ - size = feat.size() - assert len(size) == 4, "The input feature should be 4D tensor." - b, c = size[:2] - feat_var = feat.view(b, c, -1).var(dim=2) + eps - feat_std = feat_var.sqrt().view(b, c, 1, 1) - feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1) - return feat_mean, feat_std - - -def adaptive_instance_normalization(content_feat, style_feat): - """Adaptive instance normalization. - Adjust the reference features to have the similar color and illuminations - as those in the degradate features. - Args: - content_feat (Tensor): The reference feature. - style_feat (Tensor): The degradate features. - """ - size = content_feat.size() - style_mean, style_std = calc_mean_std(style_feat) - content_mean, content_std = calc_mean_std(content_feat) - normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand( - size - ) - return normalized_feat * style_std.expand(size) + style_mean.expand(size) - - -class PositionEmbeddingSine(nn.Module): - """ - This is a more standard version of the position embedding, very similar to the one - used by the Attention is all you need paper, generalized to work on images. - """ - - def __init__( - self, num_pos_feats=64, temperature=10000, normalize=False, scale=None - ): - super().__init__() - self.num_pos_feats = num_pos_feats - self.temperature = temperature - self.normalize = normalize - if scale is not None and normalize is False: - raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale - - def forward(self, x, mask=None): - if mask is None: - mask = torch.zeros( - (x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool - ) - not_mask = ~mask # pylint: disable=invalid-unary-operand-type - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) - if self.normalize: - eps = 1e-6 - y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale - x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale - - dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack( - (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos_y = torch.stack( - (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) - return pos - - -def _get_activation_fn(activation): - """Return an activation function given a string""" - if activation == "relu": - return F.relu - if activation == "gelu": - return F.gelu - if activation == "glu": - return F.glu - raise RuntimeError(f"activation should be relu/gelu, not {activation}.") - - -class TransformerSALayer(nn.Module): - def __init__( - self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu" - ): - super().__init__() - self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout) - # Implementation of Feedforward model - MLP - self.linear1 = nn.Linear(embed_dim, dim_mlp) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_mlp, embed_dim) - - self.norm1 = nn.LayerNorm(embed_dim) - self.norm2 = nn.LayerNorm(embed_dim) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - - self.activation = _get_activation_fn(activation) - - def with_pos_embed(self, tensor, pos: Optional[Tensor]): - return tensor if pos is None else tensor + pos - - def forward( - self, - tgt, - tgt_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None, - ): - # self attention - tgt2 = self.norm1(tgt) - q = k = self.with_pos_embed(tgt2, query_pos) - tgt2 = self.self_attn( - q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask - )[0] - tgt = tgt + self.dropout1(tgt2) - - # ffn - tgt2 = self.norm2(tgt) - tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) - tgt = tgt + self.dropout2(tgt2) - return tgt - - -def normalize(in_channels): - return torch.nn.GroupNorm( - num_groups=32, num_channels=in_channels, eps=1e-6, affine=True - ) - - -@torch.jit.script # type: ignore -def swish(x): - return x * torch.sigmoid(x) - - -class ResBlock(nn.Module): - def __init__(self, in_channels, out_channels=None): - super(ResBlock, self).__init__() - self.in_channels = in_channels - self.out_channels = in_channels if out_channels is None else out_channels - self.norm1 = normalize(in_channels) - self.conv1 = nn.Conv2d( - in_channels, out_channels, kernel_size=3, stride=1, padding=1 # type: ignore - ) - self.norm2 = normalize(out_channels) - self.conv2 = nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1 # type: ignore - ) - if self.in_channels != self.out_channels: - self.conv_out = nn.Conv2d( - in_channels, out_channels, kernel_size=1, stride=1, padding=0 # type: ignore - ) - - def forward(self, x_in): - x = x_in - x = self.norm1(x) - x = swish(x) - x = self.conv1(x) - x = self.norm2(x) - x = swish(x) - x = self.conv2(x) - if self.in_channels != self.out_channels: - x_in = self.conv_out(x_in) - - return x + x_in - - -class Fuse_sft_block(nn.Module): - def __init__(self, in_ch, out_ch): - super().__init__() - self.encode_enc = ResBlock(2 * in_ch, out_ch) - - self.scale = nn.Sequential( - nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), - nn.LeakyReLU(0.2, True), - nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), - ) - - self.shift = nn.Sequential( - nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1), - nn.LeakyReLU(0.2, True), - nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), - ) - - def forward(self, enc_feat, dec_feat, w=1): - enc_feat = self.encode_enc(torch.cat([enc_feat, dec_feat], dim=1)) - scale = self.scale(enc_feat) - shift = self.shift(enc_feat) - residual = w * (dec_feat * scale + shift) - out = dec_feat + residual - return out - - -class CodeFormer(VQAutoEncoder): - def __init__(self, state_dict): - dim_embd = 512 - n_head = 8 - n_layers = 9 - codebook_size = 1024 - latent_size = 256 - connect_list = ["32", "64", "128", "256"] - fix_modules = ["quantize", "generator"] - - # This is just a guess as I only have one model to look at - position_emb = state_dict["position_emb"] - dim_embd = position_emb.shape[1] - latent_size = position_emb.shape[0] - - try: - n_layers = len( - set([x.split(".")[1] for x in state_dict.keys() if "ft_layers" in x]) - ) - except: - pass - - codebook_size = state_dict["quantize.embedding.weight"].shape[0] - - # This is also just another guess - n_head_exp = ( - state_dict["ft_layers.0.self_attn.in_proj_weight"].shape[0] // dim_embd - ) - n_head = 2**n_head_exp - - in_nc = state_dict["encoder.blocks.0.weight"].shape[1] - - self.model_arch = "CodeFormer" - self.sub_type = "Face SR" - self.scale = 8 - self.in_nc = in_nc - self.out_nc = in_nc - - self.state = state_dict - - self.supports_fp16 = False - self.supports_bf16 = True - self.min_size_restriction = 16 - - super(CodeFormer, self).__init__( - 512, 64, [1, 2, 2, 4, 4, 8], "nearest", 2, [16], codebook_size - ) - - if fix_modules is not None: - for module in fix_modules: - for param in getattr(self, module).parameters(): - param.requires_grad = False - - self.connect_list = connect_list - self.n_layers = n_layers - self.dim_embd = dim_embd - self.dim_mlp = dim_embd * 2 - - self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd)) # type: ignore - self.feat_emb = nn.Linear(256, self.dim_embd) - - # transformer - self.ft_layers = nn.Sequential( - *[ - TransformerSALayer( - embed_dim=dim_embd, nhead=n_head, dim_mlp=self.dim_mlp, dropout=0.0 - ) - for _ in range(self.n_layers) - ] - ) - - # logits_predict head - self.idx_pred_layer = nn.Sequential( - nn.LayerNorm(dim_embd), nn.Linear(dim_embd, codebook_size, bias=False) - ) - - self.channels = { - "16": 512, - "32": 256, - "64": 256, - "128": 128, - "256": 128, - "512": 64, - } - - # after second residual block for > 16, before attn layer for ==16 - self.fuse_encoder_block = { - "512": 2, - "256": 5, - "128": 8, - "64": 11, - "32": 14, - "16": 18, - } - # after first residual block for > 16, before attn layer for ==16 - self.fuse_generator_block = { - "16": 6, - "32": 9, - "64": 12, - "128": 15, - "256": 18, - "512": 21, - } - - # fuse_convs_dict - self.fuse_convs_dict = nn.ModuleDict() - for f_size in self.connect_list: - in_ch = self.channels[f_size] - self.fuse_convs_dict[f_size] = Fuse_sft_block(in_ch, in_ch) - - self.load_state_dict(state_dict) - - def _init_weights(self, module): - if isinstance(module, (nn.Linear, nn.Embedding)): - module.weight.data.normal_(mean=0.0, std=0.02) - if isinstance(module, nn.Linear) and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.LayerNorm): - module.bias.data.zero_() - module.weight.data.fill_(1.0) - - def forward(self, x, weight=0.5, **kwargs): - detach_16 = True - code_only = False - adain = True - # ################### Encoder ##################### - enc_feat_dict = {} - out_list = [self.fuse_encoder_block[f_size] for f_size in self.connect_list] - for i, block in enumerate(self.encoder.blocks): - x = block(x) - if i in out_list: - enc_feat_dict[str(x.shape[-1])] = x.clone() - - lq_feat = x - # ################# Transformer ################### - # quant_feat, codebook_loss, quant_stats = self.quantize(lq_feat) - pos_emb = self.position_emb.unsqueeze(1).repeat(1, x.shape[0], 1) - # BCHW -> BC(HW) -> (HW)BC - feat_emb = self.feat_emb(lq_feat.flatten(2).permute(2, 0, 1)) - query_emb = feat_emb - # Transformer encoder - for layer in self.ft_layers: - query_emb = layer(query_emb, query_pos=pos_emb) - - # output logits - logits = self.idx_pred_layer(query_emb) # (hw)bn - logits = logits.permute(1, 0, 2) # (hw)bn -> b(hw)n - - if code_only: # for training stage II - # logits doesn't need softmax before cross_entropy loss - return logits, lq_feat - - # ################# Quantization ################### - # if self.training: - # quant_feat = torch.einsum('btn,nc->btc', [soft_one_hot, self.quantize.embedding.weight]) - # # b(hw)c -> bc(hw) -> bchw - # quant_feat = quant_feat.permute(0,2,1).view(lq_feat.shape) - # ------------ - soft_one_hot = F.softmax(logits, dim=2) - _, top_idx = torch.topk(soft_one_hot, 1, dim=2) - quant_feat = self.quantize.get_codebook_feat( - top_idx, shape=[x.shape[0], 16, 16, 256] # type: ignore - ) - # preserve gradients - # quant_feat = lq_feat + (quant_feat - lq_feat).detach() - - if detach_16: - quant_feat = quant_feat.detach() # for training stage III - if adain: - quant_feat = adaptive_instance_normalization(quant_feat, lq_feat) - - # ################## Generator #################### - x = quant_feat - fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list] - - for i, block in enumerate(self.generator.blocks): - x = block(x) - if i in fuse_list: # fuse after i-th block - f_size = str(x.shape[-1]) - if weight > 0: - x = self.fuse_convs_dict[f_size]( - enc_feat_dict[f_size].detach(), x, weight - ) - out = x - # logits doesn't need softmax before cross_entropy loss - # return out, logits, lq_feat - return out, logits diff --git a/backend/comfy_nodes/chainner_models/architecture/face/fused_act.py b/backend/comfy_nodes/chainner_models/architecture/face/fused_act.py deleted file mode 100644 index 7ed52654..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/fused_act.py +++ /dev/null @@ -1,81 +0,0 @@ -# pylint: skip-file -# type: ignore -# modify from https://github.com/rosinality/stylegan2-pytorch/blob/master/op/fused_act.py # noqa:E501 - -import torch -from torch import nn -from torch.autograd import Function - -fused_act_ext = None - - -class FusedLeakyReLUFunctionBackward(Function): - @staticmethod - def forward(ctx, grad_output, out, negative_slope, scale): - ctx.save_for_backward(out) - ctx.negative_slope = negative_slope - ctx.scale = scale - - empty = grad_output.new_empty(0) - - grad_input = fused_act_ext.fused_bias_act( - grad_output, empty, out, 3, 1, negative_slope, scale - ) - - dim = [0] - - if grad_input.ndim > 2: - dim += list(range(2, grad_input.ndim)) - - grad_bias = grad_input.sum(dim).detach() - - return grad_input, grad_bias - - @staticmethod - def backward(ctx, gradgrad_input, gradgrad_bias): - (out,) = ctx.saved_tensors - gradgrad_out = fused_act_ext.fused_bias_act( - gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale - ) - - return gradgrad_out, None, None, None - - -class FusedLeakyReLUFunction(Function): - @staticmethod - def forward(ctx, input, bias, negative_slope, scale): - empty = input.new_empty(0) - out = fused_act_ext.fused_bias_act( - input, bias, empty, 3, 0, negative_slope, scale - ) - ctx.save_for_backward(out) - ctx.negative_slope = negative_slope - ctx.scale = scale - - return out - - @staticmethod - def backward(ctx, grad_output): - (out,) = ctx.saved_tensors - - grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( - grad_output, out, ctx.negative_slope, ctx.scale - ) - - return grad_input, grad_bias, None, None - - -class FusedLeakyReLU(nn.Module): - def __init__(self, channel, negative_slope=0.2, scale=2**0.5): - super().__init__() - - self.bias = nn.Parameter(torch.zeros(channel)) - self.negative_slope = negative_slope - self.scale = scale - - def forward(self, input): - return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) - - -def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2**0.5): - return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) diff --git a/backend/comfy_nodes/chainner_models/architecture/face/gfpgan_bilinear_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/gfpgan_bilinear_arch.py deleted file mode 100644 index b6e820e0..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/gfpgan_bilinear_arch.py +++ /dev/null @@ -1,389 +0,0 @@ -# pylint: skip-file -# type: ignore -import math -import random - -import torch -from torch import nn - -from .gfpganv1_arch import ResUpBlock -from .stylegan2_bilinear_arch import ( - ConvLayer, - EqualConv2d, - EqualLinear, - ResBlock, - ScaledLeakyReLU, - StyleGAN2GeneratorBilinear, -) - - -class StyleGAN2GeneratorBilinearSFT(StyleGAN2GeneratorBilinear): - """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform). - It is the bilinear version. It does not use the complicated UpFirDnSmooth function that is not friendly for - deployment. It can be easily converted to the clean version: StyleGAN2GeneratorCSFT. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - num_mlp=8, - channel_multiplier=2, - lr_mlp=0.01, - narrow=1, - sft_half=False, - ): - super(StyleGAN2GeneratorBilinearSFT, self).__init__( - out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - lr_mlp=lr_mlp, - narrow=narrow, - ) - self.sft_half = sft_half - - def forward( - self, - styles, - conditions, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2GeneratorBilinearSFT. - Args: - styles (list[Tensor]): Sample codes of styles. - conditions (list[Tensor]): SFT conditions to generators. - input_is_latent (bool): Whether input is latent style. Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - truncation (float): The truncation ratio. Default: 1. - truncation_latent (Tensor | None): The truncation latent tensor. Default: None. - inject_index (int | None): The injection index for mixing noise. Default: None. - return_latents (bool): Whether to return style latents. Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latents with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - - # the conditions may have fewer levels - if i < len(conditions): - # SFT part to combine the conditions - if self.sft_half: # only apply SFT to half of the channels - out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1) - out_sft = out_sft * conditions[i - 1] + conditions[i] - out = torch.cat([out_same, out_sft], dim=1) - else: # apply SFT to all the channels - out = out * conditions[i - 1] + conditions[i] - - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) # feature back to the rgb space - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None - - -class GFPGANBilinear(nn.Module): - """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT. - It is the bilinear version and it does not use the complicated UpFirDnSmooth function that is not friendly for - deployment. It can be easily converted to the clean version: GFPGANv1Clean. - Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None. - fix_decoder (bool): Whether to fix the decoder. Default: True. - num_mlp (int): Layer number of MLP style layers. Default: 8. - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - input_is_latent (bool): Whether input is latent style. Default: False. - different_w (bool): Whether to use different latent w for different layers. Default: False. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - channel_multiplier=1, - decoder_load_path=None, - fix_decoder=True, - # for stylegan decoder - num_mlp=8, - lr_mlp=0.01, - input_is_latent=False, - different_w=False, - narrow=1, - sft_half=False, - ): - super(GFPGANBilinear, self).__init__() - self.input_is_latent = input_is_latent - self.different_w = different_w - self.num_style_feat = num_style_feat - self.min_size_restriction = 512 - - unet_narrow = narrow * 0.5 # by default, use a half of input channels - channels = { - "4": int(512 * unet_narrow), - "8": int(512 * unet_narrow), - "16": int(512 * unet_narrow), - "32": int(512 * unet_narrow), - "64": int(256 * channel_multiplier * unet_narrow), - "128": int(128 * channel_multiplier * unet_narrow), - "256": int(64 * channel_multiplier * unet_narrow), - "512": int(32 * channel_multiplier * unet_narrow), - "1024": int(16 * channel_multiplier * unet_narrow), - } - - self.log_size = int(math.log(out_size, 2)) - first_out_size = 2 ** (int(math.log(out_size, 2))) - - self.conv_body_first = ConvLayer( - 3, channels[f"{first_out_size}"], 1, bias=True, activate=True - ) - - # downsample - in_channels = channels[f"{first_out_size}"] - self.conv_body_down = nn.ModuleList() - for i in range(self.log_size, 2, -1): - out_channels = channels[f"{2**(i - 1)}"] - self.conv_body_down.append(ResBlock(in_channels, out_channels)) - in_channels = out_channels - - self.final_conv = ConvLayer( - in_channels, channels["4"], 3, bias=True, activate=True - ) - - # upsample - in_channels = channels["4"] - self.conv_body_up = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.conv_body_up.append(ResUpBlock(in_channels, out_channels)) - in_channels = out_channels - - # to RGB - self.toRGB = nn.ModuleList() - for i in range(3, self.log_size + 1): - self.toRGB.append( - EqualConv2d( - channels[f"{2**i}"], - 3, - 1, - stride=1, - padding=0, - bias=True, - bias_init_val=0, - ) - ) - - if different_w: - linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat - else: - linear_out_channel = num_style_feat - - self.final_linear = EqualLinear( - channels["4"] * 4 * 4, - linear_out_channel, - bias=True, - bias_init_val=0, - lr_mul=1, - activation=None, - ) - - # the decoder: stylegan2 generator with SFT modulations - self.stylegan_decoder = StyleGAN2GeneratorBilinearSFT( - out_size=out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - lr_mlp=lr_mlp, - narrow=narrow, - sft_half=sft_half, - ) - - # load pre-trained stylegan2 model if necessary - if decoder_load_path: - self.stylegan_decoder.load_state_dict( - torch.load( - decoder_load_path, map_location=lambda storage, loc: storage - )["params_ema"] - ) - # fix decoder without updating params - if fix_decoder: - for _, param in self.stylegan_decoder.named_parameters(): - param.requires_grad = False - - # for SFT modulations (scale and shift) - self.condition_scale = nn.ModuleList() - self.condition_shift = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - if sft_half: - sft_out_channels = out_channels - else: - sft_out_channels = out_channels * 2 - self.condition_scale.append( - nn.Sequential( - EqualConv2d( - out_channels, - out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ScaledLeakyReLU(0.2), - EqualConv2d( - out_channels, - sft_out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=1, - ), - ) - ) - self.condition_shift.append( - nn.Sequential( - EqualConv2d( - out_channels, - out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ScaledLeakyReLU(0.2), - EqualConv2d( - out_channels, - sft_out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ) - ) - - def forward(self, x, return_latents=False, return_rgb=True, randomize_noise=True): - """Forward function for GFPGANBilinear. - Args: - x (Tensor): Input images. - return_latents (bool): Whether to return style latents. Default: False. - return_rgb (bool): Whether return intermediate rgb images. Default: True. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - """ - conditions = [] - unet_skips = [] - out_rgbs = [] - - # encoder - feat = self.conv_body_first(x) - for i in range(self.log_size - 2): - feat = self.conv_body_down[i](feat) - unet_skips.insert(0, feat) - - feat = self.final_conv(feat) - - # style code - style_code = self.final_linear(feat.view(feat.size(0), -1)) - if self.different_w: - style_code = style_code.view(style_code.size(0), -1, self.num_style_feat) - - # decode - for i in range(self.log_size - 2): - # add unet skip - feat = feat + unet_skips[i] - # ResUpLayer - feat = self.conv_body_up[i](feat) - # generate scale and shift for SFT layers - scale = self.condition_scale[i](feat) - conditions.append(scale.clone()) - shift = self.condition_shift[i](feat) - conditions.append(shift.clone()) - # generate rgb images - if return_rgb: - out_rgbs.append(self.toRGB[i](feat)) - - # decoder - image, _ = self.stylegan_decoder( - [style_code], - conditions, - return_latents=return_latents, - input_is_latent=self.input_is_latent, - randomize_noise=randomize_noise, - ) - - return image, out_rgbs diff --git a/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_arch.py deleted file mode 100644 index 72d72fc8..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_arch.py +++ /dev/null @@ -1,566 +0,0 @@ -# pylint: skip-file -# type: ignore -import math -import random - -import torch -from torch import nn -from torch.nn import functional as F - -from .fused_act import FusedLeakyReLU -from .stylegan2_arch import ( - ConvLayer, - EqualConv2d, - EqualLinear, - ResBlock, - ScaledLeakyReLU, - StyleGAN2Generator, -) - - -class StyleGAN2GeneratorSFT(StyleGAN2Generator): - """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform). - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - resample_kernel (list[int]): A list indicating the 1D resample kernel magnitude. A cross production will be - applied to extent 1D resample kernel to 2D resample kernel. Default: (1, 3, 3, 1). - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - num_mlp=8, - channel_multiplier=2, - resample_kernel=(1, 3, 3, 1), - lr_mlp=0.01, - narrow=1, - sft_half=False, - ): - super(StyleGAN2GeneratorSFT, self).__init__( - out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - resample_kernel=resample_kernel, - lr_mlp=lr_mlp, - narrow=narrow, - ) - self.sft_half = sft_half - - def forward( - self, - styles, - conditions, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2GeneratorSFT. - Args: - styles (list[Tensor]): Sample codes of styles. - conditions (list[Tensor]): SFT conditions to generators. - input_is_latent (bool): Whether input is latent style. Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - truncation (float): The truncation ratio. Default: 1. - truncation_latent (Tensor | None): The truncation latent tensor. Default: None. - inject_index (int | None): The injection index for mixing noise. Default: None. - return_latents (bool): Whether to return style latents. Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latents with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - - # the conditions may have fewer levels - if i < len(conditions): - # SFT part to combine the conditions - if self.sft_half: # only apply SFT to half of the channels - out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1) - out_sft = out_sft * conditions[i - 1] + conditions[i] - out = torch.cat([out_same, out_sft], dim=1) - else: # apply SFT to all the channels - out = out * conditions[i - 1] + conditions[i] - - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) # feature back to the rgb space - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None - - -class ConvUpLayer(nn.Module): - """Convolutional upsampling layer. It uses bilinear upsampler + Conv. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - stride (int): Stride of the convolution. Default: 1 - padding (int): Zero-padding added to both sides of the input. Default: 0. - bias (bool): If ``True``, adds a learnable bias to the output. Default: ``True``. - bias_init_val (float): Bias initialized value. Default: 0. - activate (bool): Whether use activateion. Default: True. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - bias=True, - bias_init_val=0, - activate=True, - ): - super(ConvUpLayer, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - # self.scale is used to scale the convolution weights, which is related to the common initializations. - self.scale = 1 / math.sqrt(in_channels * kernel_size**2) - - self.weight = nn.Parameter( - torch.randn(out_channels, in_channels, kernel_size, kernel_size) - ) - - if bias and not activate: - self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val)) - else: - self.register_parameter("bias", None) - - # activation - if activate: - if bias: - self.activation = FusedLeakyReLU(out_channels) - else: - self.activation = ScaledLeakyReLU(0.2) - else: - self.activation = None - - def forward(self, x): - # bilinear upsample - out = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False) - # conv - out = F.conv2d( - out, - self.weight * self.scale, - bias=self.bias, - stride=self.stride, - padding=self.padding, - ) - # activation - if self.activation is not None: - out = self.activation(out) - return out - - -class ResUpBlock(nn.Module): - """Residual block with upsampling. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - """ - - def __init__(self, in_channels, out_channels): - super(ResUpBlock, self).__init__() - - self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True) - self.conv2 = ConvUpLayer( - in_channels, out_channels, 3, stride=1, padding=1, bias=True, activate=True - ) - self.skip = ConvUpLayer( - in_channels, out_channels, 1, bias=False, activate=False - ) - - def forward(self, x): - out = self.conv1(x) - out = self.conv2(out) - skip = self.skip(x) - out = (out + skip) / math.sqrt(2) - return out - - -class GFPGANv1(nn.Module): - """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT. - Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - resample_kernel (list[int]): A list indicating the 1D resample kernel magnitude. A cross production will be - applied to extent 1D resample kernel to 2D resample kernel. Default: (1, 3, 3, 1). - decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None. - fix_decoder (bool): Whether to fix the decoder. Default: True. - num_mlp (int): Layer number of MLP style layers. Default: 8. - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - input_is_latent (bool): Whether input is latent style. Default: False. - different_w (bool): Whether to use different latent w for different layers. Default: False. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - channel_multiplier=1, - resample_kernel=(1, 3, 3, 1), - decoder_load_path=None, - fix_decoder=True, - # for stylegan decoder - num_mlp=8, - lr_mlp=0.01, - input_is_latent=False, - different_w=False, - narrow=1, - sft_half=False, - ): - super(GFPGANv1, self).__init__() - self.input_is_latent = input_is_latent - self.different_w = different_w - self.num_style_feat = num_style_feat - - unet_narrow = narrow * 0.5 # by default, use a half of input channels - channels = { - "4": int(512 * unet_narrow), - "8": int(512 * unet_narrow), - "16": int(512 * unet_narrow), - "32": int(512 * unet_narrow), - "64": int(256 * channel_multiplier * unet_narrow), - "128": int(128 * channel_multiplier * unet_narrow), - "256": int(64 * channel_multiplier * unet_narrow), - "512": int(32 * channel_multiplier * unet_narrow), - "1024": int(16 * channel_multiplier * unet_narrow), - } - - self.log_size = int(math.log(out_size, 2)) - first_out_size = 2 ** (int(math.log(out_size, 2))) - - self.conv_body_first = ConvLayer( - 3, channels[f"{first_out_size}"], 1, bias=True, activate=True - ) - - # downsample - in_channels = channels[f"{first_out_size}"] - self.conv_body_down = nn.ModuleList() - for i in range(self.log_size, 2, -1): - out_channels = channels[f"{2**(i - 1)}"] - self.conv_body_down.append( - ResBlock(in_channels, out_channels, resample_kernel) - ) - in_channels = out_channels - - self.final_conv = ConvLayer( - in_channels, channels["4"], 3, bias=True, activate=True - ) - - # upsample - in_channels = channels["4"] - self.conv_body_up = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.conv_body_up.append(ResUpBlock(in_channels, out_channels)) - in_channels = out_channels - - # to RGB - self.toRGB = nn.ModuleList() - for i in range(3, self.log_size + 1): - self.toRGB.append( - EqualConv2d( - channels[f"{2**i}"], - 3, - 1, - stride=1, - padding=0, - bias=True, - bias_init_val=0, - ) - ) - - if different_w: - linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat - else: - linear_out_channel = num_style_feat - - self.final_linear = EqualLinear( - channels["4"] * 4 * 4, - linear_out_channel, - bias=True, - bias_init_val=0, - lr_mul=1, - activation=None, - ) - - # the decoder: stylegan2 generator with SFT modulations - self.stylegan_decoder = StyleGAN2GeneratorSFT( - out_size=out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - resample_kernel=resample_kernel, - lr_mlp=lr_mlp, - narrow=narrow, - sft_half=sft_half, - ) - - # load pre-trained stylegan2 model if necessary - if decoder_load_path: - self.stylegan_decoder.load_state_dict( - torch.load( - decoder_load_path, map_location=lambda storage, loc: storage - )["params_ema"] - ) - # fix decoder without updating params - if fix_decoder: - for _, param in self.stylegan_decoder.named_parameters(): - param.requires_grad = False - - # for SFT modulations (scale and shift) - self.condition_scale = nn.ModuleList() - self.condition_shift = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - if sft_half: - sft_out_channels = out_channels - else: - sft_out_channels = out_channels * 2 - self.condition_scale.append( - nn.Sequential( - EqualConv2d( - out_channels, - out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ScaledLeakyReLU(0.2), - EqualConv2d( - out_channels, - sft_out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=1, - ), - ) - ) - self.condition_shift.append( - nn.Sequential( - EqualConv2d( - out_channels, - out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ScaledLeakyReLU(0.2), - EqualConv2d( - out_channels, - sft_out_channels, - 3, - stride=1, - padding=1, - bias=True, - bias_init_val=0, - ), - ) - ) - - def forward( - self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs - ): - """Forward function for GFPGANv1. - Args: - x (Tensor): Input images. - return_latents (bool): Whether to return style latents. Default: False. - return_rgb (bool): Whether return intermediate rgb images. Default: True. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - """ - conditions = [] - unet_skips = [] - out_rgbs = [] - - # encoder - feat = self.conv_body_first(x) - for i in range(self.log_size - 2): - feat = self.conv_body_down[i](feat) - unet_skips.insert(0, feat) - - feat = self.final_conv(feat) - - # style code - style_code = self.final_linear(feat.view(feat.size(0), -1)) - if self.different_w: - style_code = style_code.view(style_code.size(0), -1, self.num_style_feat) - - # decode - for i in range(self.log_size - 2): - # add unet skip - feat = feat + unet_skips[i] - # ResUpLayer - feat = self.conv_body_up[i](feat) - # generate scale and shift for SFT layers - scale = self.condition_scale[i](feat) - conditions.append(scale.clone()) - shift = self.condition_shift[i](feat) - conditions.append(shift.clone()) - # generate rgb images - if return_rgb: - out_rgbs.append(self.toRGB[i](feat)) - - # decoder - image, _ = self.stylegan_decoder( - [style_code], - conditions, - return_latents=return_latents, - input_is_latent=self.input_is_latent, - randomize_noise=randomize_noise, - ) - - return image, out_rgbs - - -class FacialComponentDiscriminator(nn.Module): - """Facial component (eyes, mouth, noise) discriminator used in GFPGAN.""" - - def __init__(self): - super(FacialComponentDiscriminator, self).__init__() - # It now uses a VGG-style architectrue with fixed model size - self.conv1 = ConvLayer( - 3, - 64, - 3, - downsample=False, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ) - self.conv2 = ConvLayer( - 64, - 128, - 3, - downsample=True, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ) - self.conv3 = ConvLayer( - 128, - 128, - 3, - downsample=False, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ) - self.conv4 = ConvLayer( - 128, - 256, - 3, - downsample=True, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ) - self.conv5 = ConvLayer( - 256, - 256, - 3, - downsample=False, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ) - self.final_conv = ConvLayer(256, 1, 3, bias=True, activate=False) - - def forward(self, x, return_feats=False, **kwargs): - """Forward function for FacialComponentDiscriminator. - Args: - x (Tensor): Input images. - return_feats (bool): Whether to return intermediate features. Default: False. - """ - feat = self.conv1(x) - feat = self.conv3(self.conv2(feat)) - rlt_feats = [] - if return_feats: - rlt_feats.append(feat.clone()) - feat = self.conv5(self.conv4(feat)) - if return_feats: - rlt_feats.append(feat.clone()) - out = self.final_conv(feat) - - if return_feats: - return out, rlt_feats - else: - return out, None diff --git a/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_clean_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_clean_arch.py deleted file mode 100644 index 16470d63..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/gfpganv1_clean_arch.py +++ /dev/null @@ -1,370 +0,0 @@ -# pylint: skip-file -# type: ignore -import math -import random - -import torch -from torch import nn -from torch.nn import functional as F - -from .stylegan2_clean_arch import StyleGAN2GeneratorClean - - -class StyleGAN2GeneratorCSFT(StyleGAN2GeneratorClean): - """StyleGAN2 Generator with SFT modulation (Spatial Feature Transform). - It is the clean version without custom compiled CUDA extensions used in StyleGAN2. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - num_mlp=8, - channel_multiplier=2, - narrow=1, - sft_half=False, - ): - super(StyleGAN2GeneratorCSFT, self).__init__( - out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - narrow=narrow, - ) - self.sft_half = sft_half - - def forward( - self, - styles, - conditions, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2GeneratorCSFT. - Args: - styles (list[Tensor]): Sample codes of styles. - conditions (list[Tensor]): SFT conditions to generators. - input_is_latent (bool): Whether input is latent style. Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - truncation (float): The truncation ratio. Default: 1. - truncation_latent (Tensor | None): The truncation latent tensor. Default: None. - inject_index (int | None): The injection index for mixing noise. Default: None. - return_latents (bool): Whether to return style latents. Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latents with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - - # the conditions may have fewer levels - if i < len(conditions): - # SFT part to combine the conditions - if self.sft_half: # only apply SFT to half of the channels - out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1) - out_sft = out_sft * conditions[i - 1] + conditions[i] - out = torch.cat([out_same, out_sft], dim=1) - else: # apply SFT to all the channels - out = out * conditions[i - 1] + conditions[i] - - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) # feature back to the rgb space - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None - - -class ResBlock(nn.Module): - """Residual block with bilinear upsampling/downsampling. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - mode (str): Upsampling/downsampling mode. Options: down | up. Default: down. - """ - - def __init__(self, in_channels, out_channels, mode="down"): - super(ResBlock, self).__init__() - - self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1) - self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1) - self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False) - if mode == "down": - self.scale_factor = 0.5 - elif mode == "up": - self.scale_factor = 2 - - def forward(self, x): - out = F.leaky_relu_(self.conv1(x), negative_slope=0.2) - # upsample/downsample - out = F.interpolate( - out, scale_factor=self.scale_factor, mode="bilinear", align_corners=False - ) - out = F.leaky_relu_(self.conv2(out), negative_slope=0.2) - # skip - x = F.interpolate( - x, scale_factor=self.scale_factor, mode="bilinear", align_corners=False - ) - skip = self.skip(x) - out = out + skip - return out - - -class GFPGANv1Clean(nn.Module): - """The GFPGAN architecture: Unet + StyleGAN2 decoder with SFT. - It is the clean version without custom compiled CUDA extensions used in StyleGAN2. - Ref: GFP-GAN: Towards Real-World Blind Face Restoration with Generative Facial Prior. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - decoder_load_path (str): The path to the pre-trained decoder model (usually, the StyleGAN2). Default: None. - fix_decoder (bool): Whether to fix the decoder. Default: True. - num_mlp (int): Layer number of MLP style layers. Default: 8. - input_is_latent (bool): Whether input is latent style. Default: False. - different_w (bool): Whether to use different latent w for different layers. Default: False. - narrow (float): The narrow ratio for channels. Default: 1. - sft_half (bool): Whether to apply SFT on half of the input channels. Default: False. - """ - - def __init__( - self, - state_dict, - ): - super(GFPGANv1Clean, self).__init__() - - out_size = 512 - num_style_feat = 512 - channel_multiplier = 2 - decoder_load_path = None - fix_decoder = False - num_mlp = 8 - input_is_latent = True - different_w = True - narrow = 1 - sft_half = True - - self.model_arch = "GFPGAN" - self.sub_type = "Face SR" - self.scale = 8 - self.in_nc = 3 - self.out_nc = 3 - self.state = state_dict - - self.supports_fp16 = False - self.supports_bf16 = True - self.min_size_restriction = 512 - - self.input_is_latent = input_is_latent - self.different_w = different_w - self.num_style_feat = num_style_feat - - unet_narrow = narrow * 0.5 # by default, use a half of input channels - channels = { - "4": int(512 * unet_narrow), - "8": int(512 * unet_narrow), - "16": int(512 * unet_narrow), - "32": int(512 * unet_narrow), - "64": int(256 * channel_multiplier * unet_narrow), - "128": int(128 * channel_multiplier * unet_narrow), - "256": int(64 * channel_multiplier * unet_narrow), - "512": int(32 * channel_multiplier * unet_narrow), - "1024": int(16 * channel_multiplier * unet_narrow), - } - - self.log_size = int(math.log(out_size, 2)) - first_out_size = 2 ** (int(math.log(out_size, 2))) - - self.conv_body_first = nn.Conv2d(3, channels[f"{first_out_size}"], 1) - - # downsample - in_channels = channels[f"{first_out_size}"] - self.conv_body_down = nn.ModuleList() - for i in range(self.log_size, 2, -1): - out_channels = channels[f"{2**(i - 1)}"] - self.conv_body_down.append(ResBlock(in_channels, out_channels, mode="down")) - in_channels = out_channels - - self.final_conv = nn.Conv2d(in_channels, channels["4"], 3, 1, 1) - - # upsample - in_channels = channels["4"] - self.conv_body_up = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.conv_body_up.append(ResBlock(in_channels, out_channels, mode="up")) - in_channels = out_channels - - # to RGB - self.toRGB = nn.ModuleList() - for i in range(3, self.log_size + 1): - self.toRGB.append(nn.Conv2d(channels[f"{2**i}"], 3, 1)) - - if different_w: - linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat - else: - linear_out_channel = num_style_feat - - self.final_linear = nn.Linear(channels["4"] * 4 * 4, linear_out_channel) - - # the decoder: stylegan2 generator with SFT modulations - self.stylegan_decoder = StyleGAN2GeneratorCSFT( - out_size=out_size, - num_style_feat=num_style_feat, - num_mlp=num_mlp, - channel_multiplier=channel_multiplier, - narrow=narrow, - sft_half=sft_half, - ) - - # load pre-trained stylegan2 model if necessary - if decoder_load_path: - self.stylegan_decoder.load_state_dict( - torch.load( - decoder_load_path, map_location=lambda storage, loc: storage - )["params_ema"] - ) - # fix decoder without updating params - if fix_decoder: - for _, param in self.stylegan_decoder.named_parameters(): - param.requires_grad = False - - # for SFT modulations (scale and shift) - self.condition_scale = nn.ModuleList() - self.condition_shift = nn.ModuleList() - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - if sft_half: - sft_out_channels = out_channels - else: - sft_out_channels = out_channels * 2 - self.condition_scale.append( - nn.Sequential( - nn.Conv2d(out_channels, out_channels, 3, 1, 1), - nn.LeakyReLU(0.2, True), - nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1), - ) - ) - self.condition_shift.append( - nn.Sequential( - nn.Conv2d(out_channels, out_channels, 3, 1, 1), - nn.LeakyReLU(0.2, True), - nn.Conv2d(out_channels, sft_out_channels, 3, 1, 1), - ) - ) - self.load_state_dict(state_dict) - - def forward( - self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs - ): - """Forward function for GFPGANv1Clean. - Args: - x (Tensor): Input images. - return_latents (bool): Whether to return style latents. Default: False. - return_rgb (bool): Whether return intermediate rgb images. Default: True. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - """ - conditions = [] - unet_skips = [] - out_rgbs = [] - - # encoder - feat = F.leaky_relu_(self.conv_body_first(x), negative_slope=0.2) - for i in range(self.log_size - 2): - feat = self.conv_body_down[i](feat) - unet_skips.insert(0, feat) - feat = F.leaky_relu_(self.final_conv(feat), negative_slope=0.2) - - # style code - style_code = self.final_linear(feat.view(feat.size(0), -1)) - if self.different_w: - style_code = style_code.view(style_code.size(0), -1, self.num_style_feat) - - # decode - for i in range(self.log_size - 2): - # add unet skip - feat = feat + unet_skips[i] - # ResUpLayer - feat = self.conv_body_up[i](feat) - # generate scale and shift for SFT layers - scale = self.condition_scale[i](feat) - conditions.append(scale.clone()) - shift = self.condition_shift[i](feat) - conditions.append(shift.clone()) - # generate rgb images - if return_rgb: - out_rgbs.append(self.toRGB[i](feat)) - - # decoder - image, _ = self.stylegan_decoder( - [style_code], - conditions, - return_latents=return_latents, - input_is_latent=self.input_is_latent, - randomize_noise=randomize_noise, - ) - - return image, out_rgbs diff --git a/backend/comfy_nodes/chainner_models/architecture/face/restoreformer_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/restoreformer_arch.py deleted file mode 100644 index 44922602..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/restoreformer_arch.py +++ /dev/null @@ -1,776 +0,0 @@ -# pylint: skip-file -# type: ignore -"""Modified from https://github.com/wzhouxiff/RestoreFormer -""" -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class VectorQuantizer(nn.Module): - """ - see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py - ____________________________________________ - Discretization bottleneck part of the VQ-VAE. - Inputs: - - n_e : number of embeddings - - e_dim : dimension of embedding - - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2 - _____________________________________________ - """ - - def __init__(self, n_e, e_dim, beta): - super(VectorQuantizer, self).__init__() - self.n_e = n_e - self.e_dim = e_dim - self.beta = beta - - self.embedding = nn.Embedding(self.n_e, self.e_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - def forward(self, z): - """ - Inputs the output of the encoder network z and maps it to a discrete - one-hot vector that is the index of the closest embedding vector e_j - z (continuous) -> z_q (discrete) - z.shape = (batch, channel, height, width) - quantization pipeline: - 1. get encoder input (B,C,H,W) - 2. flatten input to (B*H*W,C) - """ - # reshape z -> (batch, height, width, channel) and flatten - z = z.permute(0, 2, 3, 1).contiguous() - z_flattened = z.view(-1, self.e_dim) - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - - d = ( - torch.sum(z_flattened**2, dim=1, keepdim=True) - + torch.sum(self.embedding.weight**2, dim=1) - - 2 * torch.matmul(z_flattened, self.embedding.weight.t()) - ) - - # could possible replace this here - # #\start... - # find closest encodings - - min_value, min_encoding_indices = torch.min(d, dim=1) - - min_encoding_indices = min_encoding_indices.unsqueeze(1) - - min_encodings = torch.zeros(min_encoding_indices.shape[0], self.n_e).to(z) - min_encodings.scatter_(1, min_encoding_indices, 1) - - # dtype min encodings: torch.float32 - # min_encodings shape: torch.Size([2048, 512]) - # min_encoding_indices.shape: torch.Size([2048, 1]) - - # get quantized latent vectors - z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape) - # .........\end - - # with: - # .........\start - # min_encoding_indices = torch.argmin(d, dim=1) - # z_q = self.embedding(min_encoding_indices) - # ......\end......... (TODO) - - # compute loss for embedding - loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean( - (z_q - z.detach()) ** 2 - ) - - # preserve gradients - z_q = z + (z_q - z).detach() - - # perplexity - - e_mean = torch.mean(min_encodings, dim=0) - perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10))) - - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices, d) - - def get_codebook_entry(self, indices, shape): - # shape specifying (batch, height, width, channel) - # TODO: check for more easy handling with nn.Embedding - min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices) - min_encodings.scatter_(1, indices[:, None], 1) - - # get quantized latent vectors - z_q = torch.matmul(min_encodings.float(), self.embedding.weight) - - 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 - - -# pytorch_diffusion + derived encoder decoder -def nonlinearity(x): - # swish - return x * torch.sigmoid(x) - - -def Normalize(in_channels): - return torch.nn.GroupNorm( - num_groups=32, num_channels=in_channels, eps=1e-6, affine=True - ) - - -class Upsample(nn.Module): - def __init__(self, in_channels, with_conv): - super().__init__() - self.with_conv = with_conv - if self.with_conv: - self.conv = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1 - ) - - def forward(self, x): - x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") - if self.with_conv: - x = self.conv(x) - return x - - -class Downsample(nn.Module): - def __init__(self, in_channels, with_conv): - super().__init__() - self.with_conv = with_conv - if self.with_conv: - # no asymmetric padding in torch conv, must do it ourselves - self.conv = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=2, padding=0 - ) - - def forward(self, x): - if self.with_conv: - pad = (0, 1, 0, 1) - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) - else: - x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) - return x - - -class ResnetBlock(nn.Module): - def __init__( - self, - *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout, - temb_channels=512 - ): - 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.norm1 = Normalize(in_channels) - self.conv1 = torch.nn.Conv2d( - in_channels, out_channels, kernel_size=3, stride=1, padding=1 - ) - if temb_channels > 0: - self.temb_proj = torch.nn.Linear(temb_channels, out_channels) - self.norm2 = Normalize(out_channels) - self.dropout = torch.nn.Dropout(dropout) - self.conv2 = torch.nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1 - ) - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = torch.nn.Conv2d( - in_channels, out_channels, kernel_size=3, stride=1, padding=1 - ) - else: - self.nin_shortcut = torch.nn.Conv2d( - in_channels, out_channels, kernel_size=1, stride=1, padding=0 - ) - - def forward(self, x, temb): - h = x - h = self.norm1(h) - h = nonlinearity(h) - h = self.conv1(h) - - if temb is not None: - h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None] - - h = self.norm2(h) - h = nonlinearity(h) - h = self.dropout(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - x = self.conv_shortcut(x) - else: - x = self.nin_shortcut(x) - - return x + h - - -class MultiHeadAttnBlock(nn.Module): - def __init__(self, in_channels, head_size=1): - super().__init__() - self.in_channels = in_channels - self.head_size = head_size - self.att_size = in_channels // head_size - assert ( - in_channels % head_size == 0 - ), "The size of head should be divided by the number of channels." - - self.norm1 = Normalize(in_channels) - self.norm2 = Normalize(in_channels) - - self.q = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.k = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.v = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.proj_out = torch.nn.Conv2d( - in_channels, in_channels, kernel_size=1, stride=1, padding=0 - ) - self.num = 0 - - def forward(self, x, y=None): - h_ = x - h_ = self.norm1(h_) - if y is None: - y = h_ - else: - y = self.norm2(y) - - q = self.q(y) - k = self.k(h_) - v = self.v(h_) - - # compute attention - b, c, h, w = q.shape - q = q.reshape(b, self.head_size, self.att_size, h * w) - q = q.permute(0, 3, 1, 2) # b, hw, head, att - - k = k.reshape(b, self.head_size, self.att_size, h * w) - k = k.permute(0, 3, 1, 2) - - v = v.reshape(b, self.head_size, self.att_size, h * w) - v = v.permute(0, 3, 1, 2) - - q = q.transpose(1, 2) - v = v.transpose(1, 2) - k = k.transpose(1, 2).transpose(2, 3) - - scale = int(self.att_size) ** (-0.5) - q.mul_(scale) - w_ = torch.matmul(q, k) - w_ = F.softmax(w_, dim=3) - - w_ = w_.matmul(v) - - w_ = w_.transpose(1, 2).contiguous() # [b, h*w, head, att] - w_ = w_.view(b, h, w, -1) - w_ = w_.permute(0, 3, 1, 2) - - w_ = self.proj_out(w_) - - return x + w_ - - -class MultiHeadEncoder(nn.Module): - def __init__( - self, - ch, - out_ch, - ch_mult=(1, 2, 4, 8), - num_res_blocks=2, - attn_resolutions=(16,), - dropout=0.0, - resamp_with_conv=True, - in_channels=3, - resolution=512, - z_channels=256, - double_z=True, - enable_mid=True, - head_size=1, - **ignore_kwargs - ): - super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.enable_mid = enable_mid - - # downsampling - self.conv_in = torch.nn.Conv2d( - in_channels, self.ch, kernel_size=3, stride=1, padding=1 - ) - - curr_res = resolution - in_ch_mult = (1,) + tuple(ch_mult) - self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch * in_ch_mult[i_level] - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append( - ResnetBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - ) - ) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(MultiHeadAttnBlock(block_in, head_size)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions - 1: - down.downsample = Downsample(block_in, resamp_with_conv) - curr_res = curr_res // 2 - self.down.append(down) - - # middle - if self.enable_mid: - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size) - self.mid.block_2 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - - # end - self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d( - block_in, - 2 * z_channels if double_z else z_channels, - kernel_size=3, - stride=1, - padding=1, - ) - - def forward(self, x): - hs = {} - # timestep embedding - temb = None - - # downsampling - h = self.conv_in(x) - hs["in"] = h - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](h, temb) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - - if i_level != self.num_resolutions - 1: - # hs.append(h) - hs["block_" + str(i_level)] = h - h = self.down[i_level].downsample(h) - - # middle - # h = hs[-1] - if self.enable_mid: - h = self.mid.block_1(h, temb) - hs["block_" + str(i_level) + "_atten"] = h - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb) - hs["mid_atten"] = h - - # end - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - # hs.append(h) - hs["out"] = h - - return hs - - -class MultiHeadDecoder(nn.Module): - def __init__( - self, - ch, - out_ch, - ch_mult=(1, 2, 4, 8), - num_res_blocks=2, - attn_resolutions=(16,), - dropout=0.0, - resamp_with_conv=True, - in_channels=3, - resolution=512, - z_channels=256, - give_pre_end=False, - enable_mid=True, - head_size=1, - **ignorekwargs - ): - super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.give_pre_end = give_pre_end - self.enable_mid = enable_mid - - # compute in_ch_mult, block_in and curr_res at lowest res - block_in = ch * ch_mult[self.num_resolutions - 1] - curr_res = resolution // 2 ** (self.num_resolutions - 1) - self.z_shape = (1, z_channels, curr_res, curr_res) - print( - "Working with z of shape {} = {} dimensions.".format( - self.z_shape, np.prod(self.z_shape) - ) - ) - - # z to block_in - self.conv_in = torch.nn.Conv2d( - z_channels, block_in, kernel_size=3, stride=1, padding=1 - ) - - # middle - if self.enable_mid: - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size) - self.mid.block_2 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - - # upsampling - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - block.append( - ResnetBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - ) - ) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(MultiHeadAttnBlock(block_in, head_size)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv) - curr_res = curr_res * 2 - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d( - block_in, out_ch, kernel_size=3, stride=1, padding=1 - ) - - def forward(self, z): - # assert z.shape[1:] == self.z_shape[1:] - self.last_z_shape = z.shape - - # timestep embedding - temb = None - - # z to block_in - h = self.conv_in(z) - - # middle - if self.enable_mid: - h = self.mid.block_1(h, temb) - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = self.up[i_level].block[i_block](h, temb) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - if i_level != 0: - h = self.up[i_level].upsample(h) - - # end - if self.give_pre_end: - return h - - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - return h - - -class MultiHeadDecoderTransformer(nn.Module): - def __init__( - self, - ch, - out_ch, - ch_mult=(1, 2, 4, 8), - num_res_blocks=2, - attn_resolutions=(16,), - dropout=0.0, - resamp_with_conv=True, - in_channels=3, - resolution=512, - z_channels=256, - give_pre_end=False, - enable_mid=True, - head_size=1, - **ignorekwargs - ): - super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.give_pre_end = give_pre_end - self.enable_mid = enable_mid - - # compute in_ch_mult, block_in and curr_res at lowest res - block_in = ch * ch_mult[self.num_resolutions - 1] - curr_res = resolution // 2 ** (self.num_resolutions - 1) - self.z_shape = (1, z_channels, curr_res, curr_res) - print( - "Working with z of shape {} = {} dimensions.".format( - self.z_shape, np.prod(self.z_shape) - ) - ) - - # z to block_in - self.conv_in = torch.nn.Conv2d( - z_channels, block_in, kernel_size=3, stride=1, padding=1 - ) - - # middle - if self.enable_mid: - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - self.mid.attn_1 = MultiHeadAttnBlock(block_in, head_size) - self.mid.block_2 = ResnetBlock( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - ) - - # upsampling - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - block.append( - ResnetBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - ) - ) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(MultiHeadAttnBlock(block_in, head_size)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv) - curr_res = curr_res * 2 - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d( - block_in, out_ch, kernel_size=3, stride=1, padding=1 - ) - - def forward(self, z, hs): - # assert z.shape[1:] == self.z_shape[1:] - # self.last_z_shape = z.shape - - # timestep embedding - temb = None - - # z to block_in - h = self.conv_in(z) - - # middle - if self.enable_mid: - h = self.mid.block_1(h, temb) - h = self.mid.attn_1(h, hs["mid_atten"]) - h = self.mid.block_2(h, temb) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = self.up[i_level].block[i_block](h, temb) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block]( - h, hs["block_" + str(i_level) + "_atten"] - ) - # hfeature = h.clone() - if i_level != 0: - h = self.up[i_level].upsample(h) - - # end - if self.give_pre_end: - return h - - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - return h - - -class RestoreFormer(nn.Module): - def __init__( - self, - state_dict, - ): - super(RestoreFormer, self).__init__() - - n_embed = 1024 - embed_dim = 256 - ch = 64 - out_ch = 3 - ch_mult = (1, 2, 2, 4, 4, 8) - num_res_blocks = 2 - attn_resolutions = (16,) - dropout = 0.0 - in_channels = 3 - resolution = 512 - z_channels = 256 - double_z = False - enable_mid = True - fix_decoder = False - fix_codebook = True - fix_encoder = False - head_size = 8 - - self.model_arch = "RestoreFormer" - self.sub_type = "Face SR" - self.scale = 8 - self.in_nc = 3 - self.out_nc = out_ch - self.state = state_dict - - self.supports_fp16 = False - self.supports_bf16 = True - self.min_size_restriction = 16 - - self.encoder = MultiHeadEncoder( - ch=ch, - out_ch=out_ch, - ch_mult=ch_mult, - num_res_blocks=num_res_blocks, - attn_resolutions=attn_resolutions, - dropout=dropout, - in_channels=in_channels, - resolution=resolution, - z_channels=z_channels, - double_z=double_z, - enable_mid=enable_mid, - head_size=head_size, - ) - self.decoder = MultiHeadDecoderTransformer( - ch=ch, - out_ch=out_ch, - ch_mult=ch_mult, - num_res_blocks=num_res_blocks, - attn_resolutions=attn_resolutions, - dropout=dropout, - in_channels=in_channels, - resolution=resolution, - z_channels=z_channels, - enable_mid=enable_mid, - head_size=head_size, - ) - - self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25) - - self.quant_conv = torch.nn.Conv2d(z_channels, embed_dim, 1) - self.post_quant_conv = torch.nn.Conv2d(embed_dim, z_channels, 1) - - if fix_decoder: - for _, param in self.decoder.named_parameters(): - param.requires_grad = False - for _, param in self.post_quant_conv.named_parameters(): - param.requires_grad = False - for _, param in self.quantize.named_parameters(): - param.requires_grad = False - elif fix_codebook: - for _, param in self.quantize.named_parameters(): - param.requires_grad = False - - if fix_encoder: - for _, param in self.encoder.named_parameters(): - param.requires_grad = False - - self.load_state_dict(state_dict) - - def encode(self, x): - hs = self.encoder(x) - h = self.quant_conv(hs["out"]) - quant, emb_loss, info = self.quantize(h) - return quant, emb_loss, info, hs - - def decode(self, quant, hs): - quant = self.post_quant_conv(quant) - dec = self.decoder(quant, hs) - - return dec - - def forward(self, input, **kwargs): - quant, diff, info, hs = self.encode(input) - dec = self.decode(quant, hs) - - return dec, None diff --git a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_arch.py deleted file mode 100644 index 1eb0e9f1..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_arch.py +++ /dev/null @@ -1,865 +0,0 @@ -# pylint: skip-file -# type: ignore -import math -import random - -import torch -from torch import nn -from torch.nn import functional as F - -from .fused_act import FusedLeakyReLU, fused_leaky_relu -from .upfirdn2d import upfirdn2d - - -class NormStyleCode(nn.Module): - def forward(self, x): - """Normalize the style codes. - - Args: - x (Tensor): Style codes with shape (b, c). - - Returns: - Tensor: Normalized tensor. - """ - return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8) - - -def make_resample_kernel(k): - """Make resampling kernel for UpFirDn. - - Args: - k (list[int]): A list indicating the 1D resample kernel magnitude. - - Returns: - Tensor: 2D resampled kernel. - """ - k = torch.tensor(k, dtype=torch.float32) - if k.ndim == 1: - k = k[None, :] * k[:, None] # to 2D kernel, outer product - # normalize - k /= k.sum() - return k - - -class UpFirDnUpsample(nn.Module): - """Upsample, FIR filter, and downsample (upsampole version). - - References: - 1. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.upfirdn.html # noqa: E501 - 2. http://www.ece.northwestern.edu/local-apps/matlabhelp/toolbox/signal/upfirdn.html # noqa: E501 - - Args: - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. - factor (int): Upsampling scale factor. Default: 2. - """ - - def __init__(self, resample_kernel, factor=2): - super(UpFirDnUpsample, self).__init__() - self.kernel = make_resample_kernel(resample_kernel) * (factor**2) - self.factor = factor - - pad = self.kernel.shape[0] - factor - self.pad = ((pad + 1) // 2 + factor - 1, pad // 2) - - def forward(self, x): - out = upfirdn2d(x, self.kernel.type_as(x), up=self.factor, down=1, pad=self.pad) - return out - - def __repr__(self): - return f"{self.__class__.__name__}(factor={self.factor})" - - -class UpFirDnDownsample(nn.Module): - """Upsample, FIR filter, and downsample (downsampole version). - - Args: - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. - factor (int): Downsampling scale factor. Default: 2. - """ - - def __init__(self, resample_kernel, factor=2): - super(UpFirDnDownsample, self).__init__() - self.kernel = make_resample_kernel(resample_kernel) - self.factor = factor - - pad = self.kernel.shape[0] - factor - self.pad = ((pad + 1) // 2, pad // 2) - - def forward(self, x): - out = upfirdn2d(x, self.kernel.type_as(x), up=1, down=self.factor, pad=self.pad) - return out - - def __repr__(self): - return f"{self.__class__.__name__}(factor={self.factor})" - - -class UpFirDnSmooth(nn.Module): - """Upsample, FIR filter, and downsample (smooth version). - - Args: - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. - upsample_factor (int): Upsampling scale factor. Default: 1. - downsample_factor (int): Downsampling scale factor. Default: 1. - kernel_size (int): Kernel size: Default: 1. - """ - - def __init__( - self, resample_kernel, upsample_factor=1, downsample_factor=1, kernel_size=1 - ): - super(UpFirDnSmooth, self).__init__() - self.upsample_factor = upsample_factor - self.downsample_factor = downsample_factor - self.kernel = make_resample_kernel(resample_kernel) - if upsample_factor > 1: - self.kernel = self.kernel * (upsample_factor**2) - - if upsample_factor > 1: - pad = (self.kernel.shape[0] - upsample_factor) - (kernel_size - 1) - self.pad = ((pad + 1) // 2 + upsample_factor - 1, pad // 2 + 1) - elif downsample_factor > 1: - pad = (self.kernel.shape[0] - downsample_factor) + (kernel_size - 1) - self.pad = ((pad + 1) // 2, pad // 2) - else: - raise NotImplementedError - - def forward(self, x): - out = upfirdn2d(x, self.kernel.type_as(x), up=1, down=1, pad=self.pad) - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(upsample_factor={self.upsample_factor}" - f", downsample_factor={self.downsample_factor})" - ) - - -class EqualLinear(nn.Module): - """Equalized Linear as StyleGAN2. - - Args: - in_channels (int): Size of each sample. - out_channels (int): Size of each output sample. - bias (bool): If set to ``False``, the layer will not learn an additive - bias. Default: ``True``. - bias_init_val (float): Bias initialized value. Default: 0. - lr_mul (float): Learning rate multiplier. Default: 1. - activation (None | str): The activation after ``linear`` operation. - Supported: 'fused_lrelu', None. Default: None. - """ - - def __init__( - self, - in_channels, - out_channels, - bias=True, - bias_init_val=0, - lr_mul=1, - activation=None, - ): - super(EqualLinear, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.lr_mul = lr_mul - self.activation = activation - if self.activation not in ["fused_lrelu", None]: - raise ValueError( - f"Wrong activation value in EqualLinear: {activation}" - "Supported ones are: ['fused_lrelu', None]." - ) - self.scale = (1 / math.sqrt(in_channels)) * lr_mul - - self.weight = nn.Parameter(torch.randn(out_channels, in_channels).div_(lr_mul)) - if bias: - self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val)) - else: - self.register_parameter("bias", None) - - def forward(self, x): - if self.bias is None: - bias = None - else: - bias = self.bias * self.lr_mul - if self.activation == "fused_lrelu": - out = F.linear(x, self.weight * self.scale) - out = fused_leaky_relu(out, bias) - else: - out = F.linear(x, self.weight * self.scale, bias=bias) - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, bias={self.bias is not None})" - ) - - -class ModulatedConv2d(nn.Module): - """Modulated Conv2d used in StyleGAN2. - - There is no bias in ModulatedConv2d. - - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether to demodulate in the conv layer. - Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. - Default: None. - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. Default: (1, 3, 3, 1). - eps (float): A value added to the denominator for numerical stability. - Default: 1e-8. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - resample_kernel=(1, 3, 3, 1), - eps=1e-8, - ): - super(ModulatedConv2d, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.demodulate = demodulate - self.sample_mode = sample_mode - self.eps = eps - - if self.sample_mode == "upsample": - self.smooth = UpFirDnSmooth( - resample_kernel, - upsample_factor=2, - downsample_factor=1, - kernel_size=kernel_size, - ) - elif self.sample_mode == "downsample": - self.smooth = UpFirDnSmooth( - resample_kernel, - upsample_factor=1, - downsample_factor=2, - kernel_size=kernel_size, - ) - elif self.sample_mode is None: - pass - else: - raise ValueError( - f"Wrong sample mode {self.sample_mode}, " - "supported ones are ['upsample', 'downsample', None]." - ) - - self.scale = 1 / math.sqrt(in_channels * kernel_size**2) - # modulation inside each modulated conv - self.modulation = EqualLinear( - num_style_feat, - in_channels, - bias=True, - bias_init_val=1, - lr_mul=1, - activation=None, - ) - - self.weight = nn.Parameter( - torch.randn(1, out_channels, in_channels, kernel_size, kernel_size) - ) - self.padding = kernel_size // 2 - - def forward(self, x, style): - """Forward function. - - Args: - x (Tensor): Tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - - Returns: - Tensor: Modulated tensor after convolution. - """ - b, c, h, w = x.shape # c = c_in - # weight modulation - style = self.modulation(style).view(b, 1, c, 1, 1) - # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1) - weight = self.scale * self.weight * style # (b, c_out, c_in, k, k) - - if self.demodulate: - demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps) - weight = weight * demod.view(b, self.out_channels, 1, 1, 1) - - weight = weight.view( - b * self.out_channels, c, self.kernel_size, self.kernel_size - ) - - if self.sample_mode == "upsample": - x = x.view(1, b * c, h, w) - weight = weight.view( - b, self.out_channels, c, self.kernel_size, self.kernel_size - ) - weight = weight.transpose(1, 2).reshape( - b * c, self.out_channels, self.kernel_size, self.kernel_size - ) - out = F.conv_transpose2d(x, weight, padding=0, stride=2, groups=b) - out = out.view(b, self.out_channels, *out.shape[2:4]) - out = self.smooth(out) - elif self.sample_mode == "downsample": - x = self.smooth(x) - x = x.view(1, b * c, *x.shape[2:4]) - out = F.conv2d(x, weight, padding=0, stride=2, groups=b) - out = out.view(b, self.out_channels, *out.shape[2:4]) - else: - x = x.view(1, b * c, h, w) - # weight: (b*c_out, c_in, k, k), groups=b - out = F.conv2d(x, weight, padding=self.padding, groups=b) - out = out.view(b, self.out_channels, *out.shape[2:4]) - - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, " - f"kernel_size={self.kernel_size}, " - f"demodulate={self.demodulate}, sample_mode={self.sample_mode})" - ) - - -class StyleConv(nn.Module): - """Style conv. - - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether demodulate in the conv layer. Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. - Default: None. - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. Default: (1, 3, 3, 1). - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - resample_kernel=(1, 3, 3, 1), - ): - super(StyleConv, self).__init__() - self.modulated_conv = ModulatedConv2d( - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=demodulate, - sample_mode=sample_mode, - resample_kernel=resample_kernel, - ) - self.weight = nn.Parameter(torch.zeros(1)) # for noise injection - self.activate = FusedLeakyReLU(out_channels) - - def forward(self, x, style, noise=None): - # modulate - out = self.modulated_conv(x, style) - # noise injection - if noise is None: - b, _, h, w = out.shape - noise = out.new_empty(b, 1, h, w).normal_() - out = out + self.weight * noise - # activation (with bias) - out = self.activate(out) - return out - - -class ToRGB(nn.Module): - """To RGB from features. - - Args: - in_channels (int): Channel number of input. - num_style_feat (int): Channel number of style features. - upsample (bool): Whether to upsample. Default: True. - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. Default: (1, 3, 3, 1). - """ - - def __init__( - self, in_channels, num_style_feat, upsample=True, resample_kernel=(1, 3, 3, 1) - ): - super(ToRGB, self).__init__() - if upsample: - self.upsample = UpFirDnUpsample(resample_kernel, factor=2) - else: - self.upsample = None - self.modulated_conv = ModulatedConv2d( - in_channels, - 3, - kernel_size=1, - num_style_feat=num_style_feat, - demodulate=False, - sample_mode=None, - ) - self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) - - def forward(self, x, style, skip=None): - """Forward function. - - Args: - x (Tensor): Feature tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - skip (Tensor): Base/skip tensor. Default: None. - - Returns: - Tensor: RGB images. - """ - out = self.modulated_conv(x, style) - out = out + self.bias - if skip is not None: - if self.upsample: - skip = self.upsample(skip) - out = out + skip - return out - - -class ConstantInput(nn.Module): - """Constant input. - - Args: - num_channel (int): Channel number of constant input. - size (int): Spatial size of constant input. - """ - - def __init__(self, num_channel, size): - super(ConstantInput, self).__init__() - self.weight = nn.Parameter(torch.randn(1, num_channel, size, size)) - - def forward(self, batch): - out = self.weight.repeat(batch, 1, 1, 1) - return out - - -class StyleGAN2Generator(nn.Module): - """StyleGAN2 Generator. - - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of - StyleGAN2. Default: 2. - resample_kernel (list[int]): A list indicating the 1D resample kernel - magnitude. A cross production will be applied to extent 1D resample - kernel to 2D resample kernel. Default: (1, 3, 3, 1). - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - narrow (float): Narrow ratio for channels. Default: 1.0. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - num_mlp=8, - channel_multiplier=2, - resample_kernel=(1, 3, 3, 1), - lr_mlp=0.01, - narrow=1, - ): - super(StyleGAN2Generator, self).__init__() - # Style MLP layers - self.num_style_feat = num_style_feat - style_mlp_layers = [NormStyleCode()] - for i in range(num_mlp): - style_mlp_layers.append( - EqualLinear( - num_style_feat, - num_style_feat, - bias=True, - bias_init_val=0, - lr_mul=lr_mlp, - activation="fused_lrelu", - ) - ) - self.style_mlp = nn.Sequential(*style_mlp_layers) - - channels = { - "4": int(512 * narrow), - "8": int(512 * narrow), - "16": int(512 * narrow), - "32": int(512 * narrow), - "64": int(256 * channel_multiplier * narrow), - "128": int(128 * channel_multiplier * narrow), - "256": int(64 * channel_multiplier * narrow), - "512": int(32 * channel_multiplier * narrow), - "1024": int(16 * channel_multiplier * narrow), - } - self.channels = channels - - self.constant_input = ConstantInput(channels["4"], size=4) - self.style_conv1 = StyleConv( - channels["4"], - channels["4"], - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - resample_kernel=resample_kernel, - ) - self.to_rgb1 = ToRGB( - channels["4"], - num_style_feat, - upsample=False, - resample_kernel=resample_kernel, - ) - - self.log_size = int(math.log(out_size, 2)) - self.num_layers = (self.log_size - 2) * 2 + 1 - self.num_latent = self.log_size * 2 - 2 - - self.style_convs = nn.ModuleList() - self.to_rgbs = nn.ModuleList() - self.noises = nn.Module() - - in_channels = channels["4"] - # noise - for layer_idx in range(self.num_layers): - resolution = 2 ** ((layer_idx + 5) // 2) - shape = [1, 1, resolution, resolution] - self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape)) - # style convs and to_rgbs - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.style_convs.append( - StyleConv( - in_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode="upsample", - resample_kernel=resample_kernel, - ) - ) - self.style_convs.append( - StyleConv( - out_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - resample_kernel=resample_kernel, - ) - ) - self.to_rgbs.append( - ToRGB( - out_channels, - num_style_feat, - upsample=True, - resample_kernel=resample_kernel, - ) - ) - in_channels = out_channels - - def make_noise(self): - """Make noise for noise injection.""" - device = self.constant_input.weight.device - noises = [torch.randn(1, 1, 4, 4, device=device)] - - for i in range(3, self.log_size + 1): - for _ in range(2): - noises.append(torch.randn(1, 1, 2**i, 2**i, device=device)) - - return noises - - def get_latent(self, x): - return self.style_mlp(x) - - def mean_latent(self, num_latent): - latent_in = torch.randn( - num_latent, self.num_style_feat, device=self.constant_input.weight.device - ) - latent = self.style_mlp(latent_in).mean(0, keepdim=True) - return latent - - def forward( - self, - styles, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2Generator. - - Args: - styles (list[Tensor]): Sample codes of styles. - input_is_latent (bool): Whether input is latent style. - Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is - False. Default: True. - truncation (float): TODO. Default: 1. - truncation_latent (Tensor | None): TODO. Default: None. - inject_index (int | None): The injection index for mixing noise. - Default: None. - return_latents (bool): Whether to return style latents. - Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latent with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None - - -class ScaledLeakyReLU(nn.Module): - """Scaled LeakyReLU. - - Args: - negative_slope (float): Negative slope. Default: 0.2. - """ - - def __init__(self, negative_slope=0.2): - super(ScaledLeakyReLU, self).__init__() - self.negative_slope = negative_slope - - def forward(self, x): - out = F.leaky_relu(x, negative_slope=self.negative_slope) - return out * math.sqrt(2) - - -class EqualConv2d(nn.Module): - """Equalized Linear as StyleGAN2. - - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - stride (int): Stride of the convolution. Default: 1 - padding (int): Zero-padding added to both sides of the input. - Default: 0. - bias (bool): If ``True``, adds a learnable bias to the output. - Default: ``True``. - bias_init_val (float): Bias initialized value. Default: 0. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - bias=True, - bias_init_val=0, - ): - super(EqualConv2d, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.scale = 1 / math.sqrt(in_channels * kernel_size**2) - - self.weight = nn.Parameter( - torch.randn(out_channels, in_channels, kernel_size, kernel_size) - ) - if bias: - self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val)) - else: - self.register_parameter("bias", None) - - def forward(self, x): - out = F.conv2d( - x, - self.weight * self.scale, - bias=self.bias, - stride=self.stride, - padding=self.padding, - ) - - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, " - f"kernel_size={self.kernel_size}," - f" stride={self.stride}, padding={self.padding}, " - f"bias={self.bias is not None})" - ) - - -class ConvLayer(nn.Sequential): - """Conv Layer used in StyleGAN2 Discriminator. - - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Kernel size. - downsample (bool): Whether downsample by a factor of 2. - Default: False. - resample_kernel (list[int]): A list indicating the 1D resample - kernel magnitude. A cross production will be applied to - extent 1D resample kernel to 2D resample kernel. - Default: (1, 3, 3, 1). - bias (bool): Whether with bias. Default: True. - activate (bool): Whether use activateion. Default: True. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - downsample=False, - resample_kernel=(1, 3, 3, 1), - bias=True, - activate=True, - ): - layers = [] - # downsample - if downsample: - layers.append( - UpFirDnSmooth( - resample_kernel, - upsample_factor=1, - downsample_factor=2, - kernel_size=kernel_size, - ) - ) - stride = 2 - self.padding = 0 - else: - stride = 1 - self.padding = kernel_size // 2 - # conv - layers.append( - EqualConv2d( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=self.padding, - bias=bias and not activate, - ) - ) - # activation - if activate: - if bias: - layers.append(FusedLeakyReLU(out_channels)) - else: - layers.append(ScaledLeakyReLU(0.2)) - - super(ConvLayer, self).__init__(*layers) - - -class ResBlock(nn.Module): - """Residual block used in StyleGAN2 Discriminator. - - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - resample_kernel (list[int]): A list indicating the 1D resample - kernel magnitude. A cross production will be applied to - extent 1D resample kernel to 2D resample kernel. - Default: (1, 3, 3, 1). - """ - - def __init__(self, in_channels, out_channels, resample_kernel=(1, 3, 3, 1)): - super(ResBlock, self).__init__() - - self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True) - self.conv2 = ConvLayer( - in_channels, - out_channels, - 3, - downsample=True, - resample_kernel=resample_kernel, - bias=True, - activate=True, - ) - self.skip = ConvLayer( - in_channels, - out_channels, - 1, - downsample=True, - resample_kernel=resample_kernel, - bias=False, - activate=False, - ) - - def forward(self, x): - out = self.conv1(x) - out = self.conv2(out) - skip = self.skip(x) - out = (out + skip) / math.sqrt(2) - return out diff --git a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_bilinear_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_bilinear_arch.py deleted file mode 100644 index 601f8cc4..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_bilinear_arch.py +++ /dev/null @@ -1,709 +0,0 @@ -# pylint: skip-file -# type: ignore -import math -import random - -import torch -from torch import nn -from torch.nn import functional as F - -from .fused_act import FusedLeakyReLU, fused_leaky_relu - - -class NormStyleCode(nn.Module): - def forward(self, x): - """Normalize the style codes. - Args: - x (Tensor): Style codes with shape (b, c). - Returns: - Tensor: Normalized tensor. - """ - return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8) - - -class EqualLinear(nn.Module): - """Equalized Linear as StyleGAN2. - Args: - in_channels (int): Size of each sample. - out_channels (int): Size of each output sample. - bias (bool): If set to ``False``, the layer will not learn an additive - bias. Default: ``True``. - bias_init_val (float): Bias initialized value. Default: 0. - lr_mul (float): Learning rate multiplier. Default: 1. - activation (None | str): The activation after ``linear`` operation. - Supported: 'fused_lrelu', None. Default: None. - """ - - def __init__( - self, - in_channels, - out_channels, - bias=True, - bias_init_val=0, - lr_mul=1, - activation=None, - ): - super(EqualLinear, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.lr_mul = lr_mul - self.activation = activation - if self.activation not in ["fused_lrelu", None]: - raise ValueError( - f"Wrong activation value in EqualLinear: {activation}" - "Supported ones are: ['fused_lrelu', None]." - ) - self.scale = (1 / math.sqrt(in_channels)) * lr_mul - - self.weight = nn.Parameter(torch.randn(out_channels, in_channels).div_(lr_mul)) - if bias: - self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val)) - else: - self.register_parameter("bias", None) - - def forward(self, x): - if self.bias is None: - bias = None - else: - bias = self.bias * self.lr_mul - if self.activation == "fused_lrelu": - out = F.linear(x, self.weight * self.scale) - out = fused_leaky_relu(out, bias) - else: - out = F.linear(x, self.weight * self.scale, bias=bias) - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, bias={self.bias is not None})" - ) - - -class ModulatedConv2d(nn.Module): - """Modulated Conv2d used in StyleGAN2. - There is no bias in ModulatedConv2d. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether to demodulate in the conv layer. - Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. - Default: None. - eps (float): A value added to the denominator for numerical stability. - Default: 1e-8. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - eps=1e-8, - interpolation_mode="bilinear", - ): - super(ModulatedConv2d, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.demodulate = demodulate - self.sample_mode = sample_mode - self.eps = eps - self.interpolation_mode = interpolation_mode - if self.interpolation_mode == "nearest": - self.align_corners = None - else: - self.align_corners = False - - self.scale = 1 / math.sqrt(in_channels * kernel_size**2) - # modulation inside each modulated conv - self.modulation = EqualLinear( - num_style_feat, - in_channels, - bias=True, - bias_init_val=1, - lr_mul=1, - activation=None, - ) - - self.weight = nn.Parameter( - torch.randn(1, out_channels, in_channels, kernel_size, kernel_size) - ) - self.padding = kernel_size // 2 - - def forward(self, x, style): - """Forward function. - Args: - x (Tensor): Tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - Returns: - Tensor: Modulated tensor after convolution. - """ - b, c, h, w = x.shape # c = c_in - # weight modulation - style = self.modulation(style).view(b, 1, c, 1, 1) - # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1) - weight = self.scale * self.weight * style # (b, c_out, c_in, k, k) - - if self.demodulate: - demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps) - weight = weight * demod.view(b, self.out_channels, 1, 1, 1) - - weight = weight.view( - b * self.out_channels, c, self.kernel_size, self.kernel_size - ) - - if self.sample_mode == "upsample": - x = F.interpolate( - x, - scale_factor=2, - mode=self.interpolation_mode, - align_corners=self.align_corners, - ) - elif self.sample_mode == "downsample": - x = F.interpolate( - x, - scale_factor=0.5, - mode=self.interpolation_mode, - align_corners=self.align_corners, - ) - - b, c, h, w = x.shape - x = x.view(1, b * c, h, w) - # weight: (b*c_out, c_in, k, k), groups=b - out = F.conv2d(x, weight, padding=self.padding, groups=b) - out = out.view(b, self.out_channels, *out.shape[2:4]) - - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, " - f"kernel_size={self.kernel_size}, " - f"demodulate={self.demodulate}, sample_mode={self.sample_mode})" - ) - - -class StyleConv(nn.Module): - """Style conv. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether demodulate in the conv layer. Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. - Default: None. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - interpolation_mode="bilinear", - ): - super(StyleConv, self).__init__() - self.modulated_conv = ModulatedConv2d( - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=demodulate, - sample_mode=sample_mode, - interpolation_mode=interpolation_mode, - ) - self.weight = nn.Parameter(torch.zeros(1)) # for noise injection - self.activate = FusedLeakyReLU(out_channels) - - def forward(self, x, style, noise=None): - # modulate - out = self.modulated_conv(x, style) - # noise injection - if noise is None: - b, _, h, w = out.shape - noise = out.new_empty(b, 1, h, w).normal_() - out = out + self.weight * noise - # activation (with bias) - out = self.activate(out) - return out - - -class ToRGB(nn.Module): - """To RGB from features. - Args: - in_channels (int): Channel number of input. - num_style_feat (int): Channel number of style features. - upsample (bool): Whether to upsample. Default: True. - """ - - def __init__( - self, in_channels, num_style_feat, upsample=True, interpolation_mode="bilinear" - ): - super(ToRGB, self).__init__() - self.upsample = upsample - self.interpolation_mode = interpolation_mode - if self.interpolation_mode == "nearest": - self.align_corners = None - else: - self.align_corners = False - self.modulated_conv = ModulatedConv2d( - in_channels, - 3, - kernel_size=1, - num_style_feat=num_style_feat, - demodulate=False, - sample_mode=None, - interpolation_mode=interpolation_mode, - ) - self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) - - def forward(self, x, style, skip=None): - """Forward function. - Args: - x (Tensor): Feature tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - skip (Tensor): Base/skip tensor. Default: None. - Returns: - Tensor: RGB images. - """ - out = self.modulated_conv(x, style) - out = out + self.bias - if skip is not None: - if self.upsample: - skip = F.interpolate( - skip, - scale_factor=2, - mode=self.interpolation_mode, - align_corners=self.align_corners, - ) - out = out + skip - return out - - -class ConstantInput(nn.Module): - """Constant input. - Args: - num_channel (int): Channel number of constant input. - size (int): Spatial size of constant input. - """ - - def __init__(self, num_channel, size): - super(ConstantInput, self).__init__() - self.weight = nn.Parameter(torch.randn(1, num_channel, size, size)) - - def forward(self, batch): - out = self.weight.repeat(batch, 1, 1, 1) - return out - - -class StyleGAN2GeneratorBilinear(nn.Module): - """StyleGAN2 Generator. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of - StyleGAN2. Default: 2. - lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01. - narrow (float): Narrow ratio for channels. Default: 1.0. - """ - - def __init__( - self, - out_size, - num_style_feat=512, - num_mlp=8, - channel_multiplier=2, - lr_mlp=0.01, - narrow=1, - interpolation_mode="bilinear", - ): - super(StyleGAN2GeneratorBilinear, self).__init__() - # Style MLP layers - self.num_style_feat = num_style_feat - style_mlp_layers = [NormStyleCode()] - for i in range(num_mlp): - style_mlp_layers.append( - EqualLinear( - num_style_feat, - num_style_feat, - bias=True, - bias_init_val=0, - lr_mul=lr_mlp, - activation="fused_lrelu", - ) - ) - self.style_mlp = nn.Sequential(*style_mlp_layers) - - channels = { - "4": int(512 * narrow), - "8": int(512 * narrow), - "16": int(512 * narrow), - "32": int(512 * narrow), - "64": int(256 * channel_multiplier * narrow), - "128": int(128 * channel_multiplier * narrow), - "256": int(64 * channel_multiplier * narrow), - "512": int(32 * channel_multiplier * narrow), - "1024": int(16 * channel_multiplier * narrow), - } - self.channels = channels - - self.constant_input = ConstantInput(channels["4"], size=4) - self.style_conv1 = StyleConv( - channels["4"], - channels["4"], - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - interpolation_mode=interpolation_mode, - ) - self.to_rgb1 = ToRGB( - channels["4"], - num_style_feat, - upsample=False, - interpolation_mode=interpolation_mode, - ) - - self.log_size = int(math.log(out_size, 2)) - self.num_layers = (self.log_size - 2) * 2 + 1 - self.num_latent = self.log_size * 2 - 2 - - self.style_convs = nn.ModuleList() - self.to_rgbs = nn.ModuleList() - self.noises = nn.Module() - - in_channels = channels["4"] - # noise - for layer_idx in range(self.num_layers): - resolution = 2 ** ((layer_idx + 5) // 2) - shape = [1, 1, resolution, resolution] - self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape)) - # style convs and to_rgbs - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.style_convs.append( - StyleConv( - in_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode="upsample", - interpolation_mode=interpolation_mode, - ) - ) - self.style_convs.append( - StyleConv( - out_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - interpolation_mode=interpolation_mode, - ) - ) - self.to_rgbs.append( - ToRGB( - out_channels, - num_style_feat, - upsample=True, - interpolation_mode=interpolation_mode, - ) - ) - in_channels = out_channels - - def make_noise(self): - """Make noise for noise injection.""" - device = self.constant_input.weight.device - noises = [torch.randn(1, 1, 4, 4, device=device)] - - for i in range(3, self.log_size + 1): - for _ in range(2): - noises.append(torch.randn(1, 1, 2**i, 2**i, device=device)) - - return noises - - def get_latent(self, x): - return self.style_mlp(x) - - def mean_latent(self, num_latent): - latent_in = torch.randn( - num_latent, self.num_style_feat, device=self.constant_input.weight.device - ) - latent = self.style_mlp(latent_in).mean(0, keepdim=True) - return latent - - def forward( - self, - styles, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2Generator. - Args: - styles (list[Tensor]): Sample codes of styles. - input_is_latent (bool): Whether input is latent style. - Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is - False. Default: True. - truncation (float): TODO. Default: 1. - truncation_latent (Tensor | None): TODO. Default: None. - inject_index (int | None): The injection index for mixing noise. - Default: None. - return_latents (bool): Whether to return style latents. - Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latent with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None - - -class ScaledLeakyReLU(nn.Module): - """Scaled LeakyReLU. - Args: - negative_slope (float): Negative slope. Default: 0.2. - """ - - def __init__(self, negative_slope=0.2): - super(ScaledLeakyReLU, self).__init__() - self.negative_slope = negative_slope - - def forward(self, x): - out = F.leaky_relu(x, negative_slope=self.negative_slope) - return out * math.sqrt(2) - - -class EqualConv2d(nn.Module): - """Equalized Linear as StyleGAN2. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - stride (int): Stride of the convolution. Default: 1 - padding (int): Zero-padding added to both sides of the input. - Default: 0. - bias (bool): If ``True``, adds a learnable bias to the output. - Default: ``True``. - bias_init_val (float): Bias initialized value. Default: 0. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - bias=True, - bias_init_val=0, - ): - super(EqualConv2d, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.scale = 1 / math.sqrt(in_channels * kernel_size**2) - - self.weight = nn.Parameter( - torch.randn(out_channels, in_channels, kernel_size, kernel_size) - ) - if bias: - self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val)) - else: - self.register_parameter("bias", None) - - def forward(self, x): - out = F.conv2d( - x, - self.weight * self.scale, - bias=self.bias, - stride=self.stride, - padding=self.padding, - ) - - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, " - f"out_channels={self.out_channels}, " - f"kernel_size={self.kernel_size}," - f" stride={self.stride}, padding={self.padding}, " - f"bias={self.bias is not None})" - ) - - -class ConvLayer(nn.Sequential): - """Conv Layer used in StyleGAN2 Discriminator. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Kernel size. - downsample (bool): Whether downsample by a factor of 2. - Default: False. - bias (bool): Whether with bias. Default: True. - activate (bool): Whether use activateion. Default: True. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - downsample=False, - bias=True, - activate=True, - interpolation_mode="bilinear", - ): - layers = [] - self.interpolation_mode = interpolation_mode - # downsample - if downsample: - if self.interpolation_mode == "nearest": - self.align_corners = None - else: - self.align_corners = False - - layers.append( - torch.nn.Upsample( - scale_factor=0.5, - mode=interpolation_mode, - align_corners=self.align_corners, - ) - ) - stride = 1 - self.padding = kernel_size // 2 - # conv - layers.append( - EqualConv2d( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=self.padding, - bias=bias and not activate, - ) - ) - # activation - if activate: - if bias: - layers.append(FusedLeakyReLU(out_channels)) - else: - layers.append(ScaledLeakyReLU(0.2)) - - super(ConvLayer, self).__init__(*layers) - - -class ResBlock(nn.Module): - """Residual block used in StyleGAN2 Discriminator. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - """ - - def __init__(self, in_channels, out_channels, interpolation_mode="bilinear"): - super(ResBlock, self).__init__() - - self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True) - self.conv2 = ConvLayer( - in_channels, - out_channels, - 3, - downsample=True, - interpolation_mode=interpolation_mode, - bias=True, - activate=True, - ) - self.skip = ConvLayer( - in_channels, - out_channels, - 1, - downsample=True, - interpolation_mode=interpolation_mode, - bias=False, - activate=False, - ) - - def forward(self, x): - out = self.conv1(x) - out = self.conv2(out) - skip = self.skip(x) - out = (out + skip) / math.sqrt(2) - return out diff --git a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_clean_arch.py b/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_clean_arch.py deleted file mode 100644 index c48de9af..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/stylegan2_clean_arch.py +++ /dev/null @@ -1,453 +0,0 @@ -# pylint: skip-file -# type: ignore -import math - -import torch -from torch import nn -from torch.nn import functional as F -from torch.nn import init -from torch.nn.modules.batchnorm import _BatchNorm - - -@torch.no_grad() -def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs): - """Initialize network weights. - Args: - module_list (list[nn.Module] | nn.Module): Modules to be initialized. - scale (float): Scale initialized weights, especially for residual - blocks. Default: 1. - bias_fill (float): The value to fill bias. Default: 0 - kwargs (dict): Other arguments for initialization function. - """ - if not isinstance(module_list, list): - module_list = [module_list] - for module in module_list: - for m in module.modules(): - if isinstance(m, nn.Conv2d): - init.kaiming_normal_(m.weight, **kwargs) - m.weight.data *= scale - if m.bias is not None: - m.bias.data.fill_(bias_fill) - elif isinstance(m, nn.Linear): - init.kaiming_normal_(m.weight, **kwargs) - m.weight.data *= scale - if m.bias is not None: - m.bias.data.fill_(bias_fill) - elif isinstance(m, _BatchNorm): - init.constant_(m.weight, 1) - if m.bias is not None: - m.bias.data.fill_(bias_fill) - - -class NormStyleCode(nn.Module): - def forward(self, x): - """Normalize the style codes. - Args: - x (Tensor): Style codes with shape (b, c). - Returns: - Tensor: Normalized tensor. - """ - return x * torch.rsqrt(torch.mean(x**2, dim=1, keepdim=True) + 1e-8) - - -class ModulatedConv2d(nn.Module): - """Modulated Conv2d used in StyleGAN2. - There is no bias in ModulatedConv2d. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether to demodulate in the conv layer. Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. Default: None. - eps (float): A value added to the denominator for numerical stability. Default: 1e-8. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - eps=1e-8, - ): - super(ModulatedConv2d, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.demodulate = demodulate - self.sample_mode = sample_mode - self.eps = eps - - # modulation inside each modulated conv - self.modulation = nn.Linear(num_style_feat, in_channels, bias=True) - # initialization - default_init_weights( - self.modulation, - scale=1, - bias_fill=1, - a=0, - mode="fan_in", - nonlinearity="linear", - ) - - self.weight = nn.Parameter( - torch.randn(1, out_channels, in_channels, kernel_size, kernel_size) - / math.sqrt(in_channels * kernel_size**2) - ) - self.padding = kernel_size // 2 - - def forward(self, x, style): - """Forward function. - Args: - x (Tensor): Tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - Returns: - Tensor: Modulated tensor after convolution. - """ - b, c, h, w = x.shape # c = c_in - # weight modulation - style = self.modulation(style).view(b, 1, c, 1, 1) - # self.weight: (1, c_out, c_in, k, k); style: (b, 1, c, 1, 1) - weight = self.weight * style # (b, c_out, c_in, k, k) - - if self.demodulate: - demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps) - weight = weight * demod.view(b, self.out_channels, 1, 1, 1) - - weight = weight.view( - b * self.out_channels, c, self.kernel_size, self.kernel_size - ) - - # upsample or downsample if necessary - if self.sample_mode == "upsample": - x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False) - elif self.sample_mode == "downsample": - x = F.interpolate(x, scale_factor=0.5, mode="bilinear", align_corners=False) - - b, c, h, w = x.shape - x = x.view(1, b * c, h, w) - # weight: (b*c_out, c_in, k, k), groups=b - out = F.conv2d(x, weight, padding=self.padding, groups=b) - out = out.view(b, self.out_channels, *out.shape[2:4]) - - return out - - def __repr__(self): - return ( - f"{self.__class__.__name__}(in_channels={self.in_channels}, out_channels={self.out_channels}, " - f"kernel_size={self.kernel_size}, demodulate={self.demodulate}, sample_mode={self.sample_mode})" - ) - - -class StyleConv(nn.Module): - """Style conv used in StyleGAN2. - Args: - in_channels (int): Channel number of the input. - out_channels (int): Channel number of the output. - kernel_size (int): Size of the convolving kernel. - num_style_feat (int): Channel number of style features. - demodulate (bool): Whether demodulate in the conv layer. Default: True. - sample_mode (str | None): Indicating 'upsample', 'downsample' or None. Default: None. - """ - - def __init__( - self, - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=True, - sample_mode=None, - ): - super(StyleConv, self).__init__() - self.modulated_conv = ModulatedConv2d( - in_channels, - out_channels, - kernel_size, - num_style_feat, - demodulate=demodulate, - sample_mode=sample_mode, - ) - self.weight = nn.Parameter(torch.zeros(1)) # for noise injection - self.bias = nn.Parameter(torch.zeros(1, out_channels, 1, 1)) - self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) - - def forward(self, x, style, noise=None): - # modulate - out = self.modulated_conv(x, style) * 2**0.5 # for conversion - # noise injection - if noise is None: - b, _, h, w = out.shape - noise = out.new_empty(b, 1, h, w).normal_() - out = out + self.weight * noise - # add bias - out = out + self.bias - # activation - out = self.activate(out) - return out - - -class ToRGB(nn.Module): - """To RGB (image space) from features. - Args: - in_channels (int): Channel number of input. - num_style_feat (int): Channel number of style features. - upsample (bool): Whether to upsample. Default: True. - """ - - def __init__(self, in_channels, num_style_feat, upsample=True): - super(ToRGB, self).__init__() - self.upsample = upsample - self.modulated_conv = ModulatedConv2d( - in_channels, - 3, - kernel_size=1, - num_style_feat=num_style_feat, - demodulate=False, - sample_mode=None, - ) - self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) - - def forward(self, x, style, skip=None): - """Forward function. - Args: - x (Tensor): Feature tensor with shape (b, c, h, w). - style (Tensor): Tensor with shape (b, num_style_feat). - skip (Tensor): Base/skip tensor. Default: None. - Returns: - Tensor: RGB images. - """ - out = self.modulated_conv(x, style) - out = out + self.bias - if skip is not None: - if self.upsample: - skip = F.interpolate( - skip, scale_factor=2, mode="bilinear", align_corners=False - ) - out = out + skip - return out - - -class ConstantInput(nn.Module): - """Constant input. - Args: - num_channel (int): Channel number of constant input. - size (int): Spatial size of constant input. - """ - - def __init__(self, num_channel, size): - super(ConstantInput, self).__init__() - self.weight = nn.Parameter(torch.randn(1, num_channel, size, size)) - - def forward(self, batch): - out = self.weight.repeat(batch, 1, 1, 1) - return out - - -class StyleGAN2GeneratorClean(nn.Module): - """Clean version of StyleGAN2 Generator. - Args: - out_size (int): The spatial size of outputs. - num_style_feat (int): Channel number of style features. Default: 512. - num_mlp (int): Layer number of MLP style layers. Default: 8. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - narrow (float): Narrow ratio for channels. Default: 1.0. - """ - - def __init__( - self, out_size, num_style_feat=512, num_mlp=8, channel_multiplier=2, narrow=1 - ): - super(StyleGAN2GeneratorClean, self).__init__() - # Style MLP layers - self.num_style_feat = num_style_feat - style_mlp_layers = [NormStyleCode()] - for i in range(num_mlp): - style_mlp_layers.extend( - [ - nn.Linear(num_style_feat, num_style_feat, bias=True), - nn.LeakyReLU(negative_slope=0.2, inplace=True), - ] - ) - self.style_mlp = nn.Sequential(*style_mlp_layers) - # initialization - default_init_weights( - self.style_mlp, - scale=1, - bias_fill=0, - a=0.2, - mode="fan_in", - nonlinearity="leaky_relu", - ) - - # channel list - channels = { - "4": int(512 * narrow), - "8": int(512 * narrow), - "16": int(512 * narrow), - "32": int(512 * narrow), - "64": int(256 * channel_multiplier * narrow), - "128": int(128 * channel_multiplier * narrow), - "256": int(64 * channel_multiplier * narrow), - "512": int(32 * channel_multiplier * narrow), - "1024": int(16 * channel_multiplier * narrow), - } - self.channels = channels - - self.constant_input = ConstantInput(channels["4"], size=4) - self.style_conv1 = StyleConv( - channels["4"], - channels["4"], - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - ) - self.to_rgb1 = ToRGB(channels["4"], num_style_feat, upsample=False) - - self.log_size = int(math.log(out_size, 2)) - self.num_layers = (self.log_size - 2) * 2 + 1 - self.num_latent = self.log_size * 2 - 2 - - self.style_convs = nn.ModuleList() - self.to_rgbs = nn.ModuleList() - self.noises = nn.Module() - - in_channels = channels["4"] - # noise - for layer_idx in range(self.num_layers): - resolution = 2 ** ((layer_idx + 5) // 2) - shape = [1, 1, resolution, resolution] - self.noises.register_buffer(f"noise{layer_idx}", torch.randn(*shape)) - # style convs and to_rgbs - for i in range(3, self.log_size + 1): - out_channels = channels[f"{2**i}"] - self.style_convs.append( - StyleConv( - in_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode="upsample", - ) - ) - self.style_convs.append( - StyleConv( - out_channels, - out_channels, - kernel_size=3, - num_style_feat=num_style_feat, - demodulate=True, - sample_mode=None, - ) - ) - self.to_rgbs.append(ToRGB(out_channels, num_style_feat, upsample=True)) - in_channels = out_channels - - def make_noise(self): - """Make noise for noise injection.""" - device = self.constant_input.weight.device - noises = [torch.randn(1, 1, 4, 4, device=device)] - - for i in range(3, self.log_size + 1): - for _ in range(2): - noises.append(torch.randn(1, 1, 2**i, 2**i, device=device)) - - return noises - - def get_latent(self, x): - return self.style_mlp(x) - - def mean_latent(self, num_latent): - latent_in = torch.randn( - num_latent, self.num_style_feat, device=self.constant_input.weight.device - ) - latent = self.style_mlp(latent_in).mean(0, keepdim=True) - return latent - - def forward( - self, - styles, - input_is_latent=False, - noise=None, - randomize_noise=True, - truncation=1, - truncation_latent=None, - inject_index=None, - return_latents=False, - ): - """Forward function for StyleGAN2GeneratorClean. - Args: - styles (list[Tensor]): Sample codes of styles. - input_is_latent (bool): Whether input is latent style. Default: False. - noise (Tensor | None): Input noise or None. Default: None. - randomize_noise (bool): Randomize noise, used when 'noise' is False. Default: True. - truncation (float): The truncation ratio. Default: 1. - truncation_latent (Tensor | None): The truncation latent tensor. Default: None. - inject_index (int | None): The injection index for mixing noise. Default: None. - return_latents (bool): Whether to return style latents. Default: False. - """ - # style codes -> latents with Style MLP layer - if not input_is_latent: - styles = [self.style_mlp(s) for s in styles] - # noises - if noise is None: - if randomize_noise: - noise = [None] * self.num_layers # for each style conv layer - else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] - # style truncation - if truncation < 1: - style_truncation = [] - for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) - styles = style_truncation - # get style latents with injection - if len(styles) == 1: - inject_index = self.num_latent - - if styles[0].ndim < 3: - # repeat latent code for all the layers - latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - else: # used for encoder with different latent code for each layer - latent = styles[0] - elif len(styles) == 2: # mixing noises - if inject_index is None: - inject_index = random.randint(1, self.num_latent - 1) - latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) - latent = torch.cat([latent1, latent2], 1) - - # main generation - out = self.constant_input(latent.shape[0]) - out = self.style_conv1(out, latent[:, 0], noise=noise[0]) - skip = self.to_rgb1(out, latent[:, 1]) - - i = 1 - for conv1, conv2, noise1, noise2, to_rgb in zip( - self.style_convs[::2], - self.style_convs[1::2], - noise[1::2], - noise[2::2], - self.to_rgbs, - ): - out = conv1(out, latent[:, i], noise=noise1) - out = conv2(out, latent[:, i + 1], noise=noise2) - skip = to_rgb(out, latent[:, i + 2], skip) # feature back to the rgb space - i += 2 - - image = skip - - if return_latents: - return image, latent - else: - return image, None diff --git a/backend/comfy_nodes/chainner_models/architecture/face/upfirdn2d.py b/backend/comfy_nodes/chainner_models/architecture/face/upfirdn2d.py deleted file mode 100644 index 4ea45415..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/face/upfirdn2d.py +++ /dev/null @@ -1,194 +0,0 @@ -# pylint: skip-file -# type: ignore -# modify from https://github.com/rosinality/stylegan2-pytorch/blob/master/op/upfirdn2d.py # noqa:E501 - -import os - -import torch -from torch.autograd import Function -from torch.nn import functional as F - -upfirdn2d_ext = None - - -class UpFirDn2dBackward(Function): - @staticmethod - def forward( - ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size - ): - up_x, up_y = up - down_x, down_y = down - g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad - - grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) - - grad_input = upfirdn2d_ext.upfirdn2d( - grad_output, - grad_kernel, - down_x, - down_y, - up_x, - up_y, - g_pad_x0, - g_pad_x1, - g_pad_y0, - g_pad_y1, - ) - grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3]) - - ctx.save_for_backward(kernel) - - pad_x0, pad_x1, pad_y0, pad_y1 = pad - - ctx.up_x = up_x - ctx.up_y = up_y - ctx.down_x = down_x - ctx.down_y = down_y - ctx.pad_x0 = pad_x0 - ctx.pad_x1 = pad_x1 - ctx.pad_y0 = pad_y0 - ctx.pad_y1 = pad_y1 - ctx.in_size = in_size - ctx.out_size = out_size - - return grad_input - - @staticmethod - def backward(ctx, gradgrad_input): - (kernel,) = ctx.saved_tensors - - gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1) - - gradgrad_out = upfirdn2d_ext.upfirdn2d( - gradgrad_input, - kernel, - ctx.up_x, - ctx.up_y, - ctx.down_x, - ctx.down_y, - ctx.pad_x0, - ctx.pad_x1, - ctx.pad_y0, - ctx.pad_y1, - ) - # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], - # ctx.out_size[1], ctx.in_size[3]) - gradgrad_out = gradgrad_out.view( - ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1] - ) - - return gradgrad_out, None, None, None, None, None, None, None, None - - -class UpFirDn2d(Function): - @staticmethod - def forward(ctx, input, kernel, up, down, pad): - up_x, up_y = up - down_x, down_y = down - pad_x0, pad_x1, pad_y0, pad_y1 = pad - - kernel_h, kernel_w = kernel.shape - _, channel, in_h, in_w = input.shape - ctx.in_size = input.shape - - input = input.reshape(-1, in_h, in_w, 1) - - ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) - - 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 - ctx.out_size = (out_h, out_w) - - ctx.up = (up_x, up_y) - ctx.down = (down_x, down_y) - ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) - - g_pad_x0 = kernel_w - pad_x0 - 1 - g_pad_y0 = kernel_h - pad_y0 - 1 - g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 - g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 - - ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) - - out = upfirdn2d_ext.upfirdn2d( - input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 - ) - # out = out.view(major, out_h, out_w, minor) - out = out.view(-1, channel, out_h, out_w) - - return out - - @staticmethod - def backward(ctx, grad_output): - kernel, grad_kernel = ctx.saved_tensors - - grad_input = UpFirDn2dBackward.apply( - grad_output, - kernel, - grad_kernel, - ctx.up, - ctx.down, - ctx.pad, - ctx.g_pad, - ctx.in_size, - ctx.out_size, - ) - - return grad_input, None, None, None, None - - -def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): - if input.device.type == "cpu": - out = upfirdn2d_native( - input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1] - ) - else: - out = UpFirDn2d.apply( - input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1]) - ) - - return out - - -def upfirdn2d_native( - input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1 -): - _, channel, in_h, in_w = input.shape - input = input.reshape(-1, in_h, in_w, 1) - - _, in_h, in_w, minor = input.shape - kernel_h, kernel_w = kernel.shape - - out = input.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[ - :, - 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) diff --git a/backend/comfy_nodes/chainner_models/architecture/timm/LICENSE b/backend/comfy_nodes/chainner_models/architecture/timm/LICENSE deleted file mode 100644 index b4e9438b..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/timm/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 2019 Ross Wightman - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. \ No newline at end of file diff --git a/backend/comfy_nodes/chainner_models/architecture/timm/drop.py b/backend/comfy_nodes/chainner_models/architecture/timm/drop.py deleted file mode 100644 index 14f0da91..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/timm/drop.py +++ /dev/null @@ -1,223 +0,0 @@ -""" DropBlock, DropPath - -PyTorch implementations of DropBlock and DropPath (Stochastic Depth) regularization layers. - -Papers: -DropBlock: A regularization method for convolutional networks (https://arxiv.org/abs/1810.12890) - -Deep Networks with Stochastic Depth (https://arxiv.org/abs/1603.09382) - -Code: -DropBlock impl inspired by two Tensorflow impl that I liked: - - https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_model.py#L74 - - https://github.com/clovaai/assembled-cnn/blob/master/nets/blocks.py - -Hacked together by / Copyright 2020 Ross Wightman -""" -import torch -import torch.nn as nn -import torch.nn.functional as F - - -def drop_block_2d( - x, - drop_prob: float = 0.1, - block_size: int = 7, - gamma_scale: float = 1.0, - with_noise: bool = False, - inplace: bool = False, - batchwise: bool = False, -): - """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf - - DropBlock with an experimental gaussian noise option. This layer has been tested on a few training - runs with success, but needs further validation and possibly optimization for lower runtime impact. - """ - _, C, H, W = x.shape - total_size = W * H - clipped_block_size = min(block_size, min(W, H)) - # seed_drop_rate, the gamma parameter - gamma = ( - gamma_scale - * drop_prob - * total_size - / clipped_block_size**2 - / ((W - block_size + 1) * (H - block_size + 1)) - ) - - # Forces the block to be inside the feature map. - w_i, h_i = torch.meshgrid( - torch.arange(W).to(x.device), torch.arange(H).to(x.device) - ) - valid_block = ( - (w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2) - ) & ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2)) - valid_block = torch.reshape(valid_block, (1, 1, H, W)).to(dtype=x.dtype) - - if batchwise: - # one mask for whole batch, quite a bit faster - uniform_noise = torch.rand((1, C, H, W), dtype=x.dtype, device=x.device) - else: - uniform_noise = torch.rand_like(x) - block_mask = ((2 - gamma - valid_block + uniform_noise) >= 1).to(dtype=x.dtype) - block_mask = -F.max_pool2d( - -block_mask, - kernel_size=clipped_block_size, # block_size, - stride=1, - padding=clipped_block_size // 2, - ) - - if with_noise: - normal_noise = ( - torch.randn((1, C, H, W), dtype=x.dtype, device=x.device) - if batchwise - else torch.randn_like(x) - ) - if inplace: - x.mul_(block_mask).add_(normal_noise * (1 - block_mask)) - else: - x = x * block_mask + normal_noise * (1 - block_mask) - else: - normalize_scale = ( - block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-7) - ).to(x.dtype) - if inplace: - x.mul_(block_mask * normalize_scale) - else: - x = x * block_mask * normalize_scale - return x - - -def drop_block_fast_2d( - x: torch.Tensor, - drop_prob: float = 0.1, - block_size: int = 7, - gamma_scale: float = 1.0, - with_noise: bool = False, - inplace: bool = False, -): - """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf - - DropBlock with an experimental gaussian noise option. Simplied from above without concern for valid - block mask at edges. - """ - _, _, H, W = x.shape - total_size = W * H - clipped_block_size = min(block_size, min(W, H)) - gamma = ( - gamma_scale - * drop_prob - * total_size - / clipped_block_size**2 - / ((W - block_size + 1) * (H - block_size + 1)) - ) - - block_mask = torch.empty_like(x).bernoulli_(gamma) - block_mask = F.max_pool2d( - block_mask.to(x.dtype), - kernel_size=clipped_block_size, - stride=1, - padding=clipped_block_size // 2, - ) - - if with_noise: - normal_noise = torch.empty_like(x).normal_() - if inplace: - x.mul_(1.0 - block_mask).add_(normal_noise * block_mask) - else: - x = x * (1.0 - block_mask) + normal_noise * block_mask - else: - block_mask = 1 - block_mask - normalize_scale = ( - block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-6) - ).to(dtype=x.dtype) - if inplace: - x.mul_(block_mask * normalize_scale) - else: - x = x * block_mask * normalize_scale - return x - - -class DropBlock2d(nn.Module): - """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf""" - - def __init__( - self, - drop_prob: float = 0.1, - block_size: int = 7, - gamma_scale: float = 1.0, - with_noise: bool = False, - inplace: bool = False, - batchwise: bool = False, - fast: bool = True, - ): - super(DropBlock2d, self).__init__() - self.drop_prob = drop_prob - self.gamma_scale = gamma_scale - self.block_size = block_size - self.with_noise = with_noise - self.inplace = inplace - self.batchwise = batchwise - self.fast = fast # FIXME finish comparisons of fast vs not - - def forward(self, x): - if not self.training or not self.drop_prob: - return x - if self.fast: - return drop_block_fast_2d( - x, - self.drop_prob, - self.block_size, - self.gamma_scale, - self.with_noise, - self.inplace, - ) - else: - return drop_block_2d( - x, - self.drop_prob, - self.block_size, - self.gamma_scale, - self.with_noise, - self.inplace, - self.batchwise, - ) - - -def drop_path( - x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True -): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - - This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, - the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... - See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for - changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use - 'survival rate' as the argument. - - """ - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * ( - x.ndim - 1 - ) # work with diff dim tensors, not just 2D ConvNets - random_tensor = x.new_empty(shape).bernoulli_(keep_prob) - if keep_prob > 0.0 and scale_by_keep: - random_tensor.div_(keep_prob) - return x * random_tensor - - -class DropPath(nn.Module): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" - - def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): - super(DropPath, self).__init__() - self.drop_prob = drop_prob - self.scale_by_keep = scale_by_keep - - def forward(self, x): - return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) - - def extra_repr(self): - return f"drop_prob={round(self.drop_prob,3):0.3f}" diff --git a/backend/comfy_nodes/chainner_models/architecture/timm/helpers.py b/backend/comfy_nodes/chainner_models/architecture/timm/helpers.py deleted file mode 100644 index cdafee07..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/timm/helpers.py +++ /dev/null @@ -1,31 +0,0 @@ -""" Layer/Module Helpers -Hacked together by / Copyright 2020 Ross Wightman -""" -import collections.abc -from itertools import repeat - - -# From PyTorch internals -def _ntuple(n): - def parse(x): - if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): - return x - return tuple(repeat(x, n)) - - return parse - - -to_1tuple = _ntuple(1) -to_2tuple = _ntuple(2) -to_3tuple = _ntuple(3) -to_4tuple = _ntuple(4) -to_ntuple = _ntuple - - -def make_divisible(v, divisor=8, min_value=None, round_limit=0.9): - min_value = min_value or divisor - new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) - # Make sure that round down does not go down by more than 10%. - if new_v < round_limit * v: - new_v += divisor - return new_v diff --git a/backend/comfy_nodes/chainner_models/architecture/timm/weight_init.py b/backend/comfy_nodes/chainner_models/architecture/timm/weight_init.py deleted file mode 100644 index b0169774..00000000 --- a/backend/comfy_nodes/chainner_models/architecture/timm/weight_init.py +++ /dev/null @@ -1,128 +0,0 @@ -import math -import warnings - -import torch -from torch.nn.init import _calculate_fan_in_and_fan_out - - -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): - warnings.warn( - "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " - "The distribution of values may be incorrect.", - stacklevel=2, - ) - - 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: torch.Tensor, mean=0.0, std=1.0, a=-2.0, b=2.0 -) -> 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`. - - NOTE: this impl is similar to the PyTorch trunc_normal_, the bounds [a, b] are - applied while sampling the normal with mean/std applied, therefore a, b args - should be adjusted to match the range of mean, std args. - - 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) - - -def trunc_normal_tf_( - tensor: torch.Tensor, mean=0.0, std=1.0, a=-2.0, b=2.0 -) -> 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`. - - NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the - bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0 - and the result is subsquently scaled and shifted by the mean and std args. - - 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) - """ - _no_grad_trunc_normal_(tensor, 0, 1.0, a, b) - with torch.no_grad(): - tensor.mul_(std).add_(mean) - return tensor - - -def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"): - fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) - if mode == "fan_in": - denom = fan_in - elif mode == "fan_out": - denom = fan_out - elif mode == "fan_avg": - denom = (fan_in + fan_out) / 2 - - variance = scale / denom # type: ignore - - if distribution == "truncated_normal": - # constant is stddev of standard normal truncated to (-2, 2) - trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978) - elif distribution == "normal": - tensor.normal_(std=math.sqrt(variance)) - elif distribution == "uniform": - bound = math.sqrt(3 * variance) - # pylint: disable=invalid-unary-operand-type - tensor.uniform_(-bound, bound) - else: - raise ValueError(f"invalid distribution {distribution}") - - -def lecun_normal_(tensor): - variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal") diff --git a/backend/comfy_nodes/chainner_models/model_loading.py b/backend/comfy_nodes/chainner_models/model_loading.py deleted file mode 100644 index e000871c..00000000 --- a/backend/comfy_nodes/chainner_models/model_loading.py +++ /dev/null @@ -1,99 +0,0 @@ -import logging as logger - -from .architecture.DAT import DAT -from .architecture.face.codeformer import CodeFormer -from .architecture.face.gfpganv1_clean_arch import GFPGANv1Clean -from .architecture.face.restoreformer_arch import RestoreFormer -from .architecture.HAT import HAT -from .architecture.LaMa import LaMa -from .architecture.OmniSR.OmniSR import OmniSR -from .architecture.RRDB import RRDBNet as ESRGAN -from .architecture.SCUNet import SCUNet -from .architecture.SPSR import SPSRNet as SPSR -from .architecture.SRVGG import SRVGGNetCompact as RealESRGANv2 -from .architecture.SwiftSRGAN import Generator as SwiftSRGAN -from .architecture.Swin2SR import Swin2SR -from .architecture.SwinIR import SwinIR -from .types import PyTorchModel - - -class UnsupportedModel(Exception): - pass - - -def load_state_dict(state_dict) -> PyTorchModel: - logger.debug(f"Loading state dict into pytorch model arch") - - state_dict_keys = list(state_dict.keys()) - - if "params_ema" in state_dict_keys: - state_dict = state_dict["params_ema"] - elif "params-ema" in state_dict_keys: - state_dict = state_dict["params-ema"] - elif "params" in state_dict_keys: - state_dict = state_dict["params"] - - state_dict_keys = list(state_dict.keys()) - # SRVGGNet Real-ESRGAN (v2) - if "body.0.weight" in state_dict_keys and "body.1.weight" in state_dict_keys: - model = RealESRGANv2(state_dict) - # SPSR (ESRGAN with lots of extra layers) - elif "f_HR_conv1.0.weight" in state_dict: - model = SPSR(state_dict) - # Swift-SRGAN - elif ( - "model" in state_dict_keys - and "initial.cnn.depthwise.weight" in state_dict["model"].keys() - ): - model = SwiftSRGAN(state_dict) - # SwinIR, Swin2SR, HAT - elif "layers.0.residual_group.blocks.0.norm1.weight" in state_dict_keys: - if ( - "layers.0.residual_group.blocks.0.conv_block.cab.0.weight" - in state_dict_keys - ): - model = HAT(state_dict) - elif "patch_embed.proj.weight" in state_dict_keys: - model = Swin2SR(state_dict) - else: - model = SwinIR(state_dict) - # GFPGAN - elif ( - "toRGB.0.weight" in state_dict_keys - and "stylegan_decoder.style_mlp.1.weight" in state_dict_keys - ): - model = GFPGANv1Clean(state_dict) - # RestoreFormer - elif ( - "encoder.conv_in.weight" in state_dict_keys - and "encoder.down.0.block.0.norm1.weight" in state_dict_keys - ): - model = RestoreFormer(state_dict) - elif ( - "encoder.blocks.0.weight" in state_dict_keys - and "quantize.embedding.weight" in state_dict_keys - ): - model = CodeFormer(state_dict) - # LaMa - elif ( - "model.model.1.bn_l.running_mean" in state_dict_keys - or "generator.model.1.bn_l.running_mean" in state_dict_keys - ): - model = LaMa(state_dict) - # Omni-SR - elif "residual_layer.0.residual_layer.0.layer.0.fn.0.weight" in state_dict_keys: - model = OmniSR(state_dict) - # SCUNet - elif "m_head.0.weight" in state_dict_keys and "m_tail.0.weight" in state_dict_keys: - model = SCUNet(state_dict) - # DAT - elif "layers.0.blocks.2.attn.attn_mask_0" in state_dict_keys: - model = DAT(state_dict) - # Regular ESRGAN, "new-arch" ESRGAN, Real-ESRGAN v1 - else: - try: - model = ESRGAN(state_dict) - except: - # pylint: disable=raise-missing-from - raise UnsupportedModel - return model diff --git a/backend/comfy_nodes/chainner_models/types.py b/backend/comfy_nodes/chainner_models/types.py deleted file mode 100644 index 193333b9..00000000 --- a/backend/comfy_nodes/chainner_models/types.py +++ /dev/null @@ -1,69 +0,0 @@ -from typing import Union - -from .architecture.DAT import DAT -from .architecture.face.codeformer import CodeFormer -from .architecture.face.gfpganv1_clean_arch import GFPGANv1Clean -from .architecture.face.restoreformer_arch import RestoreFormer -from .architecture.HAT import HAT -from .architecture.LaMa import LaMa -from .architecture.OmniSR.OmniSR import OmniSR -from .architecture.RRDB import RRDBNet as ESRGAN -from .architecture.SCUNet import SCUNet -from .architecture.SPSR import SPSRNet as SPSR -from .architecture.SRVGG import SRVGGNetCompact as RealESRGANv2 -from .architecture.SwiftSRGAN import Generator as SwiftSRGAN -from .architecture.Swin2SR import Swin2SR -from .architecture.SwinIR import SwinIR - -PyTorchSRModels = ( - RealESRGANv2, - SPSR, - SwiftSRGAN, - ESRGAN, - SwinIR, - Swin2SR, - HAT, - OmniSR, - SCUNet, - DAT, -) -PyTorchSRModel = Union[ - RealESRGANv2, - SPSR, - SwiftSRGAN, - ESRGAN, - SwinIR, - Swin2SR, - HAT, - OmniSR, - SCUNet, - DAT, -] - - -def is_pytorch_sr_model(model: object): - return isinstance(model, PyTorchSRModels) - - -PyTorchFaceModels = (GFPGANv1Clean, RestoreFormer, CodeFormer) -PyTorchFaceModel = Union[GFPGANv1Clean, RestoreFormer, CodeFormer] - - -def is_pytorch_face_model(model: object): - return isinstance(model, PyTorchFaceModels) - - -PyTorchInpaintModels = (LaMa,) -PyTorchInpaintModel = Union[LaMa] - - -def is_pytorch_inpaint_model(model: object): - return isinstance(model, PyTorchInpaintModels) - - -PyTorchModels = (*PyTorchSRModels, *PyTorchFaceModels, *PyTorchInpaintModels) -PyTorchModel = Union[PyTorchSRModel, PyTorchFaceModel, PyTorchInpaintModel] - - -def is_pytorch_model(model: object): - return isinstance(model, PyTorchModels) diff --git a/backend/comfy_nodes/nodes_canny.py b/backend/comfy_nodes/nodes_canny.py deleted file mode 100644 index 730dded5..00000000 --- a/backend/comfy_nodes/nodes_canny.py +++ /dev/null @@ -1,299 +0,0 @@ -#From https://github.com/kornia/kornia -import math - -import torch -import torch.nn.functional as F -import comfy.model_management - -def get_canny_nms_kernel(device=None, dtype=None): - """Utility function that returns 3x3 kernels for the Canny Non-maximal suppression.""" - return torch.tensor( - [ - [[[0.0, 0.0, 0.0], [0.0, 1.0, -1.0], [0.0, 0.0, 0.0]]], - [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]]], - [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, -1.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.0], [-1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]], - [[[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]], - [[[0.0, -1.0, 0.0], [0.0, 1.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]]], - ], - device=device, - dtype=dtype, - ) - - -def get_hysteresis_kernel(device=None, dtype=None): - """Utility function that returns the 3x3 kernels for the Canny hysteresis.""" - return torch.tensor( - [ - [[[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, 1.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], [1.0, 0.0, 0.0]]], - [[[0.0, 0.0, 0.0], [1.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, 1.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]]], - ], - device=device, - dtype=dtype, - ) - -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 = torch.nn.functional.pad(img, padding, mode="reflect") - img = torch.nn.functional.conv2d(img, kernel2d, groups=img.shape[-3]) - - return img - -def get_sobel_kernel2d(device=None, dtype=None): - kernel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=device, dtype=dtype) - kernel_y = kernel_x.transpose(0, 1) - return torch.stack([kernel_x, kernel_y]) - -def spatial_gradient(input, normalized: bool = True): - r"""Compute the first order image derivative in both x and y using a Sobel operator. - .. image:: _static/img/spatial_gradient.png - Args: - input: input image tensor with shape :math:`(B, C, H, W)`. - mode: derivatives modality, can be: `sobel` or `diff`. - order: the order of the derivatives. - normalized: whether the output is normalized. - Return: - the derivatives of the input feature map. with shape :math:`(B, C, 2, H, W)`. - .. note:: - See a working example `here `__. - Examples: - >>> input = torch.rand(1, 3, 4, 4) - >>> output = spatial_gradient(input) # 1x3x2x4x4 - >>> output.shape - torch.Size([1, 3, 2, 4, 4]) - """ - # KORNIA_CHECK_IS_TENSOR(input) - # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W']) - - # allocate kernel - kernel = get_sobel_kernel2d(device=input.device, dtype=input.dtype) - if normalized: - kernel = normalize_kernel2d(kernel) - - # prepare kernel - b, c, h, w = input.shape - tmp_kernel = kernel[:, None, ...] - - # Pad with "replicate for spatial dims, but with zeros for channel - spatial_pad = [kernel.size(1) // 2, kernel.size(1) // 2, kernel.size(2) // 2, kernel.size(2) // 2] - out_channels: int = 2 - padded_inp = torch.nn.functional.pad(input.reshape(b * c, 1, h, w), spatial_pad, 'replicate') - out = F.conv2d(padded_inp, tmp_kernel, groups=1, padding=0, stride=1) - return out.reshape(b, c, out_channels, h, w) - -def rgb_to_grayscale(image, rgb_weights = None): - r"""Convert a RGB image to grayscale version of image. - - .. image:: _static/img/rgb_to_grayscale.png - - The image data is assumed to be in the range of (0, 1). - - Args: - image: RGB image to be converted to grayscale with shape :math:`(*,3,H,W)`. - rgb_weights: Weights that will be applied on each channel (RGB). - The sum of the weights should add up to one. - Returns: - grayscale version of the image with shape :math:`(*,1,H,W)`. - - .. note:: - See a working example `here `__. - - Example: - >>> input = torch.rand(2, 3, 4, 5) - >>> gray = rgb_to_grayscale(input) # 2x1x4x5 - """ - - if len(image.shape) < 3 or image.shape[-3] != 3: - raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}") - - if rgb_weights is None: - # 8 bit images - if image.dtype == torch.uint8: - rgb_weights = torch.tensor([76, 150, 29], device=image.device, dtype=torch.uint8) - # floating point images - elif image.dtype in (torch.float16, torch.float32, torch.float64): - rgb_weights = torch.tensor([0.299, 0.587, 0.114], device=image.device, dtype=image.dtype) - else: - raise TypeError(f"Unknown data type: {image.dtype}") - else: - # is tensor that we make sure is in the same device/dtype - rgb_weights = rgb_weights.to(image) - - # unpack the color image channels with RGB order - r: Tensor = image[..., 0:1, :, :] - g: Tensor = image[..., 1:2, :, :] - b: Tensor = image[..., 2:3, :, :] - - w_r, w_g, w_b = rgb_weights.unbind() - return w_r * r + w_g * g + w_b * b - -def canny( - input, - low_threshold = 0.1, - high_threshold = 0.2, - kernel_size = 5, - sigma = 1, - hysteresis = True, - eps = 1e-6, -): - r"""Find edges of the input image and filters them using the Canny algorithm. - .. image:: _static/img/canny.png - Args: - input: input image tensor with shape :math:`(B,C,H,W)`. - low_threshold: lower threshold for the hysteresis procedure. - high_threshold: upper threshold for the hysteresis procedure. - kernel_size: the size of the kernel for the gaussian blur. - sigma: the standard deviation of the kernel for the gaussian blur. - hysteresis: if True, applies the hysteresis edge tracking. - Otherwise, the edges are divided between weak (0.5) and strong (1) edges. - eps: regularization number to avoid NaN during backprop. - Returns: - - the canny edge magnitudes map, shape of :math:`(B,1,H,W)`. - - the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`. - .. note:: - See a working example `here `__. - Example: - >>> input = torch.rand(5, 3, 4, 4) - >>> magnitude, edges = canny(input) # 5x3x4x4 - >>> magnitude.shape - torch.Size([5, 1, 4, 4]) - >>> edges.shape - torch.Size([5, 1, 4, 4]) - """ - # KORNIA_CHECK_IS_TENSOR(input) - # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W']) - # KORNIA_CHECK( - # low_threshold <= high_threshold, - # "Invalid input thresholds. low_threshold should be smaller than the high_threshold. Got: " - # f"{low_threshold}>{high_threshold}", - # ) - # KORNIA_CHECK(0 < low_threshold < 1, f'Invalid low threshold. Should be in range (0, 1). Got: {low_threshold}') - # KORNIA_CHECK(0 < high_threshold < 1, f'Invalid high threshold. Should be in range (0, 1). Got: {high_threshold}') - - device = input.device - dtype = input.dtype - - # To Grayscale - if input.shape[1] == 3: - input = rgb_to_grayscale(input) - - # Gaussian filter - blurred: Tensor = gaussian_blur_2d(input, kernel_size, sigma) - - # Compute the gradients - gradients: Tensor = spatial_gradient(blurred, normalized=False) - - # Unpack the edges - gx: Tensor = gradients[:, :, 0] - gy: Tensor = gradients[:, :, 1] - - # Compute gradient magnitude and angle - magnitude: Tensor = torch.sqrt(gx * gx + gy * gy + eps) - angle: Tensor = torch.atan2(gy, gx) - - # Radians to Degrees - angle = 180.0 * angle / math.pi - - # Round angle to the nearest 45 degree - angle = torch.round(angle / 45) * 45 - - # Non-maximal suppression - nms_kernels: Tensor = get_canny_nms_kernel(device, dtype) - nms_magnitude: Tensor = F.conv2d(magnitude, nms_kernels, padding=nms_kernels.shape[-1] // 2) - - # Get the indices for both directions - positive_idx: Tensor = (angle / 45) % 8 - positive_idx = positive_idx.long() - - negative_idx: Tensor = ((angle / 45) + 4) % 8 - negative_idx = negative_idx.long() - - # Apply the non-maximum suppression to the different directions - channel_select_filtered_positive: Tensor = torch.gather(nms_magnitude, 1, positive_idx) - channel_select_filtered_negative: Tensor = torch.gather(nms_magnitude, 1, negative_idx) - - channel_select_filtered: Tensor = torch.stack( - [channel_select_filtered_positive, channel_select_filtered_negative], 1 - ) - - is_max: Tensor = channel_select_filtered.min(dim=1)[0] > 0.0 - - magnitude = magnitude * is_max - - # Threshold - edges: Tensor = F.threshold(magnitude, low_threshold, 0.0) - - low: Tensor = magnitude > low_threshold - high: Tensor = magnitude > high_threshold - - edges = low * 0.5 + high * 0.5 - edges = edges.to(dtype) - - # Hysteresis - if hysteresis: - edges_old: Tensor = -torch.ones(edges.shape, device=edges.device, dtype=dtype) - hysteresis_kernels: Tensor = get_hysteresis_kernel(device, dtype) - - while ((edges_old - edges).abs() != 0).any(): - weak: Tensor = (edges == 0.5).float() - strong: Tensor = (edges == 1).float() - - hysteresis_magnitude: Tensor = F.conv2d( - edges, hysteresis_kernels, padding=hysteresis_kernels.shape[-1] // 2 - ) - hysteresis_magnitude = (hysteresis_magnitude == 1).any(1, keepdim=True).to(dtype) - hysteresis_magnitude = hysteresis_magnitude * weak + strong - - edges_old = edges.clone() - edges = hysteresis_magnitude + (hysteresis_magnitude == 0) * weak * 0.5 - - edges = hysteresis_magnitude - - return magnitude, edges - - -class Canny: - @classmethod - def INPUT_TYPES(s): - return {"required": {"image": ("IMAGE",), - "low_threshold": ("FLOAT", {"default": 0.4, "min": 0.01, "max": 0.99, "step": 0.01}), - "high_threshold": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01}) - }} - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "detect_edge" - - CATEGORY = "image/preprocessors" - - def detect_edge(self, image, low_threshold, high_threshold): - output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold) - img_out = output[1].to(comfy.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1) - return (img_out,) - -NODE_CLASS_MAPPINGS = { - "Canny": Canny, -} diff --git a/backend/comfy_nodes/nodes_clip_sdxl.py b/backend/comfy_nodes/nodes_clip_sdxl.py deleted file mode 100644 index dcf8859f..00000000 --- a/backend/comfy_nodes/nodes_clip_sdxl.py +++ /dev/null @@ -1,56 +0,0 @@ -import torch -from nodes import MAX_RESOLUTION - -class CLIPTextEncodeSDXLRefiner: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "ascore": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 1000.0, "step": 0.01}), - "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "text": ("STRING", {"multiline": True}), "clip": ("CLIP", ), - }} - RETURN_TYPES = ("CONDITIONING",) - FUNCTION = "encode" - - CATEGORY = "advanced/conditioning" - - def encode(self, clip, ascore, width, height, text): - tokens = clip.tokenize(text) - cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) - return ([[cond, {"pooled_output": pooled, "aesthetic_score": ascore, "width": width,"height": height}]], ) - -class CLIPTextEncodeSDXL: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}), - "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}), - "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), - "text_g": ("STRING", {"multiline": True, "default": "CLIP_G"}), "clip": ("CLIP", ), - "text_l": ("STRING", {"multiline": True, "default": "CLIP_L"}), "clip": ("CLIP", ), - }} - RETURN_TYPES = ("CONDITIONING",) - FUNCTION = "encode" - - CATEGORY = "advanced/conditioning" - - def encode(self, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l): - tokens = clip.tokenize(text_g) - tokens["l"] = clip.tokenize(text_l)["l"] - if len(tokens["l"]) != len(tokens["g"]): - empty = clip.tokenize("") - while len(tokens["l"]) < len(tokens["g"]): - tokens["l"] += empty["l"] - while len(tokens["l"]) > len(tokens["g"]): - tokens["g"] += empty["g"] - cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) - return ([[cond, {"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]], ) - -NODE_CLASS_MAPPINGS = { - "CLIPTextEncodeSDXLRefiner": CLIPTextEncodeSDXLRefiner, - "CLIPTextEncodeSDXL": CLIPTextEncodeSDXL, -} diff --git a/backend/comfy_nodes/nodes_compositing.py b/backend/comfy_nodes/nodes_compositing.py deleted file mode 100644 index 181b36ed..00000000 --- a/backend/comfy_nodes/nodes_compositing.py +++ /dev/null @@ -1,202 +0,0 @@ -import numpy as np -import torch -import comfy.utils -from enum import Enum - -def resize_mask(mask, shape): - return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1) - -class PorterDuffMode(Enum): - ADD = 0 - CLEAR = 1 - DARKEN = 2 - DST = 3 - DST_ATOP = 4 - DST_IN = 5 - DST_OUT = 6 - DST_OVER = 7 - LIGHTEN = 8 - MULTIPLY = 9 - OVERLAY = 10 - SCREEN = 11 - SRC = 12 - SRC_ATOP = 13 - SRC_IN = 14 - SRC_OUT = 15 - SRC_OVER = 16 - XOR = 17 - - -def porter_duff_composite(src_image: torch.Tensor, src_alpha: torch.Tensor, dst_image: torch.Tensor, dst_alpha: torch.Tensor, mode: PorterDuffMode): - if mode == PorterDuffMode.ADD: - out_alpha = torch.clamp(src_alpha + dst_alpha, 0, 1) - out_image = torch.clamp(src_image + dst_image, 0, 1) - elif mode == PorterDuffMode.CLEAR: - out_alpha = torch.zeros_like(dst_alpha) - out_image = torch.zeros_like(dst_image) - elif mode == PorterDuffMode.DARKEN: - out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha - out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.min(src_image, dst_image) - elif mode == PorterDuffMode.DST: - out_alpha = dst_alpha - out_image = dst_image - elif mode == PorterDuffMode.DST_ATOP: - out_alpha = src_alpha - out_image = src_alpha * dst_image + (1 - dst_alpha) * src_image - elif mode == PorterDuffMode.DST_IN: - out_alpha = src_alpha * dst_alpha - out_image = dst_image * src_alpha - elif mode == PorterDuffMode.DST_OUT: - out_alpha = (1 - src_alpha) * dst_alpha - out_image = (1 - src_alpha) * dst_image - elif mode == PorterDuffMode.DST_OVER: - out_alpha = dst_alpha + (1 - dst_alpha) * src_alpha - out_image = dst_image + (1 - dst_alpha) * src_image - elif mode == PorterDuffMode.LIGHTEN: - out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha - out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.max(src_image, dst_image) - elif mode == PorterDuffMode.MULTIPLY: - out_alpha = src_alpha * dst_alpha - out_image = src_image * dst_image - elif mode == PorterDuffMode.OVERLAY: - out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha - out_image = torch.where(2 * dst_image < dst_alpha, 2 * src_image * dst_image, - src_alpha * dst_alpha - 2 * (dst_alpha - src_image) * (src_alpha - dst_image)) - elif mode == PorterDuffMode.SCREEN: - out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha - out_image = src_image + dst_image - src_image * dst_image - elif mode == PorterDuffMode.SRC: - out_alpha = src_alpha - out_image = src_image - elif mode == PorterDuffMode.SRC_ATOP: - out_alpha = dst_alpha - out_image = dst_alpha * src_image + (1 - src_alpha) * dst_image - elif mode == PorterDuffMode.SRC_IN: - out_alpha = src_alpha * dst_alpha - out_image = src_image * dst_alpha - elif mode == PorterDuffMode.SRC_OUT: - out_alpha = (1 - dst_alpha) * src_alpha - out_image = (1 - dst_alpha) * src_image - elif mode == PorterDuffMode.SRC_OVER: - out_alpha = src_alpha + (1 - src_alpha) * dst_alpha - out_image = src_image + (1 - src_alpha) * dst_image - elif mode == PorterDuffMode.XOR: - out_alpha = (1 - dst_alpha) * src_alpha + (1 - src_alpha) * dst_alpha - out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image - else: - out_alpha = None - out_image = None - return out_image, out_alpha - - -class PorterDuffImageComposite: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "source": ("IMAGE",), - "source_alpha": ("MASK",), - "destination": ("IMAGE",), - "destination_alpha": ("MASK",), - "mode": ([mode.name for mode in PorterDuffMode], {"default": PorterDuffMode.DST.name}), - }, - } - - RETURN_TYPES = ("IMAGE", "MASK") - FUNCTION = "composite" - CATEGORY = "mask/compositing" - - def composite(self, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode): - batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha)) - out_images = [] - out_alphas = [] - - for i in range(batch_size): - src_image = source[i] - dst_image = destination[i] - - assert src_image.shape[2] == dst_image.shape[2] # inputs need to have same number of channels - - src_alpha = source_alpha[i].unsqueeze(2) - dst_alpha = destination_alpha[i].unsqueeze(2) - - if dst_alpha.shape[:2] != dst_image.shape[:2]: - upscale_input = dst_alpha.unsqueeze(0).permute(0, 3, 1, 2) - upscale_output = comfy.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center') - dst_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0) - if src_image.shape != dst_image.shape: - upscale_input = src_image.unsqueeze(0).permute(0, 3, 1, 2) - upscale_output = comfy.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center') - src_image = upscale_output.permute(0, 2, 3, 1).squeeze(0) - if src_alpha.shape != dst_alpha.shape: - upscale_input = src_alpha.unsqueeze(0).permute(0, 3, 1, 2) - upscale_output = comfy.utils.common_upscale(upscale_input, dst_alpha.shape[1], dst_alpha.shape[0], upscale_method='bicubic', crop='center') - src_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0) - - out_image, out_alpha = porter_duff_composite(src_image, src_alpha, dst_image, dst_alpha, PorterDuffMode[mode]) - - out_images.append(out_image) - out_alphas.append(out_alpha.squeeze(2)) - - result = (torch.stack(out_images), torch.stack(out_alphas)) - return result - - -class SplitImageWithAlpha: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - } - } - - CATEGORY = "mask/compositing" - RETURN_TYPES = ("IMAGE", "MASK") - FUNCTION = "split_image_with_alpha" - - def split_image_with_alpha(self, image: torch.Tensor): - out_images = [i[:,:,:3] for i in image] - out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image] - result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas)) - return result - - -class JoinImageWithAlpha: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "alpha": ("MASK",), - } - } - - CATEGORY = "mask/compositing" - RETURN_TYPES = ("IMAGE",) - FUNCTION = "join_image_with_alpha" - - def join_image_with_alpha(self, image: torch.Tensor, alpha: torch.Tensor): - batch_size = min(len(image), len(alpha)) - out_images = [] - - alpha = 1.0 - resize_mask(alpha, image.shape[1:]) - for i in range(batch_size): - out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2)) - - result = (torch.stack(out_images),) - return result - - -NODE_CLASS_MAPPINGS = { - "PorterDuffImageComposite": PorterDuffImageComposite, - "SplitImageWithAlpha": SplitImageWithAlpha, - "JoinImageWithAlpha": JoinImageWithAlpha, -} - - -NODE_DISPLAY_NAME_MAPPINGS = { - "PorterDuffImageComposite": "Porter-Duff Image Composite", - "SplitImageWithAlpha": "Split Image with Alpha", - "JoinImageWithAlpha": "Join Image with Alpha", -} diff --git a/backend/comfy_nodes/nodes_cond.py b/backend/comfy_nodes/nodes_cond.py deleted file mode 100644 index 646fefa1..00000000 --- a/backend/comfy_nodes/nodes_cond.py +++ /dev/null @@ -1,25 +0,0 @@ - - -class CLIPTextEncodeControlnet: - @classmethod - def INPUT_TYPES(s): - return {"required": {"clip": ("CLIP", ), "conditioning": ("CONDITIONING", ), "text": ("STRING", {"multiline": True})}} - RETURN_TYPES = ("CONDITIONING",) - FUNCTION = "encode" - - CATEGORY = "_for_testing/conditioning" - - def encode(self, clip, conditioning, text): - tokens = clip.tokenize(text) - cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) - c = [] - for t in conditioning: - n = [t[0], t[1].copy()] - n[1]['cross_attn_controlnet'] = cond - n[1]['pooled_output_controlnet'] = pooled - c.append(n) - return (c, ) - -NODE_CLASS_MAPPINGS = { - "CLIPTextEncodeControlnet": CLIPTextEncodeControlnet -} diff --git a/backend/comfy_nodes/nodes_custom_sampler.py b/backend/comfy_nodes/nodes_custom_sampler.py deleted file mode 100644 index 99f9ea7d..00000000 --- a/backend/comfy_nodes/nodes_custom_sampler.py +++ /dev/null @@ -1,295 +0,0 @@ -import comfy.samplers -import comfy.sample -from comfy.k_diffusion import sampling as k_diffusion_sampling -import latent_preview -import torch -import comfy.utils - - -class BasicScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"model": ("MODEL",), - "scheduler": (comfy.samplers.SCHEDULER_NAMES, ), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, model, scheduler, steps, denoise): - total_steps = steps - if denoise < 1.0: - total_steps = int(steps/denoise) - - comfy.model_management.load_models_gpu([model]) - sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu() - sigmas = sigmas[-(steps + 1):] - return (sigmas, ) - - -class KarrasScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, steps, sigma_max, sigma_min, rho): - sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) - return (sigmas, ) - -class ExponentialScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, steps, sigma_max, sigma_min): - sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max) - return (sigmas, ) - -class PolyexponentialScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "rho": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, steps, sigma_max, sigma_min, rho): - sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) - return (sigmas, ) - -class SDTurboScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"model": ("MODEL",), - "steps": ("INT", {"default": 1, "min": 1, "max": 10}), - "denoise": ("FLOAT", {"default": 1.0, "min": 0, "max": 1.0, "step": 0.01}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, model, steps, denoise): - start_step = 10 - int(10 * denoise) - timesteps = torch.flip(torch.arange(1, 11) * 100 - 1, (0,))[start_step:start_step + steps] - comfy.model_management.load_models_gpu([model]) - sigmas = model.model.model_sampling.sigma(timesteps) - sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) - return (sigmas, ) - -class VPScheduler: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "beta_d": ("FLOAT", {"default": 19.9, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), #TODO: fix default values - "beta_min": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), - "eps_s": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step":0.0001, "round": False}), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/schedulers" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, steps, beta_d, beta_min, eps_s): - sigmas = k_diffusion_sampling.get_sigmas_vp(n=steps, beta_d=beta_d, beta_min=beta_min, eps_s=eps_s) - return (sigmas, ) - -class SplitSigmas: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"sigmas": ("SIGMAS", ), - "step": ("INT", {"default": 0, "min": 0, "max": 10000}), - } - } - RETURN_TYPES = ("SIGMAS","SIGMAS") - CATEGORY = "sampling/custom_sampling/sigmas" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, sigmas, step): - sigmas1 = sigmas[:step + 1] - sigmas2 = sigmas[step:] - return (sigmas1, sigmas2) - -class FlipSigmas: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"sigmas": ("SIGMAS", ), - } - } - RETURN_TYPES = ("SIGMAS",) - CATEGORY = "sampling/custom_sampling/sigmas" - - FUNCTION = "get_sigmas" - - def get_sigmas(self, sigmas): - sigmas = sigmas.flip(0) - if sigmas[0] == 0: - sigmas[0] = 0.0001 - return (sigmas,) - -class KSamplerSelect: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"sampler_name": (comfy.samplers.SAMPLER_NAMES, ), - } - } - RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling/samplers" - - FUNCTION = "get_sampler" - - def get_sampler(self, sampler_name): - sampler = comfy.samplers.sampler_object(sampler_name) - return (sampler, ) - -class SamplerDPMPP_2M_SDE: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"solver_type": (['midpoint', 'heun'], ), - "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - "noise_device": (['gpu', 'cpu'], ), - } - } - RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling/samplers" - - FUNCTION = "get_sampler" - - def get_sampler(self, solver_type, eta, s_noise, noise_device): - if noise_device == 'cpu': - sampler_name = "dpmpp_2m_sde" - else: - sampler_name = "dpmpp_2m_sde_gpu" - sampler = comfy.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "solver_type": solver_type}) - return (sampler, ) - - -class SamplerDPMPP_SDE: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - "r": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 100.0, "step":0.01, "round": False}), - "noise_device": (['gpu', 'cpu'], ), - } - } - RETURN_TYPES = ("SAMPLER",) - CATEGORY = "sampling/custom_sampling/samplers" - - FUNCTION = "get_sampler" - - def get_sampler(self, eta, s_noise, r, noise_device): - if noise_device == 'cpu': - sampler_name = "dpmpp_sde" - else: - sampler_name = "dpmpp_sde_gpu" - sampler = comfy.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "r": r}) - return (sampler, ) - -class SamplerCustom: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"model": ("MODEL",), - "add_noise": ("BOOLEAN", {"default": True}), - "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), - "positive": ("CONDITIONING", ), - "negative": ("CONDITIONING", ), - "sampler": ("SAMPLER", ), - "sigmas": ("SIGMAS", ), - "latent_image": ("LATENT", ), - } - } - - RETURN_TYPES = ("LATENT","LATENT") - RETURN_NAMES = ("output", "denoised_output") - - FUNCTION = "sample" - - CATEGORY = "sampling/custom_sampling" - - def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image): - latent = latent_image - latent_image = latent["samples"] - if not add_noise: - noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") - else: - batch_inds = latent["batch_index"] if "batch_index" in latent else None - noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds) - - noise_mask = None - if "noise_mask" in latent: - noise_mask = latent["noise_mask"] - - x0_output = {} - callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output) - - disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED - samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed) - - out = latent.copy() - out["samples"] = samples - if "x0" in x0_output: - out_denoised = latent.copy() - out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu()) - else: - out_denoised = out - return (out, out_denoised) - -NODE_CLASS_MAPPINGS = { - "SamplerCustom": SamplerCustom, - "BasicScheduler": BasicScheduler, - "KarrasScheduler": KarrasScheduler, - "ExponentialScheduler": ExponentialScheduler, - "PolyexponentialScheduler": PolyexponentialScheduler, - "VPScheduler": VPScheduler, - "SDTurboScheduler": SDTurboScheduler, - "KSamplerSelect": KSamplerSelect, - "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE, - "SamplerDPMPP_SDE": SamplerDPMPP_SDE, - "SplitSigmas": SplitSigmas, - "FlipSigmas": FlipSigmas, -} diff --git a/backend/comfy_nodes/nodes_freelunch.py b/backend/comfy_nodes/nodes_freelunch.py deleted file mode 100644 index 7764aa0b..00000000 --- a/backend/comfy_nodes/nodes_freelunch.py +++ /dev/null @@ -1,113 +0,0 @@ -#code originally taken from: https://github.com/ChenyangSi/FreeU (under MIT License) - -import torch - - -def Fourier_filter(x, threshold, scale): - # FFT - x_freq = torch.fft.fftn(x.float(), dim=(-2, -1)) - x_freq = torch.fft.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 = torch.fft.ifftshift(x_freq, dim=(-2, -1)) - x_filtered = torch.fft.ifftn(x_freq, dim=(-2, -1)).real - - return x_filtered.to(x.dtype) - - -class FreeU: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.01}), - "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.01}), - "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}), - "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "model_patches" - - def patch(self, model, b1, b2, s1, s2): - model_channels = model.model.model_config.unet_config["model_channels"] - scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)} - on_cpu_devices = {} - - def output_block_patch(h, hsp, transformer_options): - scale = scale_dict.get(h.shape[1], None) - if scale is not None: - h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0] - if hsp.device not in on_cpu_devices: - try: - hsp = Fourier_filter(hsp, threshold=1, scale=scale[1]) - except: - print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.") - on_cpu_devices[hsp.device] = True - hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device) - else: - hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device) - - return h, hsp - - m = model.clone() - m.set_model_output_block_patch(output_block_patch) - return (m, ) - -class FreeU_V2: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "b1": ("FLOAT", {"default": 1.3, "min": 0.0, "max": 10.0, "step": 0.01}), - "b2": ("FLOAT", {"default": 1.4, "min": 0.0, "max": 10.0, "step": 0.01}), - "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}), - "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "model_patches" - - def patch(self, model, b1, b2, s1, s2): - model_channels = model.model.model_config.unet_config["model_channels"] - scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)} - on_cpu_devices = {} - - def output_block_patch(h, hsp, transformer_options): - scale = scale_dict.get(h.shape[1], None) - if scale is not None: - hidden_mean = h.mean(1).unsqueeze(1) - B = hidden_mean.shape[0] - hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) - hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True) - hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3) - - h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * ((scale[0] - 1 ) * hidden_mean + 1) - - if hsp.device not in on_cpu_devices: - try: - hsp = Fourier_filter(hsp, threshold=1, scale=scale[1]) - except: - print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.") - on_cpu_devices[hsp.device] = True - hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device) - else: - hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device) - - return h, hsp - - m = model.clone() - m.set_model_output_block_patch(output_block_patch) - return (m, ) - -NODE_CLASS_MAPPINGS = { - "FreeU": FreeU, - "FreeU_V2": FreeU_V2, -} diff --git a/backend/comfy_nodes/nodes_hypernetwork.py b/backend/comfy_nodes/nodes_hypernetwork.py deleted file mode 100644 index f692945a..00000000 --- a/backend/comfy_nodes/nodes_hypernetwork.py +++ /dev/null @@ -1,119 +0,0 @@ -import comfy.utils -import folder_paths -import torch - -def load_hypernetwork_patch(path, strength): - sd = comfy.utils.load_torch_file(path, safe_load=True) - activation_func = sd.get('activation_func', 'linear') - is_layer_norm = sd.get('is_layer_norm', False) - use_dropout = sd.get('use_dropout', False) - activate_output = sd.get('activate_output', False) - last_layer_dropout = sd.get('last_layer_dropout', False) - - valid_activation = { - "linear": torch.nn.Identity, - "relu": torch.nn.ReLU, - "leakyrelu": torch.nn.LeakyReLU, - "elu": torch.nn.ELU, - "swish": torch.nn.Hardswish, - "tanh": torch.nn.Tanh, - "sigmoid": torch.nn.Sigmoid, - "softsign": torch.nn.Softsign, - "mish": torch.nn.Mish, - } - - if activation_func not in valid_activation: - print("Unsupported Hypernetwork format, if you report it I might implement it.", path, " ", activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout) - return None - - out = {} - - for d in sd: - try: - dim = int(d) - except: - continue - - output = [] - for index in [0, 1]: - attn_weights = sd[dim][index] - keys = attn_weights.keys() - - linears = filter(lambda a: a.endswith(".weight"), keys) - linears = list(map(lambda a: a[:-len(".weight")], linears)) - layers = [] - - i = 0 - while i < len(linears): - lin_name = linears[i] - last_layer = (i == (len(linears) - 1)) - penultimate_layer = (i == (len(linears) - 2)) - - lin_weight = attn_weights['{}.weight'.format(lin_name)] - lin_bias = attn_weights['{}.bias'.format(lin_name)] - layer = torch.nn.Linear(lin_weight.shape[1], lin_weight.shape[0]) - layer.load_state_dict({"weight": lin_weight, "bias": lin_bias}) - layers.append(layer) - if activation_func != "linear": - if (not last_layer) or (activate_output): - layers.append(valid_activation[activation_func]()) - if is_layer_norm: - i += 1 - ln_name = linears[i] - ln_weight = attn_weights['{}.weight'.format(ln_name)] - ln_bias = attn_weights['{}.bias'.format(ln_name)] - ln = torch.nn.LayerNorm(ln_weight.shape[0]) - ln.load_state_dict({"weight": ln_weight, "bias": ln_bias}) - layers.append(ln) - if use_dropout: - if (not last_layer) and (not penultimate_layer or last_layer_dropout): - layers.append(torch.nn.Dropout(p=0.3)) - i += 1 - - output.append(torch.nn.Sequential(*layers)) - out[dim] = torch.nn.ModuleList(output) - - class hypernetwork_patch: - def __init__(self, hypernet, strength): - self.hypernet = hypernet - self.strength = strength - def __call__(self, q, k, v, extra_options): - dim = k.shape[-1] - if dim in self.hypernet: - hn = self.hypernet[dim] - k = k + hn[0](k) * self.strength - v = v + hn[1](v) * self.strength - - return q, k, v - - def to(self, device): - for d in self.hypernet.keys(): - self.hypernet[d] = self.hypernet[d].to(device) - return self - - return hypernetwork_patch(out, strength) - -class HypernetworkLoader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "hypernetwork_name": (folder_paths.get_filename_list("hypernetworks"), ), - "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_hypernetwork" - - CATEGORY = "loaders" - - def load_hypernetwork(self, model, hypernetwork_name, strength): - hypernetwork_path = folder_paths.get_full_path("hypernetworks", hypernetwork_name) - model_hypernetwork = model.clone() - patch = load_hypernetwork_patch(hypernetwork_path, strength) - if patch is not None: - model_hypernetwork.set_model_attn1_patch(patch) - model_hypernetwork.set_model_attn2_patch(patch) - return (model_hypernetwork,) - -NODE_CLASS_MAPPINGS = { - "HypernetworkLoader": HypernetworkLoader -} diff --git a/backend/comfy_nodes/nodes_hypertile.py b/backend/comfy_nodes/nodes_hypertile.py deleted file mode 100644 index ae55d23d..00000000 --- a/backend/comfy_nodes/nodes_hypertile.py +++ /dev/null @@ -1,83 +0,0 @@ -#Taken from: https://github.com/tfernd/HyperTile/ - -import math -from einops import rearrange -# Use torch rng for consistency across generations -from torch import randint - -def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int: - min_value = min(min_value, value) - - # All big divisors of value (inclusive) - divisors = [i for i in range(min_value, value + 1) if value % i == 0] - - ns = [value // i for i in divisors[:max_options]] # has at least 1 element - - if len(ns) - 1 > 0: - idx = randint(low=0, high=len(ns) - 1, size=(1,)).item() - else: - idx = 0 - - return ns[idx] - -class HyperTile: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "tile_size": ("INT", {"default": 256, "min": 1, "max": 2048}), - "swap_size": ("INT", {"default": 2, "min": 1, "max": 128}), - "max_depth": ("INT", {"default": 0, "min": 0, "max": 10}), - "scale_depth": ("BOOLEAN", {"default": False}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "model_patches" - - def patch(self, model, tile_size, swap_size, max_depth, scale_depth): - model_channels = model.model.model_config.unet_config["model_channels"] - - latent_tile_size = max(32, tile_size) // 8 - self.temp = None - - def hypertile_in(q, k, v, extra_options): - model_chans = q.shape[-2] - orig_shape = extra_options['original_shape'] - apply_to = [] - for i in range(max_depth + 1): - apply_to.append((orig_shape[-2] / (2 ** i)) * (orig_shape[-1] / (2 ** i))) - - if model_chans in apply_to: - shape = extra_options["original_shape"] - aspect_ratio = shape[-1] / shape[-2] - - hw = q.size(1) - h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio)) - - factor = (2 ** apply_to.index(model_chans)) if scale_depth else 1 - nh = random_divisor(h, latent_tile_size * factor, swap_size) - nw = random_divisor(w, latent_tile_size * factor, swap_size) - - if nh * nw > 1: - q = rearrange(q, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw) - self.temp = (nh, nw, h, w) - return q, k, v - - return q, k, v - def hypertile_out(out, extra_options): - if self.temp is not None: - nh, nw, h, w = self.temp - self.temp = None - out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw) - out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw) - return out - - - m = model.clone() - m.set_model_attn1_patch(hypertile_in) - m.set_model_attn1_output_patch(hypertile_out) - return (m, ) - -NODE_CLASS_MAPPINGS = { - "HyperTile": HyperTile, -} diff --git a/backend/comfy_nodes/nodes_images.py b/backend/comfy_nodes/nodes_images.py deleted file mode 100644 index 8f638bf8..00000000 --- a/backend/comfy_nodes/nodes_images.py +++ /dev/null @@ -1,195 +0,0 @@ -import nodes -import folder_paths -from comfy.cli_args import args - -from PIL import Image -from PIL.PngImagePlugin import PngInfo - -import numpy as np -import json -import os - -MAX_RESOLUTION = nodes.MAX_RESOLUTION - -class ImageCrop: - @classmethod - def INPUT_TYPES(s): - return {"required": { "image": ("IMAGE",), - "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - }} - RETURN_TYPES = ("IMAGE",) - FUNCTION = "crop" - - CATEGORY = "image/transform" - - def crop(self, image, width, height, x, y): - x = min(x, image.shape[2] - 1) - y = min(y, image.shape[1] - 1) - to_x = width + x - to_y = height + y - img = image[:,y:to_y, x:to_x, :] - return (img,) - -class RepeatImageBatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { "image": ("IMAGE",), - "amount": ("INT", {"default": 1, "min": 1, "max": 64}), - }} - RETURN_TYPES = ("IMAGE",) - FUNCTION = "repeat" - - CATEGORY = "image/batch" - - def repeat(self, image, amount): - s = image.repeat((amount, 1,1,1)) - return (s,) - -class ImageFromBatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { "image": ("IMAGE",), - "batch_index": ("INT", {"default": 0, "min": 0, "max": 63}), - "length": ("INT", {"default": 1, "min": 1, "max": 64}), - }} - RETURN_TYPES = ("IMAGE",) - FUNCTION = "frombatch" - - CATEGORY = "image/batch" - - def frombatch(self, image, batch_index, length): - s_in = image - batch_index = min(s_in.shape[0] - 1, batch_index) - length = min(s_in.shape[0] - batch_index, length) - s = s_in[batch_index:batch_index + length].clone() - return (s,) - -class SaveAnimatedWEBP: - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - self.type = "output" - self.prefix_append = "" - - methods = {"default": 4, "fastest": 0, "slowest": 6} - @classmethod - def INPUT_TYPES(s): - return {"required": - {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), - "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}), - "lossless": ("BOOLEAN", {"default": True}), - "quality": ("INT", {"default": 80, "min": 0, "max": 100}), - "method": (list(s.methods.keys()),), - # "num_frames": ("INT", {"default": 0, "min": 0, "max": 8192}), - }, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, - } - - RETURN_TYPES = () - FUNCTION = "save_images" - - OUTPUT_NODE = True - - CATEGORY = "image/animation" - - def save_images(self, images, fps, filename_prefix, lossless, quality, method, num_frames=0, prompt=None, extra_pnginfo=None): - method = self.methods.get(method) - filename_prefix += self.prefix_append - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) - results = list() - pil_images = [] - for image in images: - i = 255. * image.cpu().numpy() - img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - pil_images.append(img) - - metadata = pil_images[0].getexif() - if not args.disable_metadata: - if prompt is not None: - metadata[0x0110] = "prompt:{}".format(json.dumps(prompt)) - if extra_pnginfo is not None: - inital_exif = 0x010f - for x in extra_pnginfo: - metadata[inital_exif] = "{}:{}".format(x, json.dumps(extra_pnginfo[x])) - inital_exif -= 1 - - if num_frames == 0: - num_frames = len(pil_images) - - c = len(pil_images) - for i in range(0, c, num_frames): - file = f"{filename}_{counter:05}_.webp" - pil_images[i].save(os.path.join(full_output_folder, file), save_all=True, duration=int(1000.0/fps), append_images=pil_images[i + 1:i + num_frames], exif=metadata, lossless=lossless, quality=quality, method=method) - results.append({ - "filename": file, - "subfolder": subfolder, - "type": self.type - }) - counter += 1 - - animated = num_frames != 1 - return { "ui": { "images": results, "animated": (animated,) } } - -class SaveAnimatedPNG: - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - self.type = "output" - self.prefix_append = "" - - @classmethod - def INPUT_TYPES(s): - return {"required": - {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), - "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}), - "compress_level": ("INT", {"default": 4, "min": 0, "max": 9}) - }, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, - } - - RETURN_TYPES = () - FUNCTION = "save_images" - - OUTPUT_NODE = True - - CATEGORY = "image/animation" - - def save_images(self, images, fps, compress_level, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): - filename_prefix += self.prefix_append - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) - results = list() - pil_images = [] - for image in images: - i = 255. * image.cpu().numpy() - img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - pil_images.append(img) - - metadata = None - if not args.disable_metadata: - metadata = PngInfo() - if prompt is not None: - metadata.add(b"comf", "prompt".encode("latin-1", "strict") + b"\0" + json.dumps(prompt).encode("latin-1", "strict"), after_idat=True) - if extra_pnginfo is not None: - for x in extra_pnginfo: - metadata.add(b"comf", x.encode("latin-1", "strict") + b"\0" + json.dumps(extra_pnginfo[x]).encode("latin-1", "strict"), after_idat=True) - - file = f"{filename}_{counter:05}_.png" - pil_images[0].save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1:]) - results.append({ - "filename": file, - "subfolder": subfolder, - "type": self.type - }) - - return { "ui": { "images": results, "animated": (True,)} } - -NODE_CLASS_MAPPINGS = { - "ImageCrop": ImageCrop, - "RepeatImageBatch": RepeatImageBatch, - "ImageFromBatch": ImageFromBatch, - "SaveAnimatedWEBP": SaveAnimatedWEBP, - "SaveAnimatedPNG": SaveAnimatedPNG, -} diff --git a/backend/comfy_nodes/nodes_latent.py b/backend/comfy_nodes/nodes_latent.py deleted file mode 100644 index eabae088..00000000 --- a/backend/comfy_nodes/nodes_latent.py +++ /dev/null @@ -1,155 +0,0 @@ -import comfy.utils -import torch - -def reshape_latent_to(target_shape, latent): - if latent.shape[1:] != target_shape[1:]: - latent = comfy.utils.common_upscale(latent, target_shape[3], target_shape[2], "bilinear", "center") - return comfy.utils.repeat_to_batch_size(latent, target_shape[0]) - - -class LatentAdd: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "op" - - CATEGORY = "latent/advanced" - - def op(self, samples1, samples2): - samples_out = samples1.copy() - - s1 = samples1["samples"] - s2 = samples2["samples"] - - s2 = reshape_latent_to(s1.shape, s2) - samples_out["samples"] = s1 + s2 - return (samples_out,) - -class LatentSubtract: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "op" - - CATEGORY = "latent/advanced" - - def op(self, samples1, samples2): - samples_out = samples1.copy() - - s1 = samples1["samples"] - s2 = samples2["samples"] - - s2 = reshape_latent_to(s1.shape, s2) - samples_out["samples"] = s1 - s2 - return (samples_out,) - -class LatentMultiply: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples": ("LATENT",), - "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), - }} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "op" - - CATEGORY = "latent/advanced" - - def op(self, samples, multiplier): - samples_out = samples.copy() - - s1 = samples["samples"] - samples_out["samples"] = s1 * multiplier - return (samples_out,) - -class LatentInterpolate: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples1": ("LATENT",), - "samples2": ("LATENT",), - "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - }} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "op" - - CATEGORY = "latent/advanced" - - def op(self, samples1, samples2, ratio): - samples_out = samples1.copy() - - s1 = samples1["samples"] - s2 = samples2["samples"] - - s2 = reshape_latent_to(s1.shape, s2) - - m1 = torch.linalg.vector_norm(s1, dim=(1)) - m2 = torch.linalg.vector_norm(s2, dim=(1)) - - s1 = torch.nan_to_num(s1 / m1) - s2 = torch.nan_to_num(s2 / m2) - - t = (s1 * ratio + s2 * (1.0 - ratio)) - mt = torch.linalg.vector_norm(t, dim=(1)) - st = torch.nan_to_num(t / mt) - - samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio)) - return (samples_out,) - -class LatentBatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "batch" - - CATEGORY = "latent/batch" - - def batch(self, samples1, samples2): - samples_out = samples1.copy() - s1 = samples1["samples"] - s2 = samples2["samples"] - - if s1.shape[1:] != s2.shape[1:]: - s2 = comfy.utils.common_upscale(s2, s1.shape[3], s1.shape[2], "bilinear", "center") - s = torch.cat((s1, s2), dim=0) - samples_out["samples"] = s - samples_out["batch_index"] = samples1.get("batch_index", [x for x in range(0, s1.shape[0])]) + samples2.get("batch_index", [x for x in range(0, s2.shape[0])]) - return (samples_out,) - -class LatentBatchSeedBehavior: - @classmethod - def INPUT_TYPES(s): - return {"required": { "samples": ("LATENT",), - "seed_behavior": (["random", "fixed"],{"default": "fixed"}),}} - - RETURN_TYPES = ("LATENT",) - FUNCTION = "op" - - CATEGORY = "latent/advanced" - - def op(self, samples, seed_behavior): - samples_out = samples.copy() - latent = samples["samples"] - if seed_behavior == "random": - if 'batch_index' in samples_out: - samples_out.pop('batch_index') - elif seed_behavior == "fixed": - batch_number = samples_out.get("batch_index", [0])[0] - samples_out["batch_index"] = [batch_number] * latent.shape[0] - - return (samples_out,) - -NODE_CLASS_MAPPINGS = { - "LatentAdd": LatentAdd, - "LatentSubtract": LatentSubtract, - "LatentMultiply": LatentMultiply, - "LatentInterpolate": LatentInterpolate, - "LatentBatch": LatentBatch, - "LatentBatchSeedBehavior": LatentBatchSeedBehavior, -} diff --git a/backend/comfy_nodes/nodes_mask.py b/backend/comfy_nodes/nodes_mask.py deleted file mode 100644 index a7d164bf..00000000 --- a/backend/comfy_nodes/nodes_mask.py +++ /dev/null @@ -1,363 +0,0 @@ -import numpy as np -import scipy.ndimage -import torch -import comfy.utils - -from nodes import MAX_RESOLUTION - -def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False): - source = source.to(destination.device) - if resize_source: - source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear") - - source = comfy.utils.repeat_to_batch_size(source, destination.shape[0]) - - x = max(-source.shape[3] * multiplier, min(x, destination.shape[3] * multiplier)) - y = max(-source.shape[2] * multiplier, min(y, destination.shape[2] * multiplier)) - - left, top = (x // multiplier, y // multiplier) - right, bottom = (left + source.shape[3], top + source.shape[2],) - - if mask is None: - mask = torch.ones_like(source) - else: - mask = mask.to(destination.device, copy=True) - mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear") - mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0]) - - # calculate the bounds of the source that will be overlapping the destination - # this prevents the source trying to overwrite latent pixels that are out of bounds - # of the destination - visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),) - - mask = mask[:, :, :visible_height, :visible_width] - inverse_mask = torch.ones_like(mask) - mask - - source_portion = mask * source[:, :, :visible_height, :visible_width] - destination_portion = inverse_mask * destination[:, :, top:bottom, left:right] - - destination[:, :, top:bottom, left:right] = source_portion + destination_portion - return destination - -class LatentCompositeMasked: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "destination": ("LATENT",), - "source": ("LATENT",), - "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), - "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), - "resize_source": ("BOOLEAN", {"default": False}), - }, - "optional": { - "mask": ("MASK",), - } - } - RETURN_TYPES = ("LATENT",) - FUNCTION = "composite" - - CATEGORY = "latent" - - def composite(self, destination, source, x, y, resize_source, mask = None): - output = destination.copy() - destination = destination["samples"].clone() - source = source["samples"] - output["samples"] = composite(destination, source, x, y, mask, 8, resize_source) - return (output,) - -class ImageCompositeMasked: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "destination": ("IMAGE",), - "source": ("IMAGE",), - "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "resize_source": ("BOOLEAN", {"default": False}), - }, - "optional": { - "mask": ("MASK",), - } - } - RETURN_TYPES = ("IMAGE",) - FUNCTION = "composite" - - CATEGORY = "image" - - def composite(self, destination, source, x, y, resize_source, mask = None): - destination = destination.clone().movedim(-1, 1) - output = composite(destination, source.movedim(-1, 1), x, y, mask, 1, resize_source).movedim(1, -1) - return (output,) - -class MaskToImage: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "mask": ("MASK",), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "mask_to_image" - - def mask_to_image(self, mask): - result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) - return (result,) - -class ImageToMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "channel": (["red", "green", "blue", "alpha"],), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - FUNCTION = "image_to_mask" - - def image_to_mask(self, image, channel): - channels = ["red", "green", "blue", "alpha"] - mask = image[:, :, :, channels.index(channel)] - return (mask,) - -class ImageColorToMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - FUNCTION = "image_to_mask" - - def image_to_mask(self, image, color): - temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int) - temp = torch.bitwise_left_shift(temp[:,:,:,0], 16) + torch.bitwise_left_shift(temp[:,:,:,1], 8) + temp[:,:,:,2] - mask = torch.where(temp == color, 255, 0).float() - return (mask,) - -class SolidMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "solid" - - def solid(self, value, width, height): - out = torch.full((1, height, width), value, dtype=torch.float32, device="cpu") - return (out,) - -class InvertMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "invert" - - def invert(self, mask): - out = 1.0 - mask - return (out,) - -class CropMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "crop" - - def crop(self, mask, x, y, width, height): - mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) - out = mask[:, y:y + height, x:x + width] - return (out,) - -class MaskComposite: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "destination": ("MASK",), - "source": ("MASK",), - "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "operation": (["multiply", "add", "subtract", "and", "or", "xor"],), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "combine" - - def combine(self, destination, source, x, y, operation): - output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone() - source = source.reshape((-1, source.shape[-2], source.shape[-1])) - - left, top = (x, y,) - right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2])) - visible_width, visible_height = (right - left, bottom - top,) - - source_portion = source[:, :visible_height, :visible_width] - destination_portion = destination[:, top:bottom, left:right] - - if operation == "multiply": - output[:, top:bottom, left:right] = destination_portion * source_portion - elif operation == "add": - output[:, top:bottom, left:right] = destination_portion + source_portion - elif operation == "subtract": - output[:, top:bottom, left:right] = destination_portion - source_portion - elif operation == "and": - output[:, top:bottom, left:right] = torch.bitwise_and(destination_portion.round().bool(), source_portion.round().bool()).float() - elif operation == "or": - output[:, top:bottom, left:right] = torch.bitwise_or(destination_portion.round().bool(), source_portion.round().bool()).float() - elif operation == "xor": - output[:, top:bottom, left:right] = torch.bitwise_xor(destination_portion.round().bool(), source_portion.round().bool()).float() - - output = torch.clamp(output, 0.0, 1.0) - - return (output,) - -class FeatherMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), - } - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "feather" - - def feather(self, mask, left, top, right, bottom): - output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone() - - left = min(left, output.shape[-1]) - right = min(right, output.shape[-1]) - top = min(top, output.shape[-2]) - bottom = min(bottom, output.shape[-2]) - - for x in range(left): - feather_rate = (x + 1.0) / left - output[:, :, x] *= feather_rate - - for x in range(right): - feather_rate = (x + 1) / right - output[:, :, -x] *= feather_rate - - for y in range(top): - feather_rate = (y + 1) / top - output[:, y, :] *= feather_rate - - for y in range(bottom): - feather_rate = (y + 1) / bottom - output[:, -y, :] *= feather_rate - - return (output,) - -class GrowMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), - "tapered_corners": ("BOOLEAN", {"default": True}), - }, - } - - CATEGORY = "mask" - - RETURN_TYPES = ("MASK",) - - FUNCTION = "expand_mask" - - def expand_mask(self, mask, expand, tapered_corners): - c = 0 if tapered_corners else 1 - kernel = np.array([[c, 1, c], - [1, 1, 1], - [c, 1, c]]) - mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) - out = [] - for m in mask: - output = m.numpy() - for _ in range(abs(expand)): - if expand < 0: - output = scipy.ndimage.grey_erosion(output, footprint=kernel) - else: - output = scipy.ndimage.grey_dilation(output, footprint=kernel) - output = torch.from_numpy(output) - out.append(output) - return (torch.stack(out, dim=0),) - - - -NODE_CLASS_MAPPINGS = { - "LatentCompositeMasked": LatentCompositeMasked, - "ImageCompositeMasked": ImageCompositeMasked, - "MaskToImage": MaskToImage, - "ImageToMask": ImageToMask, - "ImageColorToMask": ImageColorToMask, - "SolidMask": SolidMask, - "InvertMask": InvertMask, - "CropMask": CropMask, - "MaskComposite": MaskComposite, - "FeatherMask": FeatherMask, - "GrowMask": GrowMask, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "ImageToMask": "Convert Image to Mask", - "MaskToImage": "Convert Mask to Image", -} diff --git a/backend/comfy_nodes/nodes_model_advanced.py b/backend/comfy_nodes/nodes_model_advanced.py deleted file mode 100644 index 21af4b73..00000000 --- a/backend/comfy_nodes/nodes_model_advanced.py +++ /dev/null @@ -1,217 +0,0 @@ -import folder_paths -import comfy.sd -import comfy.model_sampling -import comfy.latent_formats -import torch - -class LCM(comfy.model_sampling.EPS): - def calculate_denoised(self, sigma, model_output, model_input): - timestep = self.timestep(sigma).view(sigma.shape[:1] + (1,) * (model_output.ndim - 1)) - sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1)) - x0 = model_input - model_output * sigma - - sigma_data = 0.5 - scaled_timestep = timestep * 10.0 #timestep_scaling - - c_skip = sigma_data**2 / (scaled_timestep**2 + sigma_data**2) - c_out = scaled_timestep / (scaled_timestep**2 + sigma_data**2) ** 0.5 - - return c_out * x0 + c_skip * model_input - -class X0(comfy.model_sampling.EPS): - def calculate_denoised(self, sigma, model_output, model_input): - return model_output - -class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete): - original_timesteps = 50 - - def __init__(self, model_config=None): - super().__init__(model_config) - - self.skip_steps = self.num_timesteps // self.original_timesteps - - sigmas_valid = torch.zeros((self.original_timesteps), dtype=torch.float32) - for x in range(self.original_timesteps): - sigmas_valid[self.original_timesteps - 1 - x] = self.sigmas[self.num_timesteps - 1 - x * self.skip_steps] - - self.set_sigmas(sigmas_valid) - - def timestep(self, sigma): - log_sigma = sigma.log() - dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None] - return (dists.abs().argmin(dim=0).view(sigma.shape) * self.skip_steps + (self.skip_steps - 1)).to(sigma.device) - - def sigma(self, timestep): - t = torch.clamp(((timestep.float().to(self.log_sigmas.device) - (self.skip_steps - 1)) / self.skip_steps).float(), min=0, max=(len(self.sigmas) - 1)) - low_idx = t.floor().long() - high_idx = t.ceil().long() - w = t.frac() - log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx] - return log_sigma.exp().to(timestep.device) - - -def rescale_zero_terminal_snr_sigmas(sigmas): - alphas_cumprod = 1 / ((sigmas * sigmas) + 1) - 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_bar[-1] = 4.8973451890853435e-08 - return ((1 - alphas_bar) / alphas_bar) ** 0.5 - -class ModelSamplingDiscrete: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "sampling": (["eps", "v_prediction", "lcm", "x0"],), - "zsnr": ("BOOLEAN", {"default": False}), - }} - - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "advanced/model" - - def patch(self, model, sampling, zsnr): - m = model.clone() - - sampling_base = comfy.model_sampling.ModelSamplingDiscrete - if sampling == "eps": - sampling_type = comfy.model_sampling.EPS - elif sampling == "v_prediction": - sampling_type = comfy.model_sampling.V_PREDICTION - elif sampling == "lcm": - sampling_type = LCM - sampling_base = ModelSamplingDiscreteDistilled - elif sampling == "x0": - sampling_type = X0 - - class ModelSamplingAdvanced(sampling_base, sampling_type): - pass - - model_sampling = ModelSamplingAdvanced(model.model.model_config) - if zsnr: - model_sampling.set_sigmas(rescale_zero_terminal_snr_sigmas(model_sampling.sigmas)) - - m.add_object_patch("model_sampling", model_sampling) - return (m, ) - -class ModelSamplingStableCascade: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "shift": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 100.0, "step":0.01}), - }} - - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "advanced/model" - - def patch(self, model, shift): - m = model.clone() - - sampling_base = comfy.model_sampling.StableCascadeSampling - sampling_type = comfy.model_sampling.EPS - - class ModelSamplingAdvanced(sampling_base, sampling_type): - pass - - model_sampling = ModelSamplingAdvanced(model.model.model_config) - model_sampling.set_parameters(shift) - m.add_object_patch("model_sampling", model_sampling) - return (m, ) - -class ModelSamplingContinuousEDM: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "sampling": (["v_prediction", "edm_playground_v2.5", "eps"],), - "sigma_max": ("FLOAT", {"default": 120.0, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}), - "sigma_min": ("FLOAT", {"default": 0.002, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}), - }} - - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "advanced/model" - - def patch(self, model, sampling, sigma_max, sigma_min): - m = model.clone() - - latent_format = None - sigma_data = 1.0 - if sampling == "eps": - sampling_type = comfy.model_sampling.EPS - elif sampling == "v_prediction": - sampling_type = comfy.model_sampling.V_PREDICTION - elif sampling == "edm_playground_v2.5": - sampling_type = comfy.model_sampling.EDM - sigma_data = 0.5 - latent_format = comfy.latent_formats.SDXL_Playground_2_5() - - class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingContinuousEDM, sampling_type): - pass - - model_sampling = ModelSamplingAdvanced(model.model.model_config) - model_sampling.set_parameters(sigma_min, sigma_max, sigma_data) - m.add_object_patch("model_sampling", model_sampling) - if latent_format is not None: - m.add_object_patch("latent_format", latent_format) - return (m, ) - -class RescaleCFG: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "multiplier": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "advanced/model" - - def patch(self, model, multiplier): - def rescale_cfg(args): - cond = args["cond"] - uncond = args["uncond"] - cond_scale = args["cond_scale"] - sigma = args["sigma"] - sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1)) - x_orig = args["input"] - - #rescale cfg has to be done on v-pred model output - x = x_orig / (sigma * sigma + 1.0) - cond = ((x - (x_orig - cond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma) - uncond = ((x - (x_orig - uncond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma) - - #rescalecfg - x_cfg = uncond + cond_scale * (cond - uncond) - ro_pos = torch.std(cond, dim=(1,2,3), keepdim=True) - ro_cfg = torch.std(x_cfg, dim=(1,2,3), keepdim=True) - - x_rescaled = x_cfg * (ro_pos / ro_cfg) - x_final = multiplier * x_rescaled + (1.0 - multiplier) * x_cfg - - return x_orig - (x - x_final * sigma / (sigma * sigma + 1.0) ** 0.5) - - m = model.clone() - m.set_model_sampler_cfg_function(rescale_cfg) - return (m, ) - -NODE_CLASS_MAPPINGS = { - "ModelSamplingDiscrete": ModelSamplingDiscrete, - "ModelSamplingContinuousEDM": ModelSamplingContinuousEDM, - "ModelSamplingStableCascade": ModelSamplingStableCascade, - "RescaleCFG": RescaleCFG, -} diff --git a/backend/comfy_nodes/nodes_model_downscale.py b/backend/comfy_nodes/nodes_model_downscale.py deleted file mode 100644 index 48bcc689..00000000 --- a/backend/comfy_nodes/nodes_model_downscale.py +++ /dev/null @@ -1,53 +0,0 @@ -import torch -import comfy.utils - -class PatchModelAddDownscale: - upscale_methods = ["bicubic", "nearest-exact", "bilinear", "area", "bislerp"] - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "block_number": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1}), - "downscale_factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 9.0, "step": 0.001}), - "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), - "end_percent": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}), - "downscale_after_skip": ("BOOLEAN", {"default": True}), - "downscale_method": (s.upscale_methods,), - "upscale_method": (s.upscale_methods,), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "_for_testing" - - def patch(self, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip, downscale_method, upscale_method): - sigma_start = model.model.model_sampling.percent_to_sigma(start_percent) - sigma_end = model.model.model_sampling.percent_to_sigma(end_percent) - - def input_block_patch(h, transformer_options): - if transformer_options["block"][1] == block_number: - sigma = transformer_options["sigmas"][0].item() - if sigma <= sigma_start and sigma >= sigma_end: - h = comfy.utils.common_upscale(h, round(h.shape[-1] * (1.0 / downscale_factor)), round(h.shape[-2] * (1.0 / downscale_factor)), downscale_method, "disabled") - return h - - def output_block_patch(h, hsp, transformer_options): - if h.shape[2] != hsp.shape[2]: - h = comfy.utils.common_upscale(h, hsp.shape[-1], hsp.shape[-2], upscale_method, "disabled") - return h, hsp - - m = model.clone() - if downscale_after_skip: - m.set_model_input_block_patch_after_skip(input_block_patch) - else: - m.set_model_input_block_patch(input_block_patch) - m.set_model_output_block_patch(output_block_patch) - return (m, ) - -NODE_CLASS_MAPPINGS = { - "PatchModelAddDownscale": PatchModelAddDownscale, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - # Sampling - "PatchModelAddDownscale": "PatchModelAddDownscale (Kohya Deep Shrink)", -} diff --git a/backend/comfy_nodes/nodes_model_merging.py b/backend/comfy_nodes/nodes_model_merging.py deleted file mode 100644 index d594cf49..00000000 --- a/backend/comfy_nodes/nodes_model_merging.py +++ /dev/null @@ -1,284 +0,0 @@ -import comfy.sd -import comfy.utils -import comfy.model_base -import comfy.model_management - -import folder_paths -import json -import os - -from comfy.cli_args import args - -class ModelMergeSimple: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model1": ("MODEL",), - "model2": ("MODEL",), - "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "merge" - - CATEGORY = "advanced/model_merging" - - def merge(self, model1, model2, ratio): - m = model1.clone() - kp = model2.get_key_patches("diffusion_model.") - for k in kp: - m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) - return (m, ) - -class ModelSubtract: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model1": ("MODEL",), - "model2": ("MODEL",), - "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "merge" - - CATEGORY = "advanced/model_merging" - - def merge(self, model1, model2, multiplier): - m = model1.clone() - kp = model2.get_key_patches("diffusion_model.") - for k in kp: - m.add_patches({k: kp[k]}, - multiplier, multiplier) - return (m, ) - -class ModelAdd: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model1": ("MODEL",), - "model2": ("MODEL",), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "merge" - - CATEGORY = "advanced/model_merging" - - def merge(self, model1, model2): - m = model1.clone() - kp = model2.get_key_patches("diffusion_model.") - for k in kp: - m.add_patches({k: kp[k]}, 1.0, 1.0) - return (m, ) - - -class CLIPMergeSimple: - @classmethod - def INPUT_TYPES(s): - return {"required": { "clip1": ("CLIP",), - "clip2": ("CLIP",), - "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - }} - RETURN_TYPES = ("CLIP",) - FUNCTION = "merge" - - CATEGORY = "advanced/model_merging" - - def merge(self, clip1, clip2, ratio): - m = clip1.clone() - kp = clip2.get_key_patches() - for k in kp: - if k.endswith(".position_ids") or k.endswith(".logit_scale"): - continue - m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) - return (m, ) - -class ModelMergeBlocks: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model1": ("MODEL",), - "model2": ("MODEL",), - "input": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "middle": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}) - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "merge" - - CATEGORY = "advanced/model_merging" - - def merge(self, model1, model2, **kwargs): - m = model1.clone() - kp = model2.get_key_patches("diffusion_model.") - default_ratio = next(iter(kwargs.values())) - - for k in kp: - ratio = default_ratio - k_unet = k[len("diffusion_model."):] - - last_arg_size = 0 - for arg in kwargs: - if k_unet.startswith(arg) and last_arg_size < len(arg): - ratio = kwargs[arg] - last_arg_size = len(arg) - - m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) - return (m, ) - -def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefix=None, output_dir=None, prompt=None, extra_pnginfo=None): - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir) - prompt_info = "" - if prompt is not None: - prompt_info = json.dumps(prompt) - - metadata = {} - - enable_modelspec = True - if isinstance(model.model, comfy.model_base.SDXL): - metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base" - elif isinstance(model.model, comfy.model_base.SDXLRefiner): - metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner" - else: - enable_modelspec = False - - if enable_modelspec: - metadata["modelspec.sai_model_spec"] = "1.0.0" - metadata["modelspec.implementation"] = "sgm" - metadata["modelspec.title"] = "{} {}".format(filename, counter) - - #TODO: - # "stable-diffusion-v1", "stable-diffusion-v1-inpainting", "stable-diffusion-v2-512", - # "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h", - # "v2-inpainting" - - if model.model.model_type == comfy.model_base.ModelType.EPS: - metadata["modelspec.predict_key"] = "epsilon" - elif model.model.model_type == comfy.model_base.ModelType.V_PREDICTION: - metadata["modelspec.predict_key"] = "v" - - if not args.disable_metadata: - metadata["prompt"] = prompt_info - if extra_pnginfo is not None: - for x in extra_pnginfo: - metadata[x] = json.dumps(extra_pnginfo[x]) - - output_checkpoint = f"{filename}_{counter:05}_.safetensors" - output_checkpoint = os.path.join(full_output_folder, output_checkpoint) - - comfy.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata) - -class CheckpointSave: - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "clip": ("CLIP",), - "vae": ("VAE",), - "filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),}, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} - RETURN_TYPES = () - FUNCTION = "save" - OUTPUT_NODE = True - - CATEGORY = "advanced/model_merging" - - def save(self, model, clip, vae, filename_prefix, prompt=None, extra_pnginfo=None): - save_checkpoint(model, clip=clip, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo) - return {} - -class CLIPSave: - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - - @classmethod - def INPUT_TYPES(s): - return {"required": { "clip": ("CLIP",), - "filename_prefix": ("STRING", {"default": "clip/ComfyUI"}),}, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} - RETURN_TYPES = () - FUNCTION = "save" - OUTPUT_NODE = True - - CATEGORY = "advanced/model_merging" - - def save(self, clip, filename_prefix, prompt=None, extra_pnginfo=None): - prompt_info = "" - if prompt is not None: - prompt_info = json.dumps(prompt) - - metadata = {} - if not args.disable_metadata: - metadata["prompt"] = prompt_info - if extra_pnginfo is not None: - for x in extra_pnginfo: - metadata[x] = json.dumps(extra_pnginfo[x]) - - comfy.model_management.load_models_gpu([clip.load_model()]) - clip_sd = clip.get_sd() - - for prefix in ["clip_l.", "clip_g.", ""]: - k = list(filter(lambda a: a.startswith(prefix), clip_sd.keys())) - current_clip_sd = {} - for x in k: - current_clip_sd[x] = clip_sd.pop(x) - if len(current_clip_sd) == 0: - continue - - p = prefix[:-1] - replace_prefix = {} - filename_prefix_ = filename_prefix - if len(p) > 0: - filename_prefix_ = "{}_{}".format(filename_prefix_, p) - replace_prefix[prefix] = "" - replace_prefix["transformer."] = "" - - full_output_folder, filename, counter, subfolder, filename_prefix_ = folder_paths.get_save_image_path(filename_prefix_, self.output_dir) - - output_checkpoint = f"{filename}_{counter:05}_.safetensors" - output_checkpoint = os.path.join(full_output_folder, output_checkpoint) - - current_clip_sd = comfy.utils.state_dict_prefix_replace(current_clip_sd, replace_prefix) - - comfy.utils.save_torch_file(current_clip_sd, output_checkpoint, metadata=metadata) - return {} - -class VAESave: - def __init__(self): - self.output_dir = folder_paths.get_output_directory() - - @classmethod - def INPUT_TYPES(s): - return {"required": { "vae": ("VAE",), - "filename_prefix": ("STRING", {"default": "vae/ComfyUI_vae"}),}, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} - RETURN_TYPES = () - FUNCTION = "save" - OUTPUT_NODE = True - - CATEGORY = "advanced/model_merging" - - def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None): - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) - prompt_info = "" - if prompt is not None: - prompt_info = json.dumps(prompt) - - metadata = {} - if not args.disable_metadata: - metadata["prompt"] = prompt_info - if extra_pnginfo is not None: - for x in extra_pnginfo: - metadata[x] = json.dumps(extra_pnginfo[x]) - - output_checkpoint = f"{filename}_{counter:05}_.safetensors" - output_checkpoint = os.path.join(full_output_folder, output_checkpoint) - - comfy.utils.save_torch_file(vae.get_sd(), output_checkpoint, metadata=metadata) - return {} - -NODE_CLASS_MAPPINGS = { - "ModelMergeSimple": ModelMergeSimple, - "ModelMergeBlocks": ModelMergeBlocks, - "ModelMergeSubtract": ModelSubtract, - "ModelMergeAdd": ModelAdd, - "CheckpointSave": CheckpointSave, - "CLIPMergeSimple": CLIPMergeSimple, - "CLIPSave": CLIPSave, - "VAESave": VAESave, -} diff --git a/backend/comfy_nodes/nodes_perpneg.py b/backend/comfy_nodes/nodes_perpneg.py deleted file mode 100644 index 64bbc1dc..00000000 --- a/backend/comfy_nodes/nodes_perpneg.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch -import comfy.model_management -import comfy.sample -import comfy.samplers -import comfy.utils - - -class PerpNeg: - @classmethod - def INPUT_TYPES(s): - return {"required": {"model": ("MODEL", ), - "empty_conditioning": ("CONDITIONING", ), - "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "_for_testing" - - def patch(self, model, empty_conditioning, neg_scale): - m = model.clone() - nocond = comfy.sample.convert_cond(empty_conditioning) - - def cfg_function(args): - model = args["model"] - noise_pred_pos = args["cond_denoised"] - noise_pred_neg = args["uncond_denoised"] - cond_scale = args["cond_scale"] - x = args["input"] - sigma = args["sigma"] - model_options = args["model_options"] - nocond_processed = comfy.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative") - - (noise_pred_nocond, _) = comfy.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options) - - pos = noise_pred_pos - noise_pred_nocond - neg = noise_pred_neg - noise_pred_nocond - perp = neg - ((torch.mul(neg, pos).sum())/(torch.norm(pos)**2)) * pos - perp_neg = perp * neg_scale - cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg) - cfg_result = x - cfg_result - return cfg_result - - m.set_model_sampler_cfg_function(cfg_function) - - return (m, ) - - -NODE_CLASS_MAPPINGS = { - "PerpNeg": PerpNeg, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "PerpNeg": "Perp-Neg", -} diff --git a/backend/comfy_nodes/nodes_photomaker.py b/backend/comfy_nodes/nodes_photomaker.py deleted file mode 100644 index 90130142..00000000 --- a/backend/comfy_nodes/nodes_photomaker.py +++ /dev/null @@ -1,187 +0,0 @@ -import torch -import torch.nn as nn -import folder_paths -import comfy.clip_model -import comfy.clip_vision -import comfy.ops - -# code for model from: https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model.py under Apache License Version 2.0 -VISION_CONFIG_DICT = { - "hidden_size": 1024, - "image_size": 224, - "intermediate_size": 4096, - "num_attention_heads": 16, - "num_channels": 3, - "num_hidden_layers": 24, - "patch_size": 14, - "projection_dim": 768, - "hidden_act": "quick_gelu", -} - -class MLP(nn.Module): - def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True, operations=comfy.ops): - super().__init__() - if use_residual: - assert in_dim == out_dim - self.layernorm = operations.LayerNorm(in_dim) - self.fc1 = operations.Linear(in_dim, hidden_dim) - self.fc2 = operations.Linear(hidden_dim, out_dim) - self.use_residual = use_residual - self.act_fn = nn.GELU() - - def forward(self, x): - residual = x - x = self.layernorm(x) - x = self.fc1(x) - x = self.act_fn(x) - x = self.fc2(x) - if self.use_residual: - x = x + residual - return x - - -class FuseModule(nn.Module): - def __init__(self, embed_dim, operations): - super().__init__() - self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False, operations=operations) - self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True, operations=operations) - self.layer_norm = operations.LayerNorm(embed_dim) - - def fuse_fn(self, prompt_embeds, id_embeds): - stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1) - stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds - stacked_id_embeds = self.mlp2(stacked_id_embeds) - stacked_id_embeds = self.layer_norm(stacked_id_embeds) - return stacked_id_embeds - - def forward( - self, - prompt_embeds, - id_embeds, - class_tokens_mask, - ) -> torch.Tensor: - # id_embeds shape: [b, max_num_inputs, 1, 2048] - id_embeds = id_embeds.to(prompt_embeds.dtype) - num_inputs = class_tokens_mask.sum().unsqueeze(0) # TODO: check for training case - batch_size, max_num_inputs = id_embeds.shape[:2] - # seq_length: 77 - seq_length = prompt_embeds.shape[1] - # flat_id_embeds shape: [b*max_num_inputs, 1, 2048] - flat_id_embeds = id_embeds.view( - -1, id_embeds.shape[-2], id_embeds.shape[-1] - ) - # valid_id_mask [b*max_num_inputs] - valid_id_mask = ( - torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :] - < num_inputs[:, None] - ) - valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()] - - prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1]) - class_tokens_mask = class_tokens_mask.view(-1) - valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1]) - # slice out the image token embeddings - image_token_embeds = prompt_embeds[class_tokens_mask] - stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds) - assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}" - prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype)) - updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1) - return updated_prompt_embeds - -class PhotoMakerIDEncoder(comfy.clip_model.CLIPVisionModelProjection): - def __init__(self): - self.load_device = comfy.model_management.text_encoder_device() - offload_device = comfy.model_management.text_encoder_offload_device() - dtype = comfy.model_management.text_encoder_dtype(self.load_device) - - super().__init__(VISION_CONFIG_DICT, dtype, offload_device, comfy.ops.manual_cast) - self.visual_projection_2 = comfy.ops.manual_cast.Linear(1024, 1280, bias=False) - self.fuse_module = FuseModule(2048, comfy.ops.manual_cast) - - def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask): - b, num_inputs, c, h, w = id_pixel_values.shape - id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w) - - shared_id_embeds = self.vision_model(id_pixel_values)[2] - id_embeds = self.visual_projection(shared_id_embeds) - id_embeds_2 = self.visual_projection_2(shared_id_embeds) - - id_embeds = id_embeds.view(b, num_inputs, 1, -1) - id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1) - - id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1) - updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask) - - return updated_prompt_embeds - - -class PhotoMakerLoader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "photomaker_model_name": (folder_paths.get_filename_list("photomaker"), )}} - - RETURN_TYPES = ("PHOTOMAKER",) - FUNCTION = "load_photomaker_model" - - CATEGORY = "_for_testing/photomaker" - - def load_photomaker_model(self, photomaker_model_name): - photomaker_model_path = folder_paths.get_full_path("photomaker", photomaker_model_name) - photomaker_model = PhotoMakerIDEncoder() - data = comfy.utils.load_torch_file(photomaker_model_path, safe_load=True) - if "id_encoder" in data: - data = data["id_encoder"] - photomaker_model.load_state_dict(data) - return (photomaker_model,) - - -class PhotoMakerEncode: - @classmethod - def INPUT_TYPES(s): - return {"required": { "photomaker": ("PHOTOMAKER",), - "image": ("IMAGE",), - "clip": ("CLIP", ), - "text": ("STRING", {"multiline": True, "default": "photograph of photomaker"}), - }} - - RETURN_TYPES = ("CONDITIONING",) - FUNCTION = "apply_photomaker" - - CATEGORY = "_for_testing/photomaker" - - def apply_photomaker(self, photomaker, image, clip, text): - special_token = "photomaker" - pixel_values = comfy.clip_vision.clip_preprocess(image.to(photomaker.load_device)).float() - try: - index = text.split(" ").index(special_token) + 1 - except ValueError: - index = -1 - tokens = clip.tokenize(text, return_word_ids=True) - out_tokens = {} - for k in tokens: - out_tokens[k] = [] - for t in tokens[k]: - f = list(filter(lambda x: x[2] != index, t)) - while len(f) < len(t): - f.append(t[-1]) - out_tokens[k].append(f) - - cond, pooled = clip.encode_from_tokens(out_tokens, return_pooled=True) - - if index > 0: - token_index = index - 1 - num_id_images = 1 - class_tokens_mask = [True if token_index <= i < token_index+num_id_images else False for i in range(77)] - out = photomaker(id_pixel_values=pixel_values.unsqueeze(0), prompt_embeds=cond.to(photomaker.load_device), - class_tokens_mask=torch.tensor(class_tokens_mask, dtype=torch.bool, device=photomaker.load_device).unsqueeze(0)) - else: - out = cond - - return ([[out, {"pooled_output": pooled}]], ) - - -NODE_CLASS_MAPPINGS = { - "PhotoMakerLoader": PhotoMakerLoader, - "PhotoMakerEncode": PhotoMakerEncode, -} - diff --git a/backend/comfy_nodes/nodes_post_processing.py b/backend/comfy_nodes/nodes_post_processing.py deleted file mode 100644 index cb5c7d22..00000000 --- a/backend/comfy_nodes/nodes_post_processing.py +++ /dev/null @@ -1,276 +0,0 @@ -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image -import math - -import comfy.utils - - -class Blend: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image1": ("IMAGE",), - "image2": ("IMAGE",), - "blend_factor": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01 - }), - "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light", "difference"],), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "blend_images" - - CATEGORY = "image/postprocessing" - - def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): - image2 = image2.to(image1.device) - if image1.shape != image2.shape: - image2 = image2.permute(0, 3, 1, 2) - image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center') - image2 = image2.permute(0, 2, 3, 1) - - blended_image = self.blend_mode(image1, image2, blend_mode) - blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor - blended_image = torch.clamp(blended_image, 0, 1) - return (blended_image,) - - def blend_mode(self, img1, img2, mode): - if mode == "normal": - return img2 - elif mode == "multiply": - return img1 * img2 - elif mode == "screen": - return 1 - (1 - img1) * (1 - img2) - elif mode == "overlay": - return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2)) - elif mode == "soft_light": - return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1)) - elif mode == "difference": - return img1 - img2 - else: - raise ValueError(f"Unsupported blend mode: {mode}") - - def g(self, x): - return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) - -def gaussian_kernel(kernel_size: int, sigma: float, device=None): - x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size, device=device), torch.linspace(-1, 1, kernel_size, device=device), indexing="ij") - d = torch.sqrt(x * x + y * y) - g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) - return g / g.sum() - -class Blur: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "blur_radius": ("INT", { - "default": 1, - "min": 1, - "max": 31, - "step": 1 - }), - "sigma": ("FLOAT", { - "default": 1.0, - "min": 0.1, - "max": 10.0, - "step": 0.1 - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "blur" - - CATEGORY = "image/postprocessing" - - def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): - if blur_radius == 0: - return (image,) - - batch_size, height, width, channels = image.shape - - kernel_size = blur_radius * 2 + 1 - kernel = gaussian_kernel(kernel_size, sigma, device=image.device).repeat(channels, 1, 1).unsqueeze(1) - - image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) - padded_image = F.pad(image, (blur_radius,blur_radius,blur_radius,blur_radius), 'reflect') - blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius] - blurred = blurred.permute(0, 2, 3, 1) - - return (blurred,) - -class Quantize: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "colors": ("INT", { - "default": 256, - "min": 1, - "max": 256, - "step": 1 - }), - "dither": (["none", "floyd-steinberg", "bayer-2", "bayer-4", "bayer-8", "bayer-16"],), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "quantize" - - CATEGORY = "image/postprocessing" - - def bayer(im, pal_im, order): - def normalized_bayer_matrix(n): - if n == 0: - return np.zeros((1,1), "float32") - else: - q = 4 ** n - m = q * normalized_bayer_matrix(n - 1) - return np.bmat(((m-1.5, m+0.5), (m+1.5, m-0.5))) / q - - num_colors = len(pal_im.getpalette()) // 3 - spread = 2 * 256 / num_colors - bayer_n = int(math.log2(order)) - bayer_matrix = torch.from_numpy(spread * normalized_bayer_matrix(bayer_n) + 0.5) - - result = torch.from_numpy(np.array(im).astype(np.float32)) - tw = math.ceil(result.shape[0] / bayer_matrix.shape[0]) - th = math.ceil(result.shape[1] / bayer_matrix.shape[1]) - tiled_matrix = bayer_matrix.tile(tw, th).unsqueeze(-1) - result.add_(tiled_matrix[:result.shape[0],:result.shape[1]]).clamp_(0, 255) - result = result.to(dtype=torch.uint8) - - im = Image.fromarray(result.cpu().numpy()) - im = im.quantize(palette=pal_im, dither=Image.Dither.NONE) - return im - - def quantize(self, image: torch.Tensor, colors: int, dither: str): - batch_size, height, width, _ = image.shape - result = torch.zeros_like(image) - - for b in range(batch_size): - im = Image.fromarray((image[b] * 255).to(torch.uint8).numpy(), mode='RGB') - - pal_im = im.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 - - if dither == "none": - quantized_image = im.quantize(palette=pal_im, dither=Image.Dither.NONE) - elif dither == "floyd-steinberg": - quantized_image = im.quantize(palette=pal_im, dither=Image.Dither.FLOYDSTEINBERG) - elif dither.startswith("bayer"): - order = int(dither.split('-')[-1]) - quantized_image = Quantize.bayer(im, pal_im, order) - - quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 - result[b] = quantized_array - - return (result,) - -class Sharpen: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "sharpen_radius": ("INT", { - "default": 1, - "min": 1, - "max": 31, - "step": 1 - }), - "sigma": ("FLOAT", { - "default": 1.0, - "min": 0.1, - "max": 10.0, - "step": 0.1 - }), - "alpha": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 5.0, - "step": 0.1 - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "sharpen" - - CATEGORY = "image/postprocessing" - - def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha: float): - if sharpen_radius == 0: - return (image,) - - batch_size, height, width, channels = image.shape - - kernel_size = sharpen_radius * 2 + 1 - kernel = gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha*10) - center = kernel_size // 2 - kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0 - kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) - - tensor_image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C) - tensor_image = F.pad(tensor_image, (sharpen_radius,sharpen_radius,sharpen_radius,sharpen_radius), 'reflect') - sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)[:,:,sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius] - sharpened = sharpened.permute(0, 2, 3, 1) - - result = torch.clamp(sharpened, 0, 1) - - return (result,) - -class ImageScaleToTotalPixels: - upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] - crop_methods = ["disabled", "center"] - - @classmethod - def INPUT_TYPES(s): - return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,), - "megapixels": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01}), - }} - RETURN_TYPES = ("IMAGE",) - FUNCTION = "upscale" - - CATEGORY = "image/upscaling" - - def upscale(self, image, upscale_method, megapixels): - samples = image.movedim(-1,1) - total = int(megapixels * 1024 * 1024) - - scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) - width = round(samples.shape[3] * scale_by) - height = round(samples.shape[2] * scale_by) - - s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled") - s = s.movedim(1,-1) - return (s,) - -NODE_CLASS_MAPPINGS = { - "ImageBlend": Blend, - "ImageBlur": Blur, - "ImageQuantize": Quantize, - "ImageSharpen": Sharpen, - "ImageScaleToTotalPixels": ImageScaleToTotalPixels, -} diff --git a/backend/comfy_nodes/nodes_rebatch.py b/backend/comfy_nodes/nodes_rebatch.py deleted file mode 100644 index 3010fbd4..00000000 --- a/backend/comfy_nodes/nodes_rebatch.py +++ /dev/null @@ -1,138 +0,0 @@ -import torch - -class LatentRebatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { "latents": ("LATENT",), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - }} - RETURN_TYPES = ("LATENT",) - INPUT_IS_LIST = True - OUTPUT_IS_LIST = (True, ) - - FUNCTION = "rebatch" - - CATEGORY = "latent/batch" - - @staticmethod - def get_batch(latents, list_ind, offset): - '''prepare a batch out of the list of latents''' - samples = latents[list_ind]['samples'] - shape = samples.shape - mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu') - if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]: - torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear") - if mask.shape[0] < samples.shape[0]: - mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]] - if 'batch_index' in latents[list_ind]: - batch_inds = latents[list_ind]['batch_index'] - else: - batch_inds = [x+offset for x in range(shape[0])] - return samples, mask, batch_inds - - @staticmethod - def get_slices(indexable, num, batch_size): - '''divides an indexable object into num slices of length batch_size, and a remainder''' - slices = [] - for i in range(num): - slices.append(indexable[i*batch_size:(i+1)*batch_size]) - if num * batch_size < len(indexable): - return slices, indexable[num * batch_size:] - else: - return slices, None - - @staticmethod - def slice_batch(batch, num, batch_size): - result = [LatentRebatch.get_slices(x, num, batch_size) for x in batch] - return list(zip(*result)) - - @staticmethod - def cat_batch(batch1, batch2): - if batch1[0] is None: - return batch2 - result = [torch.cat((b1, b2)) if torch.is_tensor(b1) else b1 + b2 for b1, b2 in zip(batch1, batch2)] - return result - - def rebatch(self, latents, batch_size): - batch_size = batch_size[0] - - output_list = [] - current_batch = (None, None, None) - processed = 0 - - for i in range(len(latents)): - # fetch new entry of list - #samples, masks, indices = self.get_batch(latents, i) - next_batch = self.get_batch(latents, i, processed) - processed += len(next_batch[2]) - # set to current if current is None - if current_batch[0] is None: - current_batch = next_batch - # add previous to list if dimensions do not match - elif next_batch[0].shape[-1] != current_batch[0].shape[-1] or next_batch[0].shape[-2] != current_batch[0].shape[-2]: - sliced, _ = self.slice_batch(current_batch, 1, batch_size) - output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]}) - current_batch = next_batch - # cat if everything checks out - else: - current_batch = self.cat_batch(current_batch, next_batch) - - # add to list if dimensions gone above target batch size - if current_batch[0].shape[0] > batch_size: - num = current_batch[0].shape[0] // batch_size - sliced, remainder = self.slice_batch(current_batch, num, batch_size) - - for i in range(num): - output_list.append({'samples': sliced[0][i], 'noise_mask': sliced[1][i], 'batch_index': sliced[2][i]}) - - current_batch = remainder - - #add remainder - if current_batch[0] is not None: - sliced, _ = self.slice_batch(current_batch, 1, batch_size) - output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]}) - - #get rid of empty masks - for s in output_list: - if s['noise_mask'].mean() == 1.0: - del s['noise_mask'] - - return (output_list,) - -class ImageRebatch: - @classmethod - def INPUT_TYPES(s): - return {"required": { "images": ("IMAGE",), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - }} - RETURN_TYPES = ("IMAGE",) - INPUT_IS_LIST = True - OUTPUT_IS_LIST = (True, ) - - FUNCTION = "rebatch" - - CATEGORY = "image/batch" - - def rebatch(self, images, batch_size): - batch_size = batch_size[0] - - output_list = [] - all_images = [] - for img in images: - for i in range(img.shape[0]): - all_images.append(img[i:i+1]) - - for i in range(0, len(all_images), batch_size): - output_list.append(torch.cat(all_images[i:i+batch_size], dim=0)) - - return (output_list,) - -NODE_CLASS_MAPPINGS = { - "RebatchLatents": LatentRebatch, - "RebatchImages": ImageRebatch, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "RebatchLatents": "Rebatch Latents", - "RebatchImages": "Rebatch Images", -} diff --git a/backend/comfy_nodes/nodes_sag.py b/backend/comfy_nodes/nodes_sag.py deleted file mode 100644 index bbd38080..00000000 --- a/backend/comfy_nodes/nodes_sag.py +++ /dev/null @@ -1,170 +0,0 @@ -import torch -from torch import einsum -import torch.nn.functional as F -import math - -from einops import rearrange, repeat -import os -from comfy.ldm.modules.attention import optimized_attention, _ATTN_PRECISION -import comfy.samplers - -# from comfy/ldm/modules/attention.py -# but modified to return attention scores as well as output -def attention_basic_with_sim(q, k, v, heads, mask=None): - b, _, dim_head = q.shape - dim_head //= heads - scale = dim_head ** -0.5 - - h = heads - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, -1, heads, dim_head) - .permute(0, 2, 1, 3) - .reshape(b * heads, -1, dim_head) - .contiguous(), - (q, k, v), - ) - - # force cast to fp32 to avoid overflowing - if _ATTN_PRECISION =="fp32": - sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale - else: - sim = einsum('b i d, b j d -> b i j', q, k) * scale - - del q, k - - if mask is not None: - mask = rearrange(mask, 'b ... -> b (...)') - max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b j -> (b h) () j', h=h) - sim.masked_fill_(~mask, max_neg_value) - - # attention, what we cannot get enough of - sim = sim.softmax(dim=-1) - - out = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v) - out = ( - out.unsqueeze(0) - .reshape(b, heads, -1, dim_head) - .permute(0, 2, 1, 3) - .reshape(b, -1, heads * dim_head) - ) - return (out, sim) - -def create_blur_map(x0, attn, sigma=3.0, threshold=1.0): - # reshape and GAP the attention map - _, hw1, hw2 = attn.shape - b, _, lh, lw = x0.shape - attn = attn.reshape(b, -1, hw1, hw2) - # Global Average Pool - mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold - ratio = 2**(math.ceil(math.sqrt(lh * lw / hw1)) - 1).bit_length() - mid_shape = [math.ceil(lh / ratio), math.ceil(lw / ratio)] - - # Reshape - mask = ( - mask.reshape(b, *mid_shape) - .unsqueeze(1) - .type(attn.dtype) - ) - # Upsample - mask = F.interpolate(mask, (lh, lw)) - - blurred = gaussian_blur_2d(x0, kernel_size=9, sigma=sigma) - blurred = blurred * mask + x0 * (1 - mask) - return blurred - -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 - -class SelfAttentionGuidance: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 5.0, "step": 0.1}), - "blur_sigma": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "_for_testing" - - def patch(self, model, scale, blur_sigma): - m = model.clone() - - attn_scores = None - - # TODO: make this work properly with chunked batches - # currently, we can only save the attn from one UNet call - def attn_and_record(q, k, v, extra_options): - nonlocal attn_scores - # if uncond, save the attention scores - heads = extra_options["n_heads"] - cond_or_uncond = extra_options["cond_or_uncond"] - b = q.shape[0] // len(cond_or_uncond) - if 1 in cond_or_uncond: - uncond_index = cond_or_uncond.index(1) - # do the entire attention operation, but save the attention scores to attn_scores - (out, sim) = attention_basic_with_sim(q, k, v, heads=heads) - # when using a higher batch size, I BELIEVE the result batch dimension is [uc1, ... ucn, c1, ... cn] - n_slices = heads * b - attn_scores = sim[n_slices * uncond_index:n_slices * (uncond_index+1)] - return out - else: - return optimized_attention(q, k, v, heads=heads) - - def post_cfg_function(args): - nonlocal attn_scores - uncond_attn = attn_scores - - sag_scale = scale - sag_sigma = blur_sigma - sag_threshold = 1.0 - model = args["model"] - uncond_pred = args["uncond_denoised"] - uncond = args["uncond"] - cfg_result = args["denoised"] - sigma = args["sigma"] - model_options = args["model_options"] - x = args["input"] - if min(cfg_result.shape[2:]) <= 4: #skip when too small to add padding - return cfg_result - - # create the adversarially blurred image - degraded = create_blur_map(uncond_pred, uncond_attn, sag_sigma, sag_threshold) - degraded_noised = degraded + x - uncond_pred - # call into the UNet - (sag, _) = comfy.samplers.calc_cond_uncond_batch(model, uncond, None, degraded_noised, sigma, model_options) - return cfg_result + (degraded - sag) * sag_scale - - m.set_model_sampler_post_cfg_function(post_cfg_function, disable_cfg1_optimization=True) - - # from diffusers: - # unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch - m.set_model_attn1_replace(attn_and_record, "middle", 0, 0) - - return (m, ) - -NODE_CLASS_MAPPINGS = { - "SelfAttentionGuidance": SelfAttentionGuidance, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "SelfAttentionGuidance": "Self-Attention Guidance", -} diff --git a/backend/comfy_nodes/nodes_sdupscale.py b/backend/comfy_nodes/nodes_sdupscale.py deleted file mode 100644 index 28c1cb0f..00000000 --- a/backend/comfy_nodes/nodes_sdupscale.py +++ /dev/null @@ -1,47 +0,0 @@ -import torch -import nodes -import comfy.utils - -class SD_4XUpscale_Conditioning: - @classmethod - def INPUT_TYPES(s): - return {"required": { "images": ("IMAGE",), - "positive": ("CONDITIONING",), - "negative": ("CONDITIONING",), - "scale_ratio": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 10.0, "step": 0.01}), - "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), - }} - RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") - RETURN_NAMES = ("positive", "negative", "latent") - - FUNCTION = "encode" - - CATEGORY = "conditioning/upscale_diffusion" - - def encode(self, images, positive, negative, scale_ratio, noise_augmentation): - width = max(1, round(images.shape[-2] * scale_ratio)) - height = max(1, round(images.shape[-3] * scale_ratio)) - - pixels = comfy.utils.common_upscale((images.movedim(-1,1) * 2.0) - 1.0, width // 4, height // 4, "bilinear", "center") - - out_cp = [] - out_cn = [] - - for t in positive: - n = [t[0], t[1].copy()] - n[1]['concat_image'] = pixels - n[1]['noise_augmentation'] = noise_augmentation - out_cp.append(n) - - for t in negative: - n = [t[0], t[1].copy()] - n[1]['concat_image'] = pixels - n[1]['noise_augmentation'] = noise_augmentation - out_cn.append(n) - - latent = torch.zeros([images.shape[0], 4, height // 4, width // 4]) - return (out_cp, out_cn, {"samples":latent}) - -NODE_CLASS_MAPPINGS = { - "SD_4XUpscale_Conditioning": SD_4XUpscale_Conditioning, -} diff --git a/backend/comfy_nodes/nodes_stable3d.py b/backend/comfy_nodes/nodes_stable3d.py deleted file mode 100644 index 4375d8f9..00000000 --- a/backend/comfy_nodes/nodes_stable3d.py +++ /dev/null @@ -1,102 +0,0 @@ -import torch -import nodes -import comfy.utils - -def camera_embeddings(elevation, azimuth): - elevation = torch.as_tensor([elevation]) - azimuth = torch.as_tensor([azimuth]) - embeddings = torch.stack( - [ - torch.deg2rad( - (90 - elevation) - (90) - ), # Zero123 polar is 90-elevation - torch.sin(torch.deg2rad(azimuth)), - torch.cos(torch.deg2rad(azimuth)), - torch.deg2rad( - 90 - torch.full_like(elevation, 0) - ), - ], dim=-1).unsqueeze(1) - - return embeddings - - -class StableZero123_Conditioning: - @classmethod - def INPUT_TYPES(s): - return {"required": { "clip_vision": ("CLIP_VISION",), - "init_image": ("IMAGE",), - "vae": ("VAE",), - "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - }} - RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") - RETURN_NAMES = ("positive", "negative", "latent") - - FUNCTION = "encode" - - CATEGORY = "conditioning/3d_models" - - def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth): - output = clip_vision.encode_image(init_image) - pooled = output.image_embeds.unsqueeze(0) - pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) - encode_pixels = pixels[:,:,:,:3] - t = vae.encode(encode_pixels) - cam_embeds = camera_embeddings(elevation, azimuth) - cond = torch.cat([pooled, cam_embeds.to(pooled.device).repeat((pooled.shape[0], 1, 1))], dim=-1) - - positive = [[cond, {"concat_latent_image": t}]] - negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]] - latent = torch.zeros([batch_size, 4, height // 8, width // 8]) - return (positive, negative, {"samples":latent}) - -class StableZero123_Conditioning_Batched: - @classmethod - def INPUT_TYPES(s): - return {"required": { "clip_vision": ("CLIP_VISION",), - "init_image": ("IMAGE",), - "vae": ("VAE",), - "width": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "height": ("INT", {"default": 256, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}), - }} - RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") - RETURN_NAMES = ("positive", "negative", "latent") - - FUNCTION = "encode" - - CATEGORY = "conditioning/3d_models" - - def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth, elevation_batch_increment, azimuth_batch_increment): - output = clip_vision.encode_image(init_image) - pooled = output.image_embeds.unsqueeze(0) - pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) - encode_pixels = pixels[:,:,:,:3] - t = vae.encode(encode_pixels) - - cam_embeds = [] - for i in range(batch_size): - cam_embeds.append(camera_embeddings(elevation, azimuth)) - elevation += elevation_batch_increment - azimuth += azimuth_batch_increment - - cam_embeds = torch.cat(cam_embeds, dim=0) - cond = torch.cat([comfy.utils.repeat_to_batch_size(pooled, batch_size), cam_embeds], dim=-1) - - positive = [[cond, {"concat_latent_image": t}]] - negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]] - latent = torch.zeros([batch_size, 4, height // 8, width // 8]) - return (positive, negative, {"samples":latent, "batch_index": [0] * batch_size}) - - -NODE_CLASS_MAPPINGS = { - "StableZero123_Conditioning": StableZero123_Conditioning, - "StableZero123_Conditioning_Batched": StableZero123_Conditioning_Batched, -} diff --git a/backend/comfy_nodes/nodes_stable_cascade.py b/backend/comfy_nodes/nodes_stable_cascade.py deleted file mode 100644 index b795d008..00000000 --- a/backend/comfy_nodes/nodes_stable_cascade.py +++ /dev/null @@ -1,109 +0,0 @@ -""" - This file is part of ComfyUI. - Copyright (C) 2024 Stability AI - - 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 . -""" - -import torch -import nodes -import comfy.utils - - -class StableCascade_EmptyLatentImage: - def __init__(self, device="cpu"): - self.device = device - - @classmethod - def INPUT_TYPES(s): - return {"required": { - "width": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), - "height": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), - "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}) - }} - RETURN_TYPES = ("LATENT", "LATENT") - RETURN_NAMES = ("stage_c", "stage_b") - FUNCTION = "generate" - - CATEGORY = "_for_testing/stable_cascade" - - def generate(self, width, height, compression, batch_size=1): - c_latent = torch.zeros([batch_size, 16, height // compression, width // compression]) - b_latent = torch.zeros([batch_size, 4, height // 4, width // 4]) - return ({ - "samples": c_latent, - }, { - "samples": b_latent, - }) - -class StableCascade_StageC_VAEEncode: - def __init__(self, device="cpu"): - self.device = device - - @classmethod - def INPUT_TYPES(s): - return {"required": { - "image": ("IMAGE",), - "vae": ("VAE", ), - "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), - }} - RETURN_TYPES = ("LATENT", "LATENT") - RETURN_NAMES = ("stage_c", "stage_b") - FUNCTION = "generate" - - CATEGORY = "_for_testing/stable_cascade" - - def generate(self, image, vae, compression): - width = image.shape[-2] - height = image.shape[-3] - out_width = (width // compression) * vae.downscale_ratio - out_height = (height // compression) * vae.downscale_ratio - - s = comfy.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1) - - c_latent = vae.encode(s[:,:,:,:3]) - b_latent = torch.zeros([c_latent.shape[0], 4, height // 4, width // 4]) - return ({ - "samples": c_latent, - }, { - "samples": b_latent, - }) - -class StableCascade_StageB_Conditioning: - @classmethod - def INPUT_TYPES(s): - return {"required": { "conditioning": ("CONDITIONING",), - "stage_c": ("LATENT",), - }} - RETURN_TYPES = ("CONDITIONING",) - - FUNCTION = "set_prior" - - CATEGORY = "_for_testing/stable_cascade" - - def set_prior(self, conditioning, stage_c): - c = [] - for t in conditioning: - d = t[1].copy() - d['stable_cascade_prior'] = stage_c['samples'] - n = [t[0], d] - c.append(n) - return (c, ) - -NODE_CLASS_MAPPINGS = { - "StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage, - "StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning, - "StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode, -} diff --git a/backend/comfy_nodes/nodes_tomesd.py b/backend/comfy_nodes/nodes_tomesd.py deleted file mode 100644 index df048506..00000000 --- a/backend/comfy_nodes/nodes_tomesd.py +++ /dev/null @@ -1,177 +0,0 @@ -#Taken from: https://github.com/dbolya/tomesd - -import torch -from typing import Tuple, Callable -import math - -def do_nothing(x: torch.Tensor, mode:str=None): - return x - - -def mps_gather_workaround(input, dim, index): - if input.shape[-1] == 1: - return torch.gather( - input.unsqueeze(-1), - dim - 1 if dim < 0 else dim, - index.unsqueeze(-1) - ).squeeze(-1) - else: - return torch.gather(input, dim, index) - - -def bipartite_soft_matching_random2d(metric: torch.Tensor, - w: int, h: int, sx: int, sy: int, r: int, - no_rand: bool = False) -> Tuple[Callable, Callable]: - """ - Partitions the tokens into src and dst and merges r tokens from src to dst. - Dst tokens are partitioned by choosing one randomy in each (sx, sy) region. - Args: - - metric [B, N, C]: metric to use for similarity - - w: image width in tokens - - h: image height in tokens - - sx: stride in the x dimension for dst, must divide w - - sy: stride in the y dimension for dst, must divide h - - r: number of tokens to remove (by merging) - - no_rand: if true, disable randomness (use top left corner only) - """ - B, N, _ = metric.shape - - if r <= 0 or w == 1 or h == 1: - return do_nothing, do_nothing - - gather = mps_gather_workaround if metric.device.type == "mps" else torch.gather - - with torch.no_grad(): - - hsy, wsx = h // sy, w // sx - - # For each sy by sx kernel, randomly assign one token to be dst and the rest src - if no_rand: - rand_idx = torch.zeros(hsy, wsx, 1, device=metric.device, dtype=torch.int64) - else: - rand_idx = torch.randint(sy*sx, size=(hsy, wsx, 1), device=metric.device) - - # The image might not divide sx and sy, so we need to work on a view of the top left if the idx buffer instead - idx_buffer_view = torch.zeros(hsy, wsx, sy*sx, device=metric.device, dtype=torch.int64) - idx_buffer_view.scatter_(dim=2, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=rand_idx.dtype)) - idx_buffer_view = idx_buffer_view.view(hsy, wsx, sy, sx).transpose(1, 2).reshape(hsy * sy, wsx * sx) - - # Image is not divisible by sx or sy so we need to move it into a new buffer - if (hsy * sy) < h or (wsx * sx) < w: - idx_buffer = torch.zeros(h, w, device=metric.device, dtype=torch.int64) - idx_buffer[:(hsy * sy), :(wsx * sx)] = idx_buffer_view - else: - idx_buffer = idx_buffer_view - - # We set dst tokens to be -1 and src to be 0, so an argsort gives us dst|src indices - rand_idx = idx_buffer.reshape(1, -1, 1).argsort(dim=1) - - # We're finished with these - del idx_buffer, idx_buffer_view - - # rand_idx is currently dst|src, so split them - num_dst = hsy * wsx - a_idx = rand_idx[:, num_dst:, :] # src - b_idx = rand_idx[:, :num_dst, :] # dst - - def split(x): - C = x.shape[-1] - src = gather(x, dim=1, index=a_idx.expand(B, N - num_dst, C)) - dst = gather(x, dim=1, index=b_idx.expand(B, num_dst, C)) - return src, dst - - # Cosine similarity between A and B - metric = metric / metric.norm(dim=-1, keepdim=True) - a, b = split(metric) - scores = a @ b.transpose(-1, -2) - - # Can't reduce more than the # tokens in src - r = min(a.shape[1], r) - - # Find the most similar greedily - node_max, node_idx = scores.max(dim=-1) - edge_idx = node_max.argsort(dim=-1, descending=True)[..., None] - - unm_idx = edge_idx[..., r:, :] # Unmerged Tokens - src_idx = edge_idx[..., :r, :] # Merged Tokens - dst_idx = gather(node_idx[..., None], dim=-2, index=src_idx) - - def merge(x: torch.Tensor, mode="mean") -> torch.Tensor: - src, dst = split(x) - n, t1, c = src.shape - - unm = gather(src, dim=-2, index=unm_idx.expand(n, t1 - r, c)) - src = gather(src, dim=-2, index=src_idx.expand(n, r, c)) - dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode) - - return torch.cat([unm, dst], dim=1) - - def unmerge(x: torch.Tensor) -> torch.Tensor: - unm_len = unm_idx.shape[1] - unm, dst = x[..., :unm_len, :], x[..., unm_len:, :] - _, _, c = unm.shape - - src = gather(dst, dim=-2, index=dst_idx.expand(B, r, c)) - - # Combine back to the original shape - out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype) - out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst) - out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=unm_idx).expand(B, unm_len, c), src=unm) - out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=src_idx).expand(B, r, c), src=src) - - return out - - return merge, unmerge - - -def get_functions(x, ratio, original_shape): - b, c, original_h, original_w = original_shape - original_tokens = original_h * original_w - downsample = int(math.ceil(math.sqrt(original_tokens // x.shape[1]))) - stride_x = 2 - stride_y = 2 - max_downsample = 1 - - if downsample <= max_downsample: - w = int(math.ceil(original_w / downsample)) - h = int(math.ceil(original_h / downsample)) - r = int(x.shape[1] * ratio) - no_rand = False - m, u = bipartite_soft_matching_random2d(x, w, h, stride_x, stride_y, r, no_rand) - return m, u - - nothing = lambda y: y - return nothing, nothing - - - -class TomePatchModel: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "_for_testing" - - def patch(self, model, ratio): - self.u = None - def tomesd_m(q, k, v, extra_options): - #NOTE: In the reference code get_functions takes x (input of the transformer block) as the argument instead of q - #however from my basic testing it seems that using q instead gives better results - m, self.u = get_functions(q, ratio, extra_options["original_shape"]) - return m(q), k, v - def tomesd_u(n, extra_options): - return self.u(n) - - m = model.clone() - m.set_model_attn1_patch(tomesd_m) - m.set_model_attn1_output_patch(tomesd_u) - return (m, ) - - -NODE_CLASS_MAPPINGS = { - "TomePatchModel": TomePatchModel, -} diff --git a/backend/comfy_nodes/nodes_upscale_model.py b/backend/comfy_nodes/nodes_upscale_model.py deleted file mode 100644 index 2b5e49a5..00000000 --- a/backend/comfy_nodes/nodes_upscale_model.py +++ /dev/null @@ -1,66 +0,0 @@ -import os -from comfy_extras.chainner_models import model_loading -from comfy import model_management -import torch -import comfy.utils -import folder_paths - -class UpscaleModelLoader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ), - }} - RETURN_TYPES = ("UPSCALE_MODEL",) - FUNCTION = "load_model" - - CATEGORY = "loaders" - - def load_model(self, model_name): - model_path = folder_paths.get_full_path("upscale_models", model_name) - sd = comfy.utils.load_torch_file(model_path, safe_load=True) - if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd: - sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""}) - out = model_loading.load_state_dict(sd).eval() - return (out, ) - - -class ImageUpscaleWithModel: - @classmethod - def INPUT_TYPES(s): - return {"required": { "upscale_model": ("UPSCALE_MODEL",), - "image": ("IMAGE",), - }} - RETURN_TYPES = ("IMAGE",) - FUNCTION = "upscale" - - CATEGORY = "image/upscaling" - - def upscale(self, upscale_model, image): - device = model_management.get_torch_device() - upscale_model.to(device) - in_img = image.movedim(-1,-3).to(device) - free_memory = model_management.get_free_memory(device) - - tile = 512 - overlap = 32 - - oom = True - while oom: - try: - steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) - pbar = comfy.utils.ProgressBar(steps) - s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar) - oom = False - except model_management.OOM_EXCEPTION as e: - tile //= 2 - if tile < 128: - raise e - - upscale_model.cpu() - s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0) - return (s,) - -NODE_CLASS_MAPPINGS = { - "UpscaleModelLoader": UpscaleModelLoader, - "ImageUpscaleWithModel": ImageUpscaleWithModel -} diff --git a/backend/comfy_nodes/nodes_video_model.py b/backend/comfy_nodes/nodes_video_model.py deleted file mode 100644 index a5262565..00000000 --- a/backend/comfy_nodes/nodes_video_model.py +++ /dev/null @@ -1,106 +0,0 @@ -import nodes -import torch -import comfy.utils -import comfy.sd -import folder_paths -import comfy_extras.nodes_model_merging - - -class ImageOnlyCheckpointLoader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), - }} - RETURN_TYPES = ("MODEL", "CLIP_VISION", "VAE") - FUNCTION = "load_checkpoint" - - CATEGORY = "loaders/video_models" - - def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) - out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) - return (out[0], out[3], out[2]) - - -class SVD_img2vid_Conditioning: - @classmethod - def INPUT_TYPES(s): - return {"required": { "clip_vision": ("CLIP_VISION",), - "init_image": ("IMAGE",), - "vae": ("VAE",), - "width": ("INT", {"default": 1024, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "height": ("INT", {"default": 576, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 8}), - "video_frames": ("INT", {"default": 14, "min": 1, "max": 4096}), - "motion_bucket_id": ("INT", {"default": 127, "min": 1, "max": 1023}), - "fps": ("INT", {"default": 6, "min": 1, "max": 1024}), - "augmentation_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}) - }} - RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") - RETURN_NAMES = ("positive", "negative", "latent") - - FUNCTION = "encode" - - CATEGORY = "conditioning/video_models" - - def encode(self, clip_vision, init_image, vae, width, height, video_frames, motion_bucket_id, fps, augmentation_level): - output = clip_vision.encode_image(init_image) - pooled = output.image_embeds.unsqueeze(0) - pixels = comfy.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1) - encode_pixels = pixels[:,:,:,:3] - if augmentation_level > 0: - encode_pixels += torch.randn_like(pixels) * augmentation_level - t = vae.encode(encode_pixels) - positive = [[pooled, {"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level, "concat_latent_image": t}]] - negative = [[torch.zeros_like(pooled), {"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level, "concat_latent_image": torch.zeros_like(t)}]] - latent = torch.zeros([video_frames, 4, height // 8, width // 8]) - return (positive, negative, {"samples":latent}) - -class VideoLinearCFGGuidance: - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "min_cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "patch" - - CATEGORY = "sampling/video_models" - - def patch(self, model, min_cfg): - def linear_cfg(args): - cond = args["cond"] - uncond = args["uncond"] - cond_scale = args["cond_scale"] - - scale = torch.linspace(min_cfg, cond_scale, cond.shape[0], device=cond.device).reshape((cond.shape[0], 1, 1, 1)) - return uncond + scale * (cond - uncond) - - m = model.clone() - m.set_model_sampler_cfg_function(linear_cfg) - return (m, ) - -class ImageOnlyCheckpointSave(comfy_extras.nodes_model_merging.CheckpointSave): - CATEGORY = "_for_testing" - - @classmethod - def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "clip_vision": ("CLIP_VISION",), - "vae": ("VAE",), - "filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),}, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} - - def save(self, model, clip_vision, vae, filename_prefix, prompt=None, extra_pnginfo=None): - comfy_extras.nodes_model_merging.save_checkpoint(model, clip_vision=clip_vision, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo) - return {} - -NODE_CLASS_MAPPINGS = { - "ImageOnlyCheckpointLoader": ImageOnlyCheckpointLoader, - "SVD_img2vid_Conditioning": SVD_img2vid_Conditioning, - "VideoLinearCFGGuidance": VideoLinearCFGGuidance, - "ImageOnlyCheckpointSave": ImageOnlyCheckpointSave, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "ImageOnlyCheckpointLoader": "Image Only Checkpoint Loader (img2vid model)", -} diff --git a/backend/compose/local/django/Dockerfile b/backend/compose/local/django/Dockerfile deleted file mode 100644 index 52ec9ee1..00000000 --- a/backend/compose/local/django/Dockerfile +++ /dev/null @@ -1,91 +0,0 @@ -# define an alias for the specific python version used in this file. -FROM python:3.11.6-slim-bullseye as python - -# Python build stage -FROM python as python-build-stage - -ARG BUILD_ENVIRONMENT=local - -# Install apt packages -RUN apt-get update && apt-get install --no-install-recommends -y \ - # dependencies for building Python packages - build-essential \ - # psycopg2 dependencies - libpq-dev - -# Requirements are installed here to ensure they will be cached. -COPY ./requirements . - -# Create Python Dependency and Sub-Dependency Wheels. -RUN pip wheel --wheel-dir /usr/src/app/wheels \ - -r ${BUILD_ENVIRONMENT}.txt - - -# Python 'run' stage -FROM python as python-run-stage - -ARG BUILD_ENVIRONMENT=local -ARG APP_HOME=/app - -ENV PYTHONUNBUFFERED 1 -ENV PYTHONDONTWRITEBYTECODE 1 -ENV BUILD_ENV ${BUILD_ENVIRONMENT} - -WORKDIR ${APP_HOME} - - -# devcontainer dependencies and utils -RUN apt-get update && apt-get install --no-install-recommends -y \ - sudo git bash-completion nano ssh - -# Create devcontainer user and add it to sudoers -RUN groupadd --gid 1000 dev-user \ - && useradd --uid 1000 --gid dev-user --shell /bin/bash --create-home dev-user \ - && echo dev-user ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/dev-user \ - && chmod 0440 /etc/sudoers.d/dev-user - - -# Install required system dependencies -RUN apt-get update && apt-get install --no-install-recommends -y \ - # psycopg2 dependencies - libpq-dev \ - # Translations dependencies - gettext \ - # cleaning up unused files - && apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false \ - && rm -rf /var/lib/apt/lists/* - -# All absolute dir copies ignore workdir instruction. All relative dir copies are wrt to the workdir instruction -# copy python dependency wheels from python-build-stage -COPY --from=python-build-stage /usr/src/app/wheels /wheels/ - -# use wheels to install python dependencies -RUN pip install --no-cache-dir --no-index --find-links=/wheels/ /wheels/* \ - && rm -rf /wheels/ - -COPY ./compose/production/django/entrypoint /entrypoint -RUN sed -i 's/\r$//g' /entrypoint -RUN chmod +x /entrypoint - -COPY ./compose/local/django/start /start -RUN sed -i 's/\r$//g' /start -RUN chmod +x /start - - -COPY ./compose/local/django/celery/worker/start /start-celeryworker -RUN sed -i 's/\r$//g' /start-celeryworker -RUN chmod +x /start-celeryworker - -COPY ./compose/local/django/celery/beat/start /start-celerybeat -RUN sed -i 's/\r$//g' /start-celerybeat -RUN chmod +x /start-celerybeat - -COPY ./compose/local/django/celery/flower/start /start-flower -RUN sed -i 's/\r$//g' /start-flower -RUN chmod +x /start-flower - - -# copy application code to WORKDIR -COPY . ${APP_HOME} - -ENTRYPOINT ["/entrypoint"] diff --git a/backend/compose/local/django/celery/beat/start b/backend/compose/local/django/celery/beat/start deleted file mode 100644 index 8adc4891..00000000 --- a/backend/compose/local/django/celery/beat/start +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o nounset - - -rm -f './celerybeat.pid' -exec watchfiles --filter python celery.__main__.main --args '-A config.celery_app beat -l INFO' diff --git a/backend/compose/local/django/celery/flower/start b/backend/compose/local/django/celery/flower/start deleted file mode 100644 index b4783d2f..00000000 --- a/backend/compose/local/django/celery/flower/start +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o nounset - -exec watchfiles --filter python celery.__main__.main \ - --args \ - "-A config.celery_app -b \"${CELERY_BROKER_URL}\" flower --basic_auth=\"${CELERY_FLOWER_USER}:${CELERY_FLOWER_PASSWORD}\"" diff --git a/backend/compose/local/django/celery/worker/start b/backend/compose/local/django/celery/worker/start deleted file mode 100644 index 183a8015..00000000 --- a/backend/compose/local/django/celery/worker/start +++ /dev/null @@ -1,7 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o nounset - - -exec watchfiles --filter python celery.__main__.main --args '-A config.celery_app worker -l INFO' diff --git a/backend/compose/local/django/start b/backend/compose/local/django/start deleted file mode 100644 index 1549d133..00000000 --- a/backend/compose/local/django/start +++ /dev/null @@ -1,9 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - - -python manage.py migrate -exec uvicorn config.asgi:application --host 0.0.0.0 --reload --reload-include '*.html' diff --git a/backend/compose/local/docs/Dockerfile b/backend/compose/local/docs/Dockerfile deleted file mode 100644 index 80a086ab..00000000 --- a/backend/compose/local/docs/Dockerfile +++ /dev/null @@ -1,62 +0,0 @@ -# define an alias for the specific python version used in this file. -FROM python:3.11.6-slim-bullseye as python - - -# Python build stage -FROM python as python-build-stage - -ENV PYTHONDONTWRITEBYTECODE 1 - -RUN apt-get update && apt-get install --no-install-recommends -y \ - # dependencies for building Python packages - build-essential \ - # psycopg2 dependencies - libpq-dev \ - # cleaning up unused files - && apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false \ - && rm -rf /var/lib/apt/lists/* - -# Requirements are installed here to ensure they will be cached. -COPY ./requirements /requirements - -# create python dependency wheels -RUN pip wheel --no-cache-dir --wheel-dir /usr/src/app/wheels \ - -r /requirements/local.txt -r /requirements/production.txt \ - && rm -rf /requirements - - -# Python 'run' stage -FROM python as python-run-stage - -ARG BUILD_ENVIRONMENT -ENV PYTHONUNBUFFERED 1 -ENV PYTHONDONTWRITEBYTECODE 1 - -RUN apt-get update && apt-get install --no-install-recommends -y \ - # To run the Makefile - make \ - # psycopg2 dependencies - libpq-dev \ - # Translations dependencies - gettext \ - # Uncomment below lines to enable Sphinx output to latex and pdf - # texlive-latex-recommended \ - # texlive-fonts-recommended \ - # texlive-latex-extra \ - # latexmk \ - # cleaning up unused files - && apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false \ - && rm -rf /var/lib/apt/lists/* - -# copy python dependency wheels from python-build-stage -COPY --from=python-build-stage /usr/src/app/wheels /wheels - -# use wheels to install python dependencies -RUN pip install --no-cache /wheels/* \ - && rm -rf /wheels - -COPY ./compose/local/docs/start /start-docs -RUN sed -i 's/\r$//g' /start-docs -RUN chmod +x /start-docs - -WORKDIR /docs diff --git a/backend/compose/local/docs/start b/backend/compose/local/docs/start deleted file mode 100644 index 96a94f56..00000000 --- a/backend/compose/local/docs/start +++ /dev/null @@ -1,7 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - -exec make livehtml diff --git a/backend/compose/production/aws/Dockerfile b/backend/compose/production/aws/Dockerfile deleted file mode 100644 index 4d1ecbb2..00000000 --- a/backend/compose/production/aws/Dockerfile +++ /dev/null @@ -1,9 +0,0 @@ -FROM garland/aws-cli-docker:1.16.140 - -COPY ./compose/production/aws/maintenance /usr/local/bin/maintenance -COPY ./compose/production/postgres/maintenance/_sourced /usr/local/bin/maintenance/_sourced - -RUN chmod +x /usr/local/bin/maintenance/* - -RUN mv /usr/local/bin/maintenance/* /usr/local/bin \ - && rmdir /usr/local/bin/maintenance diff --git a/backend/compose/production/aws/maintenance/download b/backend/compose/production/aws/maintenance/download deleted file mode 100644 index 9561d917..00000000 --- a/backend/compose/production/aws/maintenance/download +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/sh - -### Download a file from your Amazon S3 bucket to the postgres /backups folder -### -### Usage: -### $ docker compose -f production.yml run --rm awscli <1> - -set -o errexit -set -o pipefail -set -o nounset - -working_dir="$(dirname ${0})" -source "${working_dir}/_sourced/constants.sh" -source "${working_dir}/_sourced/messages.sh" - -export AWS_ACCESS_KEY_ID="${DJANGO_AWS_ACCESS_KEY_ID}" -export AWS_SECRET_ACCESS_KEY="${DJANGO_AWS_SECRET_ACCESS_KEY}" -export AWS_STORAGE_BUCKET_NAME="${DJANGO_AWS_STORAGE_BUCKET_NAME}" - - -aws s3 cp s3://${AWS_STORAGE_BUCKET_NAME}${BACKUP_DIR_PATH}/${1} ${BACKUP_DIR_PATH}/${1} - -message_success "Finished downloading ${1}." diff --git a/backend/compose/production/aws/maintenance/upload b/backend/compose/production/aws/maintenance/upload deleted file mode 100644 index 73c1b9be..00000000 --- a/backend/compose/production/aws/maintenance/upload +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/sh - -### Upload the /backups folder to Amazon S3 -### -### Usage: -### $ docker compose -f production.yml run --rm awscli upload - -set -o errexit -set -o pipefail -set -o nounset - -working_dir="$(dirname ${0})" -source "${working_dir}/_sourced/constants.sh" -source "${working_dir}/_sourced/messages.sh" - -export AWS_ACCESS_KEY_ID="${DJANGO_AWS_ACCESS_KEY_ID}" -export AWS_SECRET_ACCESS_KEY="${DJANGO_AWS_SECRET_ACCESS_KEY}" -export AWS_STORAGE_BUCKET_NAME="${DJANGO_AWS_STORAGE_BUCKET_NAME}" - - -message_info "Upload the backups directory to S3 bucket {$AWS_STORAGE_BUCKET_NAME}" - -aws s3 cp ${BACKUP_DIR_PATH} s3://${AWS_STORAGE_BUCKET_NAME}${BACKUP_DIR_PATH} --recursive - -message_info "Cleaning the directory ${BACKUP_DIR_PATH}" - -rm -rf ${BACKUP_DIR_PATH}/* - -message_success "Finished uploading and cleaning." diff --git a/backend/compose/production/django/Dockerfile b/backend/compose/production/django/Dockerfile deleted file mode 100644 index 38617b78..00000000 --- a/backend/compose/production/django/Dockerfile +++ /dev/null @@ -1,96 +0,0 @@ - -# define an alias for the specific python version used in this file. -FROM python:3.11.6-slim-bullseye as python - -# Python build stage -FROM python as python-build-stage - -ARG BUILD_ENVIRONMENT=production - -# Install apt packages -RUN apt-get update && apt-get install --no-install-recommends -y \ - # dependencies for building Python packages - build-essential \ - # psycopg2 dependencies - libpq-dev - -# Requirements are installed here to ensure they will be cached. -COPY ./requirements . - -# Create Python Dependency and Sub-Dependency Wheels. -RUN pip wheel --wheel-dir /usr/src/app/wheels \ - -r ${BUILD_ENVIRONMENT}.txt - - -# Python 'run' stage -FROM python as python-run-stage - -ARG BUILD_ENVIRONMENT=production -ARG APP_HOME=/app - -ENV PYTHONUNBUFFERED 1 -ENV PYTHONDONTWRITEBYTECODE 1 -ENV BUILD_ENV ${BUILD_ENVIRONMENT} - -WORKDIR ${APP_HOME} - -RUN addgroup --system django \ - && adduser --system --ingroup django django - - -# Install required system dependencies -RUN apt-get update && apt-get install --no-install-recommends -y \ - # psycopg2 dependencies - libpq-dev \ - # Translations dependencies - gettext \ - # cleaning up unused files - && apt-get purge -y --auto-remove -o APT::AutoRemove::RecommendsImportant=false \ - && rm -rf /var/lib/apt/lists/* - -# All absolute dir copies ignore workdir instruction. All relative dir copies are wrt to the workdir instruction -# copy python dependency wheels from python-build-stage -COPY --from=python-build-stage /usr/src/app/wheels /wheels/ - -# use wheels to install python dependencies -RUN pip install --no-cache-dir --no-index --find-links=/wheels/ /wheels/* \ - && rm -rf /wheels/ - - -COPY --chown=django:django ./compose/production/django/entrypoint /entrypoint -RUN sed -i 's/\r$//g' /entrypoint -RUN chmod +x /entrypoint - - -COPY --chown=django:django ./compose/production/django/start /start -RUN sed -i 's/\r$//g' /start -RUN chmod +x /start -COPY --chown=django:django ./compose/production/django/celery/worker/start /start-celeryworker -RUN sed -i 's/\r$//g' /start-celeryworker -RUN chmod +x /start-celeryworker - - -COPY --chown=django:django ./compose/production/django/celery/beat/start /start-celerybeat -RUN sed -i 's/\r$//g' /start-celerybeat -RUN chmod +x /start-celerybeat - - -COPY --chown=django:django ./compose/production/django/celery/flower/start /start-flower -RUN sed -i 's/\r$//g' /start-flower -RUN chmod +x /start-flower - - -# copy application code to WORKDIR -COPY --chown=django:django . ${APP_HOME} - -# make django owner of the WORKDIR directory as well. -RUN chown django:django ${APP_HOME} - -USER django - -RUN DATABASE_URL="" \ - CELERY_BROKER_URL="" \ - DJANGO_SETTINGS_MODULE="config.settings.test" \ - python manage.py compilemessages - -ENTRYPOINT ["/entrypoint"] diff --git a/backend/compose/production/django/celery/beat/start b/backend/compose/production/django/celery/beat/start deleted file mode 100644 index 42ddca91..00000000 --- a/backend/compose/production/django/celery/beat/start +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - - -exec celery -A config.celery_app beat -l INFO diff --git a/backend/compose/production/django/celery/flower/start b/backend/compose/production/django/celery/flower/start deleted file mode 100644 index 4180d677..00000000 --- a/backend/compose/production/django/celery/flower/start +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o nounset - - -exec celery \ - -A config.celery_app \ - -b "${CELERY_BROKER_URL}" \ - flower \ - --basic_auth="${CELERY_FLOWER_USER}:${CELERY_FLOWER_PASSWORD}" diff --git a/backend/compose/production/django/celery/worker/start b/backend/compose/production/django/celery/worker/start deleted file mode 100644 index af0c8f7b..00000000 --- a/backend/compose/production/django/celery/worker/start +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - - -exec celery -A config.celery_app worker -l INFO diff --git a/backend/compose/production/django/entrypoint b/backend/compose/production/django/entrypoint deleted file mode 100644 index 249d8d9f..00000000 --- a/backend/compose/production/django/entrypoint +++ /dev/null @@ -1,49 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - - - -# N.B. If only .env files supported variable expansion... -export CELERY_BROKER_URL="${REDIS_URL}" - - -if [ -z "${POSTGRES_USER}" ]; then - base_postgres_image_default_user='postgres' - export POSTGRES_USER="${base_postgres_image_default_user}" -fi -export DATABASE_URL="postgres://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${POSTGRES_HOST}:${POSTGRES_PORT}/${POSTGRES_DB}" - -python << END -import sys -import time - -import psycopg - -suggest_unrecoverable_after = 30 -start = time.time() - -while True: - try: - psycopg.connect( - dbname="${POSTGRES_DB}", - user="${POSTGRES_USER}", - password="${POSTGRES_PASSWORD}", - host="${POSTGRES_HOST}", - port="${POSTGRES_PORT}", - ) - break - except psycopg.OperationalError as error: - sys.stderr.write("Waiting for PostgreSQL to become available...\n") - - if time.time() - start > suggest_unrecoverable_after: - sys.stderr.write(" This is taking longer than expected. The following exception may be indicative of an unrecoverable error: '{}'\n".format(error)) - - time.sleep(1) -END - ->&2 echo 'PostgreSQL is available' - -exec "$@" diff --git a/backend/compose/production/django/start b/backend/compose/production/django/start deleted file mode 100644 index 83ffc8f7..00000000 --- a/backend/compose/production/django/start +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash - -set -o errexit -set -o pipefail -set -o nounset - - -python /app/manage.py collectstatic --noinput - -exec /usr/local/bin/gunicorn config.asgi --bind 0.0.0.0:5000 --chdir=/app -k uvicorn.workers.UvicornWorker diff --git a/backend/compose/production/nginx/Dockerfile b/backend/compose/production/nginx/Dockerfile deleted file mode 100644 index 911b16f7..00000000 --- a/backend/compose/production/nginx/Dockerfile +++ /dev/null @@ -1,2 +0,0 @@ -FROM nginx:1.17.8-alpine -COPY ./compose/production/nginx/default.conf /etc/nginx/conf.d/default.conf diff --git a/backend/compose/production/nginx/default.conf b/backend/compose/production/nginx/default.conf deleted file mode 100644 index 562dba86..00000000 --- a/backend/compose/production/nginx/default.conf +++ /dev/null @@ -1,7 +0,0 @@ -server { - listen 80; - server_name localhost; - location /media/ { - alias /usr/share/nginx/media/; - } -} diff --git a/backend/compose/production/postgres/Dockerfile b/backend/compose/production/postgres/Dockerfile deleted file mode 100644 index 101aa812..00000000 --- a/backend/compose/production/postgres/Dockerfile +++ /dev/null @@ -1,6 +0,0 @@ -FROM postgres:14 - -COPY ./compose/production/postgres/maintenance /usr/local/bin/maintenance -RUN chmod +x /usr/local/bin/maintenance/* -RUN mv /usr/local/bin/maintenance/* /usr/local/bin \ - && rmdir /usr/local/bin/maintenance diff --git a/backend/compose/production/postgres/maintenance/_sourced/constants.sh b/backend/compose/production/postgres/maintenance/_sourced/constants.sh deleted file mode 100644 index 6ca4f0ca..00000000 --- a/backend/compose/production/postgres/maintenance/_sourced/constants.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/usr/bin/env bash - - -BACKUP_DIR_PATH='/backups' -BACKUP_FILE_PREFIX='backup' diff --git a/backend/compose/production/postgres/maintenance/_sourced/countdown.sh b/backend/compose/production/postgres/maintenance/_sourced/countdown.sh deleted file mode 100644 index e6cbfb6f..00000000 --- a/backend/compose/production/postgres/maintenance/_sourced/countdown.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/usr/bin/env bash - - -countdown() { - declare desc="A simple countdown. Source: https://superuser.com/a/611582" - local seconds="${1}" - local d=$(($(date +%s) + "${seconds}")) - while [ "$d" -ge `date +%s` ]; do - echo -ne "$(date -u --date @$(($d - `date +%s`)) +%H:%M:%S)\r"; - sleep 0.1 - done -} diff --git a/backend/compose/production/postgres/maintenance/_sourced/messages.sh b/backend/compose/production/postgres/maintenance/_sourced/messages.sh deleted file mode 100644 index f6be756e..00000000 --- a/backend/compose/production/postgres/maintenance/_sourced/messages.sh +++ /dev/null @@ -1,41 +0,0 @@ -#!/usr/bin/env bash - - -message_newline() { - echo -} - -message_debug() -{ - echo -e "DEBUG: ${@}" -} - -message_welcome() -{ - echo -e "\e[1m${@}\e[0m" -} - -message_warning() -{ - echo -e "\e[33mWARNING\e[0m: ${@}" -} - -message_error() -{ - echo -e "\e[31mERROR\e[0m: ${@}" -} - -message_info() -{ - echo -e "\e[37mINFO\e[0m: ${@}" -} - -message_suggestion() -{ - echo -e "\e[33mSUGGESTION\e[0m: ${@}" -} - -message_success() -{ - echo -e "\e[32mSUCCESS\e[0m: ${@}" -} diff --git a/backend/compose/production/postgres/maintenance/_sourced/yes_no.sh b/backend/compose/production/postgres/maintenance/_sourced/yes_no.sh deleted file mode 100644 index fd9cae16..00000000 --- a/backend/compose/production/postgres/maintenance/_sourced/yes_no.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env bash - - -yes_no() { - declare desc="Prompt for confirmation. \$\"\{1\}\": confirmation message." - local arg1="${1}" - - local response= - read -r -p "${arg1} (y/[n])? " response - if [[ "${response}" =~ ^[Yy]$ ]] - then - exit 0 - else - exit 1 - fi -} diff --git a/backend/compose/production/postgres/maintenance/backup b/backend/compose/production/postgres/maintenance/backup deleted file mode 100644 index f72304c0..00000000 --- a/backend/compose/production/postgres/maintenance/backup +++ /dev/null @@ -1,38 +0,0 @@ -#!/usr/bin/env bash - - -### Create a database backup. -### -### Usage: -### $ docker compose -f .yml (exec |run --rm) postgres backup - - -set -o errexit -set -o pipefail -set -o nounset - - -working_dir="$(dirname ${0})" -source "${working_dir}/_sourced/constants.sh" -source "${working_dir}/_sourced/messages.sh" - - -message_welcome "Backing up the '${POSTGRES_DB}' database..." - - -if [[ "${POSTGRES_USER}" == "postgres" ]]; then - message_error "Backing up as 'postgres' user is not supported. Assign 'POSTGRES_USER' env with another one and try again." - exit 1 -fi - -export PGHOST="${POSTGRES_HOST}" -export PGPORT="${POSTGRES_PORT}" -export PGUSER="${POSTGRES_USER}" -export PGPASSWORD="${POSTGRES_PASSWORD}" -export PGDATABASE="${POSTGRES_DB}" - -backup_filename="${BACKUP_FILE_PREFIX}_$(date +'%Y_%m_%dT%H_%M_%S').sql.gz" -pg_dump | gzip > "${BACKUP_DIR_PATH}/${backup_filename}" - - -message_success "'${POSTGRES_DB}' database backup '${backup_filename}' has been created and placed in '${BACKUP_DIR_PATH}'." diff --git a/backend/compose/production/postgres/maintenance/backups b/backend/compose/production/postgres/maintenance/backups deleted file mode 100644 index a18937d6..00000000 --- a/backend/compose/production/postgres/maintenance/backups +++ /dev/null @@ -1,22 +0,0 @@ -#!/usr/bin/env bash - - -### View backups. -### -### Usage: -### $ docker compose -f .yml (exec |run --rm) postgres backups - - -set -o errexit -set -o pipefail -set -o nounset - - -working_dir="$(dirname ${0})" -source "${working_dir}/_sourced/constants.sh" -source "${working_dir}/_sourced/messages.sh" - - -message_welcome "These are the backups you have got:" - -ls -lht "${BACKUP_DIR_PATH}" diff --git a/backend/compose/production/postgres/maintenance/restore b/backend/compose/production/postgres/maintenance/restore deleted file mode 100644 index c68f17d7..00000000 --- a/backend/compose/production/postgres/maintenance/restore +++ /dev/null @@ -1,55 +0,0 @@ -#!/usr/bin/env bash - - -### Restore database from a backup. -### -### Parameters: -### <1> filename of an existing backup. -### -### Usage: -### $ docker compose -f .yml (exec |run --rm) postgres restore <1> - - -set -o errexit -set -o pipefail -set -o nounset - - -working_dir="$(dirname ${0})" -source "${working_dir}/_sourced/constants.sh" -source "${working_dir}/_sourced/messages.sh" - - -if [[ -z ${1+x} ]]; then - message_error "Backup filename is not specified yet it is a required parameter. Make sure you provide one and try again." - exit 1 -fi -backup_filename="${BACKUP_DIR_PATH}/${1}" -if [[ ! -f "${backup_filename}" ]]; then - message_error "No backup with the specified filename found. Check out the 'backups' maintenance script output to see if there is one and try again." - exit 1 -fi - -message_welcome "Restoring the '${POSTGRES_DB}' database from the '${backup_filename}' backup..." - -if [[ "${POSTGRES_USER}" == "postgres" ]]; then - message_error "Restoring as 'postgres' user is not supported. Assign 'POSTGRES_USER' env with another one and try again." - exit 1 -fi - -export PGHOST="${POSTGRES_HOST}" -export PGPORT="${POSTGRES_PORT}" -export PGUSER="${POSTGRES_USER}" -export PGPASSWORD="${POSTGRES_PASSWORD}" -export PGDATABASE="${POSTGRES_DB}" - -message_info "Dropping the database..." -dropdb "${PGDATABASE}" - -message_info "Creating a new database..." -createdb --owner="${POSTGRES_USER}" - -message_info "Applying the backup to the new database..." -gunzip -c "${backup_filename}" | psql "${POSTGRES_DB}" - -message_success "The '${POSTGRES_DB}' database has been restored from the '${backup_filename}' backup." diff --git a/backend/compose/production/traefik/Dockerfile b/backend/compose/production/traefik/Dockerfile deleted file mode 100644 index e547dfbb..00000000 --- a/backend/compose/production/traefik/Dockerfile +++ /dev/null @@ -1,5 +0,0 @@ -FROM traefik:2.10.4 -RUN mkdir -p /etc/traefik/acme \ - && touch /etc/traefik/acme/acme.json \ - && chmod 600 /etc/traefik/acme/acme.json -COPY ./compose/production/traefik/traefik.yml /etc/traefik diff --git a/backend/compose/production/traefik/traefik.yml b/backend/compose/production/traefik/traefik.yml deleted file mode 100644 index 53eab040..00000000 --- a/backend/compose/production/traefik/traefik.yml +++ /dev/null @@ -1,75 +0,0 @@ -log: - level: INFO - -entryPoints: - web: - # http - address: ':80' - http: - # https://docs.traefik.io/routing/entrypoints/#entrypoint - redirections: - entryPoint: - to: web-secure - - web-secure: - # https - address: ':443' - - flower: - address: ':5555' - -certificatesResolvers: - letsencrypt: - # https://docs.traefik.io/master/https/acme/#lets-encrypt - acme: - email: 'alan@example.com' - storage: /etc/traefik/acme/acme.json - # https://docs.traefik.io/master/https/acme/#httpchallenge - httpChallenge: - entryPoint: web - -http: - routers: - web-secure-router: - rule: 'Host(`example.com`) || Host(`www.example.com`)' - entryPoints: - - web-secure - middlewares: - - csrf - service: django - tls: - # https://docs.traefik.io/master/routing/routers/#certresolver - certResolver: letsencrypt - - flower-secure-router: - rule: 'Host(`example.com`)' - entryPoints: - - flower - service: flower - tls: - # https://docs.traefik.io/master/routing/routers/#certresolver - certResolver: letsencrypt - - middlewares: - csrf: - # https://docs.traefik.io/master/middlewares/headers/#hostsproxyheaders - # https://docs.djangoproject.com/en/dev/ref/csrf/#ajax - headers: - hostsProxyHeaders: ['X-CSRFToken'] - - services: - django: - loadBalancer: - servers: - - url: http://django:5000 - - flower: - loadBalancer: - servers: - - url: http://flower:5555 - -providers: - # https://docs.traefik.io/master/providers/file/ - file: - filename: /etc/traefik/traefik.yml - watch: true diff --git a/backend/config/__init__.py b/backend/config/__init__.py deleted file mode 100644 index 10f50142..00000000 --- a/backend/config/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# This will make sure the app is always imported when -# Django starts so that shared_task will use this app. -from .celery_app import app as celery_app - -__all__ = ("celery_app",) diff --git a/backend/config/api_router.py b/backend/config/api_router.py deleted file mode 100644 index 2e49e75d..00000000 --- a/backend/config/api_router.py +++ /dev/null @@ -1,20 +0,0 @@ -from django.conf import settings -from rest_framework.routers import DefaultRouter, SimpleRouter - -from backend.users.api.views import UserViewSet -from django.urls import include, path - -if settings.DEBUG: - router = DefaultRouter() -else: - router = SimpleRouter() - -router.register("users", UserViewSet) - -sub_urls = [ - path('data-analysis/', include("data_analysis.urls")), - path('authentication/', include("users.api.urls")), -] - -app_name = "api" -urlpatterns = router.urls + sub_urls diff --git a/backend/config/asgi.py b/backend/config/asgi.py deleted file mode 100644 index 92318d15..00000000 --- a/backend/config/asgi.py +++ /dev/null @@ -1,40 +0,0 @@ -""" -ASGI config for backend project. - -It exposes the ASGI callable as a module-level variable named ``application``. - -For more information on this file, see -https://docs.djangoproject.com/en/dev/howto/deployment/asgi/ - -""" -import os -import sys -from pathlib import Path - -from django.core.asgi import get_asgi_application - -# This allows easy placement of apps within the interior -# backend directory. -BASE_DIR = Path(__file__).resolve(strict=True).parent.parent -sys.path.append(str(BASE_DIR / "backend")) - -# If DJANGO_SETTINGS_MODULE is unset, default to the local settings -os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local") - -# This application object is used by any ASGI server configured to use this file. -django_application = get_asgi_application() -# Apply ASGI middleware here. -# from helloworld.asgi import HelloWorldApplication -# application = HelloWorldApplication(application) - -# Import websocket application here, so apps from django_application are loaded first -from config.websocket import websocket_application # noqa isort:skip - - -async def application(scope, receive, send): - if scope["type"] == "http": - await django_application(scope, receive, send) - elif scope["type"] == "websocket": - await websocket_application(scope, receive, send) - else: - raise NotImplementedError(f"Unknown scope type {scope['type']}") diff --git a/backend/config/celery_app.py b/backend/config/celery_app.py deleted file mode 100644 index 55e00f1d..00000000 --- a/backend/config/celery_app.py +++ /dev/null @@ -1,17 +0,0 @@ -import os - -from celery import Celery - -# set the default Django settings module for the 'celery' program. -os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local") - -app = Celery("backend") - -# Using a string here means the worker doesn't have to serialize -# the configuration object to child processes. -# - namespace='CELERY' means all celery-related configuration keys -# should have a `CELERY_` prefix. -app.config_from_object("django.conf:settings", namespace="CELERY") - -# Load task modules from all registered Django app configs. -app.autodiscover_tasks() diff --git a/backend/config/settings/__init__.py b/backend/config/settings/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/config/settings/base.py b/backend/config/settings/base.py deleted file mode 100644 index 2c80b5b2..00000000 --- a/backend/config/settings/base.py +++ /dev/null @@ -1,369 +0,0 @@ -""" -Base settings to build other settings files upon. -""" -from pathlib import Path -import datetime - -import environ - -BASE_DIR = Path(__file__).resolve(strict=True).parent.parent.parent -# backend/ -APPS_DIR = BASE_DIR / "backend" -env = environ.Env() - -READ_DOT_ENV_FILE = env.bool("DJANGO_READ_DOT_ENV_FILE", default=True) -if READ_DOT_ENV_FILE: - # OS environment variables take precedence over variables from .env - env.read_env(str(BASE_DIR / ".env")) - -# GENERAL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#debug -DEBUG = env.bool("DJANGO_DEBUG", False) -# Local time zone. Choices are -# http://en.wikipedia.org/wiki/List_of_tz_zones_by_name -# though not all of them may be available with every OS. -# In Windows, this must be set to your system time zone. -TIME_ZONE = "UTC" -# https://docs.djangoproject.com/en/dev/ref/settings/#language-code -LANGUAGE_CODE = "en-us" -# https://docs.djangoproject.com/en/dev/ref/settings/#languages -# from django.utils.translation import gettext_lazy as _ -# LANGUAGES = [ -# ('en', _('English')), -# ('fr-fr', _('French')), -# ('pt-br', _('Portuguese')), -# ] -# https://docs.djangoproject.com/en/dev/ref/settings/#site-id -SITE_ID = 1 -# https://docs.djangoproject.com/en/dev/ref/settings/#use-i18n -USE_I18N = True -# https://docs.djangoproject.com/en/dev/ref/settings/#use-tz -USE_TZ = True -# https://docs.djangoproject.com/en/dev/ref/settings/#locale-paths -LOCALE_PATHS = [str(BASE_DIR / "locale")] - -# DATABASES -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#databases -DATABASES = {"default": env.db("DATABASE_URL")} -# https://docs.djangoproject.com/en/stable/ref/settings/#std:setting-DEFAULT_AUTO_FIELD -DEFAULT_AUTO_FIELD = "django.db.models.BigAutoField" - -# URLS -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#root-urlconf -ROOT_URLCONF = "config.urls" -# https://docs.djangoproject.com/en/dev/ref/settings/#wsgi-application -WSGI_APPLICATION = "config.wsgi.application" - -# APPS -# ------------------------------------------------------------------------------ -DJANGO_APPS = [ - "django.contrib.auth", - "django.contrib.contenttypes", - "django.contrib.sessions", - "django.contrib.sites", - "django.contrib.messages", - "django.contrib.staticfiles", - # "django.contrib.humanize", # Handy template tags - "django.contrib.admin", - "django.forms", -] -THIRD_PARTY_APPS = [ - "crispy_forms", - "crispy_bootstrap5", - "allauth", - "allauth.account", - "allauth.socialaccount", - "django_celery_beat", - "rest_framework", - "rest_framework.authtoken", - "corsheaders", - "drf_spectacular", -] - -LOCAL_APPS = [ - "backend.users", - "backend.data_analysis", - # Your stuff: custom apps go here -] -# https://docs.djangoproject.com/en/dev/ref/settings/#installed-apps -INSTALLED_APPS = DJANGO_APPS + THIRD_PARTY_APPS + LOCAL_APPS - -# MIGRATIONS -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#migration-modules -MIGRATION_MODULES = {"sites": "backend.contrib.sites.migrations"} - -# AUTHENTICATION -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#authentication-backends -AUTHENTICATION_BACKENDS = [ - "django.contrib.auth.backends.ModelBackend", - "allauth.account.auth_backends.AuthenticationBackend", -] -# https://docs.djangoproject.com/en/dev/ref/settings/#auth-user-model -AUTH_USER_MODEL = "users.User" -# https://docs.djangoproject.com/en/dev/ref/settings/#login-redirect-url -LOGIN_REDIRECT_URL = "users:redirect" -# https://docs.djangoproject.com/en/dev/ref/settings/#login-url -LOGIN_URL = "account_login" - -# PASSWORDS -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#password-hashers -PASSWORD_HASHERS = [ - # https://docs.djangoproject.com/en/dev/topics/auth/passwords/#using-argon2-with-django - "django.contrib.auth.hashers.Argon2PasswordHasher", - "django.contrib.auth.hashers.PBKDF2PasswordHasher", - "django.contrib.auth.hashers.PBKDF2SHA1PasswordHasher", - "django.contrib.auth.hashers.BCryptSHA256PasswordHasher", -] -# https://docs.djangoproject.com/en/dev/ref/settings/#auth-password-validators -AUTH_PASSWORD_VALIDATORS = [ - {"NAME": "django.contrib.auth.password_validation.UserAttributeSimilarityValidator"}, - {"NAME": "django.contrib.auth.password_validation.MinimumLengthValidator"}, - {"NAME": "django.contrib.auth.password_validation.CommonPasswordValidator"}, - {"NAME": "django.contrib.auth.password_validation.NumericPasswordValidator"}, -] - -# MIDDLEWARE -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#middleware -MIDDLEWARE = [ - "django.middleware.security.SecurityMiddleware", - "corsheaders.middleware.CorsMiddleware", - "whitenoise.middleware.WhiteNoiseMiddleware", - "django.contrib.sessions.middleware.SessionMiddleware", - "django.middleware.locale.LocaleMiddleware", - "django.middleware.common.CommonMiddleware", - "django.middleware.csrf.CsrfViewMiddleware", - "django.contrib.auth.middleware.AuthenticationMiddleware", - "django.contrib.messages.middleware.MessageMiddleware", - "django.middleware.clickjacking.XFrameOptionsMiddleware", - "allauth.account.middleware.AccountMiddleware", -] - -# STATIC -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#static-root -STATIC_ROOT = str(BASE_DIR / "staticfiles") -# https://docs.djangoproject.com/en/dev/ref/settings/#static-url -STATIC_URL = "/static/" -# https://docs.djangoproject.com/en/dev/ref/contrib/staticfiles/#std:setting-STATICFILES_DIRS -STATICFILES_DIRS = [str(APPS_DIR / "static")] -# https://docs.djangoproject.com/en/dev/ref/contrib/staticfiles/#staticfiles-finders -STATICFILES_FINDERS = [ - "django.contrib.staticfiles.finders.FileSystemFinder", - "django.contrib.staticfiles.finders.AppDirectoriesFinder", -] - -# MEDIA -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#media-root -MEDIA_ROOT = str(APPS_DIR / "media") -# https://docs.djangoproject.com/en/dev/ref/settings/#media-url -MEDIA_URL = "/media/" - -# TEMPLATES -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#templates -TEMPLATES = [ - { - # https://docs.djangoproject.com/en/dev/ref/settings/#std:setting-TEMPLATES-BACKEND - "BACKEND": "django.template.backends.django.DjangoTemplates", - # https://docs.djangoproject.com/en/dev/ref/settings/#dirs - "DIRS": [str(APPS_DIR / "templates")], - # https://docs.djangoproject.com/en/dev/ref/settings/#app-dirs - "APP_DIRS": True, - "OPTIONS": { - # https://docs.djangoproject.com/en/dev/ref/settings/#template-context-processors - "context_processors": [ - "django.template.context_processors.debug", - "django.template.context_processors.request", - "django.contrib.auth.context_processors.auth", - "django.template.context_processors.i18n", - "django.template.context_processors.media", - "django.template.context_processors.static", - "django.template.context_processors.tz", - "django.contrib.messages.context_processors.messages", - "backend.users.context_processors.allauth_settings", - ], - }, - } -] - -# https://docs.djangoproject.com/en/dev/ref/settings/#form-renderer -FORM_RENDERER = "django.forms.renderers.TemplatesSetting" - -# http://django-crispy-forms.readthedocs.io/en/latest/install.html#template-packs -CRISPY_TEMPLATE_PACK = "bootstrap5" -CRISPY_ALLOWED_TEMPLATE_PACKS = "bootstrap5" - -# FIXTURES -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#fixture-dirs -FIXTURE_DIRS = (str(APPS_DIR / "fixtures"),) - -# SECURITY -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#session-cookie-httponly -SESSION_COOKIE_HTTPONLY = True -# https://docs.djangoproject.com/en/dev/ref/settings/#csrf-cookie-httponly -CSRF_COOKIE_HTTPONLY = True -# https://docs.djangoproject.com/en/dev/ref/settings/#x-frame-options -X_FRAME_OPTIONS = "DENY" - -# EMAIL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#email-backend -EMAIL_BACKEND = env( - "DJANGO_EMAIL_BACKEND", - default="django.core.mail.backends.smtp.EmailBackend", -) -# https://docs.djangoproject.com/en/dev/ref/settings/#email-timeout -EMAIL_TIMEOUT = 5 - -# ADMIN -# ------------------------------------------------------------------------------ -# Django Admin URL. -ADMIN_URL = "admin/" -# https://docs.djangoproject.com/en/dev/ref/settings/#admins -ADMINS = [("""Alan""", "alan@example.com")] -# https://docs.djangoproject.com/en/dev/ref/settings/#managers -MANAGERS = ADMINS -# https://cookiecutter-django.readthedocs.io/en/latest/settings.html#other-environment-settings -# Force the `admin` sign in process to go through the `django-allauth` workflow -DJANGO_ADMIN_FORCE_ALLAUTH = env.bool("DJANGO_ADMIN_FORCE_ALLAUTH", default=False) - -# LOGGING -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#logging -# See https://docs.djangoproject.com/en/dev/topics/logging for -# more details on how to customize your logging configuration. -LOGGING = { - "version": 1, - "disable_existing_loggers": False, - "formatters": { - "verbose": { - "format": "%(levelname)s %(asctime)s %(module)s %(process)d %(thread)d %(message)s", - }, - }, - "handlers": { - "console": { - "level": "DEBUG", - "class": "logging.StreamHandler", - "formatter": "verbose", - } - }, - "root": {"level": "INFO", "handlers": ["console"]}, -} - -# Celery -# ------------------------------------------------------------------------------ -if USE_TZ: - # https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-timezone - CELERY_TIMEZONE = TIME_ZONE -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-broker_url -CELERY_BROKER_URL = env("CELERY_BROKER_URL") -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-result_backend -CELERY_RESULT_BACKEND = CELERY_BROKER_URL -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#result-extended -CELERY_RESULT_EXTENDED = True -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#result-backend-always-retry -# https://github.com/celery/celery/pull/6122 -CELERY_RESULT_BACKEND_ALWAYS_RETRY = True -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#result-backend-max-retries -CELERY_RESULT_BACKEND_MAX_RETRIES = 10 -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-accept_content -CELERY_ACCEPT_CONTENT = ["json"] -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-task_serializer -CELERY_TASK_SERIALIZER = "json" -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std:setting-result_serializer -CELERY_RESULT_SERIALIZER = "json" -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#task-time-limit -# TODO: set to whatever value is adequate in your circumstances -CELERY_TASK_TIME_LIMIT = 5 * 60 -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#task-soft-time-limit -# TODO: set to whatever value is adequate in your circumstances -CELERY_TASK_SOFT_TIME_LIMIT = 60 -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#beat-scheduler -CELERY_BEAT_SCHEDULER = "django_celery_beat.schedulers:DatabaseScheduler" -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#worker-send-task-events -CELERY_WORKER_SEND_TASK_EVENTS = True -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#std-setting-task_send_sent_event -CELERY_TASK_SEND_SENT_EVENT = True -# django-allauth -# ------------------------------------------------------------------------------ -ACCOUNT_ALLOW_REGISTRATION = env.bool("DJANGO_ACCOUNT_ALLOW_REGISTRATION", True) -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_AUTHENTICATION_METHOD = "email" -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_EMAIL_REQUIRED = True -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_USERNAME_REQUIRED = False -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_USER_MODEL_USERNAME_FIELD = None -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_EMAIL_VERIFICATION = "mandatory" -# https://django-allauth.readthedocs.io/en/latest/configuration.html -ACCOUNT_ADAPTER = "backend.users.adapters.AccountAdapter" -# https://django-allauth.readthedocs.io/en/latest/forms.html -ACCOUNT_FORMS = {"signup": "backend.users.forms.UserSignupForm"} -# https://django-allauth.readthedocs.io/en/latest/configuration.html -SOCIALACCOUNT_ADAPTER = "backend.users.adapters.SocialAccountAdapter" -# https://django-allauth.readthedocs.io/en/latest/forms.html -SOCIALACCOUNT_FORMS = {"signup": "backend.users.forms.UserSocialSignupForm"} - -# django-rest-framework -# ------------------------------------------------------------------------------- -# django-rest-framework - https://www.django-rest-framework.org/api-guide/settings/ -REST_FRAMEWORK = { - "DEFAULT_AUTHENTICATION_CLASSES": ( - "rest_framework_simplejwt.authentication.JWTAuthentication", - ), - "DEFAULT_PERMISSION_CLASSES": ("rest_framework.permissions.IsAuthenticated",), - "DEFAULT_SCHEMA_CLASS": "drf_spectacular.openapi.AutoSchema", -} - -# django-cors-headers - https://github.com/adamchainz/django-cors-headers#setup -CORS_URLS_REGEX = r"^/api/.*$" - -# By Default swagger ui is available only to admin user(s). You can change permission classes to change that -# See more configuration options at https://drf-spectacular.readthedocs.io/en/latest/settings.html#settings -SPECTACULAR_SETTINGS = { - "TITLE": "backend API", - "DESCRIPTION": "Documentation of API endpoints of backend", - "VERSION": "1.0.0", - "SERVE_PERMISSIONS": ["rest_framework.permissions.IsAdminUser"], -} -# Your stuff... -# ------------------------------------------------------------------------------ - -SIMPLE_JWT = { - # token expiration times - "ACCESS_TOKEN_LIFETIME": datetime.timedelta(minutes=30), - "REFRESH_TOKEN_LIFETIME": datetime.timedelta(days=1), - # token rotation options - # 'ROTATE_REFRESH_TOKENS': False, - # 'BLACKLIST_AFTER_ROTATION': False, - # user authentication options - "AUTH_HEADER_TYPES": ("Bearer",), - # 'USER_ID_FIELD': 'id', - # 'USER_ID_CLAIM': 'user_id', - # 'AUTH_TOKEN_CLASSES': ( - # 'rest_framework_simplejwt.tokens.AccessToken', - # 'rest_framework_simplejwt.tokens.RefreshToken', - # ), - # 'TOKEN_TYPE_CLAIM': 'token_type', - # token verification options - # 'ALGORITHM': 'HS256', - # 'SIGNING_KEY': None, - # 'VERIFYING_KEY': None, - # # token refresh options - # 'ALLOW_REFRESH': True, - # 'REFRESH_TOKEN_ROTATE_REFRESH_TOKENS': False, - # 'REFRESH_TOKEN_BLACKLIST_AFTER_ROTATION': False, - # 'UPDATE_LAST_LOGIN': False, -} diff --git a/backend/config/settings/local.py b/backend/config/settings/local.py deleted file mode 100644 index 7f6cf95c..00000000 --- a/backend/config/settings/local.py +++ /dev/null @@ -1,73 +0,0 @@ -from .base import * # noqa -from .base import env - -# GENERAL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#debug -DEBUG = True -# https://docs.djangoproject.com/en/dev/ref/settings/#secret-key -SECRET_KEY = env( - "DJANGO_SECRET_KEY", - default="Hno5SAr11TWUYosckJxayOIyxTgdTyhggICiGqeVGtVyEp0XCB9rACsEpoGp7ds2", -) -# https://docs.djangoproject.com/en/dev/ref/settings/#allowed-hosts -ALLOWED_HOSTS = ["localhost", "0.0.0.0", "127.0.0.1", "143.198.234.124", "metricloop.ai"] - - -CORS_ORIGIN_WHITELIST = [ - "http://localhost:3000", - "http://143.198.234.124:3000", - "http://127.0.0.1:3000", - "http://metricloop.ai:3000", -] -# CACHES -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#caches -CACHES = { - "default": { - "BACKEND": "django.core.cache.backends.locmem.LocMemCache", - "LOCATION": "", - } -} - -# EMAIL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#email-backend -EMAIL_BACKEND = env("DJANGO_EMAIL_BACKEND", default="django.core.mail.backends.console.EmailBackend") - -# WhiteNoise -# ------------------------------------------------------------------------------ -# http://whitenoise.evans.io/en/latest/django.html#using-whitenoise-in-development -INSTALLED_APPS = ["whitenoise.runserver_nostatic"] + INSTALLED_APPS # noqa: F405 - - -# django-debug-toolbar -# ------------------------------------------------------------------------------ -# https://django-debug-toolbar.readthedocs.io/en/latest/installation.html#prerequisites -INSTALLED_APPS += ["debug_toolbar"] # noqa: F405 -# https://django-debug-toolbar.readthedocs.io/en/latest/installation.html#middleware -MIDDLEWARE += ["debug_toolbar.middleware.DebugToolbarMiddleware"] # noqa: F405 -# https://django-debug-toolbar.readthedocs.io/en/latest/configuration.html#debug-toolbar-config -DEBUG_TOOLBAR_CONFIG = { - "DISABLE_PANELS": ["debug_toolbar.panels.redirects.RedirectsPanel"], - "SHOW_TEMPLATE_CONTEXT": True, -} -# https://django-debug-toolbar.readthedocs.io/en/latest/installation.html#internal-ips -INTERNAL_IPS = ["127.0.0.1", "10.0.2.2"] -if env("USE_DOCKER") == "yes": - import socket - - hostname, _, ips = socket.gethostbyname_ex(socket.gethostname()) - INTERNAL_IPS += [".".join(ip.split(".")[:-1] + ["1"]) for ip in ips] - -# django-extensions -# ------------------------------------------------------------------------------ -# https://django-extensions.readthedocs.io/en/latest/installation_instructions.html#configuration -INSTALLED_APPS += ["django_extensions"] # noqa: F405 -# Celery -# ------------------------------------------------------------------------------ - -# https://docs.celeryq.dev/en/stable/userguide/configuration.html#task-eager-propagates -CELERY_TASK_EAGER_PROPAGATES = True -# Your stuff... -# ------------------------------------------------------------------------------ diff --git a/backend/config/settings/production.py b/backend/config/settings/production.py deleted file mode 100644 index a55047a3..00000000 --- a/backend/config/settings/production.py +++ /dev/null @@ -1,173 +0,0 @@ -from .base import * # noqa -from .base import env - -# GENERAL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#secret-key -SECRET_KEY = env("DJANGO_SECRET_KEY") -# https://docs.djangoproject.com/en/dev/ref/settings/#allowed-hosts -ALLOWED_HOSTS = env.list("DJANGO_ALLOWED_HOSTS", default=["example.com"]) - -# DATABASES -# ------------------------------------------------------------------------------ -DATABASES["default"]["CONN_MAX_AGE"] = env.int("CONN_MAX_AGE", default=60) # noqa: F405 - -# CACHES -# ------------------------------------------------------------------------------ -CACHES = { - "default": { - "BACKEND": "django_redis.cache.RedisCache", - "LOCATION": env("REDIS_URL"), - "OPTIONS": { - "CLIENT_CLASS": "django_redis.client.DefaultClient", - # Mimicing memcache behavior. - # https://github.com/jazzband/django-redis#memcached-exceptions-behavior - "IGNORE_EXCEPTIONS": True, - }, - } -} - -# SECURITY -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#secure-proxy-ssl-header -SECURE_PROXY_SSL_HEADER = ("HTTP_X_FORWARDED_PROTO", "https") -# https://docs.djangoproject.com/en/dev/ref/settings/#secure-ssl-redirect -SECURE_SSL_REDIRECT = env.bool("DJANGO_SECURE_SSL_REDIRECT", default=True) -# https://docs.djangoproject.com/en/dev/ref/settings/#session-cookie-secure -SESSION_COOKIE_SECURE = True -# https://docs.djangoproject.com/en/dev/ref/settings/#csrf-cookie-secure -CSRF_COOKIE_SECURE = True -# https://docs.djangoproject.com/en/dev/topics/security/#ssl-https -# https://docs.djangoproject.com/en/dev/ref/settings/#secure-hsts-seconds -# TODO: set this to 60 seconds first and then to 518400 once you prove the former works -SECURE_HSTS_SECONDS = 60 -# https://docs.djangoproject.com/en/dev/ref/settings/#secure-hsts-include-subdomains -SECURE_HSTS_INCLUDE_SUBDOMAINS = env.bool("DJANGO_SECURE_HSTS_INCLUDE_SUBDOMAINS", default=True) -# https://docs.djangoproject.com/en/dev/ref/settings/#secure-hsts-preload -SECURE_HSTS_PRELOAD = env.bool("DJANGO_SECURE_HSTS_PRELOAD", default=True) -# https://docs.djangoproject.com/en/dev/ref/middleware/#x-content-type-options-nosniff -SECURE_CONTENT_TYPE_NOSNIFF = env.bool("DJANGO_SECURE_CONTENT_TYPE_NOSNIFF", default=True) - -# STORAGES -# ------------------------------------------------------------------------------ -# https://django-storages.readthedocs.io/en/latest/#installation -INSTALLED_APPS += ["storages"] # noqa: F405 -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_ACCESS_KEY_ID = env("DJANGO_AWS_ACCESS_KEY_ID") -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_SECRET_ACCESS_KEY = env("DJANGO_AWS_SECRET_ACCESS_KEY") -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_STORAGE_BUCKET_NAME = env("DJANGO_AWS_STORAGE_BUCKET_NAME") -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_QUERYSTRING_AUTH = False -# DO NOT change these unless you know what you're doing. -_AWS_EXPIRY = 60 * 60 * 24 * 7 -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_S3_OBJECT_PARAMETERS = { - "CacheControl": f"max-age={_AWS_EXPIRY}, s-maxage={_AWS_EXPIRY}, must-revalidate", -} -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_S3_MAX_MEMORY_SIZE = env.int( - "DJANGO_AWS_S3_MAX_MEMORY_SIZE", - default=100_000_000, # 100MB -) -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#settings -AWS_S3_REGION_NAME = env("DJANGO_AWS_S3_REGION_NAME", default=None) -# https://django-storages.readthedocs.io/en/latest/backends/amazon-S3.html#cloudfront -AWS_S3_CUSTOM_DOMAIN = env("DJANGO_AWS_S3_CUSTOM_DOMAIN", default=None) -aws_s3_domain = AWS_S3_CUSTOM_DOMAIN or f"{AWS_STORAGE_BUCKET_NAME}.s3.amazonaws.com" -# STATIC -# ------------------------ -STATICFILES_STORAGE = "whitenoise.storage.CompressedManifestStaticFilesStorage" -# MEDIA -# ------------------------------------------------------------------------------ -DEFAULT_FILE_STORAGE = "backend.utils.storages.MediaS3Storage" -MEDIA_URL = f"https://{aws_s3_domain}/media/" - -# EMAIL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#default-from-email -DEFAULT_FROM_EMAIL = env( - "DJANGO_DEFAULT_FROM_EMAIL", - default="backend ", -) -# https://docs.djangoproject.com/en/dev/ref/settings/#server-email -SERVER_EMAIL = env("DJANGO_SERVER_EMAIL", default=DEFAULT_FROM_EMAIL) -# https://docs.djangoproject.com/en/dev/ref/settings/#email-subject-prefix -EMAIL_SUBJECT_PREFIX = env( - "DJANGO_EMAIL_SUBJECT_PREFIX", - default="[backend] ", -) - -# ADMIN -# ------------------------------------------------------------------------------ -# Django Admin URL regex. -ADMIN_URL = env("DJANGO_ADMIN_URL") - -# Anymail -# ------------------------------------------------------------------------------ -# https://anymail.readthedocs.io/en/stable/installation/#installing-anymail -INSTALLED_APPS += ["anymail"] # noqa: F405 -# https://docs.djangoproject.com/en/dev/ref/settings/#email-backend -# https://anymail.readthedocs.io/en/stable/installation/#anymail-settings-reference -# https://anymail.readthedocs.io/en/stable/esps/mailgun/ -EMAIL_BACKEND = "anymail.backends.mailgun.EmailBackend" -ANYMAIL = { - "MAILGUN_API_KEY": env("MAILGUN_API_KEY"), - "MAILGUN_SENDER_DOMAIN": env("MAILGUN_DOMAIN"), - "MAILGUN_API_URL": env("MAILGUN_API_URL", default="https://api.mailgun.net/v3"), -} - - -# LOGGING -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#logging -# See https://docs.djangoproject.com/en/dev/topics/logging for -# more details on how to customize your logging configuration. -# A sample logging configuration. The only tangible logging -# performed by this configuration is to send an email to -# the site admins on every HTTP 500 error when DEBUG=False. -LOGGING = { - "version": 1, - "disable_existing_loggers": False, - "filters": {"require_debug_false": {"()": "django.utils.log.RequireDebugFalse"}}, - "formatters": { - "verbose": { - "format": "%(levelname)s %(asctime)s %(module)s %(process)d %(thread)d %(message)s", - }, - }, - "handlers": { - "mail_admins": { - "level": "ERROR", - "filters": ["require_debug_false"], - "class": "django.utils.log.AdminEmailHandler", - }, - "console": { - "level": "DEBUG", - "class": "logging.StreamHandler", - "formatter": "verbose", - }, - }, - "root": {"level": "INFO", "handlers": ["console"]}, - "loggers": { - "django.request": { - "handlers": ["mail_admins"], - "level": "ERROR", - "propagate": True, - }, - "django.security.DisallowedHost": { - "level": "ERROR", - "handlers": ["console", "mail_admins"], - "propagate": True, - }, - }, -} - -# django-rest-framework -# ------------------------------------------------------------------------------- -# Tools that generate code samples can use SERVERS to point to the correct domain -SPECTACULAR_SETTINGS["SERVERS"] = [ # noqa: F405 - {"url": "https://example.com", "description": "Production server"}, -] -# Your stuff... -# ------------------------------------------------------------------------------ diff --git a/backend/config/settings/test.py b/backend/config/settings/test.py deleted file mode 100644 index cc0913ca..00000000 --- a/backend/config/settings/test.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -With these settings, tests run faster. -""" - -from .base import * # noqa -from .base import env - -# GENERAL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#secret-key -SECRET_KEY = env( - "DJANGO_SECRET_KEY", - default="eUb3yE8YYIhmOXoQXshvrN75vBZ9RZXbDFrCC6KAqwjNSesAcavX5askNqFTeqvJ", -) -# https://docs.djangoproject.com/en/dev/ref/settings/#test-runner -TEST_RUNNER = "django.test.runner.DiscoverRunner" - -# PASSWORDS -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#password-hashers -PASSWORD_HASHERS = ["django.contrib.auth.hashers.MD5PasswordHasher"] - -# EMAIL -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#email-backend -EMAIL_BACKEND = "django.core.mail.backends.locmem.EmailBackend" - -# DEBUGGING FOR TEMPLATES -# ------------------------------------------------------------------------------ -TEMPLATES[0]["OPTIONS"]["debug"] = True # type: ignore # noqa: F405 - -# MEDIA -# ------------------------------------------------------------------------------ -# https://docs.djangoproject.com/en/dev/ref/settings/#media-url -MEDIA_URL = 'http://media.testserver' -# Your stuff... -# ------------------------------------------------------------------------------ diff --git a/backend/config/urls.py b/backend/config/urls.py deleted file mode 100644 index cbcff78c..00000000 --- a/backend/config/urls.py +++ /dev/null @@ -1,63 +0,0 @@ -from django.conf import settings -from django.conf.urls.static import static -from django.contrib import admin -from django.contrib.staticfiles.urls import staticfiles_urlpatterns -from django.urls import include, path -from django.views import defaults as default_views -from django.views.generic import TemplateView -from drf_spectacular.views import SpectacularAPIView, SpectacularSwaggerView -from rest_framework.authtoken.views import obtain_auth_token - -urlpatterns = [ - path("", TemplateView.as_view(template_name="pages/home.html"), name="home"), - path("about/", TemplateView.as_view(template_name="pages/about.html"), name="about"), - # Django Admin, use {% url 'admin:index' %} - path(settings.ADMIN_URL, admin.site.urls), - # User management - path("users/", include("backend.users.urls", namespace="users")), - path("accounts/", include("allauth.urls")), - # Your stuff: custom urls includes go here -] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT) -if settings.DEBUG: - # Static file serving when using Gunicorn + Uvicorn for local web socket development - urlpatterns += staticfiles_urlpatterns() - -# API URLS -urlpatterns += [ - # API base url - path("api/", include("config.api_router")), - # DRF auth token - path("auth-token/", obtain_auth_token), - path("api/schema/", SpectacularAPIView.as_view(), name="api-schema"), - path( - "api/docs/", - SpectacularSwaggerView.as_view(url_name="api-schema"), - name="api-docs", - ), -] - -if settings.DEBUG: - # This allows the error pages to be debugged during development, just visit - # these url in browser to see how these error pages look like. - urlpatterns += [ - path( - "400/", - default_views.bad_request, - kwargs={"exception": Exception("Bad Request!")}, - ), - path( - "403/", - default_views.permission_denied, - kwargs={"exception": Exception("Permission Denied")}, - ), - path( - "404/", - default_views.page_not_found, - kwargs={"exception": Exception("Page not Found")}, - ), - path("500/", default_views.server_error), - ] - if "debug_toolbar" in settings.INSTALLED_APPS: - import debug_toolbar - - urlpatterns = [path("__debug__/", include(debug_toolbar.urls))] + urlpatterns diff --git a/backend/config/websocket.py b/backend/config/websocket.py deleted file mode 100644 index 81adfbc6..00000000 --- a/backend/config/websocket.py +++ /dev/null @@ -1,13 +0,0 @@ -async def websocket_application(scope, receive, send): - while True: - event = await receive() - - if event["type"] == "websocket.connect": - await send({"type": "websocket.accept"}) - - if event["type"] == "websocket.disconnect": - break - - if event["type"] == "websocket.receive": - if event["text"] == "ping": - await send({"type": "websocket.send", "text": "pong!"}) diff --git a/backend/config/wsgi.py b/backend/config/wsgi.py deleted file mode 100644 index 69be6ca4..00000000 --- a/backend/config/wsgi.py +++ /dev/null @@ -1,38 +0,0 @@ -""" -WSGI config for backend project. - -This module contains the WSGI application used by Django's development server -and any production WSGI deployments. It should expose a module-level variable -named ``application``. Django's ``runserver`` and ``runfcgi`` commands discover -this application via the ``WSGI_APPLICATION`` setting. - -Usually you will have the standard Django WSGI application here, but it also -might make sense to replace the whole Django WSGI application with a custom one -that later delegates to the Django one. For example, you could introduce WSGI -middleware here, or combine a Django application with an application of another -framework. - -""" -import os -import sys -from pathlib import Path - -from django.core.wsgi import get_wsgi_application - -# This allows easy placement of apps within the interior -# backend directory. -BASE_DIR = Path(__file__).resolve(strict=True).parent.parent -sys.path.append(str(BASE_DIR / "backend")) -# We defer to a DJANGO_SETTINGS_MODULE already in the environment. This breaks -# if running multiple sites in the same mod_wsgi process. To fix this, use -# mod_wsgi daemon mode with each site in its own daemon process, or use -# os.environ["DJANGO_SETTINGS_MODULE"] = "config.settings.production" -os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.production") - -# This application object is used by any WSGI server configured to use this -# file. This includes Django's development server, if the WSGI_APPLICATION -# setting points here. -application = get_wsgi_application() -# Apply WSGI middleware here. -# from helloworld.wsgi import HelloWorldApplication -# application = HelloWorldApplication(application) diff --git a/backend/controller.py b/backend/controller.py new file mode 100644 index 00000000..15dd7d20 --- /dev/null +++ b/backend/controller.py @@ -0,0 +1,90 @@ +from sqlalchemy import select +from databases import Database +from models import FileModel, UserModel, TransactionRecordModel +import os + +database = Database(os.environ["DATABASE_URL"]) + + +async def connect_to_database(): + await database.connect() + + +async def disconnect_from_database(): + await database.disconnect() + + +async def save_file_data( + video_filename: str, + image_filename: str, + audio_filename: str, + text_data: str, + address: str, +): + try: + + await connect_to_database() + + query = FileModel.__table__.insert().values( + video_filename=video_filename, + image_filename=image_filename, + audio_filename=audio_filename, + text_data=text_data, + wallet_address=address, + ) + + return await database.execute(query) + + finally: + await disconnect_from_database() + + +async def get_file_by_id(file_id: int): + try: + await connect_to_database() + query = select(FileModel).where(FileModel.c.id == file_id) + return await database.fetch_one(query) + finally: + await disconnect_from_database() + + +async def get_all_files(): + try: + await connect_to_database() + query = select(FileModel) + return await database.fetch_all(query) + finally: + await disconnect_from_database() + + +async def save_address(address: str): + try: + await connect_to_database() + + query = select([UserModel]).where(UserModel.meta_mask_address == address) + existing_address = await database.fetch_one(query) + + if existing_address: + return "Address already exists in the database" + + insert_query = UserModel.__table__.insert().values(meta_mask_address=address) + return await database.execute(insert_query) + + finally: + await disconnect_from_database() + + +async def save_transaction( + payType: str, amount: str, destinationAddress: str, txHash: str +): + try: + await connect_to_database() + insert_query = TransactionRecordModel.__table__.insert().values( + payType=payType, + amount=amount, + destinationAddress=destinationAddress, + txHash=txHash, + ) + return await database.execute(insert_query) + finally: + await disconnect_from_database() diff --git a/backend/docs/Makefile b/backend/docs/Makefile deleted file mode 100644 index 69577002..00000000 --- a/backend/docs/Makefile +++ /dev/null @@ -1,29 +0,0 @@ -# Minimal makefile for Sphinx documentation -# - -# You can set these variables from the command line, and also -# from the environment for the first two. -SPHINXOPTS ?= -SPHINXBUILD ?= sphinx-build -SOURCEDIR = . -BUILDDIR = ./_build -APP = /app - -.PHONY: help livehtml apidocs Makefile - -# Put it first so that "make" without argument is like "make help". -help: - @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -c . - -# Build, watch and serve docs with live reload -livehtml: - sphinx-autobuild -b html --host 0.0.0.0 --port 9000 --watch $(APP) -c . $(SOURCEDIR) $(BUILDDIR)/html - -# Outputs rst files from django application code -apidocs: - sphinx-apidoc -o $(SOURCEDIR)/api $(APP) - -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). -%: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -c . diff --git a/backend/docs/__init__.py b/backend/docs/__init__.py deleted file mode 100644 index 8772c827..00000000 --- a/backend/docs/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Included so that Django's startproject comment runs against the docs directory diff --git a/backend/docs/conf.py b/backend/docs/conf.py deleted file mode 100644 index 8e54ef5a..00000000 --- a/backend/docs/conf.py +++ /dev/null @@ -1,63 +0,0 @@ -# Configuration file for the Sphinx documentation builder. -# -# This file only contains a selection of the most common options. For a full -# list see the documentation: -# https://www.sphinx-doc.org/en/master/usage/configuration.html - -# -- Path setup -------------------------------------------------------------- - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. - -import os -import sys -import django - -if os.getenv("READTHEDOCS", default=False) == "True": - sys.path.insert(0, os.path.abspath("..")) - os.environ["DJANGO_READ_DOT_ENV_FILE"] = "True" - os.environ["USE_DOCKER"] = "no" -else: - sys.path.insert(0, os.path.abspath("/app")) -os.environ["DATABASE_URL"] = "sqlite:///readthedocs.db" -os.environ["CELERY_BROKER_URL"] = os.getenv("REDIS_URL", "redis://redis:6379") -os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local") -django.setup() - -# -- Project information ----------------------------------------------------- - -project = "backend" -copyright = """2023, Alan""" -author = "Alan" - - -# -- General configuration --------------------------------------------------- - -# Add any Sphinx extension module names here, as strings. They can be -# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom -# ones. -extensions = [ - "sphinx.ext.autodoc", - "sphinx.ext.napoleon", -] - -# Add any paths that contain templates here, relative to this directory. -# templates_path = ["_templates"] - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - -# -- Options for HTML output ------------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "alabaster" - -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -# html_static_path = ["_static"] diff --git a/backend/docs/howto.rst b/backend/docs/howto.rst deleted file mode 100644 index f0d78eb4..00000000 --- a/backend/docs/howto.rst +++ /dev/null @@ -1,38 +0,0 @@ -How To - Project Documentation -====================================================================== - -Get Started ----------------------------------------------------------------------- - -Documentation can be written as rst files in `backend/docs`. - - -To build and serve docs, use the commands:: - - docker compose -f local.yml up docs - - - -Changes to files in `docs/_source` will be picked up and reloaded automatically. - -`Sphinx `_ is the tool used to build documentation. - -Docstrings to Documentation ----------------------------------------------------------------------- - -The sphinx extension `apidoc `_ is used to automatically document code using signatures and docstrings. - -Numpy or Google style docstrings will be picked up from project files and available for documentation. See the `Napoleon `_ extension for details. - -For an in-use example, see the `page source <_sources/users.rst.txt>`_ for :ref:`users`. - -To compile all docstrings automatically into documentation source files, use the command: - :: - - make apidocs - - -This can be done in the docker container: - :: - - docker run --rm docs make apidocs diff --git a/backend/docs/index.rst b/backend/docs/index.rst deleted file mode 100644 index abfe8045..00000000 --- a/backend/docs/index.rst +++ /dev/null @@ -1,23 +0,0 @@ -.. backend documentation master file, created by - sphinx-quickstart. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. - -Welcome to backend's documentation! -====================================================================== - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - howto - users - - - -Indices and tables -================== - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/backend/docs/make.bat b/backend/docs/make.bat deleted file mode 100644 index e7abb4be..00000000 --- a/backend/docs/make.bat +++ /dev/null @@ -1,46 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -c . -) -set SOURCEDIR=_source -set BUILDDIR=_build -set APP=..\backend - -if "%1" == "" goto help - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.Install sphinx-autobuild for live serving. - echo.If you don't have Sphinx installed, grab it from - echo.http://sphinx-doc.org/ - exit /b 1 -) - -%SPHINXBUILD% -b %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:livehtml -sphinx-autobuild -b html --open-browser -p 9000 --watch %APP% -c . %SOURCEDIR% %BUILDDIR%/html -GOTO :EOF - -:apidocs -sphinx-apidoc -o %SOURCEDIR%/api %APP% -GOTO :EOF - -:help -%SPHINXBUILD% -b help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/backend/docs/users.rst b/backend/docs/users.rst deleted file mode 100644 index 01c105d2..00000000 --- a/backend/docs/users.rst +++ /dev/null @@ -1,15 +0,0 @@ - .. _users: - -Users -====================================================================== - -Starting a new project, it’s highly recommended to set up a custom user model, -even if the default User model is sufficient for you. - -This model behaves identically to the default user model, -but you’ll be able to customize it in the future if the need arises. - -.. automodule:: backend.users.models - :members: - :noindex: - diff --git a/backend/embedding.py b/backend/embedding.py new file mode 100644 index 00000000..6e77a9ea --- /dev/null +++ b/backend/embedding.py @@ -0,0 +1,61 @@ +from typing import Optional + +import vertexai +from vertexai.vision_models import ( + Image, + MultiModalEmbeddingModel, + MultiModalEmbeddingResponse, + Video, + VideoSegmentConfig, +) + + + +def get_image_video_text_embeddings( + project_id: str, + location: str, + image_path: str, + video_path: str, + contextual_text: Optional[str] = None, + dimension: Optional[int] = 1408, + video_segment_config: Optional[VideoSegmentConfig] = None, +) -> MultiModalEmbeddingResponse: + """Example of how to generate multimodal embeddings from image, video, and text. + + Args: + project_id: Google Cloud Project ID, used to initialize vertexai + location: Google Cloud Region, used to initialize vertexai + image_path: Path to image (local or Google Cloud Storage) to generate embeddings for. + video_path: Path to video (local or Google Cloud Storage) to generate embeddings for. + contextual_text: Text to generate embeddings for. + dimension: Dimension for the returned embeddings. + https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-multimodal-embeddings#low-dimension + video_segment_config: Define specific segments to generate embeddings for. + https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-multimodal-embeddings#video-best-practices + """ + + vertexai.init(project=project_id, location=location) + + model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding",) + image = Image.load_from_file(image_path) + video = Video.load_from_file(video_path) + + embeddings = model.get_embeddings( + image=image, + video=video, + video_segment_config=video_segment_config, + contextual_text=contextual_text, + dimension=dimension, + ) + + print(f"Image Embedding: {embeddings.image_embedding}") + + # Video Embeddings are segmented based on the video_segment_config. + print("Video Embeddings:") + for video_embedding in embeddings.video_embeddings: + print( + f"Video Segment: {video_embedding.start_offset_sec} - {video_embedding.end_offset_sec}" + ) + print(f"Embedding: {video_embedding.embedding}") + + print(f"Text Embedding: {embeddings.text_embedding}") diff --git a/backend/local.yml b/backend/local.yml deleted file mode 100644 index 3185f1e3..00000000 --- a/backend/local.yml +++ /dev/null @@ -1,84 +0,0 @@ -version: '3' - -volumes: - backend_local_postgres_data: {} - backend_local_postgres_data_backups: {} - -services: - django: &django - build: - context: . - dockerfile: ./compose/local/django/Dockerfile - image: backend_local_django - container_name: backend_local_django - depends_on: - - postgres - - redis - volumes: - - .:/app:z - env_file: - - ./.envs/.local/.django - - ./.envs/.local/.postgres - ports: - - '8000:8000' - command: /start - - postgres: - build: - context: . - dockerfile: ./compose/production/postgres/Dockerfile - image: backend_production_postgres - container_name: backend_local_postgres - volumes: - - backend_local_postgres_data:/var/lib/postgresql/data - - backend_local_postgres_data_backups:/backups - env_file: - - ./.envs/.local/.postgres - - docs: - image: backend_local_docs - container_name: backend_local_docs - build: - context: . - dockerfile: ./compose/local/docs/Dockerfile - env_file: - - ./.envs/.local/.django - volumes: - - ./docs:/docs:z - - ./config:/app/config:z - - ./backend:/app/backend:z - ports: - - '9000:9000' - command: /start-docs - - redis: - image: redis:6 - container_name: backend_local_redis - - celeryworker: - <<: *django - image: backend_local_celeryworker - container_name: backend_local_celeryworker - depends_on: - - redis - - postgres - ports: [] - command: /start-celeryworker - - celerybeat: - <<: *django - image: backend_local_celerybeat - container_name: backend_local_celerybeat - depends_on: - - redis - - postgres - ports: [] - command: /start-celerybeat - - flower: - <<: *django - image: backend_local_flower - container_name: backend_local_flower - ports: - - '5555:5555' - command: /start-flower diff --git a/backend/locale/README.md b/backend/locale/README.md deleted file mode 100644 index 9911fec2..00000000 --- a/backend/locale/README.md +++ /dev/null @@ -1,32 +0,0 @@ -# Translations - -Start by configuring the `LANGUAGES` settings in `base.py`, by uncommenting languages you are willing to support. Then, translations strings will be placed in this folder when running: - -```bash -docker compose -f local.yml run --rm django python manage.py makemessages -all --no-location -``` - -This should generate `django.po` (stands for Portable Object) files under each locale `/LC_MESSAGES/django.po`. Each translatable string in the codebase is collected with its `msgid` and need to be translated as `msgstr`, for example: - -```po -msgid "users" -msgstr "utilisateurs" -``` - -Once all translations are done, they need to be compiled into `.mo` files (stands for Machine Object), which are the actual binary files used by the application: - -```bash -docker compose -f local.yml run --rm django python manage.py compilemessages -``` - -Note that the `.po` files are NOT used by the application directly, so if the `.mo` files are out of dates, the content won't appear as translated even if the `.po` files are up-to-date. - -## Production - -The production image runs `compilemessages` automatically at build time, so as long as your translated source files (PO) are up-to-date, you're good to go. - -## Add a new language - -1. Update the [`LANGUAGES` setting](https://docs.djangoproject.com/en/stable/ref/settings/#std-setting-LANGUAGES) to your project's base settings. -2. Create the locale folder for the language next to this file, e.g. `fr_FR` for French. Make sure the case is correct. -3. Run `makemessages` (as instructed above) to generate the PO files for the new language. diff --git a/backend/locale/en_US/LC_MESSAGES/django.po b/backend/locale/en_US/LC_MESSAGES/django.po deleted file mode 100644 index b005751d..00000000 --- a/backend/locale/en_US/LC_MESSAGES/django.po +++ /dev/null @@ -1,12 +0,0 @@ -# Translations for the backend project -# Copyright (C) 2023 Alan -# Alan , 2023. -# -#, fuzzy -msgid "" -msgstr "" -"Project-Id-Version: 0.1.0\n" -"Language: en-US\n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=UTF-8\n" -"Content-Transfer-Encoding: 8bit\n" diff --git a/backend/locale/fr_FR/LC_MESSAGES/django.po b/backend/locale/fr_FR/LC_MESSAGES/django.po deleted file mode 100644 index 9ae61f39..00000000 --- a/backend/locale/fr_FR/LC_MESSAGES/django.po +++ /dev/null @@ -1,335 +0,0 @@ -# Translations for the backend project -# Copyright (C) 2023 Alan -# Alan , 2023. -# -#, fuzzy -msgid "" -msgstr "" -"Project-Id-Version: 0.1.0\n" -"Language: fr-FR\n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=UTF-8\n" -"Content-Transfer-Encoding: 8bit\n" -"Plural-Forms: nplurals=2; plural=(n > 1);\n" -#: backend/templates/account/account_inactive.html:5 -#: backend/templates/account/account_inactive.html:8 -msgid "Account Inactive" -msgstr "Compte inactif" - -#: backend/templates/account/account_inactive.html:10 -msgid "This account is inactive." -msgstr "Ce compte est inactif." - -#: backend/templates/account/email.html:7 -msgid "Account" -msgstr "Compte" - -#: backend/templates/account/email.html:10 -msgid "E-mail Addresses" -msgstr "Adresses e-mail" - -#: backend/templates/account/email.html:13 -msgid "The following e-mail addresses are associated with your account:" -msgstr "Les adresses e-mail suivantes sont associées à votre compte :" - -#: backend/templates/account/email.html:27 -msgid "Verified" -msgstr "Vérifié" - -#: backend/templates/account/email.html:29 -msgid "Unverified" -msgstr "Non vérifié" - -#: backend/templates/account/email.html:31 -msgid "Primary" -msgstr "Primaire" - -#: backend/templates/account/email.html:37 -msgid "Make Primary" -msgstr "Changer Primaire" - -#: backend/templates/account/email.html:38 -msgid "Re-send Verification" -msgstr "Renvoyer vérification" - -#: backend/templates/account/email.html:39 -msgid "Remove" -msgstr "Supprimer" - -#: backend/templates/account/email.html:46 -msgid "Warning:" -msgstr "Avertissement:" - -#: backend/templates/account/email.html:46 -msgid "" -"You currently do not have any e-mail address set up. You should really add " -"an e-mail address so you can receive notifications, reset your password, etc." -msgstr "" -"Vous n'avez actuellement aucune adresse e-mail configurée. Vous devriez ajouter " -"une adresse e-mail pour reçevoir des notifications, réinitialiser votre mot " -"de passe, etc." - -#: backend/templates/account/email.html:51 -msgid "Add E-mail Address" -msgstr "Ajouter une adresse e-mail" - -#: backend/templates/account/email.html:56 -msgid "Add E-mail" -msgstr "Ajouter e-mail" - -#: backend/templates/account/email.html:66 -msgid "Do you really want to remove the selected e-mail address?" -msgstr "Voulez-vous vraiment supprimer l'adresse e-mail sélectionnée ?" - -#: backend/templates/account/email_confirm.html:6 -#: backend/templates/account/email_confirm.html:10 -msgid "Confirm E-mail Address" -msgstr "Confirmez votre adresse email" - -#: backend/templates/account/email_confirm.html:16 -#, python-format -msgid "" -"Please confirm that %(email)s is an e-mail " -"address for user %(user_display)s." -msgstr "" -"Veuillez confirmer que %(email)s est un e-mail " -"adresse de l'utilisateur %(user_display)s." - -#: backend/templates/account/email_confirm.html:20 -msgid "Confirm" -msgstr "Confirm" - -#: backend/templates/account/email_confirm.html:27 -#, python-format -msgid "" -"This e-mail confirmation link expired or is invalid. Please issue a new e-mail confirmation request." -msgstr "" -"Ce lien de confirmation par e-mail a expiré ou n'est pas valide. Veuillez" - "émettre une nouvelle demande de confirmation " -"par e-mail." - -#: backend/templates/account/login.html:7 -#: backend/templates/account/login.html:11 -#: backend/templates/account/login.html:56 -#: backend/templates/base.html:72 -msgid "Sign In" -msgstr "S'identifier" - -#: backend/templates/account/login.html:17 -msgid "Please sign in with one of your existing third party accounts:" -msgstr "Veuillez vous connecter avec l'un de vos comptes tiers existants :" - -#: backend/templates/account/login.html:19 -#, python-format -msgid "" -"Or, sign up for a %(site_name)s account and " -"sign in below:" -msgstr "" -"Ou, créez un compte %(site_name)s et " -"connectez-vous ci-dessous :" - -#: backend/templates/account/login.html:32 -msgid "or" -msgstr "ou" - -#: backend/templates/account/login.html:41 -#, python-format -msgid "" -"If you have not created an account yet, then please sign up first." -msgstr "" -"Si vous n'avez pas encore créé de compte, veuillez d'abord vous inscrire." - -#: backend/templates/account/login.html:55 -msgid "Forgot Password?" -msgstr "Mot de passe oublié?" - -#: backend/templates/account/logout.html:5 -#: backend/templates/account/logout.html:8 -#: backend/templates/account/logout.html:17 -#: backend/templates/base.html:61 -msgid "Sign Out" -msgstr "Se déconnecter" - -#: backend/templates/account/logout.html:10 -msgid "Are you sure you want to sign out?" -msgstr "Êtes-vous certain de vouloir vous déconnecter?" - -#: backend/templates/account/password_change.html:6 -#: backend/templates/account/password_change.html:9 -#: backend/templates/account/password_change.html:14 -#: backend/templates/account/password_reset_from_key.html:5 -#: backend/templates/account/password_reset_from_key.html:8 -#: backend/templates/account/password_reset_from_key_done.html:4 -#: backend/templates/account/password_reset_from_key_done.html:7 -msgid "Change Password" -msgstr "Changer le mot de passe" - -#: backend/templates/account/password_reset.html:7 -#: backend/templates/account/password_reset.html:11 -#: backend/templates/account/password_reset_done.html:6 -#: backend/templates/account/password_reset_done.html:9 -msgid "Password Reset" -msgstr "Réinitialisation du mot de passe" - -#: backend/templates/account/password_reset.html:16 -msgid "" -"Forgotten your password? Enter your e-mail address below, and we'll send you " -"an e-mail allowing you to reset it." -msgstr "" -"Mot de passe oublié? Entrez votre adresse e-mail ci-dessous, et nous vous " -"enverrons un e-mail vous permettant de le réinitialiser." - -#: backend/templates/account/password_reset.html:21 -msgid "Reset My Password" -msgstr "Réinitialiser mon mot de passe" - -#: backend/templates/account/password_reset.html:24 -msgid "Please contact us if you have any trouble resetting your password." -msgstr "" -"Veuillez nous contacter si vous rencontrez des difficultés pour réinitialiser" -"votre mot de passe." - -#: backend/templates/account/password_reset_done.html:15 -msgid "" -"We have sent you an e-mail. Please contact us if you do not receive it " -"within a few minutes." -msgstr "" -"Nous vous avons envoyé un e-mail. Veuillez nous contacter si vous ne le " -"recevez pas d'ici quelques minutes." - -#: backend/templates/account/password_reset_from_key.html:8 -msgid "Bad Token" -msgstr "Token Invalide" - -#: backend/templates/account/password_reset_from_key.html:12 -#, python-format -msgid "" -"The password reset link was invalid, possibly because it has already been " -"used. Please request a new password reset." -msgstr "" -"Le lien de réinitialisation du mot de passe n'était pas valide, peut-être parce " -"qu'il a déjà été utilisé. Veuillez faire une " -"nouvelle demande de réinitialisation de mot de passe." - -#: backend/templates/account/password_reset_from_key.html:18 -msgid "change password" -msgstr "changer le mot de passe" - -#: backend/templates/account/password_reset_from_key.html:21 -#: backend/templates/account/password_reset_from_key_done.html:8 -msgid "Your password is now changed." -msgstr "Votre mot de passe est maintenant modifié." - -#: backend/templates/account/password_set.html:6 -#: backend/templates/account/password_set.html:9 -#: backend/templates/account/password_set.html:14 -msgid "Set Password" -msgstr "Définir le mot de passe" - -#: backend/templates/account/signup.html:6 -msgid "Signup" -msgstr "S'inscrire" - -#: backend/templates/account/signup.html:9 -#: backend/templates/account/signup.html:19 -#: backend/templates/base.html:67 -msgid "Sign Up" -msgstr "S'inscrire" - -#: backend/templates/account/signup.html:11 -#, python-format -msgid "" -"Already have an account? Then please sign in." -msgstr "" -"Vous avez déjà un compte? Alors veuillez vous connecter." - -#: backend/templates/account/signup_closed.html:5 -#: backend/templates/account/signup_closed.html:8 -msgid "Sign Up Closed" -msgstr "Inscriptions closes" - -#: backend/templates/account/signup_closed.html:10 -msgid "We are sorry, but the sign up is currently closed." -msgstr "Désolé, mais l'inscription est actuellement fermée." - -#: backend/templates/account/verification_sent.html:5 -#: backend/templates/account/verification_sent.html:8 -#: backend/templates/account/verified_email_required.html:5 -#: backend/templates/account/verified_email_required.html:8 -msgid "Verify Your E-mail Address" -msgstr "Vérifiez votre adresse e-mail" - -#: backend/templates/account/verification_sent.html:10 -msgid "" -"We have sent an e-mail to you for verification. Follow the link provided to " -"finalize the signup process. Please contact us if you do not receive it " -"within a few minutes." -msgstr "Nous vous avons envoyé un e-mail pour vérification. Suivez le lien fourni " -"pour finalisez le processus d'inscription. Veuillez nous contacter si vous ne le " -"recevez pas d'ici quelques minutes." - -#: backend/templates/account/verified_email_required.html:12 -msgid "" -"This part of the site requires us to verify that\n" -"you are who you claim to be. For this purpose, we require that you\n" -"verify ownership of your e-mail address. " -msgstr "" -"Cette partie du site nous oblige à vérifier que\n" -"vous êtes qui vous prétendez être. Nous vous demandons donc de\n" -"vérifier la propriété de votre adresse e-mail." - -#: backend/templates/account/verified_email_required.html:16 -msgid "" -"We have sent an e-mail to you for\n" -"verification. Please click on the link inside this e-mail. Please\n" -"contact us if you do not receive it within a few minutes." -msgstr "" -"Nous vous avons envoyé un e-mail pour\n" -"vérification. Veuillez cliquer sur le lien contenu dans cet e-mail. Veuillez nous\n" -"contacter si vous ne le recevez pas d'ici quelques minutes." - -#: backend/templates/account/verified_email_required.html:20 -#, python-format -msgid "" -"Note: you can still change your e-" -"mail address." -msgstr "" -"Remarque : vous pouvez toujours changer votre e-" -"adresse e-mail." - -#: backend/templates/base.html:57 -msgid "My Profile" -msgstr "Mon Profil" - -#: backend/users/admin.py:17 -msgid "Personal info" -msgstr "Personal info" - -#: backend/users/admin.py:19 -msgid "Permissions" -msgstr "Permissions" - -#: backend/users/admin.py:30 -msgid "Important dates" -msgstr "Dates importantes" - -#: backend/users/apps.py:7 -msgid "Users" -msgstr "Utilisateurs" - -#: backend/users/forms.py:24 -#: backend/users/tests/test_forms.py:36 -msgid "This username has already been taken." -msgstr "Ce nom d'utilisateur est déjà pris." - -#: backend/users/models.py:15 -msgid "Name of User" -msgstr "Nom de l'utilisateur" - -#: backend/users/views.py:23 -msgid "Information successfully updated" -msgstr "Informations mises à jour avec succès" diff --git a/backend/locale/pt_BR/LC_MESSAGES/django.po b/backend/locale/pt_BR/LC_MESSAGES/django.po deleted file mode 100644 index 0b71ea6e..00000000 --- a/backend/locale/pt_BR/LC_MESSAGES/django.po +++ /dev/null @@ -1,315 +0,0 @@ -# Translations for the backend project -# Copyright (C) 2023 Alan -# Alan , 2023. -# -#, fuzzy -msgid "" -msgstr "" -"Project-Id-Version: 0.1.0\n" -"Language: pt-BR\n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=UTF-8\n" -"Content-Transfer-Encoding: 8bit\n" -"Plural-Forms: nplurals=2; plural=(n > 1);\n" -#: backend/templates/account/account_inactive.html:5 -#: backend/templates/account/account_inactive.html:8 -msgid "Account Inactive" -msgstr "Conta Inativa" - -#: backend/templates/account/account_inactive.html:10 -msgid "This account is inactive." -msgstr "Esta conta está inativa." - -#: backend/templates/account/email.html:7 -msgid "Account" -msgstr "Conta" - -#: backend/templates/account/email.html:10 -msgid "E-mail Addresses" -msgstr "Endereços de E-mail" - -#: backend/templates/account/email.html:13 -msgid "The following e-mail addresses are associated with your account:" -msgstr "Os seguintes endereços de e-mail estão associados à sua conta:" - -#: backend/templates/account/email.html:27 -msgid "Verified" -msgstr "Verificado" - -#: backend/templates/account/email.html:29 -msgid "Unverified" -msgstr "Não verificado" - -#: backend/templates/account/email.html:31 -msgid "Primary" -msgstr "Primário" - -#: backend/templates/account/email.html:37 -msgid "Make Primary" -msgstr "Tornar Primário" - -#: backend/templates/account/email.html:38 -msgid "Re-send Verification" -msgstr "Reenviar verificação" - -#: backend/templates/account/email.html:39 -msgid "Remove" -msgstr "Remover" - -#: backend/templates/account/email.html:46 -msgid "Warning:" -msgstr "Aviso:" - -#: backend/templates/account/email.html:46 -msgid "" -"You currently do not have any e-mail address set up. You should really add " -"an e-mail address so you can receive notifications, reset your password, etc." -msgstr "" -"No momento, você não tem nenhum endereço de e-mail configurado. Você " -"realmente deve adicionar um endereço de e-mail para receber notificações, " -"redefinir sua senha etc." - -#: backend/templates/account/email.html:51 -msgid "Add E-mail Address" -msgstr "Adicionar Endereço de E-mail" - -#: backend/templates/account/email.html:56 -msgid "Add E-mail" -msgstr "Adicionar E-mail" - -#: backend/templates/account/email.html:66 -msgid "Do you really want to remove the selected e-mail address?" -msgstr "Você realmente deseja remover o endereço de e-mail selecionado?" - -#: backend/templates/account/email_confirm.html:6 -#: backend/templates/account/email_confirm.html:10 -msgid "Confirm E-mail Address" -msgstr "Confirme o endereço de e-mail" - -#: backend/templates/account/email_confirm.html:16 -#, python-format -msgid "" -"Please confirm that %(email)s is an e-mail " -"address for user %(user_display)s." -msgstr "" -"Confirme se %(email)s é um endereço de " -"e-mail do usuário %(user_display)s." - -#: backend/templates/account/email_confirm.html:20 -msgid "Confirm" -msgstr "Confirmar" - -#: backend/templates/account/email_confirm.html:27 -#, python-format -msgid "" -"This e-mail confirmation link expired or is invalid. Please issue a new e-mail confirmation request." -msgstr "Este link de confirmação de e-mail expirou ou é inválido. " -"Por favor, emita um novo pedido de confirmação por e-mail." - -#: backend/templates/account/login.html:7 -#: backend/templates/account/login.html:11 -#: backend/templates/account/login.html:56 -#: backend/templates/base.html:72 -msgid "Sign In" -msgstr "Entrar" - -#: backend/templates/account/login.html:17 -msgid "Please sign in with one of your existing third party accounts:" -msgstr "Faça login com uma de suas contas de terceiros existentes:" - -#: backend/templates/account/login.html:19 -#, python-format -msgid "" -"Or, sign up for a %(site_name)s account and " -"sign in below:" -msgstr "Ou, cadastre-se para uma conta em %(site_name)s e entre abaixo:" - -#: backend/templates/account/login.html:32 -msgid "or" -msgstr "ou" - -#: backend/templates/account/login.html:41 -#, python-format -msgid "" -"If you have not created an account yet, then please sign up first." -msgstr "Se você ainda não criou uma conta, registre-se primeiro." - -#: backend/templates/account/login.html:55 -msgid "Forgot Password?" -msgstr "Esqueceu sua senha?" - -#: backend/templates/account/logout.html:5 -#: backend/templates/account/logout.html:8 -#: backend/templates/account/logout.html:17 -#: backend/templates/base.html:61 -msgid "Sign Out" -msgstr "Sair" - -#: backend/templates/account/logout.html:10 -msgid "Are you sure you want to sign out?" -msgstr "Você tem certeza que deseja sair?" - -#: backend/templates/account/password_change.html:6 -#: backend/templates/account/password_change.html:9 -#: backend/templates/account/password_change.html:14 -#: backend/templates/account/password_reset_from_key.html:5 -#: backend/templates/account/password_reset_from_key.html:8 -#: backend/templates/account/password_reset_from_key_done.html:4 -#: backend/templates/account/password_reset_from_key_done.html:7 -msgid "Change Password" -msgstr "Alterar Senha" - -#: backend/templates/account/password_reset.html:7 -#: backend/templates/account/password_reset.html:11 -#: backend/templates/account/password_reset_done.html:6 -#: backend/templates/account/password_reset_done.html:9 -msgid "Password Reset" -msgstr "Redefinição de senha" - -#: backend/templates/account/password_reset.html:16 -msgid "" -"Forgotten your password? Enter your e-mail address below, and we'll send you " -"an e-mail allowing you to reset it." -msgstr "Esqueceu sua senha? Digite seu endereço de e-mail abaixo e enviaremos um e-mail permitindo que você o redefina." - -#: backend/templates/account/password_reset.html:21 -msgid "Reset My Password" -msgstr "Redefinir minha senha" - -#: backend/templates/account/password_reset.html:24 -msgid "Please contact us if you have any trouble resetting your password." -msgstr "Entre em contato conosco se tiver algum problema para redefinir sua senha." - -#: backend/templates/account/password_reset_done.html:15 -msgid "" -"We have sent you an e-mail. Please contact us if you do not receive it " -"within a few minutes." -msgstr "Enviamos um e-mail para você. Entre em contato conosco se você não recebê-lo dentro de alguns minutos." - -#: backend/templates/account/password_reset_from_key.html:8 -msgid "Bad Token" -msgstr "Token Inválido" - -#: backend/templates/account/password_reset_from_key.html:12 -#, python-format -msgid "" -"The password reset link was invalid, possibly because it has already been " -"used. Please request a new password reset." -msgstr "O link de redefinição de senha era inválido, possivelmente porque já foi usado. " -"Solicite uma nova redefinição de senha." - -#: backend/templates/account/password_reset_from_key.html:18 -msgid "change password" -msgstr "alterar senha" - -#: backend/templates/account/password_reset_from_key.html:21 -#: backend/templates/account/password_reset_from_key_done.html:8 -msgid "Your password is now changed." -msgstr "Sua senha agora foi alterada." - -#: backend/templates/account/password_set.html:6 -#: backend/templates/account/password_set.html:9 -#: backend/templates/account/password_set.html:14 -msgid "Set Password" -msgstr "Definir Senha" - -#: backend/templates/account/signup.html:6 -msgid "Signup" -msgstr "Cadastro" - -#: backend/templates/account/signup.html:9 -#: backend/templates/account/signup.html:19 -#: backend/templates/base.html:67 -msgid "Sign Up" -msgstr "Cadastro" - -#: backend/templates/account/signup.html:11 -#, python-format -msgid "" -"Already have an account? Then please sign in." -msgstr "já tem uma conta? Então, por favor, faça login." - -#: backend/templates/account/signup_closed.html:5 -#: backend/templates/account/signup_closed.html:8 -msgid "Sign Up Closed" -msgstr "Inscrições encerradas" - -#: backend/templates/account/signup_closed.html:10 -msgid "We are sorry, but the sign up is currently closed." -msgstr "Lamentamos, mas as inscrições estão encerradas no momento." - -#: backend/templates/account/verification_sent.html:5 -#: backend/templates/account/verification_sent.html:8 -#: backend/templates/account/verified_email_required.html:5 -#: backend/templates/account/verified_email_required.html:8 -msgid "Verify Your E-mail Address" -msgstr "Verifique seu endereço de e-mail" - -#: backend/templates/account/verification_sent.html:10 -msgid "" -"We have sent an e-mail to you for verification. Follow the link provided to " -"finalize the signup process. Please contact us if you do not receive it " -"within a few minutes." -msgstr "Enviamos um e-mail para você para verificação. Siga o link fornecido para finalizar o processo de inscrição. Entre em contato conosco se você não recebê-lo dentro de alguns minutos." - -#: backend/templates/account/verified_email_required.html:12 -msgid "" -"This part of the site requires us to verify that\n" -"you are who you claim to be. For this purpose, we require that you\n" -"verify ownership of your e-mail address. " -msgstr "Esta parte do site exige que verifiquemos se você é quem afirma ser.\n" -"Para esse fim, exigimos que você verifique a propriedade\n" -"do seu endereço de e-mail." - -#: backend/templates/account/verified_email_required.html:16 -msgid "" -"We have sent an e-mail to you for\n" -"verification. Please click on the link inside this e-mail. Please\n" -"contact us if you do not receive it within a few minutes." -msgstr "Enviamos um e-mail para você para verificação.\n" -"Por favor, clique no link dentro deste e-mail.\n" -"Entre em contato conosco se você não recebê-lo dentro de alguns minutos." - -#: backend/templates/account/verified_email_required.html:20 -#, python-format -msgid "" -"Note: you can still change your e-" -"mail address." -msgstr "Nota: você ainda pode alterar seu endereço de e-mail." - -#: backend/templates/base.html:57 -msgid "My Profile" -msgstr "Meu perfil" - -#: backend/users/admin.py:17 -msgid "Personal info" -msgstr "Informação pessoal" - -#: backend/users/admin.py:19 -msgid "Permissions" -msgstr "Permissões" - -#: backend/users/admin.py:30 -msgid "Important dates" -msgstr "Datas importantes" - -#: backend/users/apps.py:7 -msgid "Users" -msgstr "Usuários" - -#: backend/users/forms.py:24 -#: backend/users/tests/test_forms.py:36 -msgid "This username has already been taken." -msgstr "Este nome de usuário já foi usado." - -#: backend/users/models.py:15 -msgid "Name of User" -msgstr "Nome do Usuário" - -#: backend/users/views.py:23 -msgid "Information successfully updated" -msgstr "Informação atualizada com sucesso" diff --git a/backend/main.py b/backend/main.py new file mode 100644 index 00000000..0f142b5c --- /dev/null +++ b/backend/main.py @@ -0,0 +1,161 @@ +from fastapi import FastAPI, File, UploadFile, Form, HTTPException +from fastapi.responses import JSONResponse +from controller import ( + save_address, + save_file_data, + get_file_by_id, + get_all_files, + save_transaction, +) +from databases import Database +import os +from embedding import get_image_video_text_embeddings +from dotenv import load_dotenv +from fastapi_sqlalchemy import DBSessionMiddleware, db +from fastapi.middleware.cors import CORSMiddleware +from pydantic import BaseModel + +# Load .env file +load_dotenv() + +project_id = os.getenv("PROJECT_ID") +location = os.getenv("LOCATION") +database_url = os.getenv("DATABASE_URL") + +app = FastAPI() + +database = Database(database_url) + +# Configure CORS +origins = [ + "http://localhost:3000", # Add the origins you want to allow +] + +# to avoid csrftokenError +app.add_middleware(DBSessionMiddleware, db_url=database_url) + +app.add_middleware( + CORSMiddleware, + allow_origins=origins, + allow_credentials=True, + allow_methods=["*"], # You can specify specific HTTP methods if needed + allow_headers=["*"], # You can specify specific HTTP headers if needed +) + + +@app.get("/") +def root(): + return {"message": "Hello World"} + + +@app.on_event("startup") +async def startup_db_client(): + await database.connect() + + +@app.on_event("shutdown") +async def shutdown_db_client(): + await database.disconnect() + + +@app.post("/upload") +async def upload_files( + video: UploadFile = File(...), + image: UploadFile = File(...), + audio: UploadFile = File(...), + text: str = Form(...), + address: str = Form(...), +): + # Create an "uploads" folder if it doesn't exist + upload_folder = "uploads" + os.makedirs(upload_folder, exist_ok=True) + + # Save files with full path + video_filename = os.path.join(upload_folder, video.filename) + image_filename = os.path.join(upload_folder, image.filename) + audio_filename = os.path.join(upload_folder, audio.filename) + + # Example: Get embeddings using OpenAI API + # get_image_video_text_embeddings( + # project_id, location, image_filename, video_filename, text + # ) + + # Save video file + with open(video_filename, "wb") as video_file: + video_file.write(video.file.read()) + + # Save image file + with open(image_filename, "wb") as image_file: + image_file.write(image.file.read()) + + # Save audio file + with open(audio_filename, "wb") as audio_file: + audio_file.write(audio.file.read()) + + # Process text data + processed_text = text.upper() + + # Save data to database + file_id = await save_file_data( + video_filename, image_filename, audio_filename, processed_text, address + ) + + return JSONResponse( + content={ + "message": "Files uploaded and data saved successfully", + "file_id": file_id, + } + ) + + +@app.get("/files/{file_id}") +async def get_file(file_id: int): + file_data = await get_file_by_id(file_id) + if file_data is None: + raise HTTPException(status_code=404, detail="File not found") + + return file_data + + +@app.get("/files") +async def get_all_files_endpoint(): + all_files = await get_all_files() + return all_files + + +class MetaMaskAddress(BaseModel): + address: str + + +@app.post("/saveMetamask") +async def save_metamask_address(meta_mask_address: MetaMaskAddress): + address = meta_mask_address.address + user_id = await save_address(address) + return JSONResponse( + content={ + "message": "Address saved successfully", + "user_id": user_id, + } + ) + + +class Transaction(BaseModel): + payType: str + amount: float + destinationAddress: str + txHash: str + + +@app.post("/saveTransaction") +async def save_transaction_record(transaction: Transaction): + payType = transaction.payType + amount = transaction.amount + destinationAddress = transaction.destinationAddress + txHash = transaction.txHash + user_id = await save_transaction(payType, amount, destinationAddress, txHash) + return JSONResponse( + content={ + "message": "Transaction saved successfully", + "user_id": user_id, + } + ) diff --git a/backend/manage.py b/backend/manage.py deleted file mode 100644 index 0d99e4e2..00000000 --- a/backend/manage.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python -import os -import sys -from pathlib import Path - -if __name__ == "__main__": - os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local") - - try: - from django.core.management import execute_from_command_line - except ImportError: - # The above import may fail for some other reason. Ensure that the - # issue is really that Django is missing to avoid masking other - # exceptions on Python 2. - try: - import django # noqa - except ImportError: - raise ImportError( - "Couldn't import Django. Are you sure it's installed and " - "available on your PYTHONPATH environment variable? Did you " - "forget to activate a virtual environment?" - ) - - raise - - # This allows easy placement of apps within the interior - # backend directory. - current_path = Path(__file__).parent.resolve() - sys.path.append(str(current_path / "backend")) - - execute_from_command_line(sys.argv) diff --git a/backend/merge_production_dotenvs_in_dotenv.py b/backend/merge_production_dotenvs_in_dotenv.py deleted file mode 100644 index 35139fb2..00000000 --- a/backend/merge_production_dotenvs_in_dotenv.py +++ /dev/null @@ -1,26 +0,0 @@ -import os -from collections.abc import Sequence -from pathlib import Path - -BASE_DIR = Path(__file__).parent.resolve() -PRODUCTION_DOTENVS_DIR = BASE_DIR / ".envs" / ".production" -PRODUCTION_DOTENV_FILES = [ - PRODUCTION_DOTENVS_DIR / ".django", - PRODUCTION_DOTENVS_DIR / ".postgres", -] -DOTENV_FILE = BASE_DIR / ".env" - - -def merge( - output_file: Path, - files_to_merge: Sequence[Path], -) -> None: - merged_content = "" - for merge_file in files_to_merge: - merged_content += merge_file.read_text() - merged_content += os.linesep - output_file.write_text(merged_content) - - -if __name__ == "__main__": - merge(DOTENV_FILE, PRODUCTION_DOTENV_FILES) diff --git a/backend/models.py b/backend/models.py new file mode 100644 index 00000000..927d5337 --- /dev/null +++ b/backend/models.py @@ -0,0 +1,62 @@ +from typing import Optional +from pydantic import BaseModel +from sqlalchemy import ( + Column, + Integer, + String, + Text, + Float, + create_engine, + MetaData, + DateTime, +) +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.sql import func +from datetime import datetime + +# My code +import os +from dotenv import load_dotenv + +# Load .env file +load_dotenv() + +database_url = os.getenv("DATABASE_URL") + +engine = create_engine(database_url) +metadata = MetaData() + +Base = declarative_base() + + +class FileModel(Base): + __tablename__ = "files" + id = Column(Integer, primary_key=True, index=True) + video_filename = Column(String, index=True, nullable=False) + image_filename = Column(String, index=True, nullable=False) + audio_filename = Column(String, index=True, nullable=False) + text_data = Column(Text) + # time_created = Column(DateTime(timezone=True), server_default=func.now()) + # time_updated = Column(DateTime(timezone=True), onupdate=func.now()) + # wallet_address = Column(String, index=True, nullable=False) + + +class TransactionRecordModel(Base): + __tablename__ = "transactions" + id: int = Column(Integer, primary_key=True, index=True) + payType: str = Column(String, index=True, nullable=False) + amount: float = Column(Float, index=True, nullable=False) + destinationAddress: str = Column(String, index=True, nullable=False) + txHash: str = Column(String, index=True, nullable=False) + created_at: datetime = Column(DateTime(timezone=True), server_default=func.now()) + updated_at: datetime = Column(DateTime(timezone=True), server_default=func.now()) + deleted_at: datetime = Column(DateTime(timezone=True)) + + +class UserModel(Base): + __tablename__ = "users" + id = Column(Integer, primary_key=True, index=True) + meta_mask_address = Column(String, index=True, nullable=False) + + +Base.metadata.create_all(bind=engine) diff --git a/backend/production.yml b/backend/production.yml deleted file mode 100644 index 124d2b12..00000000 --- a/backend/production.yml +++ /dev/null @@ -1,73 +0,0 @@ -version: '3' - -volumes: - production_postgres_data: {} - production_postgres_data_backups: {} - production_traefik: {} - -services: - django: &django - build: - context: . - dockerfile: ./compose/production/django/Dockerfile - - image: backend_production_django - depends_on: - - postgres - - redis - env_file: - - ./.envs/.production/.django - - ./.envs/.production/.postgres - command: /start - - postgres: - build: - context: . - dockerfile: ./compose/production/postgres/Dockerfile - image: backend_production_postgres - volumes: - - production_postgres_data:/var/lib/postgresql/data - - production_postgres_data_backups:/backups - env_file: - - ./.envs/.production/.postgres - - traefik: - build: - context: . - dockerfile: ./compose/production/traefik/Dockerfile - image: backend_production_traefik - depends_on: - - django - volumes: - - production_traefik:/etc/traefik/acme - ports: - - '0.0.0.0:80:80' - - '0.0.0.0:443:443' - - '0.0.0.0:5555:5555' - - redis: - image: redis:6 - - celeryworker: - <<: *django - image: backend_production_celeryworker - command: /start-celeryworker - - celerybeat: - <<: *django - image: backend_production_celerybeat - command: /start-celerybeat - - flower: - <<: *django - image: backend_production_flower - command: /start-flower - - awscli: - build: - context: . - dockerfile: ./compose/production/aws/Dockerfile - env_file: - - ./.envs/.production/.django - volumes: - - production_postgres_data_backups:/backups:z diff --git a/backend/pyproject.toml b/backend/pyproject.toml deleted file mode 100644 index 2ba69992..00000000 --- a/backend/pyproject.toml +++ /dev/null @@ -1,105 +0,0 @@ -# ==== pytest ==== -[tool.pytest.ini_options] -minversion = "6.0" -addopts = "--ds=config.settings.test --reuse-db" -python_files = [ - "tests.py", - "test_*.py", -] - -# ==== Coverage ==== -[tool.coverage.run] -include = ["backend/**"] -omit = ["*/migrations/*", "*/tests/*"] -plugins = ["django_coverage_plugin"] - - -# ==== black ==== -[tool.black] -line-length = 119 -target-version = ['py311'] - - -# ==== isort ==== -[tool.isort] -profile = "black" -line_length = 119 -known_first_party = [ - "backend", - "config", -] -skip = ["venv/"] -skip_glob = ["**/migrations/*.py"] - - -# ==== mypy ==== -[tool.mypy] -python_version = "3.11" -check_untyped_defs = true -ignore_missing_imports = true -warn_unused_ignores = true -warn_redundant_casts = true -warn_unused_configs = true -plugins = [ - "mypy_django_plugin.main", - "mypy_drf_plugin.main", -] - -[[tool.mypy.overrides]] -# Django migrations should not produce any errors: -module = "*.migrations.*" -ignore_errors = true - -[tool.django-stubs] -django_settings_module = "config.settings.test" - - -# ==== PyLint ==== -[tool.pylint.MASTER] -load-plugins = [ - "pylint_django", - "pylint_celery", -] -django-settings-module = "config.settings.local" - -[tool.pylint.FORMAT] -max-line-length = 119 - -[tool.pylint."MESSAGES CONTROL"] -disable = [ - "missing-docstring", - "invalid-name", -] - -[tool.pylint.DESIGN] -max-parents = 13 - -[tool.pylint.TYPECHECK] -generated-members = [ - "REQUEST", - "acl_users", - "aq_parent", - "[a-zA-Z]+_set{1,2}", - "save", - "delete", -] - - -# ==== djLint ==== -[tool.djlint] -blank_line_after_tag = "load,extends" -close_void_tags = true -format_css = true -format_js = true -# TODO: remove T002 when fixed https://github.com/Riverside-Healthcare/djLint/issues/687 -ignore = "H006,H030,H031,T002" -include = "H017,H035" -indent = 2 -max_line_length = 119 -profile = "django" - -[tool.djlint.css] -indent_size = 2 - -[tool.djlint.js] -indent_size = 2 diff --git a/backend/python_docker_image.yaml b/backend/python_docker_image.yaml deleted file mode 100644 index d9ff0a44..00000000 --- a/backend/python_docker_image.yaml +++ /dev/null @@ -1,11 +0,0 @@ -FROM python:latest - -WORKDIR /app - -COPY requirements.txt . - -RUN pip install --no-cache-dir -r requirements.txt - -COPY . /app - -CMD ["python", "your_main_script.py"] diff --git a/backend/requirements.txt b/backend/requirements.txt index 44533d7c..c088ad27 100644 Binary files a/backend/requirements.txt and b/backend/requirements.txt differ diff --git a/backend/requirements/base.txt b/backend/requirements/base.txt deleted file mode 100644 index 0721a1bf..00000000 --- a/backend/requirements/base.txt +++ /dev/null @@ -1,27 +0,0 @@ -python-slugify==8.0.1 # https://github.com/un33k/python-slugify -Pillow==10.0.1 # https://github.com/python-pillow/Pillow -argon2-cffi==23.1.0 # https://github.com/hynek/argon2_cffi -whitenoise==6.5.0 # https://github.com/evansd/whitenoise -redis==5.0.1 # https://github.com/redis/redis-py -hiredis==2.2.3 # https://github.com/redis/hiredis-py -celery==5.3.4 # pyup: < 6.0 # https://github.com/celery/celery -django-celery-beat==2.5.0 # https://github.com/celery/django-celery-beat -flower==2.0.1 # https://github.com/mher/flower -uvicorn[standard]==0.23.2 # https://github.com/encode/uvicorn - -# Django -# ------------------------------------------------------------------------------ -django==4.2.6 # pyup: < 5.0 # https://www.djangoproject.com/ -django-environ==0.11.2 # https://github.com/joke2k/django-environ -django-model-utils==4.3.1 # https://github.com/jazzband/django-model-utils -django-allauth==0.57.0 # https://github.com/pennersr/django-allauth -django-crispy-forms==2.0 # https://github.com/django-crispy-forms/django-crispy-forms -crispy-bootstrap5==0.7 # https://github.com/django-crispy-forms/crispy-bootstrap5 -django-redis==5.4.0 # https://github.com/jazzband/django-redis -# Django REST Framework -djangorestframework==3.14.0 # https://github.com/encode/django-rest-framework -django-cors-headers==4.2.0 # https://github.com/adamchainz/django-cors-headers -# DRF-spectacular for api documentation -drf-spectacular==0.26.5 # https://github.com/tfranzel/drf-spectacular - -djangorestframework-simplejwt==5.3.0 diff --git a/backend/requirements/local.txt b/backend/requirements/local.txt deleted file mode 100644 index ac1bbdbb..00000000 --- a/backend/requirements/local.txt +++ /dev/null @@ -1,42 +0,0 @@ --r base.txt - -Werkzeug[watchdog]==3.0.0 # https://github.com/pallets/werkzeug -ipdb==0.13.13 # https://github.com/gotcha/ipdb -# psycopg[c]==3.1.12 # https://github.com/psycopg/psycopg -psycopg -watchfiles==0.20.0 # https://github.com/samuelcolvin/watchfiles - -# Testing -# ------------------------------------------------------------------------------ -mypy==1.5.1 # https://github.com/python/mypy -django-stubs[compatible-mypy]==4.2.4 # https://github.com/typeddjango/django-stubs -pytest==7.4.2 # https://github.com/pytest-dev/pytest -pytest-sugar==0.9.7 # https://github.com/Frozenball/pytest-sugar -djangorestframework-stubs[compatible-mypy]==3.14.3 # https://github.com/typeddjango/djangorestframework-stubs -django -environ -celery -docker -# Documentation -# ------------------------------------------------------------------------------ -sphinx==7.2.6 # https://github.com/sphinx-doc/sphinx -sphinx-autobuild==2021.3.14 # https://github.com/GaretJax/sphinx-autobuild - -# Code quality -# ------------------------------------------------------------------------------ -flake8==6.1.0 # https://github.com/PyCQA/flake8 -flake8-isort==6.1.0 # https://github.com/gforcada/flake8-isort -coverage==7.3.2 # https://github.com/nedbat/coveragepy -black==23.9.1 # https://github.com/psf/black -djlint==1.34.0 # https://github.com/Riverside-Healthcare/djLint -pylint-django==2.5.3 # https://github.com/PyCQA/pylint-django -pylint-celery==0.3 # https://github.com/PyCQA/pylint-celery -pre-commit==3.4.0 # https://github.com/pre-commit/pre-commit -# Django -# ------------------------------------------------------------------------------ -factory-boy==3.3.0 # https://github.com/FactoryBoy/factory_boy - -django-debug-toolbar==4.2.0 # https://github.com/jazzband/django-debug-toolbar -django-extensions==3.2.3 # https://github.com/django-extensions/django-extensions -django-coverage-plugin==3.1.0 # https://github.com/nedbat/django_coverage_plugin -pytest-django==4.5.2 # https://github.com/pytest-dev/pytest-django diff --git a/backend/requirements/production.txt b/backend/requirements/production.txt deleted file mode 100644 index 9cd3bf97..00000000 --- a/backend/requirements/production.txt +++ /dev/null @@ -1,11 +0,0 @@ -# PRECAUTION: avoid production dependencies that aren't in development - --r base.txt - -gunicorn==21.2.0 # https://github.com/benoitc/gunicorn -psycopg[c]==3.1.12 # https://github.com/psycopg/psycopg - -# Django -# ------------------------------------------------------------------------------ -django-storages[s3]==1.14.2 # https://github.com/jschneier/django-storages -django-anymail[mailgun]==10.1 # https://github.com/anymail/django-anymail diff --git a/backend/schema.py b/backend/schema.py new file mode 100644 index 00000000..64b1284b --- /dev/null +++ b/backend/schema.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel + +class FileModel(BaseModel): + video_filename: str + image_filename: str + audio_filename: str + text_data: str + time_created: str + time_updated: str diff --git a/backend/setup.cfg b/backend/setup.cfg deleted file mode 100644 index 2412f174..00000000 --- a/backend/setup.cfg +++ /dev/null @@ -1,11 +0,0 @@ -# flake8 and pycodestyle don't support pyproject.toml -# https://github.com/PyCQA/flake8/issues/234 -# https://github.com/PyCQA/pycodestyle/issues/813 -[flake8] -max-line-length = 119 -exclude = .tox,.git,*/migrations/*,*/static/CACHE/*,docs,node_modules,venv,.venv -extend-ignore = E203 - -[pycodestyle] -max-line-length = 119 -exclude = .tox,.git,*/migrations/*,*/static/CACHE/*,docs,node_modules,venv,.venv diff --git a/backend/tests/test_merge_production_dotenvs_in_dotenv.py b/backend/tests/test_merge_production_dotenvs_in_dotenv.py deleted file mode 100644 index c0e68f60..00000000 --- a/backend/tests/test_merge_production_dotenvs_in_dotenv.py +++ /dev/null @@ -1,34 +0,0 @@ -from pathlib import Path - -import pytest - -from merge_production_dotenvs_in_dotenv import merge - - -@pytest.mark.parametrize( - ("input_contents", "expected_output"), - [ - ([], ""), - ([""], "\n"), - (["JANE=doe"], "JANE=doe\n"), - (["SEP=true", "AR=ator"], "SEP=true\nAR=ator\n"), - (["A=0", "B=1", "C=2"], "A=0\nB=1\nC=2\n"), - (["X=x\n", "Y=y", "Z=z\n"], "X=x\n\nY=y\nZ=z\n\n"), - ], -) -def test_merge( - tmp_path: Path, - input_contents: list[str], - expected_output: str, -): - output_file = tmp_path / ".env" - - files_to_merge = [] - for num, input_content in enumerate(input_contents, start=1): - merge_file = tmp_path / f".service{num}" - merge_file.write_text(input_content) - files_to_merge.append(merge_file) - - merge(output_file, files_to_merge) - - assert output_file.read_text() == expected_output diff --git a/frontend/mdx-components.tsx b/frontend/mdx-components.tsx index 792fe383..d721029a 100644 --- a/frontend/mdx-components.tsx +++ b/frontend/mdx-components.tsx @@ -1,6 +1,6 @@ +import type { MDXComponents } from 'mdx/types'; import classnames from 'classnames'; import classes from './mdx-components.module.css'; -import type { MDXComponents } from 'mdx/types'; export function useMDXComponents(components: MDXComponents): MDXComponents { diff --git a/frontend/modules.json b/frontend/modules.json new file mode 100644 index 00000000..9bf17a6f --- /dev/null +++ b/frontend/modules.json @@ -0,0 +1,102 @@ +[ + { + "name": "Modules", + "url": "/modules", + "image": "/img/frontpage/comchat.png", + "description": "The hub for Commune Ai modules", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "com", + "hub" + ] + }, + { + "name": "Telemetry", + "url": "/telemetry", + "image": "/img/frontpage/comai-logo.png", + "description": "Telemetry of Commune Ai", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "stats", + "staking", + "wallet" + ] + }, + { + "name": "ComChat", + "url": "https://comchat.io/", + "image": "/img/frontpage/comchat.png", + "description": "LLM GPT aggregator built on Commune Ai", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "chat", + "GPT", + "Ai" + ] + }, + { + "name": "ComfyUILauncher", + "url": "https://huggingface.co/spaces/subbytech/comfyui-launcher/", + "image": "/img/frontpage/comai-logo.png", + "description": "ComfyUILauncher of Commune Ai", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "Docs", + "AI" + ] + }, + { + "name": "ComWallet", + "url": "https://comwallet.io/", + "image": "/img/frontpage/combridge-logo.png", + "description": "ComWallet of Commune Ai", + "registerKey": "5CXiWwsS76H2vwroWu4SvdAS3kxprb7aFtqWLxxZC5FNhYri", + "verified": true, + "tags": [ + "staking", + "wallet" + ], + "disabled": false + }, + { + "name": "Bittensor", + "url": "/bittensor", + "image": "/img/frontpage/comai-webp.webp", + "description": "Bittensor of Commune Ai", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": true, + "tags": [ + "AI" + ] + }, + { + "name": "Comscanner", + "url": "https://comscanner.vercel.app/app/welcome", + "image": "/img/frontpage/comai-logo.png", + "description": "coming soon", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "wallet", + "predict", + "explorer", + "AI" + ] + }, + { + "name": "Staking", + "url": "/staking", + "image": "/img/frontpage/comai-logo.png", + "description": "coming soon", + "registerKey": "5EJ9AUpSGafWeagdP5nwc5AwcYBkagYSZyx2BmLKWJrGBZUZ", + "verified": false, + "tags": [ + "wallet", + "bridge" + ] + } +] \ No newline at end of file diff --git a/frontend/package.json b/frontend/package.json index 644d2e6f..51b22dd5 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -30,25 +30,33 @@ "@typescript-eslint/eslint-plugin": "^7.1.1", "@web3modal/wagmi": "3.2.1", "antd": "^5.12.8", + "antd-style": "^3.6.1", "axios": "^1.6.5", - "classnames": "^2.3.2", + "classnames": "^2.5.1", + "daisyui": "^4.7.3", "eslint": "^8.57.0", "eslint-config-next": "^14.1.3", "eslint-plugin-react": "^7.34.0", "ethers": "^6.11.1", + "fuse.js": "^7.0.0", + "i": "^0.3.7", "moralis": "^2.23.2", "next": "14.0.1", "next-auth": "^4.24.5", + "npm": "^10.5.0", "react": "^18", "react-dom": "^18", "react-flow-renderer": "^10.3.17", "react-github-login": "^1.0.3", "react-google-login": "^5.2.2", + "react-hamburger-menus": "^1.0.16", "react-hook-form": "^7.49.3", "react-icons": "^5.0.1", "react-notifications": "^1.7.4", "react-paginate": "8.1.3", "react-redux": "^9.1.0", + "react-responsive-modal": "^6.4.2", + "react-toastify": "^10.0.5", "reactflow": "^11.10.4", "redux": "^5.0.1", "redux-devtools-extension": "^2.13.9", diff --git a/frontend/public/gif/cubes/commune-single-block_blue.webp b/frontend/public/gif/cubes/commune-single-block_blue.webp new file mode 100644 index 00000000..8e493b0d Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_blue.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_gray.webp b/frontend/public/gif/cubes/commune-single-block_gray.webp new file mode 100644 index 00000000..74dafdc6 Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_gray.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_green.webp b/frontend/public/gif/cubes/commune-single-block_green.webp new file mode 100644 index 00000000..153be794 Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_green.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_purple.webp b/frontend/public/gif/cubes/commune-single-block_purple.webp new file mode 100644 index 00000000..e3ba409a Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_purple.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_red.webp b/frontend/public/gif/cubes/commune-single-block_red.webp new file mode 100644 index 00000000..e3489d20 Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_red.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_white.webp b/frontend/public/gif/cubes/commune-single-block_white.webp new file mode 100644 index 00000000..7c9202fa Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_white.webp differ diff --git a/frontend/public/gif/cubes/commune-single-block_yellow.webp b/frontend/public/gif/cubes/commune-single-block_yellow.webp new file mode 100644 index 00000000..6d64e87f Binary files /dev/null and b/frontend/public/gif/cubes/commune-single-block_yellow.webp differ diff --git a/frontend/public/gif/logo/CubesShufflingGIF.webp b/frontend/public/gif/logo/CubesShufflingGIF.webp new file mode 100644 index 00000000..448a1700 Binary files /dev/null and b/frontend/public/gif/logo/CubesShufflingGIF.webp differ diff --git a/frontend/public/gif/logo/commune.webp b/frontend/public/gif/logo/commune.webp new file mode 100644 index 00000000..cbb127fc Binary files /dev/null and b/frontend/public/gif/logo/commune.webp differ diff --git a/frontend/public/img/frontpage/comai-logo.png b/frontend/public/img/frontpage/comai-logo.png new file mode 100644 index 00000000..84d78be7 Binary files /dev/null and b/frontend/public/img/frontpage/comai-logo.png differ diff --git a/frontend/public/img/frontpage/comai-webp.webp b/frontend/public/img/frontpage/comai-webp.webp new file mode 100644 index 00000000..ebd19a94 Binary files /dev/null and b/frontend/public/img/frontpage/comai-webp.webp differ diff --git a/frontend/public/img/frontpage/combridge-logo.png b/frontend/public/img/frontpage/combridge-logo.png new file mode 100644 index 00000000..ce645ea4 Binary files /dev/null and b/frontend/public/img/frontpage/combridge-logo.png differ diff --git a/frontend/public/img/frontpage/comchat.png b/frontend/public/img/frontpage/comchat.png new file mode 100644 index 00000000..0e22cfac Binary files /dev/null and b/frontend/public/img/frontpage/comchat.png differ diff --git a/frontend/public/img/frontpage/comstats-webp.webp b/frontend/public/img/frontpage/comstats-webp.webp new file mode 100644 index 00000000..6793cdbf Binary files /dev/null and b/frontend/public/img/frontpage/comstats-webp.webp differ diff --git a/frontend/src/app/bittensor/page.tsx b/frontend/src/app/bittensor/page.tsx index 3fab08c9..07c70a64 100644 --- a/frontend/src/app/bittensor/page.tsx +++ b/frontend/src/app/bittensor/page.tsx @@ -2,10 +2,10 @@ import BittensorItem from "@/components/molecules/bittensor/item"; import { items } from "@/components/molecules/bittensor/item-date"; -const BittensorPage = () => { +export default function BittensorPage() { return (
-

+

Bittensor Subnets

@@ -18,5 +18,3 @@ const BittensorPage = () => {
) } - -export default BittensorPage diff --git a/frontend/src/app/cog-modules/page.tsx b/frontend/src/app/cog-modules/page.tsx index 1d1cb612..89c8af2e 100644 --- a/frontend/src/app/cog-modules/page.tsx +++ b/frontend/src/app/cog-modules/page.tsx @@ -23,16 +23,16 @@ const CogModulePage = () => { const [searchString, setSearchString] = useState(""); const [currentPage, setCurrentPage] = useState(1); const itemsPerPage = 6; - const [loadedModules, setLoadedModules] = useState([]); - const [displayedModules, setDisplayedModules] = useState([]); - const [filteredModules, setFilteredModules] = useState([]); + const [loadedModules, setLoadedModules] = useState(Array); + const [displayedModules, setDisplayedModules] = useState(Array); + const [filteredModules, setFilteredModules] = useState(Array); const [isShowPolkadotWalletModalOpen, setIsShowPolkadotWalletModalOpen] = useState(false); - const [replicateData] = useState([]); + const [replicateData] = useState(Array); useEffect(() => { const filtered = searchString ? loadedModules.filter((module) => - module?.data.description ? module.data.description.toLowerCase().includes(searchString.toLowerCase()) : false + module?.description ? module?.description.toLowerCase().includes(searchString.toLowerCase()) : false ) : loadedModules; setFilteredModules(filtered); @@ -71,7 +71,6 @@ const CogModulePage = () => { fetchModules(''); }, []); - const handlePageChange = (selectedItem: { selected: number; }) => { @@ -82,14 +81,12 @@ const CogModulePage = () => { const handleModulesFetched = (modules: string[]) => { const formattedModules = modules.map((moduleName: string, index: number) => ({ id: index, - data: { - url: 'string', - cover_image_url: 'string', - owner: 'string', - name: moduleName, - description: 'string', - run_count: 'number' - } + url: 'string', + cover_image_url: 'string', + owner: 'string', + name: moduleName, + description: 'string', + run_count: 'number' })); setLoadedModules(formattedModules); updateDisplayedModules(formattedModules, currentPage); @@ -123,9 +120,17 @@ const CogModulePage = () => { searchString={searchString} /> {displayedModules && displayedModules.length > 0 ? ( -
+
{displayedModules.map((item, idx) => ( - + ))}
) : ( @@ -137,8 +142,8 @@ const CogModulePage = () => { pageCount={pageCount} onPageChange={handlePageChange} forcePage={currentPage - 1} - containerClassName="flex justify-center items-center space-x-3 my-4 text-lg dark:text-white" - pageLinkClassName="px-5 text-lg border rounded hover:bg-gray-200 transition-colors duration-200 py-3" + containerClassName="flex justify-center items-center space-x-1 my-4 text-lg dark:text-white" + pageLinkClassName="px-3 text-lg border rounded hover:bg-gray-200 transition-colors duration-200 py-3" activeClassName="bg-blue-500 text-white py-3 rounded" previousLabel={"previous"} nextLabel={"next"} diff --git a/frontend/src/app/docs/bittensor/page.tsx b/frontend/src/app/docs/bittensor/page.tsx index 9a0c24a9..d5a25ead 100644 --- a/frontend/src/app/docs/bittensor/page.tsx +++ b/frontend/src/app/docs/bittensor/page.tsx @@ -2,7 +2,7 @@ import BittensorItem from "@/components/molecules/bittensor/item"; import { items } from "@/components/molecules/bittensor/item-date"; -export default function () { +export default function DocsBittensor() { return (

@@ -17,4 +17,4 @@ export default function () {

) -} \ No newline at end of file +} diff --git a/frontend/src/app/globals.css b/frontend/src/app/globals.css index dcd61f90..8039123f 100644 --- a/frontend/src/app/globals.css +++ b/frontend/src/app/globals.css @@ -8,8 +8,8 @@ --ifm-color-primary: #25c2a0; --ifm-color-white: #fff; --foreground-rgb: 0, 0, 0; - --background-start-rgb: 214, 219, 220; - --background-end-rgb: 255, 255, 255; + --background-start-rgb: 14, 23, 42; + --background-end-rgb: 8, 8, 8; --ifm-font-family-base: system-ui, -apple-system, Segoe UI, Roboto, Ubuntu, Cantarell, Noto Sans, sans-serif, BlinkMacSystemFont, 'Segoe UI', Helvetica, Arial, sans-serif, 'Apple Color Emoji', 'Segoe UI Emoji', 'Segoe UI Symbol'; --ifm-h1-font-size: 2rem; --ifm-heading-color: inherit; diff --git a/frontend/src/app/layout.tsx b/frontend/src/app/layout.tsx index ca35f8b9..2cdb72cc 100644 --- a/frontend/src/app/layout.tsx +++ b/frontend/src/app/layout.tsx @@ -111,12 +111,12 @@ export default function RootLayout({ - - {/* */} - - {children} -