mirror of
https://github.com/deepseek-ai/Janus.git
synced 2025-02-22 13:48:57 -05:00
commit
This commit is contained in:
wuchengyue
commit
659982f413
8
.gitattributes
vendored
Normal file
8
.gitattributes
vendored
Normal file
@ -0,0 +1,8 @@
|
|||||||
|
* text eol=lf
|
||||||
|
*.ipynb linguist-detectable=false
|
||||||
|
|
||||||
|
*.png binary
|
||||||
|
*.jpg binary
|
||||||
|
*.jpeg binary
|
||||||
|
*.gif binary
|
||||||
|
*.pdf binary
|
||||||
421
.gitignore
vendored
Normal file
421
.gitignore
vendored
Normal file
@ -0,0 +1,421 @@
|
|||||||
|
##### Python.gitignore #####
|
||||||
|
# Byte-compiled / optimized / DLL files
|
||||||
|
**/__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
*.pyd
|
||||||
|
*.py[cod]
|
||||||
|
*$py.class
|
||||||
|
|
||||||
|
# C extensions
|
||||||
|
*.so
|
||||||
|
|
||||||
|
# Distribution / packaging
|
||||||
|
.Python
|
||||||
|
build/
|
||||||
|
develop-eggs/
|
||||||
|
dist/
|
||||||
|
downloads/
|
||||||
|
eggs/
|
||||||
|
.eggs/
|
||||||
|
lib/
|
||||||
|
lib64/
|
||||||
|
parts/
|
||||||
|
sdist/
|
||||||
|
var/
|
||||||
|
wheels/
|
||||||
|
wheelhouse/
|
||||||
|
share/python-wheels/
|
||||||
|
*.egg-info/
|
||||||
|
.installed.cfg
|
||||||
|
*.egg
|
||||||
|
MANIFEST
|
||||||
|
*.whl
|
||||||
|
|
||||||
|
# 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/
|
||||||
|
.nox/
|
||||||
|
.coverage
|
||||||
|
.coverage.*
|
||||||
|
.cache
|
||||||
|
nosetests.xml
|
||||||
|
coverage.xml
|
||||||
|
*.cover
|
||||||
|
*.py,cover
|
||||||
|
.hypothesis/
|
||||||
|
.pytest_cache/
|
||||||
|
cover/
|
||||||
|
|
||||||
|
# Translations
|
||||||
|
*.mo
|
||||||
|
*.pot
|
||||||
|
|
||||||
|
# Django stuff:
|
||||||
|
*.log
|
||||||
|
local_settings.py
|
||||||
|
db.sqlite3
|
||||||
|
db.sqlite3-journal
|
||||||
|
|
||||||
|
# Flask stuff:
|
||||||
|
instance/
|
||||||
|
.webassets-cache
|
||||||
|
|
||||||
|
# Scrapy stuff:
|
||||||
|
.scrapy
|
||||||
|
|
||||||
|
# Sphinx documentation
|
||||||
|
docs/_build/
|
||||||
|
docs/source/_build/
|
||||||
|
_autosummary/
|
||||||
|
|
||||||
|
# PyBuilder
|
||||||
|
.pybuilder/
|
||||||
|
target/
|
||||||
|
|
||||||
|
# Jupyter Notebook
|
||||||
|
.ipynb_checkpoints
|
||||||
|
|
||||||
|
# IPython
|
||||||
|
profile_default/
|
||||||
|
ipython_config.py
|
||||||
|
|
||||||
|
# pyenv
|
||||||
|
# For a library or package, you might want to ignore these files since the code is
|
||||||
|
# intended to run in multiple environments; otherwise, check them in:
|
||||||
|
.python-version
|
||||||
|
|
||||||
|
# pipenv
|
||||||
|
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||||
|
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||||
|
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||||
|
# install all needed dependencies.
|
||||||
|
#Pipfile.lock
|
||||||
|
|
||||||
|
# poetry
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||||
|
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||||
|
# commonly ignored for libraries.
|
||||||
|
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||||
|
#poetry.lock
|
||||||
|
|
||||||
|
# pdm
|
||||||
|
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||||
|
#pdm.lock
|
||||||
|
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||||
|
# in version control.
|
||||||
|
# https://pdm.fming.dev/#use-with-ide
|
||||||
|
.pdm.toml
|
||||||
|
|
||||||
|
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||||
|
__pypackages__/
|
||||||
|
|
||||||
|
# Celery stuff
|
||||||
|
celerybeat-schedule
|
||||||
|
celerybeat.pid
|
||||||
|
|
||||||
|
# SageMath parsed files
|
||||||
|
*.sage.py
|
||||||
|
|
||||||
|
# Environments
|
||||||
|
.env
|
||||||
|
.venv
|
||||||
|
env/
|
||||||
|
venv/
|
||||||
|
ENV/
|
||||||
|
env.bak/
|
||||||
|
venv.bak/
|
||||||
|
|
||||||
|
# Spyder project settings
|
||||||
|
.spyderproject
|
||||||
|
.spyproject
|
||||||
|
|
||||||
|
# Rope project settings
|
||||||
|
.ropeproject
|
||||||
|
|
||||||
|
# mkdocs documentation
|
||||||
|
/site
|
||||||
|
|
||||||
|
# ruff
|
||||||
|
.ruff_cache/
|
||||||
|
|
||||||
|
# mypy
|
||||||
|
.mypy_cache/
|
||||||
|
.dmypy.json
|
||||||
|
dmypy.json
|
||||||
|
|
||||||
|
# Pyre type checker
|
||||||
|
.pyre/
|
||||||
|
|
||||||
|
# pytype static type analyzer
|
||||||
|
.pytype/
|
||||||
|
|
||||||
|
# Cython debug symbols
|
||||||
|
cython_debug/
|
||||||
|
|
||||||
|
# PyCharm
|
||||||
|
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||||
|
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||||
|
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||||
|
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||||
|
.idea/
|
||||||
|
|
||||||
|
|
||||||
|
##### macOS.gitignore #####
|
||||||
|
# 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
|
||||||
|
|
||||||
|
|
||||||
|
##### Linux.gitignore #####
|
||||||
|
*~
|
||||||
|
|
||||||
|
# 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*
|
||||||
|
|
||||||
|
|
||||||
|
##### Windows.gitignore #####
|
||||||
|
# Windows thumbnail cache files
|
||||||
|
Thumbs.db
|
||||||
|
Thumbs.db:encryptable
|
||||||
|
ehthumbs.db
|
||||||
|
ehthumbs_vista.db
|
||||||
|
|
||||||
|
# Dump file
|
||||||
|
*.stackdump
|
||||||
|
|
||||||
|
# Folder config file
|
||||||
|
[Dd]esktop.ini
|
||||||
|
|
||||||
|
# Recycle Bin used on file shares
|
||||||
|
$RECYCLE.BIN/
|
||||||
|
|
||||||
|
# Windows Installer files
|
||||||
|
*.cab
|
||||||
|
*.msi
|
||||||
|
*.msix
|
||||||
|
*.msm
|
||||||
|
*.msp
|
||||||
|
|
||||||
|
# Windows shortcuts
|
||||||
|
*.lnk
|
||||||
|
|
||||||
|
|
||||||
|
##### Archives.gitignore #####
|
||||||
|
# It's better to unpack these files and commit the raw source because
|
||||||
|
# git has its own built in compression methods.
|
||||||
|
*.7z
|
||||||
|
*.jar
|
||||||
|
*.rar
|
||||||
|
*.zip
|
||||||
|
*.gz
|
||||||
|
*.gzip
|
||||||
|
*.tgz
|
||||||
|
*.bzip
|
||||||
|
*.bzip2
|
||||||
|
*.bz2
|
||||||
|
*.xz
|
||||||
|
*.lzma
|
||||||
|
*.cab
|
||||||
|
*.xar
|
||||||
|
|
||||||
|
# Packing-only formats
|
||||||
|
*.iso
|
||||||
|
*.tar
|
||||||
|
|
||||||
|
# Package management formats
|
||||||
|
*.dmg
|
||||||
|
*.xpi
|
||||||
|
*.gem
|
||||||
|
*.egg
|
||||||
|
*.deb
|
||||||
|
*.rpm
|
||||||
|
*.msi
|
||||||
|
*.msm
|
||||||
|
*.msp
|
||||||
|
*.txz
|
||||||
|
|
||||||
|
|
||||||
|
##### Xcode.gitignore #####
|
||||||
|
# Xcode
|
||||||
|
#
|
||||||
|
# gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore
|
||||||
|
|
||||||
|
## User settings
|
||||||
|
xcuserdata/
|
||||||
|
|
||||||
|
## Compatibility with Xcode 8 and earlier (ignoring not required starting Xcode 9)
|
||||||
|
*.xcscmblueprint
|
||||||
|
*.xccheckout
|
||||||
|
|
||||||
|
## Compatibility with Xcode 3 and earlier (ignoring not required starting Xcode 4)
|
||||||
|
build/
|
||||||
|
DerivedData/
|
||||||
|
*.moved-aside
|
||||||
|
*.pbxuser
|
||||||
|
!default.pbxuser
|
||||||
|
*.mode1v3
|
||||||
|
!default.mode1v3
|
||||||
|
*.mode2v3
|
||||||
|
!default.mode2v3
|
||||||
|
*.perspectivev3
|
||||||
|
!default.perspectivev3
|
||||||
|
|
||||||
|
## Gcc Patch
|
||||||
|
/*.gcno
|
||||||
|
|
||||||
|
|
||||||
|
##### JetBrains.gitignore #####
|
||||||
|
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm
|
||||||
|
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||||
|
|
||||||
|
# User settings
|
||||||
|
.idea/*
|
||||||
|
|
||||||
|
# User-specific stuff
|
||||||
|
.idea/**/workspace.xml
|
||||||
|
.idea/**/tasks.xml
|
||||||
|
.idea/**/usage.statistics.xml
|
||||||
|
.idea/**/dictionaries
|
||||||
|
.idea/**/shelf
|
||||||
|
|
||||||
|
# Generated files
|
||||||
|
.idea/**/contentModel.xml
|
||||||
|
|
||||||
|
# Sensitive or high-churn files
|
||||||
|
.idea/**/dataSources/
|
||||||
|
.idea/**/dataSources.ids
|
||||||
|
.idea/**/dataSources.local.xml
|
||||||
|
.idea/**/sqlDataSources.xml
|
||||||
|
.idea/**/dynamic.xml
|
||||||
|
.idea/**/uiDesigner.xml
|
||||||
|
.idea/**/dbnavigator.xml
|
||||||
|
|
||||||
|
# Gradle
|
||||||
|
.idea/**/gradle.xml
|
||||||
|
.idea/**/libraries
|
||||||
|
|
||||||
|
# Gradle and Maven with auto-import
|
||||||
|
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||||
|
# since they will be recreated, and may cause churn. Uncomment if using
|
||||||
|
# auto-import.
|
||||||
|
# .idea/artifacts
|
||||||
|
# .idea/compiler.xml
|
||||||
|
# .idea/jarRepositories.xml
|
||||||
|
# .idea/modules.xml
|
||||||
|
# .idea/*.iml
|
||||||
|
# .idea/modules
|
||||||
|
# *.iml
|
||||||
|
# *.ipr
|
||||||
|
|
||||||
|
# CMake
|
||||||
|
cmake-build-*/
|
||||||
|
|
||||||
|
# Mongo Explorer plugin
|
||||||
|
.idea/**/mongoSettings.xml
|
||||||
|
|
||||||
|
# File-based project format
|
||||||
|
*.iws
|
||||||
|
|
||||||
|
# IntelliJ
|
||||||
|
out/
|
||||||
|
|
||||||
|
# mpeltonen/sbt-idea plugin
|
||||||
|
.idea_modules/
|
||||||
|
|
||||||
|
# JIRA plugin
|
||||||
|
atlassian-ide-plugin.xml
|
||||||
|
|
||||||
|
# Cursive Clojure plugin
|
||||||
|
.idea/replstate.xml
|
||||||
|
|
||||||
|
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||||
|
com_crashlytics_export_strings.xml
|
||||||
|
crashlytics.properties
|
||||||
|
crashlytics-build.properties
|
||||||
|
fabric.properties
|
||||||
|
|
||||||
|
# Editor-based Rest Client
|
||||||
|
.idea/httpRequests
|
||||||
|
|
||||||
|
# Android studio 3.1+ serialized cache file
|
||||||
|
.idea/caches/build_file_checksums.ser
|
||||||
|
|
||||||
|
|
||||||
|
##### VisualStudioCode.gitignore #####
|
||||||
|
.vscode/*
|
||||||
|
# !.vscode/settings.json
|
||||||
|
# !.vscode/tasks.json
|
||||||
|
# !.vscode/launch.json
|
||||||
|
!.vscode/extensions.json
|
||||||
|
*.code-workspace
|
||||||
|
|
||||||
|
# Local History for Visual Studio Code
|
||||||
|
.history/
|
||||||
|
|
||||||
|
|
||||||
|
##### Vim.gitignore #####
|
||||||
|
# Swap
|
||||||
|
.*.s[a-v][a-z]
|
||||||
|
!*.svg # comment out if you don't need vector files
|
||||||
|
.*.sw[a-p]
|
||||||
|
.s[a-rt-v][a-z]
|
||||||
|
.ss[a-gi-z]
|
||||||
|
.sw[a-p]
|
||||||
|
|
||||||
|
# Session
|
||||||
|
Session.vim
|
||||||
|
Sessionx.vim
|
||||||
|
|
||||||
|
# Temporary
|
||||||
|
.netrwhist
|
||||||
|
*~
|
||||||
|
# Auto-generated tag files
|
||||||
|
tags
|
||||||
|
# Persistent undo
|
||||||
|
[._]*.un~
|
||||||
|
.vscode
|
||||||
|
.github
|
||||||
|
generated_samples/
|
||||||
21
LICENSE-CODE
Normal file
21
LICENSE-CODE
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2023 DeepSeek
|
||||||
|
|
||||||
|
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.
|
||||||
91
LICENSE-MODEL
Normal file
91
LICENSE-MODEL
Normal file
@ -0,0 +1,91 @@
|
|||||||
|
DEEPSEEK LICENSE AGREEMENT
|
||||||
|
|
||||||
|
Version 1.0, 23 October 2023
|
||||||
|
|
||||||
|
Copyright (c) 2023 DeepSeek
|
||||||
|
|
||||||
|
Section I: PREAMBLE
|
||||||
|
|
||||||
|
Large generative models are being widely adopted and used, and have the potential to transform the way individuals conceive and benefit from AI or ML technologies.
|
||||||
|
|
||||||
|
Notwithstanding the current and potential benefits that these artifacts can bring to society at large, there are also concerns about potential misuses of them, either due to their technical limitations or ethical considerations.
|
||||||
|
|
||||||
|
In short, this license strives for both the open and responsible downstream use of the accompanying model. When it comes to the open character, we took inspiration from open source permissive licenses regarding the grant of IP rights. Referring to the downstream responsible use, we added use-based restrictions not permitting the use of the model in very specific scenarios, in order for the licensor to be able to enforce the license in case potential misuses of the Model may occur. At the same time, we strive to promote open and responsible research on generative models for content generation.
|
||||||
|
|
||||||
|
Even though downstream derivative versions of the model could be released under different licensing terms, the latter will always have to include - at minimum - the same use-based restrictions as the ones in the original license (this license). We believe in the intersection between open and responsible AI development; thus, this agreement aims to strike a balance between both in order to enable responsible open-science in the field of AI.
|
||||||
|
|
||||||
|
This License governs the use of the model (and its derivatives) and is informed by the model card associated with the model.
|
||||||
|
|
||||||
|
NOW THEREFORE, You and DeepSeek agree as follows:
|
||||||
|
|
||||||
|
1. Definitions
|
||||||
|
"License" means the terms and conditions for use, reproduction, and Distribution as defined in this document.
|
||||||
|
"Data" means a collection of information and/or content extracted from the dataset used with the Model, including to train, pretrain, or otherwise evaluate the Model. The Data is not licensed under this License.
|
||||||
|
"Output" means the results of operating a Model as embodied in informational content resulting therefrom.
|
||||||
|
"Model" means any accompanying machine-learning based assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Complementary Material, that have been trained or tuned, in whole or in part on the Data, using the Complementary Material.
|
||||||
|
"Derivatives of the Model" means all modifications to the Model, works based on the Model, or any other model which is created or initialized by transfer of patterns of the weights, parameters, activations or output of the Model, to the other model, in order to cause the other model to perform similarly to the Model, including - but not limited to - distillation methods entailing the use of intermediate data representations or methods based on the generation of synthetic data by the Model for training the other model.
|
||||||
|
"Complementary Material" means the accompanying source code and scripts used to define, run, load, benchmark or evaluate the Model, and used to prepare data for training or evaluation, if any. This includes any accompanying documentation, tutorials, examples, etc, if any.
|
||||||
|
"Distribution" means any transmission, reproduction, publication or other sharing of the Model or Derivatives of the Model to a third party, including providing the Model as a hosted service made available by electronic or other remote means - e.g. API-based or web access.
|
||||||
|
"DeepSeek" (or "we") means Beijing DeepSeek Artificial Intelligence Fundamental Technology Research Co., Ltd., Hangzhou DeepSeek Artificial Intelligence Fundamental Technology Research Co., Ltd. and/or any of their affiliates.
|
||||||
|
"You" (or "Your") means an individual or Legal Entity exercising permissions granted by this License and/or making use of the Model for whichever purpose and in any field of use, including usage of the Model in an end-use application - e.g. chatbot, translator, etc.
|
||||||
|
"Third Parties" means individuals or legal entities that are not under common control with DeepSeek or You.
|
||||||
|
|
||||||
|
Section II: INTELLECTUAL PROPERTY RIGHTS
|
||||||
|
|
||||||
|
Both copyright and patent grants apply to the Model, Derivatives of the Model and Complementary Material. The Model and Derivatives of the Model are subject to additional terms as described in Section III.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of this License, DeepSeek hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare, publicly display, publicly perform, sublicense, and distribute the Complementary Material, the Model, and Derivatives of the Model.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of this License and where and as applicable, DeepSeek hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Model and the Complementary Material, where such license applies only to those patent claims licensable by DeepSeek that are necessarily infringed by its contribution(s). If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Model and/or Complementary Material constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for the Model and/or works shall terminate as of the date such litigation is asserted or filed.
|
||||||
|
|
||||||
|
|
||||||
|
Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION
|
||||||
|
|
||||||
|
4. Distribution and Redistribution. You may host for Third Party remote access purposes (e.g. software-as-a-service), reproduce and distribute copies of the Model or Derivatives of the Model thereof in any medium, with or without modifications, provided that You meet the following conditions:
|
||||||
|
a. Use-based restrictions as referenced in paragraph 5 MUST be included as an enforceable provision by You in any type of legal agreement (e.g. a license) governing the use and/or distribution of the Model or Derivatives of the Model, and You shall give notice to subsequent users You Distribute to, that the Model or Derivatives of the Model are subject to paragraph 5. This provision does not apply to the use of Complementary Material.
|
||||||
|
b. You must give any Third Party recipients of the Model or Derivatives of the Model a copy of this License;
|
||||||
|
c. You must cause any modified files to carry prominent notices stating that You changed the files;
|
||||||
|
d. You must retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Model, Derivatives of the Model.
|
||||||
|
e. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions - respecting paragraph 4.a. – for use, reproduction, or Distribution of Your modifications, or for any such Derivatives of the Model as a whole, provided Your use, reproduction, and Distribution of the Model otherwise complies with the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Use-based restrictions. The restrictions set forth in Attachment A are considered Use-based restrictions. Therefore You cannot use the Model and the Derivatives of the Model for the specified restricted uses. You may use the Model subject to this License, including only for lawful purposes and in accordance with the License. Use may include creating any content with, finetuning, updating, running, training, evaluating and/or reparametrizing the Model. You shall require all of Your users who use the Model or a Derivative of the Model to comply with the terms of this paragraph (paragraph 5).
|
||||||
|
|
||||||
|
6. The Output You Generate. Except as set forth herein, DeepSeek claims no rights in the Output You generate using the Model. You are accountable for the Output you generate and its subsequent uses. No use of the output can contravene any provision as stated in the License.
|
||||||
|
|
||||||
|
Section IV: OTHER PROVISIONS
|
||||||
|
|
||||||
|
7. Updates and Runtime Restrictions. To the maximum extent permitted by law, DeepSeek reserves the right to restrict (remotely or otherwise) usage of the Model in violation of this License.
|
||||||
|
|
||||||
|
8. Trademarks and related. Nothing in this License permits You to make use of DeepSeek’ trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between the parties; and any rights not expressly granted herein are reserved by DeepSeek.
|
||||||
|
|
||||||
|
9. Personal information, IP rights and related. This Model may contain personal information and works with IP rights. You commit to complying with applicable laws and regulations in the handling of personal information and the use of such works. Please note that DeepSeek's license granted to you to use the Model does not imply that you have obtained a legitimate basis for processing the related information or works. As an independent personal information processor and IP rights user, you need to ensure full compliance with relevant legal and regulatory requirements when handling personal information and works with IP rights that may be contained in the Model, and are willing to assume solely any risks and consequences that may arise from that.
|
||||||
|
|
||||||
|
10. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, DeepSeek provides the Model and the Complementary Material 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 Model, Derivatives of the Model, and the Complementary Material and assume any risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
11. 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 DeepSeek 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 Model and the Complementary Material (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 DeepSeek has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
12. Accepting Warranty or Additional Liability. While redistributing the Model, Derivatives of the Model and the Complementary Material 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 DeepSeek, and only if You agree to indemnify, defend, and hold DeepSeek harmless for any liability incurred by, or claims asserted against, DeepSeek by reason of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
13. If any provision of this License is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
|
||||||
|
|
||||||
|
14. Governing Law and Jurisdiction. This agreement will be governed and construed under PRC laws without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this agreement. The courts located in the domicile of Hangzhou DeepSeek Artificial Intelligence Fundamental Technology Research Co., Ltd. shall have exclusive jurisdiction of any dispute arising out of this agreement.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
Attachment A
|
||||||
|
|
||||||
|
Use Restrictions
|
||||||
|
|
||||||
|
You agree not to use the Model or Derivatives of the Model:
|
||||||
|
|
||||||
|
- In any way that violates any applicable national or international law or regulation or infringes upon the lawful rights and interests of any third party;
|
||||||
|
- For military use in any way;
|
||||||
|
- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way;
|
||||||
|
- To generate or disseminate verifiably false information and/or content with the purpose of harming others;
|
||||||
|
- To generate or disseminate inappropriate content subject to applicable regulatory requirements;
|
||||||
|
- To generate or disseminate personal identifiable information without due authorization or for unreasonable use;
|
||||||
|
- To defame, disparage or otherwise harass others;
|
||||||
|
- For fully automated decision making that adversely impacts an individual’s legal rights or otherwise creates or modifies a binding, enforceable obligation;
|
||||||
|
- For any use intended to or which has the effect of discriminating against or harming individuals or groups based on online or offline social behavior or known or predicted personal or personality characteristics;
|
||||||
|
- To exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm;
|
||||||
|
- For any use intended to or which has the effect of discriminating against individuals or groups based on legally protected characteristics or categories.
|
||||||
99
Makefile
Normal file
99
Makefile
Normal file
@ -0,0 +1,99 @@
|
|||||||
|
print-% : ; @echo $* = $($*)
|
||||||
|
PROJECT_NAME = Janus
|
||||||
|
COPYRIGHT = "DeepSeek."
|
||||||
|
PROJECT_PATH = janus
|
||||||
|
SHELL = /bin/bash
|
||||||
|
SOURCE_FOLDERS = janus
|
||||||
|
PYTHON_FILES = $(shell find $(SOURCE_FOLDERS) -type f -name "*.py" -o -name "*.pyi") inference.py
|
||||||
|
COMMIT_HASH = $(shell git log -1 --format=%h)
|
||||||
|
PATH := $(HOME)/go/bin:$(PATH)
|
||||||
|
PYTHON ?= $(shell command -v python3 || command -v python)
|
||||||
|
PYTESTOPTS ?=
|
||||||
|
|
||||||
|
.PHONY: default
|
||||||
|
default: install
|
||||||
|
|
||||||
|
# Tools Installation
|
||||||
|
|
||||||
|
check_pip_install = $(PYTHON) -m pip show $(1) &>/dev/null || (cd && $(PYTHON) -m pip install $(1) --upgrade)
|
||||||
|
check_pip_install_extra = $(PYTHON) -m pip show $(1) &>/dev/null || (cd && $(PYTHON) -m pip install $(2) --upgrade)
|
||||||
|
|
||||||
|
pylint-install:
|
||||||
|
$(call check_pip_install_extra,pylint,pylint[spelling])
|
||||||
|
$(call check_pip_install,pyenchant)
|
||||||
|
|
||||||
|
flake8-install:
|
||||||
|
$(call check_pip_install,flake8)
|
||||||
|
$(call check_pip_install,flake8-bugbear)
|
||||||
|
$(call check_pip_install,flake8-comprehensions)
|
||||||
|
$(call check_pip_install,flake8-docstrings)
|
||||||
|
$(call check_pip_install,flake8-pyi)
|
||||||
|
$(call check_pip_install,flake8-simplify)
|
||||||
|
|
||||||
|
py-format-install:
|
||||||
|
$(call check_pip_install,isort)
|
||||||
|
$(call check_pip_install_extra,black,black[jupyter])
|
||||||
|
|
||||||
|
ruff-install:
|
||||||
|
$(call check_pip_install,ruff)
|
||||||
|
|
||||||
|
mypy-install:
|
||||||
|
$(call check_pip_install,mypy)
|
||||||
|
|
||||||
|
pre-commit-install:
|
||||||
|
$(call check_pip_install,pre-commit)
|
||||||
|
$(PYTHON) -m pre_commit install --install-hooks
|
||||||
|
|
||||||
|
go-install:
|
||||||
|
# requires go >= 1.16
|
||||||
|
command -v go || (sudo apt-get install -y golang && sudo ln -sf /usr/lib/go/bin/go /usr/bin/go)
|
||||||
|
|
||||||
|
addlicense-install: go-install
|
||||||
|
command -v addlicense || go install github.com/google/addlicense@latest
|
||||||
|
|
||||||
|
addlicense: addlicense-install
|
||||||
|
addlicense -c $(COPYRIGHT) -ignore tests/coverage.xml -l mit -y 2023-$(shell date +"%Y") -check $(SOURCE_FOLDERS)
|
||||||
|
|
||||||
|
# Python linters
|
||||||
|
|
||||||
|
pylint: pylint-install
|
||||||
|
$(PYTHON) -m pylint $(PROJECT_PATH)
|
||||||
|
|
||||||
|
flake8: flake8-install
|
||||||
|
$(PYTHON) -m flake8 --count --show-source --statistics
|
||||||
|
|
||||||
|
py-format: py-format-install
|
||||||
|
$(PYTHON) -m isort --project $(PROJECT_PATH) --check $(PYTHON_FILES) && \
|
||||||
|
$(PYTHON) -m black --check $(PYTHON_FILES)
|
||||||
|
|
||||||
|
black-format: py-format-install
|
||||||
|
$(PYTHON) -m black --check $(PYTHON_FILES)
|
||||||
|
|
||||||
|
ruff: ruff-install
|
||||||
|
$(PYTHON) -m ruff check .
|
||||||
|
|
||||||
|
ruff-fix: ruff-install
|
||||||
|
$(PYTHON) -m ruff check . --fix --exit-non-zero-on-fix
|
||||||
|
|
||||||
|
mypy: mypy-install
|
||||||
|
$(PYTHON) -m mypy $(PROJECT_PATH) --install-types --non-interactive
|
||||||
|
|
||||||
|
pre-commit: pre-commit-install
|
||||||
|
$(PYTHON) -m pre_commit run --all-files
|
||||||
|
|
||||||
|
# Utility functions
|
||||||
|
|
||||||
|
lint: ruff flake8 py-format mypy pylint addlicense
|
||||||
|
|
||||||
|
format: py-format-install ruff-install addlicense-install
|
||||||
|
$(PYTHON) -m isort --project $(PROJECT_PATH) $(PYTHON_FILES)
|
||||||
|
$(PYTHON) -m black $(PYTHON_FILES)
|
||||||
|
addlicense -c $(COPYRIGHT) -ignore tests/coverage.xml -l mit -y 2023-$(shell date +"%Y") $(SOURCE_FOLDERS) inference.py
|
||||||
|
|
||||||
|
clean-py:
|
||||||
|
find . -type f -name '*.py[co]' -delete
|
||||||
|
find . -depth -type d -name "__pycache__" -exec rm -r "{}" +
|
||||||
|
find . -depth -type d -name ".ruff_cache" -exec rm -r "{}" +
|
||||||
|
find . -depth -type d -name ".mypy_cache" -exec rm -r "{}" +
|
||||||
|
|
||||||
|
clean: clean-py
|
||||||
266
README.md
Normal file
266
README.md
Normal file
@ -0,0 +1,266 @@
|
|||||||
|
<!-- markdownlint-disable first-line-h1 -->
|
||||||
|
<!-- markdownlint-disable html -->
|
||||||
|
<!-- markdownlint-disable no-duplicate-header -->
|
||||||
|
|
||||||
|
<div align="center">
|
||||||
|
<img src="images/logo.svg" width="60%" alt="DeepSeek LLM" />
|
||||||
|
</div>
|
||||||
|
<hr>
|
||||||
|
<div align="center">
|
||||||
|
|
||||||
|
<a href="https://www.deepseek.com/" target="_blank">
|
||||||
|
<img alt="Homepage" src="images/badge.svg" />
|
||||||
|
</a>
|
||||||
|
</a>
|
||||||
|
<a href="https://huggingface.co/deepseek-ai" target="_blank">
|
||||||
|
<img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DeepSeek%20AI-ffc107?color=ffc107&logoColor=white" />
|
||||||
|
</a>
|
||||||
|
|
||||||
|
</div>
|
||||||
|
|
||||||
|
|
||||||
|
<div align="center">
|
||||||
|
|
||||||
|
<!-- <a href="https://discord.gg/Tc7c45Zzu5" target="_blank">
|
||||||
|
<img alt="Discord" src="https://img.shields.io/badge/Discord-DeepSeek%20AI-7289da?logo=discord&logoColor=white&color=7289da" />
|
||||||
|
</a> -->
|
||||||
|
<!-- <a href="images/qr.jpeg" target="_blank">
|
||||||
|
<img alt="Wechat" src="https://img.shields.io/badge/WeChat-DeepSeek%20AI-brightgreen?logo=wechat&logoColor=white" />
|
||||||
|
</a> -->
|
||||||
|
<!-- <a href="https://twitter.com/deepseek_ai" target="_blank">
|
||||||
|
<img alt="Twitter Follow" src="https://img.shields.io/badge/Twitter-deepseek_ai-white?logo=x&logoColor=white" />
|
||||||
|
</a> -->
|
||||||
|
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div align="center">
|
||||||
|
|
||||||
|
<a href="LICENSE-CODE">
|
||||||
|
<img alt="Code License" src="https://img.shields.io/badge/Code_License-MIT-f5de53?&color=f5de53">
|
||||||
|
</a>
|
||||||
|
<a href="LICENSE-MODEL">
|
||||||
|
<img alt="Model License" src="https://img.shields.io/badge/Model_License-Model_Agreement-f5de53?&color=f5de53">
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
|
||||||
|
<p align="center">
|
||||||
|
<a href="#2-model-download"><b>📥 Model Download</b></a> |
|
||||||
|
<a href="#3-quick-start"><b>⚡ Quick Start</b></a> |
|
||||||
|
<a href="#4-license"><b>📜 License</b></a> |
|
||||||
|
<a href="#5-citation"><b>📖 Citation</b></a> <br>
|
||||||
|
<a href="https://arxiv.org/abs/2410.13848"><b>📄 Paper Link</b></a> |
|
||||||
|
</p>
|
||||||
|
|
||||||
|
|
||||||
|
## 1. Introduction
|
||||||
|
|
||||||
|
Janus is a novel autoregressive framework that unifies multimodal understanding and generation. It addresses the limitations of previous approaches by decoupling visual encoding into separate pathways, while still utilizing a single, unified transformer architecture for processing. The decoupling not only alleviates the conflict between the visual encoder’s roles in understanding and generation, but also enhances the framework’s flexibility. Janus surpasses previous unified model and matches or exceeds the performance of task-specific models. The simplicity, high flexibility, and effectiveness of Janus make it a strong candidate for next-generation unified multimodal models.
|
||||||
|
|
||||||
|
<div align="center">
|
||||||
|
<img alt="image" src="images/teaser.png" style="width:90%;">
|
||||||
|
</div>
|
||||||
|
|
||||||
|
|
||||||
|
## 2. Model Download
|
||||||
|
|
||||||
|
We release Janus to the public to support a broader and more diverse range of research within both academic and commercial communities.
|
||||||
|
Please note that the use of this model is subject to the terms outlined in [License section](#4-license). Commercial usage is
|
||||||
|
permitted under these terms.
|
||||||
|
|
||||||
|
### Huggingface
|
||||||
|
|
||||||
|
| Model | Sequence Length | Download |
|
||||||
|
|-----------------------|-----------------|-----------------------------------------------------------------------------|
|
||||||
|
| Janus-1.3B | 4096 | [🤗 Hugging Face](https://huggingface.co/deepseek-ai/Janus-1.3B) |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## 3. Quick Start
|
||||||
|
|
||||||
|
### Installation
|
||||||
|
|
||||||
|
On the basis of `Python >= 3.8` environment, install the necessary dependencies by running the following command:
|
||||||
|
|
||||||
|
```shell
|
||||||
|
pip install -e .
|
||||||
|
```
|
||||||
|
|
||||||
|
### Simple Inference Example
|
||||||
|
|
||||||
|
#### Multimodal Understanding
|
||||||
|
```python
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from transformers import AutoModelForCausalLM
|
||||||
|
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||||
|
from janus.utils.io import load_pil_images
|
||||||
|
|
||||||
|
# specify the path to the model
|
||||||
|
model_path = "deepseek-ai/Janus-1.3B"
|
||||||
|
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
|
||||||
|
tokenizer = vl_chat_processor.tokenizer
|
||||||
|
|
||||||
|
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
|
||||||
|
model_path, trust_remote_code=True
|
||||||
|
)
|
||||||
|
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
|
||||||
|
|
||||||
|
conversation = [
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "<image_placeholder>\nConvert the formula into latex code.",
|
||||||
|
"images": ["images/equation.png"],
|
||||||
|
},
|
||||||
|
{"role": "Assistant", "content": ""},
|
||||||
|
]
|
||||||
|
|
||||||
|
# load images and prepare for inputs
|
||||||
|
pil_images = load_pil_images(conversation)
|
||||||
|
prepare_inputs = vl_chat_processor(
|
||||||
|
conversations=conversation, images=pil_images, force_batchify=True
|
||||||
|
).to(vl_gpt.device)
|
||||||
|
|
||||||
|
# # run image encoder to get the image embeddings
|
||||||
|
inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
|
||||||
|
|
||||||
|
# # run the model to get the response
|
||||||
|
outputs = vl_gpt.language_model.generate(
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
attention_mask=prepare_inputs.attention_mask,
|
||||||
|
pad_token_id=tokenizer.eos_token_id,
|
||||||
|
bos_token_id=tokenizer.bos_token_id,
|
||||||
|
eos_token_id=tokenizer.eos_token_id,
|
||||||
|
max_new_tokens=512,
|
||||||
|
do_sample=False,
|
||||||
|
use_cache=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
|
||||||
|
print(f"{prepare_inputs['sft_format'][0]}", answer)
|
||||||
|
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Text-to-Image Generation
|
||||||
|
```python
|
||||||
|
import os
|
||||||
|
import PIL.Image
|
||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
from transformers import AutoModelForCausalLM
|
||||||
|
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||||
|
|
||||||
|
|
||||||
|
# specify the path to the model
|
||||||
|
model_path = "deepseek-ai/Janus-1.3B"
|
||||||
|
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
|
||||||
|
tokenizer = vl_chat_processor.tokenizer
|
||||||
|
|
||||||
|
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
|
||||||
|
model_path, trust_remote_code=True
|
||||||
|
)
|
||||||
|
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
|
||||||
|
|
||||||
|
conversation = [
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "A stunning princess from kabul in red, white traditional clothing, blue eyes, brown hair",
|
||||||
|
},
|
||||||
|
{"role": "Assistant", "content": ""},
|
||||||
|
]
|
||||||
|
|
||||||
|
sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
|
||||||
|
conversations=conversation,
|
||||||
|
sft_format=vl_chat_processor.sft_format,
|
||||||
|
system_prompt="",
|
||||||
|
)
|
||||||
|
prompt = sft_format + vl_chat_processor.image_start_tag
|
||||||
|
|
||||||
|
|
||||||
|
@torch.inference_mode()
|
||||||
|
def generate(
|
||||||
|
mmgpt: MultiModalityCausalLM,
|
||||||
|
vl_chat_processor: VLChatProcessor,
|
||||||
|
prompt: str,
|
||||||
|
temperature: float = 1,
|
||||||
|
parallel_size: int = 16,
|
||||||
|
cfg_weight: float = 5,
|
||||||
|
image_token_num_per_image: int = 576,
|
||||||
|
img_size: int = 384,
|
||||||
|
patch_size: int = 16,
|
||||||
|
):
|
||||||
|
input_ids = vl_chat_processor.tokenizer.encode(prompt)
|
||||||
|
input_ids = torch.LongTensor(input_ids)
|
||||||
|
|
||||||
|
tokens = torch.zeros((parallel_size*2, len(input_ids)), dtype=torch.int).cuda()
|
||||||
|
for i in range(parallel_size*2):
|
||||||
|
tokens[i, :] = input_ids
|
||||||
|
if i % 2 != 0:
|
||||||
|
tokens[i, 1:-1] = vl_chat_processor.pad_id
|
||||||
|
|
||||||
|
inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)
|
||||||
|
|
||||||
|
generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda()
|
||||||
|
|
||||||
|
for i in range(image_token_num_per_image):
|
||||||
|
outputs = mmgpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=outputs.past_key_values if i != 0 else None)
|
||||||
|
hidden_states = outputs.last_hidden_state
|
||||||
|
|
||||||
|
logits = mmgpt.gen_head(hidden_states[:, -1, :])
|
||||||
|
logit_cond = logits[0::2, :]
|
||||||
|
logit_uncond = logits[1::2, :]
|
||||||
|
|
||||||
|
logits = logit_uncond + cfg_weight * (logit_cond-logit_uncond)
|
||||||
|
probs = torch.softmax(logits / temperature, dim=-1)
|
||||||
|
|
||||||
|
next_token = torch.multinomial(probs, num_samples=1)
|
||||||
|
generated_tokens[:, i] = next_token.squeeze(dim=-1)
|
||||||
|
|
||||||
|
next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
|
||||||
|
img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
|
||||||
|
inputs_embeds = img_embeds.unsqueeze(dim=1)
|
||||||
|
|
||||||
|
|
||||||
|
dec = mmgpt.gen_vision_model.decode_code(generated_tokens.to(dtype=torch.int), shape=[parallel_size, 8, img_size//patch_size, img_size//patch_size])
|
||||||
|
dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
|
||||||
|
|
||||||
|
dec = np.clip((dec + 1) / 2 * 255, 0, 255)
|
||||||
|
|
||||||
|
visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
|
||||||
|
visual_img[:, :, :] = dec
|
||||||
|
|
||||||
|
os.makedirs('generated_samples', exist_ok=True)
|
||||||
|
for i in range(parallel_size):
|
||||||
|
save_path = os.path.join('generated_samples', "img_{}.jpg".format(i))
|
||||||
|
PIL.Image.fromarray(visual_img[i]).save(save_path)
|
||||||
|
|
||||||
|
|
||||||
|
generate(
|
||||||
|
vl_gpt,
|
||||||
|
vl_chat_processor,
|
||||||
|
prompt,
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## 4. License
|
||||||
|
|
||||||
|
This code repository is licensed under [the MIT License](https://github.com/deepseek-ai/DeepSeek-LLM/blob/HEAD/LICENSE-CODE). The use of Janus models is subject to [DeepSeek Model License](https://github.com/deepseek-ai/DeepSeek-LLM/blob/HEAD/LICENSE-MODEL).
|
||||||
|
|
||||||
|
## 5. Citation
|
||||||
|
|
||||||
|
```
|
||||||
|
@misc{wu2024janus,
|
||||||
|
title={Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation},
|
||||||
|
author={Chengyue Wu and Xiaokang Chen and Zhiyu Wu and Yiyang Ma and Xingchao Liu and Zizheng Pan and Wen Liu and Zhenda Xie and Xingkai Yu and Chong Ruan and Ping Luo},
|
||||||
|
year={2024},
|
||||||
|
eprint={2410.13848},
|
||||||
|
archivePrefix={arXiv},
|
||||||
|
primaryClass={cs.CV},
|
||||||
|
url={https://arxiv.org/abs/2410.13848},
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## 6. Contact
|
||||||
|
|
||||||
|
If you have any questions, please raise an issue or contact us at [service@deepseek.com](mailto:service@deepseek.com).
|
||||||
116
generation_inference.py
Normal file
116
generation_inference.py
Normal file
@ -0,0 +1,116 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from transformers import AutoModelForCausalLM
|
||||||
|
|
||||||
|
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||||
|
import numpy as np
|
||||||
|
import os
|
||||||
|
import PIL.Image
|
||||||
|
|
||||||
|
# specify the path to the model
|
||||||
|
model_path = "deepseek-ai/Janus-1.3B"
|
||||||
|
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
|
||||||
|
tokenizer = vl_chat_processor.tokenizer
|
||||||
|
|
||||||
|
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
|
||||||
|
model_path, trust_remote_code=True
|
||||||
|
)
|
||||||
|
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
|
||||||
|
|
||||||
|
conversation = [
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "A stunning princess from kabul in red, white traditional clothing, blue eyes, brown hair",
|
||||||
|
},
|
||||||
|
{"role": "Assistant", "content": ""},
|
||||||
|
]
|
||||||
|
|
||||||
|
sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
|
||||||
|
conversations=conversation,
|
||||||
|
sft_format=vl_chat_processor.sft_format,
|
||||||
|
system_prompt="",
|
||||||
|
)
|
||||||
|
prompt = sft_format + vl_chat_processor.image_start_tag
|
||||||
|
|
||||||
|
|
||||||
|
@torch.inference_mode()
|
||||||
|
def generate(
|
||||||
|
mmgpt: MultiModalityCausalLM,
|
||||||
|
vl_chat_processor: VLChatProcessor,
|
||||||
|
prompt: str,
|
||||||
|
temperature: float = 1,
|
||||||
|
parallel_size: int = 16,
|
||||||
|
cfg_weight: float = 5,
|
||||||
|
image_token_num_per_image: int = 576,
|
||||||
|
img_size: int = 384,
|
||||||
|
patch_size: int = 16,
|
||||||
|
):
|
||||||
|
input_ids = vl_chat_processor.tokenizer.encode(prompt)
|
||||||
|
input_ids = torch.LongTensor(input_ids)
|
||||||
|
|
||||||
|
tokens = torch.zeros((parallel_size*2, len(input_ids)), dtype=torch.int).cuda()
|
||||||
|
for i in range(parallel_size*2):
|
||||||
|
tokens[i, :] = input_ids
|
||||||
|
if i % 2 != 0:
|
||||||
|
tokens[i, 1:-1] = vl_chat_processor.pad_id
|
||||||
|
|
||||||
|
inputs_embeds = mmgpt.language_model.get_input_embeddings()(tokens)
|
||||||
|
|
||||||
|
generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda()
|
||||||
|
|
||||||
|
for i in range(image_token_num_per_image):
|
||||||
|
outputs = mmgpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=outputs.past_key_values if i != 0 else None)
|
||||||
|
hidden_states = outputs.last_hidden_state
|
||||||
|
|
||||||
|
logits = mmgpt.gen_head(hidden_states[:, -1, :])
|
||||||
|
logit_cond = logits[0::2, :]
|
||||||
|
logit_uncond = logits[1::2, :]
|
||||||
|
|
||||||
|
logits = logit_uncond + cfg_weight * (logit_cond-logit_uncond)
|
||||||
|
probs = torch.softmax(logits / temperature, dim=-1)
|
||||||
|
|
||||||
|
next_token = torch.multinomial(probs, num_samples=1)
|
||||||
|
generated_tokens[:, i] = next_token.squeeze(dim=-1)
|
||||||
|
|
||||||
|
next_token = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
|
||||||
|
img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
|
||||||
|
inputs_embeds = img_embeds.unsqueeze(dim=1)
|
||||||
|
|
||||||
|
|
||||||
|
dec = mmgpt.gen_vision_model.decode_code(generated_tokens.to(dtype=torch.int), shape=[parallel_size, 8, img_size//patch_size, img_size//patch_size])
|
||||||
|
dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
|
||||||
|
|
||||||
|
dec = np.clip((dec + 1) / 2 * 255, 0, 255)
|
||||||
|
|
||||||
|
visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
|
||||||
|
visual_img[:, :, :] = dec
|
||||||
|
|
||||||
|
os.makedirs('generated_samples', exist_ok=True)
|
||||||
|
for i in range(parallel_size):
|
||||||
|
save_path = os.path.join('generated_samples', "img_{}.jpg".format(i))
|
||||||
|
PIL.Image.fromarray(visual_img[i]).save(save_path)
|
||||||
|
|
||||||
|
|
||||||
|
generate(
|
||||||
|
vl_gpt,
|
||||||
|
vl_chat_processor,
|
||||||
|
prompt,
|
||||||
|
)
|
||||||
1
images/badge.svg
Normal file
1
images/badge.svg
Normal file
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 6.0 KiB |
BIN
images/equation.png
Normal file
BIN
images/equation.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 31 KiB |
BIN
images/logo.png
Normal file
BIN
images/logo.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 8.5 KiB |
22
images/logo.svg
Normal file
22
images/logo.svg
Normal file
@ -0,0 +1,22 @@
|
|||||||
|
<svg width="195.000000" height="41.359375" viewBox="0 0 195 41.3594" fill="none" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
|
||||||
|
<desc>
|
||||||
|
Created with Pixso.
|
||||||
|
</desc>
|
||||||
|
<defs>
|
||||||
|
<clipPath id="clip30_2029">
|
||||||
|
<rect id="_图层_1" width="134.577469" height="25.511124" transform="translate(60.422485 10.022217)" fill="white"/>
|
||||||
|
</clipPath>
|
||||||
|
</defs>
|
||||||
|
<g clip-path="url(#clip30_2029)">
|
||||||
|
<path id="path" d="M119.508 30.113L117.562 30.113L117.562 27.0967L119.508 27.0967C120.713 27.0967 121.931 26.7961 122.715 25.9614C123.5 25.1265 123.796 23.8464 123.796 22.5664C123.796 21.2864 123.512 20.0063 122.715 19.1716C121.919 18.3369 120.713 18.0364 119.508 18.0364C118.302 18.0364 117.085 18.3369 116.3 19.1716C115.515 20.0063 115.219 21.2864 115.219 22.5664L115.219 34.9551L111.806 34.9551L111.806 15.031L115.219 15.031L115.219 16.2998L115.845 16.2998C115.913 16.2219 115.981 16.1553 116.049 16.0884C116.903 15.3093 118.211 15.031 119.496 15.031C121.51 15.031 123.523 15.532 124.843 16.9233C126.162 18.3145 126.629 20.4517 126.629 22.5776C126.629 24.7036 126.151 26.8296 124.843 28.2319C123.535 29.6345 121.51 30.113 119.508 30.113Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M67.5664 15.5654L69.5117 15.5654L69.5117 18.5818L67.5664 18.5818C66.3606 18.5818 65.1434 18.8823 64.3585 19.717C63.5736 20.552 63.2778 21.832 63.2778 23.1121C63.2778 24.3921 63.5623 25.6721 64.3585 26.5068C65.1548 27.3418 66.3606 27.6423 67.5664 27.6423C68.7722 27.6423 69.9895 27.3418 70.7744 26.5068C71.5593 25.6721 71.8551 24.3921 71.8551 23.1121L71.8551 10.7124L75.2677 10.7124L75.2677 30.6475L71.8551 30.6475L71.8551 29.3787L71.2294 29.3787C71.1611 29.4565 71.0929 29.5234 71.0247 29.5901C70.1715 30.3691 68.8633 30.6475 67.5779 30.6475C65.5643 30.6475 63.5509 30.1467 62.2313 28.7554C60.9117 27.364 60.4453 25.2268 60.4453 23.1008C60.4453 20.9749 60.9231 18.8489 62.2313 17.4465C63.5509 16.0552 65.5643 15.5654 67.5664 15.5654Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M92.3881 22.845L92.3881 24.0581L83.299 24.0581L83.299 21.6428L89.328 21.6428C89.1914 20.7634 88.8729 19.9397 88.3042 19.3386C87.4851 18.4705 86.2224 18.1589 84.9711 18.1589C83.7198 18.1589 82.4572 18.4705 81.6381 19.3386C80.819 20.2068 80.5232 21.5315 80.5232 22.845C80.5232 24.1582 80.819 25.4939 81.6381 26.3511C82.4572 27.208 83.7198 27.531 84.9711 27.531C86.2224 27.531 87.4851 27.2192 88.3042 26.3511C88.418 26.2285 88.5203 26.095 88.6227 25.9614L91.9899 25.9614C91.6941 27.0078 91.2277 27.9539 90.5225 28.6885C89.1573 30.1243 87.0529 30.6475 84.9711 30.6475C82.8894 30.6475 80.7849 30.1355 79.4198 28.6885C78.0547 27.2415 77.5542 25.0376 77.5542 22.845C77.5542 20.6521 78.0433 18.437 79.4198 17.0012C80.7963 15.5654 82.8894 15.0422 84.9711 15.0422C87.0529 15.0422 89.1573 15.5542 90.5225 17.0012C91.8988 18.4482 92.3881 20.6521 92.3881 22.845Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M109.52 22.845L109.52 24.0581L100.431 24.0581L100.431 21.6428L106.46 21.6428C106.323 20.7634 106.005 19.9397 105.436 19.3386C104.617 18.4705 103.354 18.1589 102.103 18.1589C100.852 18.1589 99.5889 18.4705 98.7698 19.3386C97.9507 20.2068 97.6549 21.5315 97.6549 22.845C97.6549 24.1582 97.9507 25.4939 98.7698 26.3511C99.5889 27.208 100.852 27.531 102.103 27.531C103.354 27.531 104.617 27.2192 105.436 26.3511C105.55 26.2285 105.652 26.095 105.754 25.9614L109.122 25.9614C108.826 27.0078 108.359 27.9539 107.654 28.6885C106.289 30.1243 104.185 30.6475 102.103 30.6475C100.021 30.6475 97.9166 30.1355 96.5515 28.6885C95.1864 27.2415 94.6859 25.0376 94.6859 22.845C94.6859 20.6521 95.175 18.437 96.5515 17.0012C97.928 15.5654 100.021 15.0422 102.103 15.0422C104.185 15.0422 106.289 15.5542 107.654 17.0012C109.031 18.4482 109.52 20.6521 109.52 22.845Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M136.355 30.6475C138.437 30.6475 140.541 30.3469 141.906 29.49C143.271 28.6328 143.772 27.3306 143.772 26.0393C143.772 24.7483 143.282 23.4348 141.906 22.5889C140.541 21.7429 138.437 21.4312 136.355 21.4312C135.467 21.4312 134.648 21.3088 134.068 20.9861C133.488 20.6521 133.272 20.1511 133.272 19.6504C133.272 19.1494 133.477 18.6375 134.068 18.3147C134.648 17.9807 135.547 17.8694 136.434 17.8694C137.322 17.8694 138.22 17.9919 138.801 18.3147C139.381 18.6487 139.597 19.1494 139.597 19.6504L143.066 19.6504C143.066 18.3591 142.623 17.0457 141.383 16.2C140.143 15.354 138.243 15.0422 136.355 15.0422C134.466 15.0422 132.567 15.3428 131.327 16.2C130.087 17.0569 129.643 18.3591 129.643 19.6504C129.643 20.9414 130.087 22.2549 131.327 23.1008C132.567 23.9468 134.466 24.2585 136.355 24.2585C137.333 24.2585 138.414 24.3809 139.062 24.7036C139.711 25.0266 139.938 25.5386 139.938 26.0393C139.938 26.5403 139.711 27.0522 139.062 27.375C138.414 27.6978 137.424 27.8203 136.446 27.8203C135.467 27.8203 134.466 27.6978 133.829 27.375C133.192 27.0522 132.953 26.5403 132.953 26.0393L128.949 26.0393C128.949 27.3306 129.438 28.644 130.815 29.49C132.191 30.3359 134.273 30.6475 136.355 30.6475Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M160.903 22.845L160.903 24.0581L151.814 24.0581L151.814 21.6428L157.843 21.6428C157.707 20.7634 157.388 19.9397 156.82 19.3386C156 18.4705 154.738 18.1589 153.486 18.1589C152.235 18.1589 150.972 18.4705 150.153 19.3386C149.334 20.2068 149.039 21.5315 149.039 22.845C149.039 24.1582 149.334 25.4939 150.153 26.3511C150.972 27.208 152.235 27.531 153.486 27.531C154.738 27.531 156 27.2192 156.82 26.3511C156.933 26.2285 157.036 26.095 157.138 25.9614L160.505 25.9614C160.209 27.0078 159.743 27.9539 159.038 28.6885C157.673 30.1243 155.568 30.6475 153.486 30.6475C151.405 30.6475 149.3 30.1355 147.935 28.6885C146.57 27.2415 146.07 25.0376 146.07 22.845C146.07 20.6521 146.559 18.437 147.935 17.0012C149.312 15.5654 151.405 15.0422 153.486 15.0422C155.568 15.0422 157.673 15.5542 159.038 17.0012C160.414 18.4482 160.903 20.6521 160.903 22.845Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<path id="path" d="M178.035 22.845L178.035 24.0581L168.946 24.0581L168.946 21.6428L174.975 21.6428C174.839 20.7634 174.52 19.9397 173.951 19.3386C173.132 18.4705 171.87 18.1589 170.618 18.1589C169.367 18.1589 168.104 18.4705 167.285 19.3386C166.466 20.2068 166.17 21.5315 166.17 22.845C166.17 24.1582 166.466 25.4939 167.285 26.3511C168.104 27.208 169.367 27.531 170.618 27.531C171.87 27.531 173.132 27.2192 173.951 26.3511C174.065 26.2285 174.167 26.095 174.27 25.9614L177.637 25.9614C177.341 27.0078 176.875 27.9539 176.17 28.6885C174.804 30.1243 172.7 30.6475 170.618 30.6475C168.536 30.6475 166.432 30.1355 165.067 28.6885C163.702 27.2415 163.201 25.0376 163.201 22.845C163.201 20.6521 163.69 18.437 165.067 17.0012C166.443 15.5654 168.536 15.0422 170.618 15.0422C172.7 15.0422 174.804 15.5542 176.17 17.0012C177.546 18.4482 178.035 20.6521 178.035 22.845Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
<rect id="rect" x="180.321533" y="10.022217" width="3.412687" height="20.625223" fill="#4D6BFE"/>
|
||||||
|
<path id="polygon" d="M189.559 22.3772L195.155 30.6475L190.935 30.6475L185.338 22.3772L190.935 15.7322L195.155 15.7322L189.559 22.3772Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
</g>
|
||||||
|
<path id="path" d="M55.6128 3.47119C55.0175 3.17944 54.7611 3.73535 54.413 4.01782C54.2939 4.10889 54.1932 4.22729 54.0924 4.33667C53.2223 5.26587 52.2057 5.87646 50.8776 5.80347C48.9359 5.69409 47.2781 6.30469 45.8126 7.78979C45.5012 5.9585 44.4663 4.86499 42.8909 4.16357C42.0667 3.79907 41.2332 3.43457 40.6561 2.64185C40.2532 2.07715 40.1432 1.44849 39.9418 0.828857C39.8135 0.455322 39.6853 0.0725098 39.2548 0.00878906C38.7877 -0.0639648 38.6045 0.327637 38.4213 0.655762C37.6886 1.99512 37.4047 3.47119 37.4321 4.96533C37.4962 8.32739 38.9159 11.0059 41.7369 12.9102C42.0575 13.1289 42.1399 13.3474 42.0392 13.6665C41.8468 14.3225 41.6178 14.9602 41.4164 15.6162C41.2881 16.0354 41.0957 16.1265 40.647 15.9441C39.0991 15.2974 37.7618 14.3406 36.5803 13.1836C34.5745 11.2429 32.761 9.10181 30.4988 7.42529C29.9675 7.03345 29.4363 6.66919 28.8867 6.32275C26.5786 4.08154 29.189 2.24097 29.7935 2.02246C30.4254 1.79468 30.0133 1.01099 27.9708 1.02026C25.9283 1.0293 24.0599 1.71265 21.6786 2.62378C21.3306 2.7605 20.9641 2.8606 20.5886 2.94263C18.4271 2.53271 16.1831 2.44141 13.8384 2.70581C9.42371 3.19775 5.89758 5.28418 3.30554 8.84668C0.191406 13.1289 -0.54126 17.9941 0.356323 23.0691C1.29968 28.4172 4.02905 32.8452 8.22388 36.3076C12.5745 39.8972 17.5845 41.6558 23.2997 41.3186C26.771 41.1182 30.6361 40.6536 34.9958 36.9636C36.0948 37.5103 37.2489 37.7288 39.1632 37.8928C40.6378 38.0295 42.0575 37.8201 43.1565 37.5923C44.8784 37.2278 44.7594 35.6333 44.1366 35.3418C39.09 32.9912 40.1981 33.9478 39.1907 33.1733C41.7552 30.1394 45.6204 26.9868 47.1316 16.7732C47.2506 15.9624 47.1499 15.4521 47.1316 14.7961C47.1224 14.3953 47.214 14.2405 47.672 14.1948C48.9359 14.0491 50.1632 13.7029 51.2898 13.0833C54.5596 11.2976 55.8784 8.36377 56.1898 4.84692C56.2357 4.30933 56.1807 3.75342 55.6128 3.47119ZM27.119 35.123C22.2281 31.2783 19.856 30.0117 18.8759 30.0664C17.96 30.1211 18.1249 31.1689 18.3263 31.8523C18.537 32.5264 18.8118 32.9912 19.1964 33.5833C19.462 33.9751 19.6453 34.5581 18.9309 34.9956C17.3555 35.9705 14.6169 34.6675 14.4886 34.6038C11.3014 32.7268 8.63611 30.2485 6.75842 26.8594C4.94495 23.5974 3.89172 20.0989 3.71765 16.3633C3.67188 15.4614 3.9375 15.1423 4.83508 14.9785C6.0166 14.7598 7.23474 14.7141 8.41626 14.8872C13.408 15.6162 17.6577 17.8484 21.2206 21.3835C23.2539 23.397 24.7926 25.8025 26.3772 28.1531C28.0624 30.6494 29.8759 33.0276 32.184 34.9773C32.9991 35.6606 33.6494 36.1799 34.2722 36.5627C32.3947 36.7722 29.2622 36.8179 27.119 35.123ZM29.4637 20.0442C29.4637 19.6433 29.7843 19.3245 30.1874 19.3245C30.2789 19.3245 30.3613 19.3425 30.4346 19.3699C30.5354 19.4065 30.627 19.4612 30.7002 19.543C30.8285 19.6707 30.9017 19.8528 30.9017 20.0442C30.9017 20.4451 30.5812 20.7639 30.1782 20.7639C29.7751 20.7639 29.4637 20.4451 29.4637 20.0442ZM36.7452 23.7798C36.2781 23.9712 35.811 24.135 35.3622 24.1533C34.6661 24.1897 33.9059 23.9072 33.4938 23.561C32.8527 23.0234 32.3947 22.7229 32.2023 21.7844C32.1199 21.3835 32.1656 20.7639 32.239 20.4087C32.4038 19.6433 32.2206 19.1514 31.6803 18.7048C31.2406 18.3403 30.6819 18.2402 30.0682 18.2402C29.8392 18.2402 29.6287 18.1399 29.4729 18.0579C29.2164 17.9304 29.0059 17.6116 29.2073 17.2197C29.2714 17.0923 29.5829 16.7825 29.6561 16.7278C30.4896 16.2539 31.4513 16.4089 32.3397 16.7642C33.1641 17.1013 33.7869 17.7209 34.6844 18.5955C35.6003 19.6523 35.7651 19.9441 36.2872 20.7366C36.6995 21.3562 37.075 21.9939 37.3314 22.7229C37.4871 23.1785 37.2856 23.552 36.7452 23.7798Z" fill-rule="nonzero" fill="#4D6BFE"/>
|
||||||
|
</svg>
|
||||||
|
After Width: | Height: | Size: 10 KiB |
BIN
images/pie_chart.png
Normal file
BIN
images/pie_chart.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 57 KiB |
BIN
images/teaser.png
Normal file
BIN
images/teaser.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 561 KiB |
BIN
images/ve.png
Normal file
BIN
images/ve.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 269 KiB |
67
inference.py
Normal file
67
inference.py
Normal file
@ -0,0 +1,67 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from transformers import AutoModelForCausalLM
|
||||||
|
|
||||||
|
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||||
|
from janus.utils.io import load_pil_images
|
||||||
|
|
||||||
|
# specify the path to the model
|
||||||
|
model_path = "deepseek-ai/Janus-1.3B"
|
||||||
|
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
|
||||||
|
tokenizer = vl_chat_processor.tokenizer
|
||||||
|
|
||||||
|
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
|
||||||
|
model_path, trust_remote_code=True
|
||||||
|
)
|
||||||
|
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
|
||||||
|
|
||||||
|
conversation = [
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "<image_placeholder>\nConvert the formula into latex code.",
|
||||||
|
"images": ["images/equation.png"],
|
||||||
|
},
|
||||||
|
{"role": "Assistant", "content": ""},
|
||||||
|
]
|
||||||
|
|
||||||
|
# load images and prepare for inputs
|
||||||
|
pil_images = load_pil_images(conversation)
|
||||||
|
prepare_inputs = vl_chat_processor(
|
||||||
|
conversations=conversation, images=pil_images, force_batchify=True
|
||||||
|
).to(vl_gpt.device)
|
||||||
|
|
||||||
|
# # run image encoder to get the image embeddings
|
||||||
|
inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
|
||||||
|
|
||||||
|
# # run the model to get the response
|
||||||
|
outputs = vl_gpt.language_model.generate(
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
attention_mask=prepare_inputs.attention_mask,
|
||||||
|
pad_token_id=tokenizer.eos_token_id,
|
||||||
|
bos_token_id=tokenizer.bos_token_id,
|
||||||
|
eos_token_id=tokenizer.eos_token_id,
|
||||||
|
max_new_tokens=512,
|
||||||
|
do_sample=False,
|
||||||
|
use_cache=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
|
||||||
|
print(f"{prepare_inputs['sft_format'][0]}", answer)
|
||||||
31
janus/__init__.py
Normal file
31
janus/__init__.py
Normal file
@ -0,0 +1,31 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
|
||||||
|
# check if python version is above 3.10
|
||||||
|
import sys
|
||||||
|
|
||||||
|
if sys.version_info >= (3, 10):
|
||||||
|
print("Python version is above 3.10, patching the collections module.")
|
||||||
|
# Monkey patch collections
|
||||||
|
import collections
|
||||||
|
import collections.abc
|
||||||
|
|
||||||
|
for type_name in collections.abc.__all__:
|
||||||
|
setattr(collections, type_name, getattr(collections.abc, type_name))
|
||||||
28
janus/models/__init__.py
Normal file
28
janus/models/__init__.py
Normal file
@ -0,0 +1,28 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from .image_processing_vlm import VLMImageProcessor
|
||||||
|
from .modeling_vlm import MultiModalityCausalLM
|
||||||
|
from .processing_vlm import VLChatProcessor
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VLMImageProcessor",
|
||||||
|
"VLChatProcessor",
|
||||||
|
"MultiModalityCausalLM",
|
||||||
|
]
|
||||||
122
janus/models/clip_encoder.py
Normal file
122
janus/models/clip_encoder.py
Normal file
@ -0,0 +1,122 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from typing import Dict, List, Literal, Optional, Tuple, Union
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torchvision.transforms
|
||||||
|
from einops import rearrange
|
||||||
|
|
||||||
|
from janus.models.siglip_vit import create_siglip_vit
|
||||||
|
|
||||||
|
|
||||||
|
class CLIPVisionTower(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
model_name: str = "siglip_large_patch16_384",
|
||||||
|
image_size: Union[Tuple[int, int], int] = 336,
|
||||||
|
select_feature: str = "patch",
|
||||||
|
select_layer: int = -2,
|
||||||
|
select_layers: list = None,
|
||||||
|
ckpt_path: str = "",
|
||||||
|
pixel_mean: Optional[List[float]] = None,
|
||||||
|
pixel_std: Optional[List[float]] = None,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.model_name = model_name
|
||||||
|
self.select_feature = select_feature
|
||||||
|
self.select_layer = select_layer
|
||||||
|
self.select_layers = select_layers
|
||||||
|
|
||||||
|
vision_tower_params = {
|
||||||
|
"model_name": model_name,
|
||||||
|
"image_size": image_size,
|
||||||
|
"ckpt_path": ckpt_path,
|
||||||
|
"select_layer": select_layer,
|
||||||
|
}
|
||||||
|
vision_tower_params.update(kwargs)
|
||||||
|
self.vision_tower, self.forward_kwargs = self.build_vision_tower(
|
||||||
|
vision_tower_params
|
||||||
|
)
|
||||||
|
|
||||||
|
if pixel_mean is not None and pixel_std is not None:
|
||||||
|
image_norm = torchvision.transforms.Normalize(
|
||||||
|
mean=pixel_mean, std=pixel_std
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
image_norm = None
|
||||||
|
|
||||||
|
self.image_norm = image_norm
|
||||||
|
|
||||||
|
def build_vision_tower(self, vision_tower_params):
|
||||||
|
if self.model_name.startswith("siglip"):
|
||||||
|
self.select_feature = "same"
|
||||||
|
vision_tower = create_siglip_vit(**vision_tower_params)
|
||||||
|
forward_kwargs = dict()
|
||||||
|
|
||||||
|
elif self.model_name.startswith("sam"):
|
||||||
|
vision_tower = create_sam_vit(**vision_tower_params)
|
||||||
|
forward_kwargs = dict()
|
||||||
|
|
||||||
|
else: # huggingface
|
||||||
|
from transformers import CLIPVisionModel
|
||||||
|
|
||||||
|
vision_tower = CLIPVisionModel.from_pretrained(**vision_tower_params)
|
||||||
|
forward_kwargs = dict(output_hidden_states=True)
|
||||||
|
|
||||||
|
return vision_tower, forward_kwargs
|
||||||
|
|
||||||
|
def feature_select(self, image_forward_outs):
|
||||||
|
if isinstance(image_forward_outs, torch.Tensor):
|
||||||
|
# the output has been the self.select_layer"s features
|
||||||
|
image_features = image_forward_outs
|
||||||
|
else:
|
||||||
|
image_features = image_forward_outs.hidden_states[self.select_layer]
|
||||||
|
|
||||||
|
if self.select_feature == "patch":
|
||||||
|
# if the output has cls_token
|
||||||
|
image_features = image_features[:, 1:]
|
||||||
|
elif self.select_feature == "cls_patch":
|
||||||
|
image_features = image_features
|
||||||
|
elif self.select_feature == "same":
|
||||||
|
image_features = image_features
|
||||||
|
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unexpected select feature: {self.select_feature}")
|
||||||
|
return image_features
|
||||||
|
|
||||||
|
def forward(self, images):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
images (torch.Tensor): [b, 3, H, W]
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
image_features (torch.Tensor): [b, n_patch, d]
|
||||||
|
"""
|
||||||
|
|
||||||
|
if self.image_norm is not None:
|
||||||
|
images = self.image_norm(images)
|
||||||
|
|
||||||
|
image_forward_outs = self.vision_tower(images, **self.forward_kwargs)
|
||||||
|
image_features = self.feature_select(image_forward_outs)
|
||||||
|
return image_features
|
||||||
208
janus/models/image_processing_vlm.py
Normal file
208
janus/models/image_processing_vlm.py
Normal file
@ -0,0 +1,208 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from typing import List, Tuple, Union
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torchvision
|
||||||
|
import torchvision.transforms.functional
|
||||||
|
from PIL import Image
|
||||||
|
from transformers import AutoImageProcessor, PretrainedConfig
|
||||||
|
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
|
||||||
|
from transformers.image_utils import to_numpy_array
|
||||||
|
from transformers.utils import logging
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
ImageType = Union[np.ndarray, torch.Tensor, Image.Image]
|
||||||
|
IMAGENET_MEAN = (0.48145466, 0.4578275, 0.40821073)
|
||||||
|
IMAGENET_STD = (0.26862954, 0.26130258, 0.27577711)
|
||||||
|
IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5)
|
||||||
|
IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)
|
||||||
|
|
||||||
|
|
||||||
|
def expand2square(pil_img, background_color):
|
||||||
|
width, height = pil_img.size
|
||||||
|
if width == height:
|
||||||
|
return pil_img
|
||||||
|
elif width > height:
|
||||||
|
result = Image.new(pil_img.mode, (width, width), background_color)
|
||||||
|
result.paste(pil_img, (0, (width - height) // 2))
|
||||||
|
return result
|
||||||
|
else:
|
||||||
|
result = Image.new(pil_img.mode, (height, height), background_color)
|
||||||
|
result.paste(pil_img, ((height - width) // 2, 0))
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
class VLMImageProcessorConfig(PretrainedConfig):
|
||||||
|
model_type = "deepseek_vlm"
|
||||||
|
image_size: int
|
||||||
|
min_size: int
|
||||||
|
image_mean: Union[Tuple[float, float, float], List[float]]
|
||||||
|
image_std: Union[Tuple[float, float, float], List[float]]
|
||||||
|
rescale_factor: float
|
||||||
|
do_normalize: bool
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
image_size: int,
|
||||||
|
min_size: int = 14,
|
||||||
|
image_mean: Union[Tuple[float, float, float], List[float]] = (
|
||||||
|
0.48145466,
|
||||||
|
0.4578275,
|
||||||
|
0.40821073,
|
||||||
|
),
|
||||||
|
image_std: Union[Tuple[float, float, float], List[float]] = (
|
||||||
|
0.26862954,
|
||||||
|
0.26130258,
|
||||||
|
0.27577711,
|
||||||
|
),
|
||||||
|
rescale_factor: float = 1.0 / 255.0,
|
||||||
|
do_normalize: bool = True,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
self.image_size = image_size
|
||||||
|
self.min_size = min_size
|
||||||
|
self.image_mean = image_mean
|
||||||
|
self.image_std = image_std
|
||||||
|
self.rescale_factor = rescale_factor
|
||||||
|
self.do_normalize = do_normalize
|
||||||
|
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class VLMImageProcessor(BaseImageProcessor):
|
||||||
|
model_input_names = ["pixel_values"]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
image_size: int,
|
||||||
|
min_size: int = 14,
|
||||||
|
image_mean: Union[Tuple[float, float, float], List[float]] = (
|
||||||
|
0.48145466,
|
||||||
|
0.4578275,
|
||||||
|
0.40821073,
|
||||||
|
),
|
||||||
|
image_std: Union[Tuple[float, float, float], List[float]] = (
|
||||||
|
0.26862954,
|
||||||
|
0.26130258,
|
||||||
|
0.27577711,
|
||||||
|
),
|
||||||
|
rescale_factor: float = 1.0 / 255.0,
|
||||||
|
do_normalize: bool = True,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.image_size = image_size
|
||||||
|
self.rescale_factor = rescale_factor
|
||||||
|
self.image_mean = image_mean
|
||||||
|
self.image_std = image_std
|
||||||
|
self.min_size = min_size
|
||||||
|
self.do_normalize = do_normalize
|
||||||
|
|
||||||
|
if image_mean is None:
|
||||||
|
self.background_color = (127, 127, 127)
|
||||||
|
else:
|
||||||
|
self.background_color = tuple([int(x * 255) for x in image_mean])
|
||||||
|
|
||||||
|
def resize(self, pil_img: Image) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pil_img (PIL.Image): [H, W, 3] in PIL.Image in RGB
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
x (np.ndarray): [3, self.image_size, self.image_size]
|
||||||
|
"""
|
||||||
|
|
||||||
|
width, height = pil_img.size
|
||||||
|
max_size = max(width, height)
|
||||||
|
|
||||||
|
size = [
|
||||||
|
max(int(height / max_size * self.image_size), self.min_size),
|
||||||
|
max(int(width / max_size * self.image_size), self.min_size),
|
||||||
|
]
|
||||||
|
|
||||||
|
if width <= 0 or height <= 0 or size[0] <= 0 or size[1] <= 0:
|
||||||
|
print(f"orig size = {pil_img.size}, new size = {size}")
|
||||||
|
raise ValueError("Invalid size!")
|
||||||
|
|
||||||
|
pil_img = torchvision.transforms.functional.resize(
|
||||||
|
pil_img,
|
||||||
|
size,
|
||||||
|
interpolation=torchvision.transforms.functional.InterpolationMode.BICUBIC,
|
||||||
|
antialias=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
pil_img = expand2square(pil_img, self.background_color)
|
||||||
|
x = to_numpy_array(pil_img)
|
||||||
|
|
||||||
|
# [H, W, 3] -> [3, H, W]
|
||||||
|
x = np.transpose(x, (2, 0, 1))
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
def preprocess(self, images, return_tensors: str = "pt", **kwargs) -> BatchFeature:
|
||||||
|
# resize and pad to [self.image_size, self.image_size]
|
||||||
|
# then convert from [H, W, 3] to [3, H, W]
|
||||||
|
images: List[np.ndarray] = [self.resize(image) for image in images]
|
||||||
|
|
||||||
|
# resacle from [0, 255] -> [0, 1]
|
||||||
|
images = [
|
||||||
|
self.rescale(
|
||||||
|
image=image,
|
||||||
|
scale=self.rescale_factor,
|
||||||
|
input_data_format="channels_first",
|
||||||
|
)
|
||||||
|
for image in images
|
||||||
|
]
|
||||||
|
|
||||||
|
# normalize
|
||||||
|
if self.do_normalize:
|
||||||
|
images = [
|
||||||
|
self.normalize(
|
||||||
|
image=image,
|
||||||
|
mean=self.image_mean,
|
||||||
|
std=self.image_std,
|
||||||
|
input_data_format="channels_first",
|
||||||
|
)
|
||||||
|
for image in images
|
||||||
|
]
|
||||||
|
|
||||||
|
data = {"pixel_values": images}
|
||||||
|
return BatchFeature(data=data, tensor_type=return_tensors)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def default_shape(self):
|
||||||
|
return [3, self.image_size, self.image_size]
|
||||||
|
|
||||||
|
|
||||||
|
AutoImageProcessor.register(VLMImageProcessorConfig, VLMImageProcessor)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
image_processor = VLMImageProcessor(
|
||||||
|
image_size=1024,
|
||||||
|
image_mean=IMAGENET_INCEPTION_MEAN,
|
||||||
|
image_std=IMAGENET_INCEPTION_STD,
|
||||||
|
do_normalize=True,
|
||||||
|
)
|
||||||
272
janus/models/modeling_vlm.py
Normal file
272
janus/models/modeling_vlm.py
Normal file
@ -0,0 +1,272 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from attrdict import AttrDict
|
||||||
|
from einops import rearrange
|
||||||
|
from transformers import (
|
||||||
|
AutoConfig,
|
||||||
|
AutoModelForCausalLM,
|
||||||
|
LlamaConfig,
|
||||||
|
LlamaForCausalLM,
|
||||||
|
PreTrainedModel,
|
||||||
|
)
|
||||||
|
from transformers.configuration_utils import PretrainedConfig
|
||||||
|
|
||||||
|
from janus.models.clip_encoder import CLIPVisionTower
|
||||||
|
from janus.models.projector import MlpProjector
|
||||||
|
|
||||||
|
|
||||||
|
class vision_head(torch.nn.Module):
|
||||||
|
def __init__(self, params):
|
||||||
|
super().__init__()
|
||||||
|
self.output_mlp_projector = torch.nn.Linear(
|
||||||
|
params.n_embed, params.image_token_embed
|
||||||
|
)
|
||||||
|
self.vision_activation = torch.nn.GELU()
|
||||||
|
self.vision_head = torch.nn.Linear(
|
||||||
|
params.image_token_embed, params.image_token_size
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.output_mlp_projector(x)
|
||||||
|
x = self.vision_activation(x)
|
||||||
|
x = self.vision_head(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
def model_name_to_cls(cls_name):
|
||||||
|
if "MlpProjector" in cls_name:
|
||||||
|
cls = MlpProjector
|
||||||
|
|
||||||
|
elif "CLIPVisionTower" in cls_name:
|
||||||
|
cls = CLIPVisionTower
|
||||||
|
|
||||||
|
elif "VQ" in cls_name:
|
||||||
|
from janus.models.vq_model import VQ_models
|
||||||
|
|
||||||
|
cls = VQ_models[cls_name]
|
||||||
|
elif "vision_head" in cls_name:
|
||||||
|
cls = vision_head
|
||||||
|
else:
|
||||||
|
raise ValueError(f"class_name {cls_name} is invalid.")
|
||||||
|
|
||||||
|
return cls
|
||||||
|
|
||||||
|
|
||||||
|
class VisionConfig(PretrainedConfig):
|
||||||
|
model_type = "vision"
|
||||||
|
cls: str = ""
|
||||||
|
params: AttrDict = {}
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.cls = kwargs.get("cls", "")
|
||||||
|
if not isinstance(self.cls, str):
|
||||||
|
self.cls = self.cls.__name__
|
||||||
|
|
||||||
|
self.params = AttrDict(kwargs.get("params", {}))
|
||||||
|
|
||||||
|
|
||||||
|
class AlignerConfig(PretrainedConfig):
|
||||||
|
model_type = "aligner"
|
||||||
|
cls: str = ""
|
||||||
|
params: AttrDict = {}
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.cls = kwargs.get("cls", "")
|
||||||
|
if not isinstance(self.cls, str):
|
||||||
|
self.cls = self.cls.__name__
|
||||||
|
|
||||||
|
self.params = AttrDict(kwargs.get("params", {}))
|
||||||
|
|
||||||
|
|
||||||
|
class GenVisionConfig(PretrainedConfig):
|
||||||
|
model_type = "gen_vision"
|
||||||
|
cls: str = ""
|
||||||
|
params: AttrDict = {}
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.cls = kwargs.get("cls", "")
|
||||||
|
if not isinstance(self.cls, str):
|
||||||
|
self.cls = self.cls.__name__
|
||||||
|
|
||||||
|
self.params = AttrDict(kwargs.get("params", {}))
|
||||||
|
|
||||||
|
|
||||||
|
class GenAlignerConfig(PretrainedConfig):
|
||||||
|
model_type = "gen_aligner"
|
||||||
|
cls: str = ""
|
||||||
|
params: AttrDict = {}
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.cls = kwargs.get("cls", "")
|
||||||
|
if not isinstance(self.cls, str):
|
||||||
|
self.cls = self.cls.__name__
|
||||||
|
|
||||||
|
self.params = AttrDict(kwargs.get("params", {}))
|
||||||
|
|
||||||
|
|
||||||
|
class GenHeadConfig(PretrainedConfig):
|
||||||
|
model_type = "gen_head"
|
||||||
|
cls: str = ""
|
||||||
|
params: AttrDict = {}
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
self.cls = kwargs.get("cls", "")
|
||||||
|
if not isinstance(self.cls, str):
|
||||||
|
self.cls = self.cls.__name__
|
||||||
|
|
||||||
|
self.params = AttrDict(kwargs.get("params", {}))
|
||||||
|
|
||||||
|
|
||||||
|
class MultiModalityConfig(PretrainedConfig):
|
||||||
|
model_type = "multi_modality"
|
||||||
|
vision_config: VisionConfig
|
||||||
|
aligner_config: AlignerConfig
|
||||||
|
|
||||||
|
gen_vision_config: GenVisionConfig
|
||||||
|
gen_aligner_config: GenAlignerConfig
|
||||||
|
gen_head_config: GenHeadConfig
|
||||||
|
|
||||||
|
language_config: LlamaConfig
|
||||||
|
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
vision_config = kwargs.get("vision_config", {})
|
||||||
|
self.vision_config = VisionConfig(**vision_config)
|
||||||
|
|
||||||
|
aligner_config = kwargs.get("aligner_config", {})
|
||||||
|
self.aligner_config = AlignerConfig(**aligner_config)
|
||||||
|
|
||||||
|
gen_vision_config = kwargs.get("gen_vision_config", {})
|
||||||
|
self.gen_vision_config = GenVisionConfig(**gen_vision_config)
|
||||||
|
|
||||||
|
gen_aligner_config = kwargs.get("gen_aligner_config", {})
|
||||||
|
self.gen_aligner_config = GenAlignerConfig(**gen_aligner_config)
|
||||||
|
|
||||||
|
gen_head_config = kwargs.get("gen_head_config", {})
|
||||||
|
self.gen_head_config = GenHeadConfig(**gen_head_config)
|
||||||
|
|
||||||
|
language_config = kwargs.get("language_config", {})
|
||||||
|
if isinstance(language_config, LlamaConfig):
|
||||||
|
self.language_config = language_config
|
||||||
|
else:
|
||||||
|
self.language_config = LlamaConfig(**language_config)
|
||||||
|
|
||||||
|
|
||||||
|
class MultiModalityPreTrainedModel(PreTrainedModel):
|
||||||
|
config_class = MultiModalityConfig
|
||||||
|
base_model_prefix = "multi_modality"
|
||||||
|
_no_split_modules = []
|
||||||
|
_skip_keys_device_placement = "past_key_values"
|
||||||
|
|
||||||
|
|
||||||
|
class MultiModalityCausalLM(MultiModalityPreTrainedModel):
|
||||||
|
def __init__(self, config: MultiModalityConfig):
|
||||||
|
super().__init__(config)
|
||||||
|
|
||||||
|
vision_config = config.vision_config
|
||||||
|
vision_cls = model_name_to_cls(vision_config.cls)
|
||||||
|
self.vision_model = vision_cls(**vision_config.params)
|
||||||
|
|
||||||
|
aligner_config = config.aligner_config
|
||||||
|
aligner_cls = model_name_to_cls(aligner_config.cls)
|
||||||
|
self.aligner = aligner_cls(aligner_config.params)
|
||||||
|
|
||||||
|
gen_vision_config = config.gen_vision_config
|
||||||
|
gen_vision_cls = model_name_to_cls(gen_vision_config.cls)
|
||||||
|
self.gen_vision_model = gen_vision_cls()
|
||||||
|
|
||||||
|
gen_aligner_config = config.gen_aligner_config
|
||||||
|
gen_aligner_cls = model_name_to_cls(gen_aligner_config.cls)
|
||||||
|
self.gen_aligner = gen_aligner_cls(gen_aligner_config.params)
|
||||||
|
|
||||||
|
gen_head_config = config.gen_head_config
|
||||||
|
gen_head_cls = model_name_to_cls(gen_head_config.cls)
|
||||||
|
self.gen_head = gen_head_cls(gen_head_config.params)
|
||||||
|
|
||||||
|
self.gen_embed = torch.nn.Embedding(
|
||||||
|
gen_vision_config.params.image_token_size, gen_vision_config.params.n_embed
|
||||||
|
)
|
||||||
|
|
||||||
|
language_config = config.language_config
|
||||||
|
self.language_model = LlamaForCausalLM(language_config)
|
||||||
|
|
||||||
|
def prepare_inputs_embeds(
|
||||||
|
self,
|
||||||
|
input_ids: torch.LongTensor,
|
||||||
|
pixel_values: torch.FloatTensor,
|
||||||
|
images_seq_mask: torch.LongTensor,
|
||||||
|
images_emb_mask: torch.LongTensor,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_ids (torch.LongTensor): [b, T]
|
||||||
|
pixel_values (torch.FloatTensor): [b, n_images, 3, h, w]
|
||||||
|
images_seq_mask (torch.BoolTensor): [b, T]
|
||||||
|
images_emb_mask (torch.BoolTensor): [b, n_images, n_image_tokens]
|
||||||
|
|
||||||
|
assert torch.sum(images_seq_mask) == torch.sum(images_emb_mask)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
input_embeds (torch.Tensor): [b, T, D]
|
||||||
|
"""
|
||||||
|
|
||||||
|
bs, n = pixel_values.shape[0:2]
|
||||||
|
images = rearrange(pixel_values, "b n c h w -> (b n) c h w")
|
||||||
|
# [b x n, T2, D]
|
||||||
|
images_embeds = self.aligner(self.vision_model(images))
|
||||||
|
|
||||||
|
# [b x n, T2, D] -> [b, n x T2, D]
|
||||||
|
images_embeds = rearrange(images_embeds, "(b n) t d -> b (n t) d", b=bs, n=n)
|
||||||
|
# [b, n, T2] -> [b, n x T2]
|
||||||
|
images_emb_mask = rearrange(images_emb_mask, "b n t -> b (n t)")
|
||||||
|
|
||||||
|
# [b, T, D]
|
||||||
|
input_ids[input_ids < 0] = 0 # ignore the image embeddings
|
||||||
|
inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
|
||||||
|
|
||||||
|
# replace with the image embeddings
|
||||||
|
inputs_embeds[images_seq_mask] = images_embeds[images_emb_mask]
|
||||||
|
|
||||||
|
return inputs_embeds
|
||||||
|
|
||||||
|
def prepare_gen_img_embeds(self, image_ids: torch.LongTensor):
|
||||||
|
return self.gen_aligner(self.gen_embed(image_ids))
|
||||||
|
|
||||||
|
|
||||||
|
AutoConfig.register("vision", VisionConfig)
|
||||||
|
AutoConfig.register("aligner", AlignerConfig)
|
||||||
|
AutoConfig.register("gen_vision", GenVisionConfig)
|
||||||
|
AutoConfig.register("gen_aligner", GenAlignerConfig)
|
||||||
|
AutoConfig.register("gen_head", GenHeadConfig)
|
||||||
|
AutoConfig.register("multi_modality", MultiModalityConfig)
|
||||||
|
AutoModelForCausalLM.register(MultiModalityConfig, MultiModalityCausalLM)
|
||||||
415
janus/models/processing_vlm.py
Normal file
415
janus/models/processing_vlm.py
Normal file
@ -0,0 +1,415 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from PIL.Image import Image
|
||||||
|
from transformers import LlamaTokenizerFast
|
||||||
|
from transformers.processing_utils import ProcessorMixin
|
||||||
|
|
||||||
|
from janus.models.image_processing_vlm import VLMImageProcessor
|
||||||
|
from janus.utils.conversation import get_conv_template
|
||||||
|
|
||||||
|
|
||||||
|
class DictOutput(object):
|
||||||
|
def keys(self):
|
||||||
|
return self.__dict__.keys()
|
||||||
|
|
||||||
|
def __getitem__(self, item):
|
||||||
|
return self.__dict__[item]
|
||||||
|
|
||||||
|
def __setitem__(self, key, value):
|
||||||
|
self.__dict__[key] = value
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class VLChatProcessorOutput(DictOutput):
|
||||||
|
sft_format: str
|
||||||
|
input_ids: torch.Tensor
|
||||||
|
pixel_values: torch.Tensor
|
||||||
|
num_image_tokens: torch.IntTensor
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
return len(self.input_ids)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class BatchedVLChatProcessorOutput(DictOutput):
|
||||||
|
sft_format: List[str]
|
||||||
|
input_ids: torch.Tensor
|
||||||
|
pixel_values: torch.Tensor
|
||||||
|
attention_mask: torch.Tensor
|
||||||
|
images_seq_mask: torch.BoolTensor
|
||||||
|
images_emb_mask: torch.BoolTensor
|
||||||
|
|
||||||
|
def to(self, device, dtype=torch.bfloat16):
|
||||||
|
self.input_ids = self.input_ids.to(device)
|
||||||
|
self.attention_mask = self.attention_mask.to(device)
|
||||||
|
self.images_seq_mask = self.images_seq_mask.to(device)
|
||||||
|
self.images_emb_mask = self.images_emb_mask.to(device)
|
||||||
|
self.pixel_values = self.pixel_values.to(device=device, dtype=dtype)
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
class VLChatProcessor(ProcessorMixin):
|
||||||
|
image_processor_class = "AutoImageProcessor"
|
||||||
|
tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast")
|
||||||
|
|
||||||
|
attributes = ["image_processor", "tokenizer"]
|
||||||
|
|
||||||
|
system_prompt = (
|
||||||
|
"You are a helpful language and vision assistant. "
|
||||||
|
"You are able to understand the visual content that the user provides, "
|
||||||
|
"and assist the user with a variety of tasks using natural language."
|
||||||
|
)
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
image_processor: VLMImageProcessor,
|
||||||
|
tokenizer: LlamaTokenizerFast,
|
||||||
|
image_tag: str = "<image_placeholder>",
|
||||||
|
image_start_tag: str = "<begin_of_image>",
|
||||||
|
image_end_tag: str = "<end_of_image>",
|
||||||
|
num_image_tokens: int = 576,
|
||||||
|
add_special_token: bool = False,
|
||||||
|
sft_format: str = "deepseek",
|
||||||
|
mask_prompt: bool = True,
|
||||||
|
ignore_id: int = -100,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
self.image_processor = image_processor
|
||||||
|
self.tokenizer = tokenizer
|
||||||
|
|
||||||
|
image_id = self.tokenizer.vocab.get(image_tag)
|
||||||
|
if image_id is None:
|
||||||
|
special_tokens = [image_tag]
|
||||||
|
special_tokens_dict = {"additional_special_tokens": special_tokens}
|
||||||
|
self.tokenizer.add_special_tokens(special_tokens_dict)
|
||||||
|
print(f"Add image tag = {image_tag} to the tokenizer")
|
||||||
|
|
||||||
|
self.image_tag = image_tag
|
||||||
|
self.image_start_tag = image_start_tag
|
||||||
|
self.image_end_tag = image_end_tag
|
||||||
|
|
||||||
|
self.num_image_tokens = num_image_tokens
|
||||||
|
self.add_special_token = add_special_token
|
||||||
|
self.sft_format = sft_format
|
||||||
|
self.mask_prompt = mask_prompt
|
||||||
|
self.ignore_id = ignore_id
|
||||||
|
|
||||||
|
super().__init__(
|
||||||
|
image_processor,
|
||||||
|
tokenizer,
|
||||||
|
image_tag,
|
||||||
|
num_image_tokens,
|
||||||
|
add_special_token,
|
||||||
|
sft_format,
|
||||||
|
mask_prompt,
|
||||||
|
ignore_id,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
def new_chat_template(self):
|
||||||
|
conv = get_conv_template(self.sft_format)
|
||||||
|
conv.set_system_message(self.system_prompt)
|
||||||
|
return conv
|
||||||
|
|
||||||
|
def apply_sft_template_for_multi_turn_prompts(
|
||||||
|
self,
|
||||||
|
conversations: List[Dict[str, str]],
|
||||||
|
sft_format: str = "deepseek",
|
||||||
|
system_prompt: str = "",
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Applies the SFT template to conversation.
|
||||||
|
|
||||||
|
An example of conversation:
|
||||||
|
conversation = [
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "<image_placeholder> is Figure 1.\n<image_placeholder> is Figure 2.\nWhich image is brighter?",
|
||||||
|
"images": [
|
||||||
|
"./multi-images/attribute_comparison_1.png",
|
||||||
|
"./multi-images/attribute_comparison_2.png"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"role": "Assistant",
|
||||||
|
"content": ""
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
Args:
|
||||||
|
conversations (List[Dict]): A conversation with a List of Dict[str, str] text.
|
||||||
|
sft_format (str, optional): The format of the SFT template to use. Defaults to "deepseek".
|
||||||
|
system_prompt (str, optional): The system prompt to use in the SFT template. Defaults to "".
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
sft_prompt (str): The formatted text.
|
||||||
|
"""
|
||||||
|
|
||||||
|
conv = get_conv_template(sft_format)
|
||||||
|
conv.set_system_message(system_prompt)
|
||||||
|
for message in conversations:
|
||||||
|
conv.append_message(message["role"], message["content"].strip())
|
||||||
|
sft_prompt = conv.get_prompt().strip()
|
||||||
|
|
||||||
|
return sft_prompt
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_token(self):
|
||||||
|
return self.image_tag
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_id(self):
|
||||||
|
image_id = self.tokenizer.vocab.get(self.image_tag)
|
||||||
|
return image_id
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_start_id(self):
|
||||||
|
image_start_id = self.tokenizer.vocab.get(self.image_start_tag)
|
||||||
|
return image_start_id
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_end_id(self):
|
||||||
|
image_end_id = self.tokenizer.vocab.get(self.image_end_tag)
|
||||||
|
return image_end_id
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_start_token(self):
|
||||||
|
return self.image_start_tag
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_end_token(self):
|
||||||
|
return self.image_end_tag
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pad_id(self):
|
||||||
|
pad_id = self.tokenizer.pad_token_id
|
||||||
|
if pad_id is None:
|
||||||
|
pad_id = self.tokenizer.eos_token_id
|
||||||
|
|
||||||
|
return pad_id
|
||||||
|
|
||||||
|
def add_image_token(
|
||||||
|
self,
|
||||||
|
image_indices: List[int],
|
||||||
|
input_ids: torch.LongTensor,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image_indices (List[int]): [index_0, index_1, ..., index_j]
|
||||||
|
input_ids (torch.LongTensor): [N]
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
input_ids (torch.LongTensor): [N + image tokens]
|
||||||
|
num_image_tokens (torch.IntTensor): [n_images]
|
||||||
|
"""
|
||||||
|
|
||||||
|
input_slices = []
|
||||||
|
|
||||||
|
start = 0
|
||||||
|
for index in image_indices:
|
||||||
|
if self.add_special_token:
|
||||||
|
end = index + 1
|
||||||
|
else:
|
||||||
|
end = index
|
||||||
|
|
||||||
|
# original text tokens
|
||||||
|
input_slices.append(input_ids[start:end])
|
||||||
|
|
||||||
|
# add boi, image tokens, eoi and set the mask as False
|
||||||
|
input_slices.append(self.image_start_id * torch.ones((1), dtype=torch.long))
|
||||||
|
input_slices.append(
|
||||||
|
self.image_id * torch.ones((self.num_image_tokens,), dtype=torch.long)
|
||||||
|
)
|
||||||
|
input_slices.append(self.image_end_id * torch.ones((1), dtype=torch.long))
|
||||||
|
start = index + 1
|
||||||
|
|
||||||
|
# the left part
|
||||||
|
input_slices.append(input_ids[start:])
|
||||||
|
|
||||||
|
# concat all slices
|
||||||
|
input_ids = torch.cat(input_slices, dim=0)
|
||||||
|
num_image_tokens = torch.IntTensor([self.num_image_tokens] * len(image_indices))
|
||||||
|
|
||||||
|
return input_ids, num_image_tokens
|
||||||
|
|
||||||
|
def process_one(
|
||||||
|
self,
|
||||||
|
prompt: str = None,
|
||||||
|
conversations: List[Dict[str, str]] = None,
|
||||||
|
images: List[Image] = None,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompt (str): the formatted prompt;
|
||||||
|
conversations (List[Dict]): conversations with a list of messages;
|
||||||
|
images (List[ImageType]): the list of images;
|
||||||
|
**kwargs:
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
outputs (BaseProcessorOutput): the output of the processor,
|
||||||
|
- input_ids (torch.LongTensor): [N + image tokens]
|
||||||
|
- target_ids (torch.LongTensor): [N + image tokens]
|
||||||
|
- images (torch.FloatTensor): [n_images, 3, H, W]
|
||||||
|
- image_id (int): the id of the image token
|
||||||
|
- num_image_tokens (List[int]): the number of image tokens
|
||||||
|
"""
|
||||||
|
|
||||||
|
assert (
|
||||||
|
prompt is None or conversations is None
|
||||||
|
), "prompt and conversations cannot be used at the same time."
|
||||||
|
|
||||||
|
if prompt is None:
|
||||||
|
# apply sft format
|
||||||
|
sft_format = self.apply_sft_template_for_multi_turn_prompts(
|
||||||
|
conversations=conversations,
|
||||||
|
sft_format=self.sft_format,
|
||||||
|
system_prompt=self.system_prompt,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
sft_format = prompt
|
||||||
|
|
||||||
|
# tokenize
|
||||||
|
input_ids = self.tokenizer.encode(sft_format)
|
||||||
|
input_ids = torch.LongTensor(input_ids)
|
||||||
|
|
||||||
|
# add image tokens to the input_ids
|
||||||
|
image_token_mask: torch.BoolTensor = input_ids == self.image_id
|
||||||
|
image_indices = image_token_mask.nonzero()
|
||||||
|
input_ids, num_image_tokens = self.add_image_token(
|
||||||
|
image_indices=image_indices,
|
||||||
|
input_ids=input_ids,
|
||||||
|
)
|
||||||
|
|
||||||
|
# load images
|
||||||
|
images_outputs = self.image_processor(images, return_tensors="pt")
|
||||||
|
|
||||||
|
prepare = VLChatProcessorOutput(
|
||||||
|
sft_format=sft_format,
|
||||||
|
input_ids=input_ids,
|
||||||
|
pixel_values=images_outputs.pixel_values,
|
||||||
|
num_image_tokens=num_image_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
return prepare
|
||||||
|
|
||||||
|
def __call__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
prompt: str = None,
|
||||||
|
conversations: List[Dict[str, str]] = None,
|
||||||
|
images: List[Image] = None,
|
||||||
|
force_batchify: bool = True,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompt (str): the formatted prompt;
|
||||||
|
conversations (List[Dict]): conversations with a list of messages;
|
||||||
|
images (List[ImageType]): the list of images;
|
||||||
|
force_batchify (bool): force batchify the inputs;
|
||||||
|
**kwargs:
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
outputs (BaseProcessorOutput): the output of the processor,
|
||||||
|
- input_ids (torch.LongTensor): [N + image tokens]
|
||||||
|
- images (torch.FloatTensor): [n_images, 3, H, W]
|
||||||
|
- image_id (int): the id of the image token
|
||||||
|
- num_image_tokens (List[int]): the number of image tokens
|
||||||
|
"""
|
||||||
|
|
||||||
|
prepare = self.process_one(
|
||||||
|
prompt=prompt, conversations=conversations, images=images
|
||||||
|
)
|
||||||
|
|
||||||
|
if force_batchify:
|
||||||
|
prepare = self.batchify([prepare])
|
||||||
|
|
||||||
|
return prepare
|
||||||
|
|
||||||
|
def batchify(
|
||||||
|
self, prepare_list: List[VLChatProcessorOutput]
|
||||||
|
) -> BatchedVLChatProcessorOutput:
|
||||||
|
"""
|
||||||
|
Preprocesses the inputs for multimodal inference.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prepare_list (List[VLChatProcessorOutput]): A list of VLChatProcessorOutput.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
BatchedVLChatProcessorOutput: A dictionary of the inputs to use for multimodal inference.
|
||||||
|
"""
|
||||||
|
|
||||||
|
batch_size = len(prepare_list)
|
||||||
|
sft_format = []
|
||||||
|
n_images = []
|
||||||
|
seq_lens = []
|
||||||
|
for prepare in prepare_list:
|
||||||
|
n_images.append(len(prepare.num_image_tokens))
|
||||||
|
seq_lens.append(len(prepare))
|
||||||
|
|
||||||
|
input_token_max_len = max(seq_lens)
|
||||||
|
max_n_images = max(1, max(n_images))
|
||||||
|
|
||||||
|
batched_input_ids = torch.full(
|
||||||
|
(batch_size, input_token_max_len), self.pad_id
|
||||||
|
).long() # FIXME
|
||||||
|
batched_attention_mask = torch.zeros((batch_size, input_token_max_len)).long()
|
||||||
|
batched_pixel_values = torch.zeros(
|
||||||
|
(batch_size, max_n_images, *self.image_processor.default_shape)
|
||||||
|
).float()
|
||||||
|
batched_images_seq_mask = torch.zeros((batch_size, input_token_max_len)).bool()
|
||||||
|
batched_images_emb_mask = torch.zeros(
|
||||||
|
(batch_size, max_n_images, self.num_image_tokens)
|
||||||
|
).bool()
|
||||||
|
|
||||||
|
for i, prepare in enumerate(prepare_list):
|
||||||
|
input_ids = prepare.input_ids
|
||||||
|
seq_len = len(prepare)
|
||||||
|
n_image = len(prepare.num_image_tokens)
|
||||||
|
# left-padding
|
||||||
|
batched_attention_mask[i, -seq_len:] = 1
|
||||||
|
batched_input_ids[i, -seq_len:] = torch.LongTensor(input_ids)
|
||||||
|
batched_images_seq_mask[i, -seq_len:] = input_ids == self.image_id
|
||||||
|
|
||||||
|
if n_image > 0:
|
||||||
|
batched_pixel_values[i, :n_image] = prepare.pixel_values
|
||||||
|
for j, n_image_tokens in enumerate(prepare.num_image_tokens):
|
||||||
|
batched_images_emb_mask[i, j, :n_image_tokens] = True
|
||||||
|
|
||||||
|
sft_format.append(prepare.sft_format)
|
||||||
|
|
||||||
|
batched_prepares = BatchedVLChatProcessorOutput(
|
||||||
|
input_ids=batched_input_ids,
|
||||||
|
attention_mask=batched_attention_mask,
|
||||||
|
pixel_values=batched_pixel_values,
|
||||||
|
images_seq_mask=batched_images_seq_mask,
|
||||||
|
images_emb_mask=batched_images_emb_mask,
|
||||||
|
sft_format=sft_format,
|
||||||
|
)
|
||||||
|
|
||||||
|
return batched_prepares
|
||||||
100
janus/models/projector.py
Normal file
100
janus/models/projector.py
Normal file
@ -0,0 +1,100 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from typing import Tuple, Union
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from attrdict import AttrDict
|
||||||
|
|
||||||
|
|
||||||
|
class MlpProjector(nn.Module):
|
||||||
|
def __init__(self, cfg):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.cfg = cfg
|
||||||
|
|
||||||
|
if cfg.projector_type == "identity":
|
||||||
|
modules = nn.Identity()
|
||||||
|
|
||||||
|
elif cfg.projector_type == "linear":
|
||||||
|
modules = nn.Linear(cfg.input_dim, cfg.n_embed)
|
||||||
|
|
||||||
|
elif cfg.projector_type == "mlp_gelu":
|
||||||
|
mlp_depth = cfg.get("depth", 1)
|
||||||
|
modules = [nn.Linear(cfg.input_dim, cfg.n_embed)]
|
||||||
|
for _ in range(1, mlp_depth):
|
||||||
|
modules.append(nn.GELU())
|
||||||
|
modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))
|
||||||
|
modules = nn.Sequential(*modules)
|
||||||
|
|
||||||
|
elif cfg.projector_type == "low_high_hybrid_split_mlp_gelu":
|
||||||
|
mlp_depth = cfg.get("depth", 1)
|
||||||
|
self.high_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)
|
||||||
|
self.low_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)
|
||||||
|
|
||||||
|
modules = []
|
||||||
|
for _ in range(1, mlp_depth):
|
||||||
|
modules.append(nn.GELU())
|
||||||
|
modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))
|
||||||
|
modules = nn.Sequential(*modules)
|
||||||
|
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unknown projector type: {cfg.projector_type}")
|
||||||
|
|
||||||
|
self.layers = modules
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self, x_or_tuple: Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
|
||||||
|
Args:
|
||||||
|
x_or_tuple (Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: if it is a tuple of torch.Tensor,
|
||||||
|
then it comes from the hybrid vision encoder, and x = high_res_x, low_res_x);
|
||||||
|
otherwise it is the feature from the single vision encoder.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
x (torch.Tensor): [b, s, c]
|
||||||
|
"""
|
||||||
|
|
||||||
|
if isinstance(x_or_tuple, tuple):
|
||||||
|
# self.cfg.projector_type == "low_high_hybrid_split_mlp_gelu":
|
||||||
|
high_x, low_x = x_or_tuple
|
||||||
|
high_x = self.high_up_proj(high_x)
|
||||||
|
low_x = self.low_up_proj(low_x)
|
||||||
|
x = torch.concat([high_x, low_x], dim=-1)
|
||||||
|
else:
|
||||||
|
x = x_or_tuple
|
||||||
|
|
||||||
|
return self.layers(x)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
cfg = AttrDict(
|
||||||
|
input_dim=1024,
|
||||||
|
n_embed=2048,
|
||||||
|
depth=2,
|
||||||
|
projector_type="low_high_hybrid_split_mlp_gelu",
|
||||||
|
)
|
||||||
|
inputs = (torch.rand(4, 576, 1024), torch.rand(4, 576, 1024))
|
||||||
|
|
||||||
|
m = MlpProjector(cfg)
|
||||||
|
out = m(inputs)
|
||||||
|
print(out.shape)
|
||||||
681
janus/models/siglip_vit.py
Normal file
681
janus/models/siglip_vit.py
Normal file
@ -0,0 +1,681 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
# https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py
|
||||||
|
import math
|
||||||
|
import warnings
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from functools import partial
|
||||||
|
from typing import (
|
||||||
|
Callable,
|
||||||
|
Dict,
|
||||||
|
Final,
|
||||||
|
List,
|
||||||
|
Literal,
|
||||||
|
Optional,
|
||||||
|
Sequence,
|
||||||
|
Set,
|
||||||
|
Tuple,
|
||||||
|
Type,
|
||||||
|
Union,
|
||||||
|
)
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from timm.layers import (
|
||||||
|
AttentionPoolLatent,
|
||||||
|
DropPath,
|
||||||
|
LayerType,
|
||||||
|
Mlp,
|
||||||
|
PatchDropout,
|
||||||
|
PatchEmbed,
|
||||||
|
resample_abs_pos_embed,
|
||||||
|
)
|
||||||
|
from timm.models._manipulate import checkpoint_seq, named_apply
|
||||||
|
|
||||||
|
|
||||||
|
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) # noqa: E741
|
||||||
|
u = norm_cdf((b - mean) / std)
|
||||||
|
|
||||||
|
# Uniformly fill tensor with values from [l, u], then translate to
|
||||||
|
# [2l-1, 2u-1].
|
||||||
|
tensor.uniform_(2 * l - 1, 2 * u - 1)
|
||||||
|
|
||||||
|
# Use inverse cdf transform for normal distribution to get truncated
|
||||||
|
# standard normal
|
||||||
|
tensor.erfinv_()
|
||||||
|
|
||||||
|
# Transform to proper mean, std
|
||||||
|
tensor.mul_(std * math.sqrt(2.0))
|
||||||
|
tensor.add_(mean)
|
||||||
|
|
||||||
|
# Clamp to ensure it's in the proper range
|
||||||
|
tensor.clamp_(min=a, max=b)
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
|
||||||
|
def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
|
||||||
|
# type: (torch.Tensor, float, float, float, float) -> torch.Tensor
|
||||||
|
r"""The original timm.models.layers.weight_init.trunc_normal_ can not handle bfloat16 yet, here we first
|
||||||
|
convert the tensor to float32, apply the trunc_normal_() in float32, and then convert it back to its orignal dtype.
|
||||||
|
Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn
|
||||||
|
from the normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
|
||||||
|
with values outside :math:`[a, b]` redrawn until they are within
|
||||||
|
the bounds. The method used for generating the random values works
|
||||||
|
best when :math:`a \leq \text{mean} \leq b`.
|
||||||
|
Args:
|
||||||
|
tensor: an n-dimensional `torch.Tensor`
|
||||||
|
mean: the mean of the normal distribution
|
||||||
|
std: the standard deviation of the normal distribution
|
||||||
|
a: the minimum cutoff value
|
||||||
|
b: the maximum cutoff value
|
||||||
|
Examples:
|
||||||
|
>>> w = torch.empty(3, 5)
|
||||||
|
>>> nn.init.trunc_normal_(w)
|
||||||
|
"""
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
dtype = tensor.dtype
|
||||||
|
tensor_fp32 = tensor.float()
|
||||||
|
tensor_fp32 = _no_grad_trunc_normal_(tensor_fp32, mean, std, a, b)
|
||||||
|
tensor_dtype = tensor_fp32.to(dtype=dtype)
|
||||||
|
tensor.copy_(tensor_dtype)
|
||||||
|
|
||||||
|
|
||||||
|
def init_weights(self):
|
||||||
|
if self.pos_embed is not None:
|
||||||
|
trunc_normal_(self.pos_embed, std=self.pos_embed.shape[1] ** -0.5)
|
||||||
|
trunc_normal_(self.latent, std=self.latent_dim**-0.5)
|
||||||
|
|
||||||
|
|
||||||
|
def init_weights_vit_timm(module: nn.Module, name: str = "") -> None:
|
||||||
|
"""ViT weight initialization, original timm impl (for reproducibility)"""
|
||||||
|
if isinstance(module, nn.Linear):
|
||||||
|
trunc_normal_(module.weight, std=0.02)
|
||||||
|
if module.bias is not None:
|
||||||
|
nn.init.zeros_(module.bias)
|
||||||
|
elif hasattr(module, "init_weights"):
|
||||||
|
module.init_weights()
|
||||||
|
|
||||||
|
|
||||||
|
class Attention(nn.Module):
|
||||||
|
fused_attn: Final[bool]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
num_heads: int = 8,
|
||||||
|
qkv_bias: bool = False,
|
||||||
|
qk_norm: bool = False,
|
||||||
|
attn_drop: float = 0.0,
|
||||||
|
proj_drop: float = 0.0,
|
||||||
|
norm_layer: nn.Module = nn.LayerNorm,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
assert dim % num_heads == 0, "dim should be divisible by num_heads"
|
||||||
|
self.num_heads = num_heads
|
||||||
|
self.head_dim = dim // num_heads
|
||||||
|
self.scale = self.head_dim**-0.5
|
||||||
|
# self.fused_attn = use_fused_attn()
|
||||||
|
self.fused_attn = True
|
||||||
|
|
||||||
|
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||||
|
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
||||||
|
self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
||||||
|
self.attn_drop = nn.Dropout(attn_drop)
|
||||||
|
self.proj = nn.Linear(dim, dim)
|
||||||
|
self.proj_drop = nn.Dropout(proj_drop) if proj_drop > 0.0 else nn.Identity()
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
B, N, C = x.shape
|
||||||
|
qkv = (
|
||||||
|
self.qkv(x)
|
||||||
|
.reshape(B, N, 3, self.num_heads, self.head_dim)
|
||||||
|
.permute(2, 0, 3, 1, 4)
|
||||||
|
)
|
||||||
|
q, k, v = qkv.unbind(0)
|
||||||
|
q, k = self.q_norm(q), self.k_norm(k)
|
||||||
|
|
||||||
|
if self.fused_attn:
|
||||||
|
x = F.scaled_dot_product_attention(
|
||||||
|
q,
|
||||||
|
k,
|
||||||
|
v,
|
||||||
|
dropout_p=self.attn_drop.p if self.training else 0.0,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
q = q * self.scale
|
||||||
|
attn = q @ k.transpose(-2, -1)
|
||||||
|
attn = attn.softmax(dim=-1)
|
||||||
|
attn = self.attn_drop(attn)
|
||||||
|
x = attn @ v
|
||||||
|
|
||||||
|
x = x.transpose(1, 2).reshape(B, N, C)
|
||||||
|
x = self.proj(x)
|
||||||
|
x = self.proj_drop(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class LayerScale(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
init_values: float = 1e-5,
|
||||||
|
inplace: bool = False,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.inplace = inplace
|
||||||
|
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
||||||
|
|
||||||
|
|
||||||
|
class Block(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim: int,
|
||||||
|
num_heads: int,
|
||||||
|
mlp_ratio: float = 4.0,
|
||||||
|
qkv_bias: bool = False,
|
||||||
|
qk_norm: bool = False,
|
||||||
|
proj_drop: float = 0.0,
|
||||||
|
attn_drop: float = 0.0,
|
||||||
|
init_values: Optional[float] = None,
|
||||||
|
drop_path: float = 0.0,
|
||||||
|
act_layer: nn.Module = nn.GELU,
|
||||||
|
norm_layer: nn.Module = nn.LayerNorm,
|
||||||
|
mlp_layer: nn.Module = Mlp,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.norm1 = norm_layer(dim)
|
||||||
|
self.attn = Attention(
|
||||||
|
dim,
|
||||||
|
num_heads=num_heads,
|
||||||
|
qkv_bias=qkv_bias,
|
||||||
|
qk_norm=qk_norm,
|
||||||
|
attn_drop=attn_drop,
|
||||||
|
proj_drop=proj_drop,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
)
|
||||||
|
self.ls1 = (
|
||||||
|
LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||||
|
)
|
||||||
|
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||||
|
|
||||||
|
self.norm2 = norm_layer(dim)
|
||||||
|
self.mlp = mlp_layer(
|
||||||
|
in_features=dim,
|
||||||
|
hidden_features=int(dim * mlp_ratio),
|
||||||
|
act_layer=act_layer,
|
||||||
|
drop=proj_drop,
|
||||||
|
)
|
||||||
|
self.ls2 = (
|
||||||
|
LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||||
|
)
|
||||||
|
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
|
||||||
|
x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class VisionTransformer(nn.Module):
|
||||||
|
"""Vision Transformer
|
||||||
|
|
||||||
|
A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
|
||||||
|
- https://arxiv.org/abs/2010.11929
|
||||||
|
"""
|
||||||
|
|
||||||
|
dynamic_img_size: Final[bool]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
img_size: Union[int, Tuple[int, int]] = 224,
|
||||||
|
patch_size: Union[int, Tuple[int, int]] = 16,
|
||||||
|
in_chans: int = 3,
|
||||||
|
num_classes: int = 1000,
|
||||||
|
global_pool: Literal["", "avg", "token", "map"] = "token",
|
||||||
|
embed_dim: int = 768,
|
||||||
|
depth: int = 12,
|
||||||
|
num_heads: int = 12,
|
||||||
|
mlp_ratio: float = 4.0,
|
||||||
|
qkv_bias: bool = True,
|
||||||
|
qk_norm: bool = False,
|
||||||
|
init_values: Optional[float] = None,
|
||||||
|
class_token: bool = True,
|
||||||
|
no_embed_class: bool = False,
|
||||||
|
reg_tokens: int = 0,
|
||||||
|
pre_norm: bool = False,
|
||||||
|
fc_norm: Optional[bool] = None,
|
||||||
|
dynamic_img_size: bool = False,
|
||||||
|
dynamic_img_pad: bool = False,
|
||||||
|
drop_rate: float = 0.0,
|
||||||
|
pos_drop_rate: float = 0.0,
|
||||||
|
patch_drop_rate: float = 0.0,
|
||||||
|
proj_drop_rate: float = 0.0,
|
||||||
|
attn_drop_rate: float = 0.0,
|
||||||
|
drop_path_rate: float = 0.0,
|
||||||
|
weight_init: Literal["skip", "jax", "jax_nlhb", "moco", ""] = "",
|
||||||
|
embed_layer: Callable = PatchEmbed,
|
||||||
|
norm_layer: Optional[LayerType] = None,
|
||||||
|
act_layer: Optional[LayerType] = None,
|
||||||
|
block_fn: Type[nn.Module] = Block,
|
||||||
|
mlp_layer: Type[nn.Module] = Mlp,
|
||||||
|
ignore_head: bool = False,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
img_size: Input image size.
|
||||||
|
patch_size: Patch size.
|
||||||
|
in_chans: Number of image input channels.
|
||||||
|
num_classes: Mumber of classes for classification head.
|
||||||
|
global_pool: Type of global pooling for final sequence (default: 'token').
|
||||||
|
embed_dim: Transformer embedding dimension.
|
||||||
|
depth: Depth of transformer.
|
||||||
|
num_heads: Number of attention heads.
|
||||||
|
mlp_ratio: Ratio of mlp hidden dim to embedding dim.
|
||||||
|
qkv_bias: Enable bias for qkv projections if True.
|
||||||
|
init_values: Layer-scale init values (layer-scale enabled if not None).
|
||||||
|
class_token: Use class token.
|
||||||
|
no_embed_class: Don't include position embeddings for class (or reg) tokens.
|
||||||
|
reg_tokens: Number of register tokens.
|
||||||
|
fc_norm: Pre head norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
|
||||||
|
drop_rate: Head dropout rate.
|
||||||
|
pos_drop_rate: Position embedding dropout rate.
|
||||||
|
attn_drop_rate: Attention dropout rate.
|
||||||
|
drop_path_rate: Stochastic depth rate.
|
||||||
|
weight_init: Weight initialization scheme.
|
||||||
|
embed_layer: Patch embedding layer.
|
||||||
|
norm_layer: Normalization layer.
|
||||||
|
act_layer: MLP activation layer.
|
||||||
|
block_fn: Transformer block layer.
|
||||||
|
"""
|
||||||
|
super().__init__()
|
||||||
|
assert global_pool in ("", "avg", "token", "map")
|
||||||
|
assert class_token or global_pool != "token"
|
||||||
|
use_fc_norm = global_pool == "avg" if fc_norm is None else fc_norm
|
||||||
|
# norm_layer = get_norm_layer(norm_layer) or partial(nn.LayerNorm, eps=1e-6)
|
||||||
|
# act_layer = get_act_layer(act_layer) or nn.GELU
|
||||||
|
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||||
|
act_layer = nn.GELU
|
||||||
|
|
||||||
|
self.num_classes = num_classes
|
||||||
|
self.global_pool = global_pool
|
||||||
|
self.num_features = self.embed_dim = (
|
||||||
|
embed_dim # num_features for consistency with other models
|
||||||
|
)
|
||||||
|
self.num_prefix_tokens = 1 if class_token else 0
|
||||||
|
self.num_prefix_tokens += reg_tokens
|
||||||
|
self.num_reg_tokens = reg_tokens
|
||||||
|
self.has_class_token = class_token
|
||||||
|
self.no_embed_class = (
|
||||||
|
no_embed_class # don't embed prefix positions (includes reg)
|
||||||
|
)
|
||||||
|
self.dynamic_img_size = dynamic_img_size
|
||||||
|
self.grad_checkpointing = False
|
||||||
|
self.ignore_head = ignore_head
|
||||||
|
|
||||||
|
embed_args = {}
|
||||||
|
if dynamic_img_size:
|
||||||
|
# flatten deferred until after pos embed
|
||||||
|
embed_args.update(dict(strict_img_size=False, output_fmt="NHWC"))
|
||||||
|
self.patch_embed = embed_layer(
|
||||||
|
img_size=img_size,
|
||||||
|
patch_size=patch_size,
|
||||||
|
in_chans=in_chans,
|
||||||
|
embed_dim=embed_dim,
|
||||||
|
bias=not pre_norm, # disable bias if pre-norm is used (e.g. CLIP)
|
||||||
|
dynamic_img_pad=dynamic_img_pad,
|
||||||
|
**embed_args,
|
||||||
|
)
|
||||||
|
num_patches = self.patch_embed.num_patches
|
||||||
|
|
||||||
|
self.cls_token = (
|
||||||
|
nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
|
||||||
|
)
|
||||||
|
self.reg_token = (
|
||||||
|
nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None
|
||||||
|
)
|
||||||
|
embed_len = (
|
||||||
|
num_patches if no_embed_class else num_patches + self.num_prefix_tokens
|
||||||
|
)
|
||||||
|
self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02)
|
||||||
|
self.pos_drop = nn.Dropout(p=pos_drop_rate)
|
||||||
|
if patch_drop_rate > 0:
|
||||||
|
self.patch_drop = PatchDropout(
|
||||||
|
patch_drop_rate,
|
||||||
|
num_prefix_tokens=self.num_prefix_tokens,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.patch_drop = nn.Identity()
|
||||||
|
self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
|
||||||
|
|
||||||
|
dpr = [
|
||||||
|
x.item() for x in torch.linspace(0, drop_path_rate, depth)
|
||||||
|
] # stochastic depth decay rule
|
||||||
|
self.blocks = nn.Sequential(
|
||||||
|
*[
|
||||||
|
block_fn(
|
||||||
|
dim=embed_dim,
|
||||||
|
num_heads=num_heads,
|
||||||
|
mlp_ratio=mlp_ratio,
|
||||||
|
qkv_bias=qkv_bias,
|
||||||
|
qk_norm=qk_norm,
|
||||||
|
init_values=init_values,
|
||||||
|
proj_drop=proj_drop_rate,
|
||||||
|
attn_drop=attn_drop_rate,
|
||||||
|
drop_path=dpr[i],
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
act_layer=act_layer,
|
||||||
|
mlp_layer=mlp_layer,
|
||||||
|
)
|
||||||
|
for i in range(depth)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
self.norm = norm_layer(embed_dim) if not use_fc_norm else nn.Identity()
|
||||||
|
|
||||||
|
# Classifier Head
|
||||||
|
if global_pool == "map":
|
||||||
|
AttentionPoolLatent.init_weights = init_weights
|
||||||
|
self.attn_pool = AttentionPoolLatent(
|
||||||
|
self.embed_dim,
|
||||||
|
num_heads=num_heads,
|
||||||
|
mlp_ratio=mlp_ratio,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.attn_pool = None
|
||||||
|
self.fc_norm = norm_layer(embed_dim) if use_fc_norm else nn.Identity()
|
||||||
|
self.head_drop = nn.Dropout(drop_rate)
|
||||||
|
self.head = (
|
||||||
|
nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
|
||||||
|
)
|
||||||
|
|
||||||
|
if weight_init != "skip":
|
||||||
|
self.init_weights(weight_init)
|
||||||
|
|
||||||
|
def init_weights(self, mode: Literal["jax", "jax_nlhb", "moco", ""] = "") -> None:
|
||||||
|
assert mode in ("jax", "jax_nlhb", "moco", "")
|
||||||
|
# head_bias = -math.log(self.num_classes) if "nlhb" in mode else 0.0
|
||||||
|
trunc_normal_(self.pos_embed, std=0.02)
|
||||||
|
if self.cls_token is not None:
|
||||||
|
nn.init.normal_(self.cls_token, std=1e-6)
|
||||||
|
named_apply(init_weights_vit_timm, self)
|
||||||
|
|
||||||
|
@torch.jit.ignore
|
||||||
|
def no_weight_decay(self) -> Set:
|
||||||
|
return {"pos_embed", "cls_token", "dist_token"}
|
||||||
|
|
||||||
|
@torch.jit.ignore
|
||||||
|
def group_matcher(self, coarse: bool = False) -> Dict:
|
||||||
|
return dict(
|
||||||
|
stem=r"^cls_token|pos_embed|patch_embed", # stem and embed
|
||||||
|
blocks=[(r"^blocks\.(\d+)", None), (r"^norm", (99999,))],
|
||||||
|
)
|
||||||
|
|
||||||
|
@torch.jit.ignore
|
||||||
|
def set_grad_checkpointing(self, enable: bool = True) -> None:
|
||||||
|
self.grad_checkpointing = enable
|
||||||
|
|
||||||
|
@torch.jit.ignore
|
||||||
|
def get_classifier(self) -> nn.Module:
|
||||||
|
return self.head
|
||||||
|
|
||||||
|
def reset_classifier(self, num_classes: int, global_pool=None) -> None:
|
||||||
|
self.num_classes = num_classes
|
||||||
|
if global_pool is not None:
|
||||||
|
assert global_pool in ("", "avg", "token", "map")
|
||||||
|
if global_pool == "map" and self.attn_pool is None:
|
||||||
|
assert (
|
||||||
|
False
|
||||||
|
), "Cannot currently add attention pooling in reset_classifier()."
|
||||||
|
elif global_pool != "map " and self.attn_pool is not None:
|
||||||
|
self.attn_pool = None # remove attention pooling
|
||||||
|
self.global_pool = global_pool
|
||||||
|
self.head = (
|
||||||
|
nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
|
||||||
|
)
|
||||||
|
|
||||||
|
def _pos_embed(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
if self.dynamic_img_size:
|
||||||
|
B, H, W, C = x.shape
|
||||||
|
pos_embed = resample_abs_pos_embed(
|
||||||
|
self.pos_embed,
|
||||||
|
(H, W),
|
||||||
|
num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens,
|
||||||
|
)
|
||||||
|
x = x.view(B, -1, C)
|
||||||
|
else:
|
||||||
|
pos_embed = self.pos_embed
|
||||||
|
|
||||||
|
to_cat = []
|
||||||
|
if self.cls_token is not None:
|
||||||
|
to_cat.append(self.cls_token.expand(x.shape[0], -1, -1))
|
||||||
|
if self.reg_token is not None:
|
||||||
|
to_cat.append(self.reg_token.expand(x.shape[0], -1, -1))
|
||||||
|
|
||||||
|
if self.no_embed_class:
|
||||||
|
# deit-3, updated JAX (big vision)
|
||||||
|
# position embedding does not overlap with class token, add then concat
|
||||||
|
x = x + pos_embed
|
||||||
|
if to_cat:
|
||||||
|
x = torch.cat(to_cat + [x], dim=1)
|
||||||
|
else:
|
||||||
|
# original timm, JAX, and deit vit impl
|
||||||
|
# pos_embed has entry for class token, concat then add
|
||||||
|
if to_cat:
|
||||||
|
x = torch.cat(to_cat + [x], dim=1)
|
||||||
|
x = x + pos_embed
|
||||||
|
|
||||||
|
return self.pos_drop(x)
|
||||||
|
|
||||||
|
def _intermediate_layers(
|
||||||
|
self,
|
||||||
|
x: torch.Tensor,
|
||||||
|
n: Union[int, Sequence] = 1,
|
||||||
|
) -> List[torch.Tensor]:
|
||||||
|
outputs, num_blocks = [], len(self.blocks)
|
||||||
|
take_indices = set(
|
||||||
|
range(num_blocks - n, num_blocks) if isinstance(n, int) else n
|
||||||
|
)
|
||||||
|
|
||||||
|
# forward pass
|
||||||
|
x = self.patch_embed(x)
|
||||||
|
x = self._pos_embed(x)
|
||||||
|
x = self.patch_drop(x)
|
||||||
|
x = self.norm_pre(x)
|
||||||
|
for i, blk in enumerate(self.blocks):
|
||||||
|
x = blk(x)
|
||||||
|
if i in take_indices:
|
||||||
|
outputs.append(x)
|
||||||
|
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
def get_intermediate_layers(
|
||||||
|
self,
|
||||||
|
x: torch.Tensor,
|
||||||
|
n: Union[int, Sequence] = 1,
|
||||||
|
reshape: bool = False,
|
||||||
|
return_prefix_tokens: bool = False,
|
||||||
|
norm: bool = False,
|
||||||
|
) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
|
||||||
|
"""Intermediate layer accessor (NOTE: This is a WIP experiment).
|
||||||
|
Inspired by DINO / DINOv2 interface
|
||||||
|
"""
|
||||||
|
# take last n blocks if n is an int, if in is a sequence, select by matching indices
|
||||||
|
outputs = self._intermediate_layers(x, n)
|
||||||
|
if norm:
|
||||||
|
outputs = [self.norm(out) for out in outputs]
|
||||||
|
prefix_tokens = [out[:, 0 : self.num_prefix_tokens] for out in outputs]
|
||||||
|
outputs = [out[:, self.num_prefix_tokens :] for out in outputs]
|
||||||
|
|
||||||
|
if reshape:
|
||||||
|
grid_size = self.patch_embed.grid_size
|
||||||
|
outputs = [
|
||||||
|
out.reshape(x.shape[0], grid_size[0], grid_size[1], -1)
|
||||||
|
.permute(0, 3, 1, 2)
|
||||||
|
.contiguous()
|
||||||
|
for out in outputs
|
||||||
|
]
|
||||||
|
|
||||||
|
if return_prefix_tokens:
|
||||||
|
return tuple(zip(outputs, prefix_tokens))
|
||||||
|
return tuple(outputs)
|
||||||
|
|
||||||
|
def forward_features(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
x = self.patch_embed(x)
|
||||||
|
x = self._pos_embed(x)
|
||||||
|
x = self.patch_drop(x)
|
||||||
|
x = self.norm_pre(x)
|
||||||
|
if self.grad_checkpointing and not torch.jit.is_scripting():
|
||||||
|
x = checkpoint_seq(self.blocks, x)
|
||||||
|
else:
|
||||||
|
x = self.blocks(x)
|
||||||
|
x = self.norm(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
|
||||||
|
if self.attn_pool is not None:
|
||||||
|
x = self.attn_pool(x)
|
||||||
|
elif self.global_pool == "avg":
|
||||||
|
x = x[:, self.num_prefix_tokens :].mean(dim=1)
|
||||||
|
elif self.global_pool:
|
||||||
|
x = x[:, 0] # class token
|
||||||
|
x = self.fc_norm(x)
|
||||||
|
x = self.head_drop(x)
|
||||||
|
return x if pre_logits else self.head(x)
|
||||||
|
|
||||||
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
x = self.forward_features(x)
|
||||||
|
if not self.ignore_head:
|
||||||
|
x = self.forward_head(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SigLIPVisionCfg:
|
||||||
|
width: int = 1152
|
||||||
|
layers: Union[Tuple[int, int, int, int], int] = 27
|
||||||
|
heads: int = 16
|
||||||
|
patch_size: int = 14
|
||||||
|
image_size: Union[Tuple[int, int], int] = 336
|
||||||
|
global_pool: str = "map"
|
||||||
|
mlp_ratio: float = 3.7362
|
||||||
|
class_token: bool = False
|
||||||
|
num_classes: int = 0
|
||||||
|
use_checkpoint: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
SigLIP_MODEL_CONFIG = {
|
||||||
|
"siglip_so400m_patch14_384": {
|
||||||
|
"image_size": 336,
|
||||||
|
"patch_size": 14,
|
||||||
|
"width": 1152,
|
||||||
|
"layers": 27,
|
||||||
|
"heads": 16,
|
||||||
|
"mlp_ratio": 3.7362,
|
||||||
|
"global_pool": "map",
|
||||||
|
"use_checkpoint": False,
|
||||||
|
},
|
||||||
|
"siglip_so400m_patch14_224": {
|
||||||
|
"image_size": 224,
|
||||||
|
"patch_size": 14,
|
||||||
|
"width": 1152,
|
||||||
|
"layers": 27,
|
||||||
|
"heads": 16,
|
||||||
|
"mlp_ratio": 3.7362,
|
||||||
|
"global_pool": "map",
|
||||||
|
"use_checkpoint": False,
|
||||||
|
},
|
||||||
|
"siglip_large_patch16_384": {
|
||||||
|
"image_size": 384,
|
||||||
|
"patch_size": 16,
|
||||||
|
"width": 1024,
|
||||||
|
"layers": 24,
|
||||||
|
"heads": 16,
|
||||||
|
"mlp_ratio": 4,
|
||||||
|
"global_pool": "map",
|
||||||
|
"use_checkpoint": False,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def create_siglip_vit(
|
||||||
|
model_name: str = "siglip_so400m_patch14_384",
|
||||||
|
image_size: int = 384,
|
||||||
|
select_layer: int = -1,
|
||||||
|
ckpt_path: str = "",
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
assert (
|
||||||
|
model_name in SigLIP_MODEL_CONFIG.keys()
|
||||||
|
), f"model name should be in {SigLIP_MODEL_CONFIG.keys()}"
|
||||||
|
|
||||||
|
vision_cfg = SigLIPVisionCfg(**SigLIP_MODEL_CONFIG[model_name])
|
||||||
|
|
||||||
|
if select_layer <= 0:
|
||||||
|
layers = min(vision_cfg.layers, vision_cfg.layers + select_layer + 1)
|
||||||
|
else:
|
||||||
|
layers = min(vision_cfg.layers, select_layer)
|
||||||
|
|
||||||
|
model = VisionTransformer(
|
||||||
|
img_size=image_size,
|
||||||
|
patch_size=vision_cfg.patch_size,
|
||||||
|
embed_dim=vision_cfg.width,
|
||||||
|
depth=layers,
|
||||||
|
num_heads=vision_cfg.heads,
|
||||||
|
mlp_ratio=vision_cfg.mlp_ratio,
|
||||||
|
class_token=vision_cfg.class_token,
|
||||||
|
global_pool=vision_cfg.global_pool,
|
||||||
|
ignore_head=kwargs.get("ignore_head", True),
|
||||||
|
weight_init=kwargs.get("weight_init", "skip"),
|
||||||
|
num_classes=0,
|
||||||
|
)
|
||||||
|
|
||||||
|
if ckpt_path:
|
||||||
|
state_dict = torch.load(ckpt_path, map_location="cpu")
|
||||||
|
|
||||||
|
incompatible_keys = model.load_state_dict(state_dict, strict=False)
|
||||||
|
print(
|
||||||
|
f"SigLIP-ViT restores from {ckpt_path},\n"
|
||||||
|
f"\tincompatible_keys:', {incompatible_keys}."
|
||||||
|
)
|
||||||
|
|
||||||
|
return model
|
||||||
527
janus/models/vq_model.py
Executable file
527
janus/models/vq_model.py
Executable file
@ -0,0 +1,527 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
from functools import partial
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ModelArgs:
|
||||||
|
codebook_size: int = 16384
|
||||||
|
codebook_embed_dim: int = 8
|
||||||
|
codebook_l2_norm: bool = True
|
||||||
|
codebook_show_usage: bool = True
|
||||||
|
commit_loss_beta: float = 0.25
|
||||||
|
entropy_loss_ratio: float = 0.0
|
||||||
|
|
||||||
|
encoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
|
||||||
|
decoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
|
||||||
|
z_channels: int = 256
|
||||||
|
dropout_p: float = 0.0
|
||||||
|
|
||||||
|
|
||||||
|
class Encoder(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels=3,
|
||||||
|
ch=128,
|
||||||
|
ch_mult=(1, 1, 2, 2, 4),
|
||||||
|
num_res_blocks=2,
|
||||||
|
norm_type="group",
|
||||||
|
dropout=0.0,
|
||||||
|
resamp_with_conv=True,
|
||||||
|
z_channels=256,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.num_resolutions = len(ch_mult)
|
||||||
|
self.num_res_blocks = num_res_blocks
|
||||||
|
self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
|
||||||
|
|
||||||
|
# downsampling
|
||||||
|
in_ch_mult = (1,) + tuple(ch_mult)
|
||||||
|
self.conv_blocks = nn.ModuleList()
|
||||||
|
for i_level in range(self.num_resolutions):
|
||||||
|
conv_block = nn.Module()
|
||||||
|
# res & attn
|
||||||
|
res_block = nn.ModuleList()
|
||||||
|
attn_block = nn.ModuleList()
|
||||||
|
block_in = ch * in_ch_mult[i_level]
|
||||||
|
block_out = ch * ch_mult[i_level]
|
||||||
|
for _ in range(self.num_res_blocks):
|
||||||
|
res_block.append(
|
||||||
|
ResnetBlock(
|
||||||
|
block_in, block_out, dropout=dropout, norm_type=norm_type
|
||||||
|
)
|
||||||
|
)
|
||||||
|
block_in = block_out
|
||||||
|
if i_level == self.num_resolutions - 1:
|
||||||
|
attn_block.append(AttnBlock(block_in, norm_type))
|
||||||
|
conv_block.res = res_block
|
||||||
|
conv_block.attn = attn_block
|
||||||
|
# downsample
|
||||||
|
if i_level != self.num_resolutions - 1:
|
||||||
|
conv_block.downsample = Downsample(block_in, resamp_with_conv)
|
||||||
|
self.conv_blocks.append(conv_block)
|
||||||
|
|
||||||
|
# middle
|
||||||
|
self.mid = nn.ModuleList()
|
||||||
|
self.mid.append(
|
||||||
|
ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)
|
||||||
|
)
|
||||||
|
self.mid.append(AttnBlock(block_in, norm_type=norm_type))
|
||||||
|
self.mid.append(
|
||||||
|
ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)
|
||||||
|
)
|
||||||
|
|
||||||
|
# end
|
||||||
|
self.norm_out = Normalize(block_in, norm_type)
|
||||||
|
self.conv_out = nn.Conv2d(
|
||||||
|
block_in, z_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
h = self.conv_in(x)
|
||||||
|
# downsampling
|
||||||
|
for i_level, block in enumerate(self.conv_blocks):
|
||||||
|
for i_block in range(self.num_res_blocks):
|
||||||
|
h = block.res[i_block](h)
|
||||||
|
if len(block.attn) > 0:
|
||||||
|
h = block.attn[i_block](h)
|
||||||
|
if i_level != self.num_resolutions - 1:
|
||||||
|
h = block.downsample(h)
|
||||||
|
|
||||||
|
# middle
|
||||||
|
for mid_block in self.mid:
|
||||||
|
h = mid_block(h)
|
||||||
|
|
||||||
|
# end
|
||||||
|
h = self.norm_out(h)
|
||||||
|
h = nonlinearity(h)
|
||||||
|
h = self.conv_out(h)
|
||||||
|
return h
|
||||||
|
|
||||||
|
|
||||||
|
class Decoder(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
z_channels=256,
|
||||||
|
ch=128,
|
||||||
|
ch_mult=(1, 1, 2, 2, 4),
|
||||||
|
num_res_blocks=2,
|
||||||
|
norm_type="group",
|
||||||
|
dropout=0.0,
|
||||||
|
resamp_with_conv=True,
|
||||||
|
out_channels=3,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.num_resolutions = len(ch_mult)
|
||||||
|
self.num_res_blocks = num_res_blocks
|
||||||
|
|
||||||
|
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||||
|
# z to block_in
|
||||||
|
self.conv_in = nn.Conv2d(
|
||||||
|
z_channels, block_in, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
|
||||||
|
# middle
|
||||||
|
self.mid = nn.ModuleList()
|
||||||
|
self.mid.append(
|
||||||
|
ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)
|
||||||
|
)
|
||||||
|
self.mid.append(AttnBlock(block_in, norm_type=norm_type))
|
||||||
|
self.mid.append(
|
||||||
|
ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)
|
||||||
|
)
|
||||||
|
|
||||||
|
# upsampling
|
||||||
|
self.conv_blocks = nn.ModuleList()
|
||||||
|
for i_level in reversed(range(self.num_resolutions)):
|
||||||
|
conv_block = nn.Module()
|
||||||
|
# res & attn
|
||||||
|
res_block = nn.ModuleList()
|
||||||
|
attn_block = nn.ModuleList()
|
||||||
|
block_out = ch * ch_mult[i_level]
|
||||||
|
for _ in range(self.num_res_blocks + 1):
|
||||||
|
res_block.append(
|
||||||
|
ResnetBlock(
|
||||||
|
block_in, block_out, dropout=dropout, norm_type=norm_type
|
||||||
|
)
|
||||||
|
)
|
||||||
|
block_in = block_out
|
||||||
|
if i_level == self.num_resolutions - 1:
|
||||||
|
attn_block.append(AttnBlock(block_in, norm_type))
|
||||||
|
conv_block.res = res_block
|
||||||
|
conv_block.attn = attn_block
|
||||||
|
# downsample
|
||||||
|
if i_level != 0:
|
||||||
|
conv_block.upsample = Upsample(block_in, resamp_with_conv)
|
||||||
|
self.conv_blocks.append(conv_block)
|
||||||
|
|
||||||
|
# end
|
||||||
|
self.norm_out = Normalize(block_in, norm_type)
|
||||||
|
self.conv_out = nn.Conv2d(
|
||||||
|
block_in, out_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def last_layer(self):
|
||||||
|
return self.conv_out.weight
|
||||||
|
|
||||||
|
def forward(self, z):
|
||||||
|
# z to block_in
|
||||||
|
h = self.conv_in(z)
|
||||||
|
|
||||||
|
# middle
|
||||||
|
for mid_block in self.mid:
|
||||||
|
h = mid_block(h)
|
||||||
|
|
||||||
|
# upsampling
|
||||||
|
for i_level, block in enumerate(self.conv_blocks):
|
||||||
|
for i_block in range(self.num_res_blocks + 1):
|
||||||
|
h = block.res[i_block](h)
|
||||||
|
if len(block.attn) > 0:
|
||||||
|
h = block.attn[i_block](h)
|
||||||
|
if i_level != self.num_resolutions - 1:
|
||||||
|
h = block.upsample(h)
|
||||||
|
|
||||||
|
# end
|
||||||
|
h = self.norm_out(h)
|
||||||
|
h = nonlinearity(h)
|
||||||
|
h = self.conv_out(h)
|
||||||
|
return h
|
||||||
|
|
||||||
|
|
||||||
|
class VectorQuantizer(nn.Module):
|
||||||
|
def __init__(self, n_e, e_dim, beta, entropy_loss_ratio, l2_norm, show_usage):
|
||||||
|
super().__init__()
|
||||||
|
self.n_e = n_e
|
||||||
|
self.e_dim = e_dim
|
||||||
|
self.beta = beta
|
||||||
|
self.entropy_loss_ratio = entropy_loss_ratio
|
||||||
|
self.l2_norm = l2_norm
|
||||||
|
self.show_usage = show_usage
|
||||||
|
|
||||||
|
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)
|
||||||
|
if self.l2_norm:
|
||||||
|
self.embedding.weight.data = F.normalize(
|
||||||
|
self.embedding.weight.data, p=2, dim=-1
|
||||||
|
)
|
||||||
|
if self.show_usage:
|
||||||
|
self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536)))
|
||||||
|
|
||||||
|
def forward(self, z):
|
||||||
|
# reshape z -> (batch, height, width, channel) and flatten
|
||||||
|
z = torch.einsum("b c h w -> b h w c", z).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
|
||||||
|
|
||||||
|
if self.l2_norm:
|
||||||
|
z = F.normalize(z, p=2, dim=-1)
|
||||||
|
z_flattened = F.normalize(z_flattened, p=2, dim=-1)
|
||||||
|
embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
|
||||||
|
else:
|
||||||
|
embedding = self.embedding.weight
|
||||||
|
|
||||||
|
d = (
|
||||||
|
torch.sum(z_flattened**2, dim=1, keepdim=True)
|
||||||
|
+ torch.sum(embedding**2, dim=1)
|
||||||
|
- 2
|
||||||
|
* torch.einsum(
|
||||||
|
"bd,dn->bn", z_flattened, torch.einsum("n d -> d n", embedding)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
min_encoding_indices = torch.argmin(d, dim=1)
|
||||||
|
z_q = embedding[min_encoding_indices].view(z.shape)
|
||||||
|
perplexity = None
|
||||||
|
min_encodings = None
|
||||||
|
vq_loss = None
|
||||||
|
commit_loss = None
|
||||||
|
entropy_loss = None
|
||||||
|
|
||||||
|
# compute loss for embedding
|
||||||
|
if self.training:
|
||||||
|
vq_loss = torch.mean((z_q - z.detach()) ** 2)
|
||||||
|
commit_loss = self.beta * torch.mean((z_q.detach() - z) ** 2)
|
||||||
|
entropy_loss = self.entropy_loss_ratio * compute_entropy_loss(-d)
|
||||||
|
|
||||||
|
# preserve gradients
|
||||||
|
z_q = z + (z_q - z).detach()
|
||||||
|
|
||||||
|
# reshape back to match original input shape
|
||||||
|
z_q = torch.einsum("b h w c -> b c h w", z_q)
|
||||||
|
|
||||||
|
return (
|
||||||
|
z_q,
|
||||||
|
(vq_loss, commit_loss, entropy_loss),
|
||||||
|
(perplexity, min_encodings, min_encoding_indices),
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_codebook_entry(self, indices, shape=None, channel_first=True):
|
||||||
|
# shape = (batch, channel, height, width) if channel_first else (batch, height, width, channel)
|
||||||
|
if self.l2_norm:
|
||||||
|
embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
|
||||||
|
else:
|
||||||
|
embedding = self.embedding.weight
|
||||||
|
z_q = embedding[indices] # (b*h*w, c)
|
||||||
|
|
||||||
|
if shape is not None:
|
||||||
|
if channel_first:
|
||||||
|
z_q = z_q.reshape(shape[0], shape[2], shape[3], shape[1])
|
||||||
|
# reshape back to match original input shape
|
||||||
|
z_q = z_q.permute(0, 3, 1, 2).contiguous()
|
||||||
|
else:
|
||||||
|
z_q = z_q.view(shape)
|
||||||
|
return z_q
|
||||||
|
|
||||||
|
|
||||||
|
class ResnetBlock(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels=None,
|
||||||
|
conv_shortcut=False,
|
||||||
|
dropout=0.0,
|
||||||
|
norm_type="group",
|
||||||
|
):
|
||||||
|
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, norm_type)
|
||||||
|
self.conv1 = nn.Conv2d(
|
||||||
|
in_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
self.norm2 = Normalize(out_channels, norm_type)
|
||||||
|
self.dropout = nn.Dropout(dropout)
|
||||||
|
self.conv2 = nn.Conv2d(
|
||||||
|
out_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.in_channels != self.out_channels:
|
||||||
|
if self.use_conv_shortcut:
|
||||||
|
self.conv_shortcut = nn.Conv2d(
|
||||||
|
in_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.nin_shortcut = nn.Conv2d(
|
||||||
|
in_channels, out_channels, kernel_size=1, stride=1, padding=0
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
h = x
|
||||||
|
h = self.norm1(h)
|
||||||
|
h = nonlinearity(h)
|
||||||
|
h = self.conv1(h)
|
||||||
|
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 AttnBlock(nn.Module):
|
||||||
|
def __init__(self, in_channels, norm_type="group"):
|
||||||
|
super().__init__()
|
||||||
|
self.norm = Normalize(in_channels, norm_type)
|
||||||
|
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||||
|
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||||
|
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
|
||||||
|
self.proj_out = nn.Conv2d(
|
||||||
|
in_channels, in_channels, kernel_size=1, stride=1, padding=0
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
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) # b,hw,c
|
||||||
|
k = k.reshape(b, c, h * w) # b,c,hw
|
||||||
|
w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
||||||
|
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) # b,hw,hw (first hw of k, second of q)
|
||||||
|
h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
||||||
|
h_ = h_.reshape(b, c, h, w)
|
||||||
|
|
||||||
|
h_ = self.proj_out(h_)
|
||||||
|
|
||||||
|
return x + h_
|
||||||
|
|
||||||
|
|
||||||
|
def nonlinearity(x):
|
||||||
|
# swish
|
||||||
|
return x * torch.sigmoid(x)
|
||||||
|
|
||||||
|
|
||||||
|
def Normalize(in_channels, norm_type="group"):
|
||||||
|
assert norm_type in ["group", "batch"]
|
||||||
|
if norm_type == "group":
|
||||||
|
return nn.GroupNorm(
|
||||||
|
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
|
||||||
|
)
|
||||||
|
elif norm_type == "batch":
|
||||||
|
return nn.SyncBatchNorm(in_channels)
|
||||||
|
|
||||||
|
|
||||||
|
class Upsample(nn.Module):
|
||||||
|
def __init__(self, in_channels, with_conv):
|
||||||
|
super().__init__()
|
||||||
|
self.with_conv = with_conv
|
||||||
|
if self.with_conv:
|
||||||
|
self.conv = nn.Conv2d(
|
||||||
|
in_channels, in_channels, kernel_size=3, stride=1, padding=1
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if x.dtype != torch.float32:
|
||||||
|
x = F.interpolate(x.to(torch.float), scale_factor=2.0, mode="nearest").to(
|
||||||
|
torch.bfloat16
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
x = F.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 = 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 = F.pad(x, pad, mode="constant", value=0)
|
||||||
|
x = self.conv(x)
|
||||||
|
else:
|
||||||
|
x = F.avg_pool2d(x, kernel_size=2, stride=2)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
def compute_entropy_loss(affinity, loss_type="softmax", temperature=0.01):
|
||||||
|
flat_affinity = affinity.reshape(-1, affinity.shape[-1])
|
||||||
|
flat_affinity /= temperature
|
||||||
|
probs = F.softmax(flat_affinity, dim=-1)
|
||||||
|
log_probs = F.log_softmax(flat_affinity + 1e-5, dim=-1)
|
||||||
|
if loss_type == "softmax":
|
||||||
|
target_probs = probs
|
||||||
|
else:
|
||||||
|
raise ValueError("Entropy loss {} not supported".format(loss_type))
|
||||||
|
avg_probs = torch.mean(target_probs, dim=0)
|
||||||
|
avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + 1e-5))
|
||||||
|
sample_entropy = -torch.mean(torch.sum(target_probs * log_probs, dim=-1))
|
||||||
|
loss = sample_entropy - avg_entropy
|
||||||
|
return loss
|
||||||
|
|
||||||
|
|
||||||
|
class VQModel(nn.Module):
|
||||||
|
def __init__(self, config: ModelArgs):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.encoder = Encoder(
|
||||||
|
ch_mult=config.encoder_ch_mult,
|
||||||
|
z_channels=config.z_channels,
|
||||||
|
dropout=config.dropout_p,
|
||||||
|
)
|
||||||
|
self.decoder = Decoder(
|
||||||
|
ch_mult=config.decoder_ch_mult,
|
||||||
|
z_channels=config.z_channels,
|
||||||
|
dropout=config.dropout_p,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.quantize = VectorQuantizer(
|
||||||
|
config.codebook_size,
|
||||||
|
config.codebook_embed_dim,
|
||||||
|
config.commit_loss_beta,
|
||||||
|
config.entropy_loss_ratio,
|
||||||
|
config.codebook_l2_norm,
|
||||||
|
config.codebook_show_usage,
|
||||||
|
)
|
||||||
|
self.quant_conv = nn.Conv2d(config.z_channels, config.codebook_embed_dim, 1)
|
||||||
|
self.post_quant_conv = nn.Conv2d(
|
||||||
|
config.codebook_embed_dim, config.z_channels, 1
|
||||||
|
)
|
||||||
|
|
||||||
|
def encode(self, x):
|
||||||
|
h = self.encoder(x)
|
||||||
|
h = self.quant_conv(h)
|
||||||
|
quant, emb_loss, info = self.quantize(h)
|
||||||
|
return quant, emb_loss, info
|
||||||
|
|
||||||
|
def decode(self, quant):
|
||||||
|
quant = self.post_quant_conv(quant)
|
||||||
|
dec = self.decoder(quant)
|
||||||
|
return dec
|
||||||
|
|
||||||
|
def decode_code(self, code_b, shape=None, channel_first=True):
|
||||||
|
quant_b = self.quantize.get_codebook_entry(code_b, shape, channel_first)
|
||||||
|
dec = self.decode(quant_b)
|
||||||
|
return dec
|
||||||
|
|
||||||
|
def forward(self, input):
|
||||||
|
quant, diff, _ = self.encode(input)
|
||||||
|
dec = self.decode(quant)
|
||||||
|
return dec, diff
|
||||||
|
|
||||||
|
|
||||||
|
#################################################################################
|
||||||
|
# VQ Model Configs #
|
||||||
|
#################################################################################
|
||||||
|
def VQ_16(**kwargs):
|
||||||
|
return VQModel(
|
||||||
|
ModelArgs(
|
||||||
|
encoder_ch_mult=[1, 1, 2, 2, 4], decoder_ch_mult=[1, 1, 2, 2, 4], **kwargs
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
VQ_models = {"VQ-16": VQ_16}
|
||||||
18
janus/utils/__init__.py
Normal file
18
janus/utils/__init__.py
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
348
janus/utils/conversation.py
Normal file
348
janus/utils/conversation.py
Normal file
@ -0,0 +1,348 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
"""
|
||||||
|
From https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
|
||||||
|
"""
|
||||||
|
|
||||||
|
import dataclasses
|
||||||
|
from enum import IntEnum, auto
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
|
||||||
|
class SeparatorStyle(IntEnum):
|
||||||
|
"""Separator styles."""
|
||||||
|
|
||||||
|
ADD_COLON_SINGLE = auto()
|
||||||
|
ADD_COLON_TWO = auto()
|
||||||
|
ADD_COLON_SPACE_SINGLE = auto()
|
||||||
|
NO_COLON_SINGLE = auto()
|
||||||
|
NO_COLON_TWO = auto()
|
||||||
|
ADD_NEW_LINE_SINGLE = auto()
|
||||||
|
LLAMA2 = auto()
|
||||||
|
CHATGLM = auto()
|
||||||
|
CHATML = auto()
|
||||||
|
CHATINTERN = auto()
|
||||||
|
DOLLY = auto()
|
||||||
|
RWKV = auto()
|
||||||
|
PHOENIX = auto()
|
||||||
|
ROBIN = auto()
|
||||||
|
DeepSeek = auto()
|
||||||
|
PLAIN = auto()
|
||||||
|
ALIGNMENT = auto()
|
||||||
|
|
||||||
|
|
||||||
|
@dataclasses.dataclass
|
||||||
|
class Conversation:
|
||||||
|
"""A class that manages prompt templates and keeps all conversation history."""
|
||||||
|
|
||||||
|
# The name of this template
|
||||||
|
name: str
|
||||||
|
# The template of the system prompt
|
||||||
|
system_template: str = "{system_message}"
|
||||||
|
# The system message
|
||||||
|
system_message: str = ""
|
||||||
|
# The names of two roles
|
||||||
|
roles: List[str] = (("USER", "ASSISTANT"),)
|
||||||
|
# All messages. Each item is (role, message).
|
||||||
|
messages: List[List[str]] = ()
|
||||||
|
# The number of few shot examples
|
||||||
|
offset: int = 0
|
||||||
|
# The separator style and configurations
|
||||||
|
sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
|
||||||
|
sep: str = "\n"
|
||||||
|
sep2: str = None
|
||||||
|
# Stop criteria (the default one is EOS token)
|
||||||
|
stop_str: str = None
|
||||||
|
# Stops generation if meeting any token in this list
|
||||||
|
stop_token_ids: List[int] = None
|
||||||
|
|
||||||
|
def get_prompt(self) -> str:
|
||||||
|
"""Get the prompt for generation."""
|
||||||
|
system_prompt = self.system_template.format(system_message=self.system_message)
|
||||||
|
|
||||||
|
if self.sep_style == SeparatorStyle.DeepSeek:
|
||||||
|
seps = [self.sep, self.sep2]
|
||||||
|
if system_prompt == "" or system_prompt is None:
|
||||||
|
ret = ""
|
||||||
|
else:
|
||||||
|
ret = system_prompt + seps[0]
|
||||||
|
for i, (role, message) in enumerate(self.messages):
|
||||||
|
if message:
|
||||||
|
ret += role + ": " + message + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += role + ":"
|
||||||
|
return ret
|
||||||
|
elif self.sep_style == SeparatorStyle.LLAMA2:
|
||||||
|
seps = [self.sep, self.sep2]
|
||||||
|
if self.system_message:
|
||||||
|
ret = system_prompt
|
||||||
|
else:
|
||||||
|
ret = "[INST] "
|
||||||
|
for i, (role, message) in enumerate(self.messages):
|
||||||
|
tag = self.roles[i % 2]
|
||||||
|
if message:
|
||||||
|
if type(message) is tuple: # multimodal message
|
||||||
|
message, _ = message
|
||||||
|
if i == 0:
|
||||||
|
ret += message + " "
|
||||||
|
else:
|
||||||
|
ret += tag + " " + message + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += tag
|
||||||
|
return ret
|
||||||
|
elif self.sep_style == SeparatorStyle.PLAIN:
|
||||||
|
seps = [self.sep, self.sep2]
|
||||||
|
ret = ""
|
||||||
|
for i, (role, message) in enumerate(self.messages):
|
||||||
|
if message:
|
||||||
|
if type(message) is tuple:
|
||||||
|
message, _, _ = message
|
||||||
|
if i % 2 == 0:
|
||||||
|
ret += message + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += message + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += ""
|
||||||
|
return ret
|
||||||
|
elif self.sep_style == SeparatorStyle.ALIGNMENT:
|
||||||
|
seps = [self.sep, self.sep2]
|
||||||
|
ret = ""
|
||||||
|
for i, (role, message) in enumerate(self.messages):
|
||||||
|
if message:
|
||||||
|
if type(message) is tuple:
|
||||||
|
message, _, _ = message
|
||||||
|
if i % 2 == 0:
|
||||||
|
ret += "<image>\n" + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += message + seps[i % 2]
|
||||||
|
else:
|
||||||
|
ret += ""
|
||||||
|
return ret
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Invalid style: {self.sep_style}")
|
||||||
|
|
||||||
|
def get_prompt_for_current_round(self, content=None):
|
||||||
|
"""Get current round formatted question prompt during sft training"""
|
||||||
|
if self.sep_style == SeparatorStyle.PLAIN:
|
||||||
|
formatted_question = "<image>\n"
|
||||||
|
elif self.sep_style == SeparatorStyle.DeepSeek:
|
||||||
|
formatted_question = (
|
||||||
|
f"{self.roles[0]}: " + content.strip() + self.sep + f"{self.roles[1]}:"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported sep_style: {self.sep_style}")
|
||||||
|
return formatted_question
|
||||||
|
|
||||||
|
def set_system_message(self, system_message: str):
|
||||||
|
"""Set the system message."""
|
||||||
|
self.system_message = system_message
|
||||||
|
|
||||||
|
def append_message(self, role: str, message: str):
|
||||||
|
"""Append a new message."""
|
||||||
|
self.messages.append([role, message])
|
||||||
|
|
||||||
|
def reset_message(self):
|
||||||
|
"""Reset a new message."""
|
||||||
|
self.messages = []
|
||||||
|
|
||||||
|
def update_last_message(self, message: str):
|
||||||
|
"""Update the last output.
|
||||||
|
|
||||||
|
The last message is typically set to be None when constructing the prompt,
|
||||||
|
so we need to update it in-place after getting the response from a model.
|
||||||
|
"""
|
||||||
|
self.messages[-1][1] = message
|
||||||
|
|
||||||
|
def to_gradio_chatbot(self):
|
||||||
|
"""Convert the conversation to gradio chatbot format."""
|
||||||
|
ret = []
|
||||||
|
for i, (role, msg) in enumerate(self.messages[self.offset :]):
|
||||||
|
if i % 2 == 0:
|
||||||
|
ret.append([msg, None])
|
||||||
|
else:
|
||||||
|
ret[-1][-1] = msg
|
||||||
|
return ret
|
||||||
|
|
||||||
|
def to_openai_api_messages(self):
|
||||||
|
"""Convert the conversation to OpenAI chat completion format."""
|
||||||
|
system_prompt = self.system_template.format(system_message=self.system_message)
|
||||||
|
ret = [{"role": "system", "content": system_prompt}]
|
||||||
|
|
||||||
|
for i, (_, msg) in enumerate(self.messages[self.offset :]):
|
||||||
|
if i % 2 == 0:
|
||||||
|
ret.append({"role": "user", "content": msg})
|
||||||
|
else:
|
||||||
|
if msg is not None:
|
||||||
|
ret.append({"role": "assistant", "content": msg})
|
||||||
|
return ret
|
||||||
|
|
||||||
|
def copy(self):
|
||||||
|
return Conversation(
|
||||||
|
name=self.name,
|
||||||
|
system_template=self.system_template,
|
||||||
|
system_message=self.system_message,
|
||||||
|
roles=self.roles,
|
||||||
|
messages=[[x, y] for x, y in self.messages],
|
||||||
|
offset=self.offset,
|
||||||
|
sep_style=self.sep_style,
|
||||||
|
sep=self.sep,
|
||||||
|
sep2=self.sep2,
|
||||||
|
stop_str=self.stop_str,
|
||||||
|
stop_token_ids=self.stop_token_ids,
|
||||||
|
)
|
||||||
|
|
||||||
|
def dict(self):
|
||||||
|
return {
|
||||||
|
"template_name": self.name,
|
||||||
|
"system_message": self.system_message,
|
||||||
|
"roles": self.roles,
|
||||||
|
"messages": self.messages,
|
||||||
|
"offset": self.offset,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# A global registry for all conversation templates
|
||||||
|
conv_templates: Dict[str, Conversation] = {}
|
||||||
|
|
||||||
|
|
||||||
|
def register_conv_template(template: Conversation, override: bool = False):
|
||||||
|
"""Register a new conversation template."""
|
||||||
|
if not override:
|
||||||
|
assert (
|
||||||
|
template.name not in conv_templates
|
||||||
|
), f"{template.name} has been registered."
|
||||||
|
|
||||||
|
conv_templates[template.name] = template
|
||||||
|
|
||||||
|
|
||||||
|
def get_conv_template(name: str) -> Conversation:
|
||||||
|
"""Get a conversation template."""
|
||||||
|
return conv_templates[name].copy()
|
||||||
|
|
||||||
|
|
||||||
|
# llava_llama2 template
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="llava_llama2",
|
||||||
|
system_message="You are a helpful language and vision assistant. "
|
||||||
|
"You are able to understand the visual content that the user provides, "
|
||||||
|
"and assist the user with a variety of tasks using natural language.",
|
||||||
|
system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
|
||||||
|
roles=("[INST]", "[/INST]"),
|
||||||
|
messages=(),
|
||||||
|
offset=0,
|
||||||
|
sep_style=SeparatorStyle.LLAMA2,
|
||||||
|
sep=" ",
|
||||||
|
sep2=" </s><s>",
|
||||||
|
stop_token_ids=[2],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# llama2 template
|
||||||
|
# reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="llama-2",
|
||||||
|
system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
|
||||||
|
roles=("[INST]", "[/INST]"),
|
||||||
|
messages=(),
|
||||||
|
offset=0,
|
||||||
|
sep_style=SeparatorStyle.LLAMA2,
|
||||||
|
sep=" ",
|
||||||
|
sep2=" </s><s>",
|
||||||
|
stop_token_ids=[2],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# deepseek template
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="deepseek",
|
||||||
|
system_template="{system_message}",
|
||||||
|
# system_message="You are a helpful assistant. Please answer truthfully and write out your "
|
||||||
|
# "thinking step by step to be sure you get the right answer.",
|
||||||
|
system_message="",
|
||||||
|
roles=("User", "Assistant"),
|
||||||
|
messages=(),
|
||||||
|
offset=0,
|
||||||
|
sep_style=SeparatorStyle.DeepSeek,
|
||||||
|
sep="\n\n",
|
||||||
|
sep2="<｜end▁of▁sentence｜>",
|
||||||
|
stop_token_ids=[100001],
|
||||||
|
stop_str=["User:", "<｜end▁of▁sentence｜>"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="plain",
|
||||||
|
system_template="",
|
||||||
|
system_message="",
|
||||||
|
roles=("", ""),
|
||||||
|
messages=(),
|
||||||
|
offset=0,
|
||||||
|
sep_style=SeparatorStyle.PLAIN,
|
||||||
|
sep="",
|
||||||
|
sep2="",
|
||||||
|
stop_token_ids=[2],
|
||||||
|
stop_str=["</s>"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="alignment",
|
||||||
|
system_template="",
|
||||||
|
system_message="",
|
||||||
|
roles=("", ""),
|
||||||
|
messages=(),
|
||||||
|
offset=0,
|
||||||
|
sep_style=SeparatorStyle.ALIGNMENT,
|
||||||
|
sep="",
|
||||||
|
sep2="",
|
||||||
|
stop_token_ids=[2],
|
||||||
|
stop_str=["</s>"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# print("Llama-2 template:")
|
||||||
|
# conv = get_conv_template("llama-2")
|
||||||
|
# conv.set_system_message("You are a helpful, respectful and honest assistant.")
|
||||||
|
# conv.append_message(conv.roles[0], "Hello!")
|
||||||
|
# conv.append_message(conv.roles[1], "Hi!")
|
||||||
|
# conv.append_message(conv.roles[0], "How are you?")
|
||||||
|
# conv.append_message(conv.roles[1], None)
|
||||||
|
# print(conv.get_prompt())
|
||||||
|
|
||||||
|
# print("\n")
|
||||||
|
|
||||||
|
print("deepseek template:")
|
||||||
|
conv = get_conv_template("deepseek")
|
||||||
|
conv.append_message(conv.roles[0], "Hello!")
|
||||||
|
conv.append_message(conv.roles[1], "Hi! This is Tony.")
|
||||||
|
conv.append_message(conv.roles[0], "Who are you?")
|
||||||
|
conv.append_message(conv.roles[1], "I am a helpful assistant.")
|
||||||
|
conv.append_message(conv.roles[0], "How are you?")
|
||||||
|
conv.append_message(conv.roles[1], None)
|
||||||
|
print(conv.get_prompt())
|
||||||
89
janus/utils/io.py
Normal file
89
janus/utils/io.py
Normal file
@ -0,0 +1,89 @@
|
|||||||
|
# Copyright (c) 2023-2024 DeepSeek.
|
||||||
|
#
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import Dict, List
|
||||||
|
|
||||||
|
import PIL.Image
|
||||||
|
import torch
|
||||||
|
import base64
|
||||||
|
import io
|
||||||
|
from transformers import AutoModelForCausalLM
|
||||||
|
|
||||||
|
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||||
|
|
||||||
|
|
||||||
|
def load_pretrained_model(model_path: str):
|
||||||
|
vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
|
||||||
|
tokenizer = vl_chat_processor.tokenizer
|
||||||
|
|
||||||
|
vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
|
||||||
|
model_path, trust_remote_code=True
|
||||||
|
)
|
||||||
|
vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
|
||||||
|
|
||||||
|
return tokenizer, vl_chat_processor, vl_gpt
|
||||||
|
|
||||||
|
|
||||||
|
def load_pil_images(conversations: List[Dict[str, str]]) -> List[PIL.Image.Image]:
|
||||||
|
"""
|
||||||
|
|
||||||
|
Support file path or base64 images.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
conversations (List[Dict[str, str]]): the conversations with a list of messages. An example is :
|
||||||
|
[
|
||||||
|
{
|
||||||
|
"role": "User",
|
||||||
|
"content": "<image_placeholder>\nExtract all information from this image and convert them into markdown format.",
|
||||||
|
"images": ["./examples/table_datasets.png"]
|
||||||
|
},
|
||||||
|
{"role": "Assistant", "content": ""},
|
||||||
|
]
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
pil_images (List[PIL.Image.Image]): the list of PIL images.
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
pil_images = []
|
||||||
|
|
||||||
|
for message in conversations:
|
||||||
|
if "images" not in message:
|
||||||
|
continue
|
||||||
|
|
||||||
|
for image_data in message["images"]:
|
||||||
|
if image_data.startswith("data:image"):
|
||||||
|
# Image data is in base64 format
|
||||||
|
_, image_data = image_data.split(",", 1)
|
||||||
|
image_bytes = base64.b64decode(image_data)
|
||||||
|
pil_img = PIL.Image.open(io.BytesIO(image_bytes))
|
||||||
|
else:
|
||||||
|
# Image data is a file path
|
||||||
|
pil_img = PIL.Image.open(image_data)
|
||||||
|
pil_img = pil_img.convert("RGB")
|
||||||
|
pil_images.append(pil_img)
|
||||||
|
|
||||||
|
return pil_images
|
||||||
|
|
||||||
|
|
||||||
|
def load_json(filepath):
|
||||||
|
with open(filepath, "r") as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return data
|
||||||
53
pyproject.toml
Normal file
53
pyproject.toml
Normal file
@ -0,0 +1,53 @@
|
|||||||
|
[build-system]
|
||||||
|
requires = ["setuptools>=40.6.0", "wheel"]
|
||||||
|
build-backend = "setuptools.build_meta"
|
||||||
|
|
||||||
|
[project]
|
||||||
|
name = "janus"
|
||||||
|
version = "1.0.0"
|
||||||
|
description = "Janus"
|
||||||
|
authors = [{name = "DeepSeek-AI"}]
|
||||||
|
license = {file = "LICENSE-CODE"}
|
||||||
|
urls = {homepage = "https://github.com/deepseek-ai/Janus"}
|
||||||
|
readme = "README.md"
|
||||||
|
requires-python = ">=3.8"
|
||||||
|
dependencies = [
|
||||||
|
"torch>=2.0.1",
|
||||||
|
"transformers>=4.38.2",
|
||||||
|
"timm>=0.9.16",
|
||||||
|
"accelerate",
|
||||||
|
"sentencepiece",
|
||||||
|
"attrdict",
|
||||||
|
"einops",
|
||||||
|
]
|
||||||
|
|
||||||
|
[project.optional-dependencies]
|
||||||
|
gradio = [
|
||||||
|
"gradio==3.48.0",
|
||||||
|
"gradio-client==0.6.1",
|
||||||
|
"mdtex2html==1.3.0",
|
||||||
|
"pypinyin==0.50.0",
|
||||||
|
"tiktoken==0.5.2",
|
||||||
|
"tqdm==4.64.0",
|
||||||
|
"colorama==0.4.5",
|
||||||
|
"Pygments==2.12.0",
|
||||||
|
"markdown==3.4.1",
|
||||||
|
"SentencePiece==0.1.96"
|
||||||
|
]
|
||||||
|
lint = [
|
||||||
|
"isort",
|
||||||
|
"black[jupyter] >= 22.6.0",
|
||||||
|
"pylint[spelling] >= 2.15.0",
|
||||||
|
"flake8",
|
||||||
|
"flake8-bugbear",
|
||||||
|
"flake8-comprehensions",
|
||||||
|
"flake8-docstrings",
|
||||||
|
"flake8-pyi",
|
||||||
|
"flake8-simplify",
|
||||||
|
"ruff",
|
||||||
|
"pyenchant",
|
||||||
|
"pre-commit",
|
||||||
|
]
|
||||||
|
|
||||||
|
[tool.setuptools]
|
||||||
|
packages = {find = {exclude = ["images"]}}
|
||||||
19
requirements.txt
Normal file
19
requirements.txt
Normal file
@ -0,0 +1,19 @@
|
|||||||
|
torch==2.0.1
|
||||||
|
transformers>=4.38.2
|
||||||
|
timm>=0.9.16
|
||||||
|
accelerate
|
||||||
|
sentencepiece
|
||||||
|
attrdict
|
||||||
|
einops
|
||||||
|
|
||||||
|
# for gradio demo
|
||||||
|
gradio==3.48.0
|
||||||
|
gradio-client==0.6.1
|
||||||
|
mdtex2html==1.3.0
|
||||||
|
pypinyin==0.50.0
|
||||||
|
tiktoken==0.5.2
|
||||||
|
tqdm==4.64.0
|
||||||
|
colorama==0.4.5
|
||||||
|
Pygments==2.12.0
|
||||||
|
markdown==3.4.1
|
||||||
|
SentencePiece==0.1.96
|
||||||
Loading…
Reference in New Issue
Block a user