all
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- .github/CODEOWNERS +1 -0
- .github/ISSUE_TEMPLATE/bug_report.md +18 -0
- .github/ISSUE_TEMPLATE/feature_request.md +14 -0
- .gitignore +52 -0
- LICENSE +674 -0
- args_manager.py +38 -0
- auth-example.json +6 -0
- build_launcher.py +26 -0
- css/style.css +198 -0
- entry_with_update.py +46 -0
- environment.yaml +7 -0
- experiments_expansion.py +8 -0
- experiments_face.py +7 -0
- experiments_interrogate.py +8 -0
- extras/BLIP/configs/bert_config.json +21 -0
- extras/BLIP/configs/caption_coco.yaml +33 -0
- extras/BLIP/configs/med_config.json +21 -0
- extras/BLIP/configs/nlvr.yaml +21 -0
- extras/BLIP/configs/nocaps.yaml +15 -0
- extras/BLIP/configs/pretrain.yaml +27 -0
- extras/BLIP/configs/retrieval_coco.yaml +34 -0
- extras/BLIP/configs/retrieval_flickr.yaml +34 -0
- extras/BLIP/configs/retrieval_msrvtt.yaml +12 -0
- extras/BLIP/configs/vqa.yaml +25 -0
- extras/BLIP/models/bert_tokenizer/config.json +23 -0
- extras/BLIP/models/bert_tokenizer/tokenizer.json +0 -0
- extras/BLIP/models/bert_tokenizer/tokenizer_config.json +3 -0
- extras/BLIP/models/bert_tokenizer/vocab.txt +0 -0
- extras/BLIP/models/blip.py +239 -0
- extras/BLIP/models/blip_itm.py +76 -0
- extras/BLIP/models/blip_nlvr.py +105 -0
- extras/BLIP/models/blip_pretrain.py +339 -0
- extras/BLIP/models/blip_retrieval.py +319 -0
- extras/BLIP/models/blip_vqa.py +186 -0
- extras/BLIP/models/med.py +955 -0
- extras/BLIP/models/nlvr_encoder.py +843 -0
- extras/BLIP/models/vit.py +308 -0
- extras/expansion.py +126 -0
- extras/face_crop.py +50 -0
- extras/facexlib/detection/__init__.py +31 -0
- extras/facexlib/detection/align_trans.py +219 -0
- extras/facexlib/detection/matlab_cp2tform.py +317 -0
- extras/facexlib/detection/retinaface.py +366 -0
- extras/facexlib/detection/retinaface_net.py +196 -0
- extras/facexlib/detection/retinaface_utils.py +421 -0
- extras/facexlib/parsing/__init__.py +24 -0
- extras/facexlib/parsing/bisenet.py +140 -0
- extras/facexlib/parsing/parsenet.py +194 -0
- extras/facexlib/parsing/resnet.py +69 -0
- extras/facexlib/utils/__init__.py +7 -0
.github/CODEOWNERS
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
* @lllyasviel
|
.github/ISSUE_TEMPLATE/bug_report.md
ADDED
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
name: Bug report
|
3 |
+
about: Describe a problem
|
4 |
+
title: ''
|
5 |
+
labels: ''
|
6 |
+
assignees: ''
|
7 |
+
|
8 |
+
---
|
9 |
+
|
10 |
+
**Read Troubleshoot**
|
11 |
+
|
12 |
+
[x] I admit that I have read the [Troubleshoot](https://github.com/lllyasviel/Fooocus/blob/main/troubleshoot.md) before making this issue.
|
13 |
+
|
14 |
+
**Describe the problem**
|
15 |
+
A clear and concise description of what the bug is.
|
16 |
+
|
17 |
+
**Full Console Log**
|
18 |
+
Paste **full** console log here. You will make our job easier if you give a **full** log.
|
.github/ISSUE_TEMPLATE/feature_request.md
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
name: Feature request
|
3 |
+
about: Suggest an idea for this project
|
4 |
+
title: ''
|
5 |
+
labels: ''
|
6 |
+
assignees: ''
|
7 |
+
|
8 |
+
---
|
9 |
+
|
10 |
+
**Is your feature request related to a problem? Please describe.**
|
11 |
+
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
|
12 |
+
|
13 |
+
**Describe the idea you'd like**
|
14 |
+
A clear and concise description of what you want to happen.
|
.gitignore
ADDED
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
__pycache__
|
2 |
+
*.ckpt
|
3 |
+
*.safetensors
|
4 |
+
*.pth
|
5 |
+
*.pt
|
6 |
+
*.bin
|
7 |
+
*.patch
|
8 |
+
*.backup
|
9 |
+
*.corrupted
|
10 |
+
*.partial
|
11 |
+
*.onnx
|
12 |
+
sorted_styles.json
|
13 |
+
/input
|
14 |
+
/cache
|
15 |
+
/language/default.json
|
16 |
+
/test_imgs
|
17 |
+
config.txt
|
18 |
+
config_modification_tutorial.txt
|
19 |
+
user_path_config.txt
|
20 |
+
user_path_config-deprecated.txt
|
21 |
+
/modules/*.png
|
22 |
+
/repositories
|
23 |
+
/venv
|
24 |
+
/tmp
|
25 |
+
/ui-config.json
|
26 |
+
/outputs
|
27 |
+
/config.json
|
28 |
+
/log
|
29 |
+
/webui.settings.bat
|
30 |
+
/embeddings
|
31 |
+
/styles.csv
|
32 |
+
/params.txt
|
33 |
+
/styles.csv.bak
|
34 |
+
/webui-user.bat
|
35 |
+
/webui-user.sh
|
36 |
+
/interrogate
|
37 |
+
/user.css
|
38 |
+
/.idea
|
39 |
+
/notification.ogg
|
40 |
+
/notification.mp3
|
41 |
+
/SwinIR
|
42 |
+
/textual_inversion
|
43 |
+
.vscode
|
44 |
+
/extensions
|
45 |
+
/test/stdout.txt
|
46 |
+
/test/stderr.txt
|
47 |
+
/cache.json*
|
48 |
+
/config_states/
|
49 |
+
/node_modules
|
50 |
+
/package-lock.json
|
51 |
+
/.coverage*
|
52 |
+
/auth.json
|
LICENSE
ADDED
@@ -0,0 +1,674 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
GNU GENERAL PUBLIC LICENSE
|
2 |
+
Version 3, 29 June 2007
|
3 |
+
|
4 |
+
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
|
5 |
+
Everyone is permitted to copy and distribute verbatim copies
|
6 |
+
of this license document, but changing it is not allowed.
|
7 |
+
|
8 |
+
Preamble
|
9 |
+
|
10 |
+
The GNU General Public License is a free, copyleft license for
|
11 |
+
software and other kinds of works.
|
12 |
+
|
13 |
+
The licenses for most software and other practical works are designed
|
14 |
+
to take away your freedom to share and change the works. By contrast,
|
15 |
+
the GNU General Public License is intended to guarantee your freedom to
|
16 |
+
share and change all versions of a program--to make sure it remains free
|
17 |
+
software for all its users. We, the Free Software Foundation, use the
|
18 |
+
GNU General Public License for most of our software; it applies also to
|
19 |
+
any other work released this way by its authors. You can apply it to
|
20 |
+
your programs, too.
|
21 |
+
|
22 |
+
When we speak of free software, we are referring to freedom, not
|
23 |
+
price. Our General Public Licenses are designed to make sure that you
|
24 |
+
have the freedom to distribute copies of free software (and charge for
|
25 |
+
them if you wish), that you receive source code or can get it if you
|
26 |
+
want it, that you can change the software or use pieces of it in new
|
27 |
+
free programs, and that you know you can do these things.
|
28 |
+
|
29 |
+
To protect your rights, we need to prevent others from denying you
|
30 |
+
these rights or asking you to surrender the rights. Therefore, you have
|
31 |
+
certain responsibilities if you distribute copies of the software, or if
|
32 |
+
you modify it: responsibilities to respect the freedom of others.
|
33 |
+
|
34 |
+
For example, if you distribute copies of such a program, whether
|
35 |
+
gratis or for a fee, you must pass on to the recipients the same
|
36 |
+
freedoms that you received. You must make sure that they, too, receive
|
37 |
+
or can get the source code. And you must show them these terms so they
|
38 |
+
know their rights.
|
39 |
+
|
40 |
+
Developers that use the GNU GPL protect your rights with two steps:
|
41 |
+
(1) assert copyright on the software, and (2) offer you this License
|
42 |
+
giving you legal permission to copy, distribute and/or modify it.
|
43 |
+
|
44 |
+
For the developers' and authors' protection, the GPL clearly explains
|
45 |
+
that there is no warranty for this free software. For both users' and
|
46 |
+
authors' sake, the GPL requires that modified versions be marked as
|
47 |
+
changed, so that their problems will not be attributed erroneously to
|
48 |
+
authors of previous versions.
|
49 |
+
|
50 |
+
Some devices are designed to deny users access to install or run
|
51 |
+
modified versions of the software inside them, although the manufacturer
|
52 |
+
can do so. This is fundamentally incompatible with the aim of
|
53 |
+
protecting users' freedom to change the software. The systematic
|
54 |
+
pattern of such abuse occurs in the area of products for individuals to
|
55 |
+
use, which is precisely where it is most unacceptable. Therefore, we
|
56 |
+
have designed this version of the GPL to prohibit the practice for those
|
57 |
+
products. If such problems arise substantially in other domains, we
|
58 |
+
stand ready to extend this provision to those domains in future versions
|
59 |
+
of the GPL, as needed to protect the freedom of users.
|
60 |
+
|
61 |
+
Finally, every program is threatened constantly by software patents.
|
62 |
+
States should not allow patents to restrict development and use of
|
63 |
+
software on general-purpose computers, but in those that do, we wish to
|
64 |
+
avoid the special danger that patents applied to a free program could
|
65 |
+
make it effectively proprietary. To prevent this, the GPL assures that
|
66 |
+
patents cannot be used to render the program non-free.
|
67 |
+
|
68 |
+
The precise terms and conditions for copying, distribution and
|
69 |
+
modification follow.
|
70 |
+
|
71 |
+
TERMS AND CONDITIONS
|
72 |
+
|
73 |
+
0. Definitions.
|
74 |
+
|
75 |
+
"This License" refers to version 3 of the GNU General Public License.
|
76 |
+
|
77 |
+
"Copyright" also means copyright-like laws that apply to other kinds of
|
78 |
+
works, such as semiconductor masks.
|
79 |
+
|
80 |
+
"The Program" refers to any copyrightable work licensed under this
|
81 |
+
License. Each licensee is addressed as "you". "Licensees" and
|
82 |
+
"recipients" may be individuals or organizations.
|
83 |
+
|
84 |
+
To "modify" a work means to copy from or adapt all or part of the work
|
85 |
+
in a fashion requiring copyright permission, other than the making of an
|
86 |
+
exact copy. The resulting work is called a "modified version" of the
|
87 |
+
earlier work or a work "based on" the earlier work.
|
88 |
+
|
89 |
+
A "covered work" means either the unmodified Program or a work based
|
90 |
+
on the Program.
|
91 |
+
|
92 |
+
To "propagate" a work means to do anything with it that, without
|
93 |
+
permission, would make you directly or secondarily liable for
|
94 |
+
infringement under applicable copyright law, except executing it on a
|
95 |
+
computer or modifying a private copy. Propagation includes copying,
|
96 |
+
distribution (with or without modification), making available to the
|
97 |
+
public, and in some countries other activities as well.
|
98 |
+
|
99 |
+
To "convey" a work means any kind of propagation that enables other
|
100 |
+
parties to make or receive copies. Mere interaction with a user through
|
101 |
+
a computer network, with no transfer of a copy, is not conveying.
|
102 |
+
|
103 |
+
An interactive user interface displays "Appropriate Legal Notices"
|
104 |
+
to the extent that it includes a convenient and prominently visible
|
105 |
+
feature that (1) displays an appropriate copyright notice, and (2)
|
106 |
+
tells the user that there is no warranty for the work (except to the
|
107 |
+
extent that warranties are provided), that licensees may convey the
|
108 |
+
work under this License, and how to view a copy of this License. If
|
109 |
+
the interface presents a list of user commands or options, such as a
|
110 |
+
menu, a prominent item in the list meets this criterion.
|
111 |
+
|
112 |
+
1. Source Code.
|
113 |
+
|
114 |
+
The "source code" for a work means the preferred form of the work
|
115 |
+
for making modifications to it. "Object code" means any non-source
|
116 |
+
form of a work.
|
117 |
+
|
118 |
+
A "Standard Interface" means an interface that either is an official
|
119 |
+
standard defined by a recognized standards body, or, in the case of
|
120 |
+
interfaces specified for a particular programming language, one that
|
121 |
+
is widely used among developers working in that language.
|
122 |
+
|
123 |
+
The "System Libraries" of an executable work include anything, other
|
124 |
+
than the work as a whole, that (a) is included in the normal form of
|
125 |
+
packaging a Major Component, but which is not part of that Major
|
126 |
+
Component, and (b) serves only to enable use of the work with that
|
127 |
+
Major Component, or to implement a Standard Interface for which an
|
128 |
+
implementation is available to the public in source code form. A
|
129 |
+
"Major Component", in this context, means a major essential component
|
130 |
+
(kernel, window system, and so on) of the specific operating system
|
131 |
+
(if any) on which the executable work runs, or a compiler used to
|
132 |
+
produce the work, or an object code interpreter used to run it.
|
133 |
+
|
134 |
+
The "Corresponding Source" for a work in object code form means all
|
135 |
+
the source code needed to generate, install, and (for an executable
|
136 |
+
work) run the object code and to modify the work, including scripts to
|
137 |
+
control those activities. However, it does not include the work's
|
138 |
+
System Libraries, or general-purpose tools or generally available free
|
139 |
+
programs which are used unmodified in performing those activities but
|
140 |
+
which are not part of the work. For example, Corresponding Source
|
141 |
+
includes interface definition files associated with source files for
|
142 |
+
the work, and the source code for shared libraries and dynamically
|
143 |
+
linked subprograms that the work is specifically designed to require,
|
144 |
+
such as by intimate data communication or control flow between those
|
145 |
+
subprograms and other parts of the work.
|
146 |
+
|
147 |
+
The Corresponding Source need not include anything that users
|
148 |
+
can regenerate automatically from other parts of the Corresponding
|
149 |
+
Source.
|
150 |
+
|
151 |
+
The Corresponding Source for a work in source code form is that
|
152 |
+
same work.
|
153 |
+
|
154 |
+
2. Basic Permissions.
|
155 |
+
|
156 |
+
All rights granted under this License are granted for the term of
|
157 |
+
copyright on the Program, and are irrevocable provided the stated
|
158 |
+
conditions are met. This License explicitly affirms your unlimited
|
159 |
+
permission to run the unmodified Program. The output from running a
|
160 |
+
covered work is covered by this License only if the output, given its
|
161 |
+
content, constitutes a covered work. This License acknowledges your
|
162 |
+
rights of fair use or other equivalent, as provided by copyright law.
|
163 |
+
|
164 |
+
You may make, run and propagate covered works that you do not
|
165 |
+
convey, without conditions so long as your license otherwise remains
|
166 |
+
in force. You may convey covered works to others for the sole purpose
|
167 |
+
of having them make modifications exclusively for you, or provide you
|
168 |
+
with facilities for running those works, provided that you comply with
|
169 |
+
the terms of this License in conveying all material for which you do
|
170 |
+
not control copyright. Those thus making or running the covered works
|
171 |
+
for you must do so exclusively on your behalf, under your direction
|
172 |
+
and control, on terms that prohibit them from making any copies of
|
173 |
+
your copyrighted material outside their relationship with you.
|
174 |
+
|
175 |
+
Conveying under any other circumstances is permitted solely under
|
176 |
+
the conditions stated below. Sublicensing is not allowed; section 10
|
177 |
+
makes it unnecessary.
|
178 |
+
|
179 |
+
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
|
180 |
+
|
181 |
+
No covered work shall be deemed part of an effective technological
|
182 |
+
measure under any applicable law fulfilling obligations under article
|
183 |
+
11 of the WIPO copyright treaty adopted on 20 December 1996, or
|
184 |
+
similar laws prohibiting or restricting circumvention of such
|
185 |
+
measures.
|
186 |
+
|
187 |
+
When you convey a covered work, you waive any legal power to forbid
|
188 |
+
circumvention of technological measures to the extent such circumvention
|
189 |
+
is effected by exercising rights under this License with respect to
|
190 |
+
the covered work, and you disclaim any intention to limit operation or
|
191 |
+
modification of the work as a means of enforcing, against the work's
|
192 |
+
users, your or third parties' legal rights to forbid circumvention of
|
193 |
+
technological measures.
|
194 |
+
|
195 |
+
4. Conveying Verbatim Copies.
|
196 |
+
|
197 |
+
You may convey verbatim copies of the Program's source code as you
|
198 |
+
receive it, in any medium, provided that you conspicuously and
|
199 |
+
appropriately publish on each copy an appropriate copyright notice;
|
200 |
+
keep intact all notices stating that this License and any
|
201 |
+
non-permissive terms added in accord with section 7 apply to the code;
|
202 |
+
keep intact all notices of the absence of any warranty; and give all
|
203 |
+
recipients a copy of this License along with the Program.
|
204 |
+
|
205 |
+
You may charge any price or no price for each copy that you convey,
|
206 |
+
and you may offer support or warranty protection for a fee.
|
207 |
+
|
208 |
+
5. Conveying Modified Source Versions.
|
209 |
+
|
210 |
+
You may convey a work based on the Program, or the modifications to
|
211 |
+
produce it from the Program, in the form of source code under the
|
212 |
+
terms of section 4, provided that you also meet all of these conditions:
|
213 |
+
|
214 |
+
a) The work must carry prominent notices stating that you modified
|
215 |
+
it, and giving a relevant date.
|
216 |
+
|
217 |
+
b) The work must carry prominent notices stating that it is
|
218 |
+
released under this License and any conditions added under section
|
219 |
+
7. This requirement modifies the requirement in section 4 to
|
220 |
+
"keep intact all notices".
|
221 |
+
|
222 |
+
c) You must license the entire work, as a whole, under this
|
223 |
+
License to anyone who comes into possession of a copy. This
|
224 |
+
License will therefore apply, along with any applicable section 7
|
225 |
+
additional terms, to the whole of the work, and all its parts,
|
226 |
+
regardless of how they are packaged. This License gives no
|
227 |
+
permission to license the work in any other way, but it does not
|
228 |
+
invalidate such permission if you have separately received it.
|
229 |
+
|
230 |
+
d) If the work has interactive user interfaces, each must display
|
231 |
+
Appropriate Legal Notices; however, if the Program has interactive
|
232 |
+
interfaces that do not display Appropriate Legal Notices, your
|
233 |
+
work need not make them do so.
|
234 |
+
|
235 |
+
A compilation of a covered work with other separate and independent
|
236 |
+
works, which are not by their nature extensions of the covered work,
|
237 |
+
and which are not combined with it such as to form a larger program,
|
238 |
+
in or on a volume of a storage or distribution medium, is called an
|
239 |
+
"aggregate" if the compilation and its resulting copyright are not
|
240 |
+
used to limit the access or legal rights of the compilation's users
|
241 |
+
beyond what the individual works permit. Inclusion of a covered work
|
242 |
+
in an aggregate does not cause this License to apply to the other
|
243 |
+
parts of the aggregate.
|
244 |
+
|
245 |
+
6. Conveying Non-Source Forms.
|
246 |
+
|
247 |
+
You may convey a covered work in object code form under the terms
|
248 |
+
of sections 4 and 5, provided that you also convey the
|
249 |
+
machine-readable Corresponding Source under the terms of this License,
|
250 |
+
in one of these ways:
|
251 |
+
|
252 |
+
a) Convey the object code in, or embodied in, a physical product
|
253 |
+
(including a physical distribution medium), accompanied by the
|
254 |
+
Corresponding Source fixed on a durable physical medium
|
255 |
+
customarily used for software interchange.
|
256 |
+
|
257 |
+
b) Convey the object code in, or embodied in, a physical product
|
258 |
+
(including a physical distribution medium), accompanied by a
|
259 |
+
written offer, valid for at least three years and valid for as
|
260 |
+
long as you offer spare parts or customer support for that product
|
261 |
+
model, to give anyone who possesses the object code either (1) a
|
262 |
+
copy of the Corresponding Source for all the software in the
|
263 |
+
product that is covered by this License, on a durable physical
|
264 |
+
medium customarily used for software interchange, for a price no
|
265 |
+
more than your reasonable cost of physically performing this
|
266 |
+
conveying of source, or (2) access to copy the
|
267 |
+
Corresponding Source from a network server at no charge.
|
268 |
+
|
269 |
+
c) Convey individual copies of the object code with a copy of the
|
270 |
+
written offer to provide the Corresponding Source. This
|
271 |
+
alternative is allowed only occasionally and noncommercially, and
|
272 |
+
only if you received the object code with such an offer, in accord
|
273 |
+
with subsection 6b.
|
274 |
+
|
275 |
+
d) Convey the object code by offering access from a designated
|
276 |
+
place (gratis or for a charge), and offer equivalent access to the
|
277 |
+
Corresponding Source in the same way through the same place at no
|
278 |
+
further charge. You need not require recipients to copy the
|
279 |
+
Corresponding Source along with the object code. If the place to
|
280 |
+
copy the object code is a network server, the Corresponding Source
|
281 |
+
may be on a different server (operated by you or a third party)
|
282 |
+
that supports equivalent copying facilities, provided you maintain
|
283 |
+
clear directions next to the object code saying where to find the
|
284 |
+
Corresponding Source. Regardless of what server hosts the
|
285 |
+
Corresponding Source, you remain obligated to ensure that it is
|
286 |
+
available for as long as needed to satisfy these requirements.
|
287 |
+
|
288 |
+
e) Convey the object code using peer-to-peer transmission, provided
|
289 |
+
you inform other peers where the object code and Corresponding
|
290 |
+
Source of the work are being offered to the general public at no
|
291 |
+
charge under subsection 6d.
|
292 |
+
|
293 |
+
A separable portion of the object code, whose source code is excluded
|
294 |
+
from the Corresponding Source as a System Library, need not be
|
295 |
+
included in conveying the object code work.
|
296 |
+
|
297 |
+
A "User Product" is either (1) a "consumer product", which means any
|
298 |
+
tangible personal property which is normally used for personal, family,
|
299 |
+
or household purposes, or (2) anything designed or sold for incorporation
|
300 |
+
into a dwelling. In determining whether a product is a consumer product,
|
301 |
+
doubtful cases shall be resolved in favor of coverage. For a particular
|
302 |
+
product received by a particular user, "normally used" refers to a
|
303 |
+
typical or common use of that class of product, regardless of the status
|
304 |
+
of the particular user or of the way in which the particular user
|
305 |
+
actually uses, or expects or is expected to use, the product. A product
|
306 |
+
is a consumer product regardless of whether the product has substantial
|
307 |
+
commercial, industrial or non-consumer uses, unless such uses represent
|
308 |
+
the only significant mode of use of the product.
|
309 |
+
|
310 |
+
"Installation Information" for a User Product means any methods,
|
311 |
+
procedures, authorization keys, or other information required to install
|
312 |
+
and execute modified versions of a covered work in that User Product from
|
313 |
+
a modified version of its Corresponding Source. The information must
|
314 |
+
suffice to ensure that the continued functioning of the modified object
|
315 |
+
code is in no case prevented or interfered with solely because
|
316 |
+
modification has been made.
|
317 |
+
|
318 |
+
If you convey an object code work under this section in, or with, or
|
319 |
+
specifically for use in, a User Product, and the conveying occurs as
|
320 |
+
part of a transaction in which the right of possession and use of the
|
321 |
+
User Product is transferred to the recipient in perpetuity or for a
|
322 |
+
fixed term (regardless of how the transaction is characterized), the
|
323 |
+
Corresponding Source conveyed under this section must be accompanied
|
324 |
+
by the Installation Information. But this requirement does not apply
|
325 |
+
if neither you nor any third party retains the ability to install
|
326 |
+
modified object code on the User Product (for example, the work has
|
327 |
+
been installed in ROM).
|
328 |
+
|
329 |
+
The requirement to provide Installation Information does not include a
|
330 |
+
requirement to continue to provide support service, warranty, or updates
|
331 |
+
for a work that has been modified or installed by the recipient, or for
|
332 |
+
the User Product in which it has been modified or installed. Access to a
|
333 |
+
network may be denied when the modification itself materially and
|
334 |
+
adversely affects the operation of the network or violates the rules and
|
335 |
+
protocols for communication across the network.
|
336 |
+
|
337 |
+
Corresponding Source conveyed, and Installation Information provided,
|
338 |
+
in accord with this section must be in a format that is publicly
|
339 |
+
documented (and with an implementation available to the public in
|
340 |
+
source code form), and must require no special password or key for
|
341 |
+
unpacking, reading or copying.
|
342 |
+
|
343 |
+
7. Additional Terms.
|
344 |
+
|
345 |
+
"Additional permissions" are terms that supplement the terms of this
|
346 |
+
License by making exceptions from one or more of its conditions.
|
347 |
+
Additional permissions that are applicable to the entire Program shall
|
348 |
+
be treated as though they were included in this License, to the extent
|
349 |
+
that they are valid under applicable law. If additional permissions
|
350 |
+
apply only to part of the Program, that part may be used separately
|
351 |
+
under those permissions, but the entire Program remains governed by
|
352 |
+
this License without regard to the additional permissions.
|
353 |
+
|
354 |
+
When you convey a copy of a covered work, you may at your option
|
355 |
+
remove any additional permissions from that copy, or from any part of
|
356 |
+
it. (Additional permissions may be written to require their own
|
357 |
+
removal in certain cases when you modify the work.) You may place
|
358 |
+
additional permissions on material, added by you to a covered work,
|
359 |
+
for which you have or can give appropriate copyright permission.
|
360 |
+
|
361 |
+
Notwithstanding any other provision of this License, for material you
|
362 |
+
add to a covered work, you may (if authorized by the copyright holders of
|
363 |
+
that material) supplement the terms of this License with terms:
|
364 |
+
|
365 |
+
a) Disclaiming warranty or limiting liability differently from the
|
366 |
+
terms of sections 15 and 16 of this License; or
|
367 |
+
|
368 |
+
b) Requiring preservation of specified reasonable legal notices or
|
369 |
+
author attributions in that material or in the Appropriate Legal
|
370 |
+
Notices displayed by works containing it; or
|
371 |
+
|
372 |
+
c) Prohibiting misrepresentation of the origin of that material, or
|
373 |
+
requiring that modified versions of such material be marked in
|
374 |
+
reasonable ways as different from the original version; or
|
375 |
+
|
376 |
+
d) Limiting the use for publicity purposes of names of licensors or
|
377 |
+
authors of the material; or
|
378 |
+
|
379 |
+
e) Declining to grant rights under trademark law for use of some
|
380 |
+
trade names, trademarks, or service marks; or
|
381 |
+
|
382 |
+
f) Requiring indemnification of licensors and authors of that
|
383 |
+
material by anyone who conveys the material (or modified versions of
|
384 |
+
it) with contractual assumptions of liability to the recipient, for
|
385 |
+
any liability that these contractual assumptions directly impose on
|
386 |
+
those licensors and authors.
|
387 |
+
|
388 |
+
All other non-permissive additional terms are considered "further
|
389 |
+
restrictions" within the meaning of section 10. If the Program as you
|
390 |
+
received it, or any part of it, contains a notice stating that it is
|
391 |
+
governed by this License along with a term that is a further
|
392 |
+
restriction, you may remove that term. If a license document contains
|
393 |
+
a further restriction but permits relicensing or conveying under this
|
394 |
+
License, you may add to a covered work material governed by the terms
|
395 |
+
of that license document, provided that the further restriction does
|
396 |
+
not survive such relicensing or conveying.
|
397 |
+
|
398 |
+
If you add terms to a covered work in accord with this section, you
|
399 |
+
must place, in the relevant source files, a statement of the
|
400 |
+
additional terms that apply to those files, or a notice indicating
|
401 |
+
where to find the applicable terms.
|
402 |
+
|
403 |
+
Additional terms, permissive or non-permissive, may be stated in the
|
404 |
+
form of a separately written license, or stated as exceptions;
|
405 |
+
the above requirements apply either way.
|
406 |
+
|
407 |
+
8. Termination.
|
408 |
+
|
409 |
+
You may not propagate or modify a covered work except as expressly
|
410 |
+
provided under this License. Any attempt otherwise to propagate or
|
411 |
+
modify it is void, and will automatically terminate your rights under
|
412 |
+
this License (including any patent licenses granted under the third
|
413 |
+
paragraph of section 11).
|
414 |
+
|
415 |
+
However, if you cease all violation of this License, then your
|
416 |
+
license from a particular copyright holder is reinstated (a)
|
417 |
+
provisionally, unless and until the copyright holder explicitly and
|
418 |
+
finally terminates your license, and (b) permanently, if the copyright
|
419 |
+
holder fails to notify you of the violation by some reasonable means
|
420 |
+
prior to 60 days after the cessation.
|
421 |
+
|
422 |
+
Moreover, your license from a particular copyright holder is
|
423 |
+
reinstated permanently if the copyright holder notifies you of the
|
424 |
+
violation by some reasonable means, this is the first time you have
|
425 |
+
received notice of violation of this License (for any work) from that
|
426 |
+
copyright holder, and you cure the violation prior to 30 days after
|
427 |
+
your receipt of the notice.
|
428 |
+
|
429 |
+
Termination of your rights under this section does not terminate the
|
430 |
+
licenses of parties who have received copies or rights from you under
|
431 |
+
this License. If your rights have been terminated and not permanently
|
432 |
+
reinstated, you do not qualify to receive new licenses for the same
|
433 |
+
material under section 10.
|
434 |
+
|
435 |
+
9. Acceptance Not Required for Having Copies.
|
436 |
+
|
437 |
+
You are not required to accept this License in order to receive or
|
438 |
+
run a copy of the Program. Ancillary propagation of a covered work
|
439 |
+
occurring solely as a consequence of using peer-to-peer transmission
|
440 |
+
to receive a copy likewise does not require acceptance. However,
|
441 |
+
nothing other than this License grants you permission to propagate or
|
442 |
+
modify any covered work. These actions infringe copyright if you do
|
443 |
+
not accept this License. Therefore, by modifying or propagating a
|
444 |
+
covered work, you indicate your acceptance of this License to do so.
|
445 |
+
|
446 |
+
10. Automatic Licensing of Downstream Recipients.
|
447 |
+
|
448 |
+
Each time you convey a covered work, the recipient automatically
|
449 |
+
receives a license from the original licensors, to run, modify and
|
450 |
+
propagate that work, subject to this License. You are not responsible
|
451 |
+
for enforcing compliance by third parties with this License.
|
452 |
+
|
453 |
+
An "entity transaction" is a transaction transferring control of an
|
454 |
+
organization, or substantially all assets of one, or subdividing an
|
455 |
+
organization, or merging organizations. If propagation of a covered
|
456 |
+
work results from an entity transaction, each party to that
|
457 |
+
transaction who receives a copy of the work also receives whatever
|
458 |
+
licenses to the work the party's predecessor in interest had or could
|
459 |
+
give under the previous paragraph, plus a right to possession of the
|
460 |
+
Corresponding Source of the work from the predecessor in interest, if
|
461 |
+
the predecessor has it or can get it with reasonable efforts.
|
462 |
+
|
463 |
+
You may not impose any further restrictions on the exercise of the
|
464 |
+
rights granted or affirmed under this License. For example, you may
|
465 |
+
not impose a license fee, royalty, or other charge for exercise of
|
466 |
+
rights granted under this License, and you may not initiate litigation
|
467 |
+
(including a cross-claim or counterclaim in a lawsuit) alleging that
|
468 |
+
any patent claim is infringed by making, using, selling, offering for
|
469 |
+
sale, or importing the Program or any portion of it.
|
470 |
+
|
471 |
+
11. Patents.
|
472 |
+
|
473 |
+
A "contributor" is a copyright holder who authorizes use under this
|
474 |
+
License of the Program or a work on which the Program is based. The
|
475 |
+
work thus licensed is called the contributor's "contributor version".
|
476 |
+
|
477 |
+
A contributor's "essential patent claims" are all patent claims
|
478 |
+
owned or controlled by the contributor, whether already acquired or
|
479 |
+
hereafter acquired, that would be infringed by some manner, permitted
|
480 |
+
by this License, of making, using, or selling its contributor version,
|
481 |
+
but do not include claims that would be infringed only as a
|
482 |
+
consequence of further modification of the contributor version. For
|
483 |
+
purposes of this definition, "control" includes the right to grant
|
484 |
+
patent sublicenses in a manner consistent with the requirements of
|
485 |
+
this License.
|
486 |
+
|
487 |
+
Each contributor grants you a non-exclusive, worldwide, royalty-free
|
488 |
+
patent license under the contributor's essential patent claims, to
|
489 |
+
make, use, sell, offer for sale, import and otherwise run, modify and
|
490 |
+
propagate the contents of its contributor version.
|
491 |
+
|
492 |
+
In the following three paragraphs, a "patent license" is any express
|
493 |
+
agreement or commitment, however denominated, not to enforce a patent
|
494 |
+
(such as an express permission to practice a patent or covenant not to
|
495 |
+
sue for patent infringement). To "grant" such a patent license to a
|
496 |
+
party means to make such an agreement or commitment not to enforce a
|
497 |
+
patent against the party.
|
498 |
+
|
499 |
+
If you convey a covered work, knowingly relying on a patent license,
|
500 |
+
and the Corresponding Source of the work is not available for anyone
|
501 |
+
to copy, free of charge and under the terms of this License, through a
|
502 |
+
publicly available network server or other readily accessible means,
|
503 |
+
then you must either (1) cause the Corresponding Source to be so
|
504 |
+
available, or (2) arrange to deprive yourself of the benefit of the
|
505 |
+
patent license for this particular work, or (3) arrange, in a manner
|
506 |
+
consistent with the requirements of this License, to extend the patent
|
507 |
+
license to downstream recipients. "Knowingly relying" means you have
|
508 |
+
actual knowledge that, but for the patent license, your conveying the
|
509 |
+
covered work in a country, or your recipient's use of the covered work
|
510 |
+
in a country, would infringe one or more identifiable patents in that
|
511 |
+
country that you have reason to believe are valid.
|
512 |
+
|
513 |
+
If, pursuant to or in connection with a single transaction or
|
514 |
+
arrangement, you convey, or propagate by procuring conveyance of, a
|
515 |
+
covered work, and grant a patent license to some of the parties
|
516 |
+
receiving the covered work authorizing them to use, propagate, modify
|
517 |
+
or convey a specific copy of the covered work, then the patent license
|
518 |
+
you grant is automatically extended to all recipients of the covered
|
519 |
+
work and works based on it.
|
520 |
+
|
521 |
+
A patent license is "discriminatory" if it does not include within
|
522 |
+
the scope of its coverage, prohibits the exercise of, or is
|
523 |
+
conditioned on the non-exercise of one or more of the rights that are
|
524 |
+
specifically granted under this License. You may not convey a covered
|
525 |
+
work if you are a party to an arrangement with a third party that is
|
526 |
+
in the business of distributing software, under which you make payment
|
527 |
+
to the third party based on the extent of your activity of conveying
|
528 |
+
the work, and under which the third party grants, to any of the
|
529 |
+
parties who would receive the covered work from you, a discriminatory
|
530 |
+
patent license (a) in connection with copies of the covered work
|
531 |
+
conveyed by you (or copies made from those copies), or (b) primarily
|
532 |
+
for and in connection with specific products or compilations that
|
533 |
+
contain the covered work, unless you entered into that arrangement,
|
534 |
+
or that patent license was granted, prior to 28 March 2007.
|
535 |
+
|
536 |
+
Nothing in this License shall be construed as excluding or limiting
|
537 |
+
any implied license or other defenses to infringement that may
|
538 |
+
otherwise be available to you under applicable patent law.
|
539 |
+
|
540 |
+
12. No Surrender of Others' Freedom.
|
541 |
+
|
542 |
+
If conditions are imposed on you (whether by court order, agreement or
|
543 |
+
otherwise) that contradict the conditions of this License, they do not
|
544 |
+
excuse you from the conditions of this License. If you cannot convey a
|
545 |
+
covered work so as to satisfy simultaneously your obligations under this
|
546 |
+
License and any other pertinent obligations, then as a consequence you may
|
547 |
+
not convey it at all. For example, if you agree to terms that obligate you
|
548 |
+
to collect a royalty for further conveying from those to whom you convey
|
549 |
+
the Program, the only way you could satisfy both those terms and this
|
550 |
+
License would be to refrain entirely from conveying the Program.
|
551 |
+
|
552 |
+
13. Use with the GNU Affero General Public License.
|
553 |
+
|
554 |
+
Notwithstanding any other provision of this License, you have
|
555 |
+
permission to link or combine any covered work with a work licensed
|
556 |
+
under version 3 of the GNU Affero General Public License into a single
|
557 |
+
combined work, and to convey the resulting work. The terms of this
|
558 |
+
License will continue to apply to the part which is the covered work,
|
559 |
+
but the special requirements of the GNU Affero General Public License,
|
560 |
+
section 13, concerning interaction through a network will apply to the
|
561 |
+
combination as such.
|
562 |
+
|
563 |
+
14. Revised Versions of this License.
|
564 |
+
|
565 |
+
The Free Software Foundation may publish revised and/or new versions of
|
566 |
+
the GNU General Public License from time to time. Such new versions will
|
567 |
+
be similar in spirit to the present version, but may differ in detail to
|
568 |
+
address new problems or concerns.
|
569 |
+
|
570 |
+
Each version is given a distinguishing version number. If the
|
571 |
+
Program specifies that a certain numbered version of the GNU General
|
572 |
+
Public License "or any later version" applies to it, you have the
|
573 |
+
option of following the terms and conditions either of that numbered
|
574 |
+
version or of any later version published by the Free Software
|
575 |
+
Foundation. If the Program does not specify a version number of the
|
576 |
+
GNU General Public License, you may choose any version ever published
|
577 |
+
by the Free Software Foundation.
|
578 |
+
|
579 |
+
If the Program specifies that a proxy can decide which future
|
580 |
+
versions of the GNU General Public License can be used, that proxy's
|
581 |
+
public statement of acceptance of a version permanently authorizes you
|
582 |
+
to choose that version for the Program.
|
583 |
+
|
584 |
+
Later license versions may give you additional or different
|
585 |
+
permissions. However, no additional obligations are imposed on any
|
586 |
+
author or copyright holder as a result of your choosing to follow a
|
587 |
+
later version.
|
588 |
+
|
589 |
+
15. Disclaimer of Warranty.
|
590 |
+
|
591 |
+
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
|
592 |
+
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
|
593 |
+
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
|
594 |
+
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
|
595 |
+
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
|
596 |
+
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
|
597 |
+
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
|
598 |
+
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
|
599 |
+
|
600 |
+
16. Limitation of Liability.
|
601 |
+
|
602 |
+
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
|
603 |
+
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
|
604 |
+
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
|
605 |
+
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
|
606 |
+
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
|
607 |
+
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
|
608 |
+
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
|
609 |
+
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
|
610 |
+
SUCH DAMAGES.
|
611 |
+
|
612 |
+
17. Interpretation of Sections 15 and 16.
|
613 |
+
|
614 |
+
If the disclaimer of warranty and limitation of liability provided
|
615 |
+
above cannot be given local legal effect according to their terms,
|
616 |
+
reviewing courts shall apply local law that most closely approximates
|
617 |
+
an absolute waiver of all civil liability in connection with the
|
618 |
+
Program, unless a warranty or assumption of liability accompanies a
|
619 |
+
copy of the Program in return for a fee.
|
620 |
+
|
621 |
+
END OF TERMS AND CONDITIONS
|
622 |
+
|
623 |
+
How to Apply These Terms to Your New Programs
|
624 |
+
|
625 |
+
If you develop a new program, and you want it to be of the greatest
|
626 |
+
possible use to the public, the best way to achieve this is to make it
|
627 |
+
free software which everyone can redistribute and change under these terms.
|
628 |
+
|
629 |
+
To do so, attach the following notices to the program. It is safest
|
630 |
+
to attach them to the start of each source file to most effectively
|
631 |
+
state the exclusion of warranty; and each file should have at least
|
632 |
+
the "copyright" line and a pointer to where the full notice is found.
|
633 |
+
|
634 |
+
<one line to give the program's name and a brief idea of what it does.>
|
635 |
+
Copyright (C) <year> <name of author>
|
636 |
+
|
637 |
+
This program is free software: you can redistribute it and/or modify
|
638 |
+
it under the terms of the GNU General Public License as published by
|
639 |
+
the Free Software Foundation, either version 3 of the License, or
|
640 |
+
(at your option) any later version.
|
641 |
+
|
642 |
+
This program is distributed in the hope that it will be useful,
|
643 |
+
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
644 |
+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
645 |
+
GNU General Public License for more details.
|
646 |
+
|
647 |
+
You should have received a copy of the GNU General Public License
|
648 |
+
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
649 |
+
|
650 |
+
Also add information on how to contact you by electronic and paper mail.
|
651 |
+
|
652 |
+
If the program does terminal interaction, make it output a short
|
653 |
+
notice like this when it starts in an interactive mode:
|
654 |
+
|
655 |
+
<program> Copyright (C) <year> <name of author>
|
656 |
+
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
|
657 |
+
This is free software, and you are welcome to redistribute it
|
658 |
+
under certain conditions; type `show c' for details.
|
659 |
+
|
660 |
+
The hypothetical commands `show w' and `show c' should show the appropriate
|
661 |
+
parts of the General Public License. Of course, your program's commands
|
662 |
+
might be different; for a GUI interface, you would use an "about box".
|
663 |
+
|
664 |
+
You should also get your employer (if you work as a programmer) or school,
|
665 |
+
if any, to sign a "copyright disclaimer" for the program, if necessary.
|
666 |
+
For more information on this, and how to apply and follow the GNU GPL, see
|
667 |
+
<https://www.gnu.org/licenses/>.
|
668 |
+
|
669 |
+
The GNU General Public License does not permit incorporating your program
|
670 |
+
into proprietary programs. If your program is a subroutine library, you
|
671 |
+
may consider it more useful to permit linking proprietary applications with
|
672 |
+
the library. If this is what you want to do, use the GNU Lesser General
|
673 |
+
Public License instead of this License. But first, please read
|
674 |
+
<https://www.gnu.org/licenses/why-not-lgpl.html>.
|
args_manager.py
ADDED
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import ldm_patched.modules.args_parser as args_parser
|
2 |
+
|
3 |
+
|
4 |
+
args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.")
|
5 |
+
args_parser.parser.add_argument("--preset", type=str, default=None, help="Apply specified UI preset.")
|
6 |
+
|
7 |
+
args_parser.parser.add_argument("--language", type=str, default='default',
|
8 |
+
help="Translate UI using json files in [language] folder. "
|
9 |
+
"For example, [--language example] will use [language/example.json] for translation.")
|
10 |
+
|
11 |
+
# For example, https://github.com/lllyasviel/Fooocus/issues/849
|
12 |
+
args_parser.parser.add_argument("--disable-offload-from-vram", action="store_true",
|
13 |
+
help="Force loading models to vram when the unload can be avoided. "
|
14 |
+
"Some Mac users may need this.")
|
15 |
+
|
16 |
+
args_parser.parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)
|
17 |
+
args_parser.parser.add_argument("--disable-image-log", action='store_true',
|
18 |
+
help="Prevent writing images and logs to hard drive.")
|
19 |
+
|
20 |
+
args_parser.parser.add_argument("--disable-analytics", action='store_true',
|
21 |
+
help="Disables analytics for Gradio", default=False)
|
22 |
+
|
23 |
+
args_parser.parser.set_defaults(
|
24 |
+
disable_cuda_malloc=True,
|
25 |
+
in_browser=True,
|
26 |
+
port=None
|
27 |
+
)
|
28 |
+
|
29 |
+
args_parser.args = args_parser.parser.parse_args()
|
30 |
+
|
31 |
+
# (Disable by default because of issues like https://github.com/lllyasviel/Fooocus/issues/724)
|
32 |
+
args_parser.args.always_offload_from_vram = not args_parser.args.disable_offload_from_vram
|
33 |
+
|
34 |
+
if args_parser.args.disable_analytics:
|
35 |
+
import os
|
36 |
+
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
|
37 |
+
|
38 |
+
args = args_parser.args
|
auth-example.json
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
[
|
2 |
+
{
|
3 |
+
"user": "sitting-duck-1",
|
4 |
+
"pass": "very-bad-publicly-known-password-change-it"
|
5 |
+
}
|
6 |
+
]
|
build_launcher.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
|
3 |
+
win32_root = os.path.dirname(os.path.dirname(__file__))
|
4 |
+
python_embeded_path = os.path.join(win32_root, 'python_embeded')
|
5 |
+
|
6 |
+
is_win32_standalone_build = os.path.exists(python_embeded_path) and os.path.isdir(python_embeded_path)
|
7 |
+
|
8 |
+
win32_cmd = '''
|
9 |
+
.\python_embeded\python.exe -s Fooocus\entry_with_update.py {cmds} %*
|
10 |
+
pause
|
11 |
+
'''
|
12 |
+
|
13 |
+
|
14 |
+
def build_launcher():
|
15 |
+
if not is_win32_standalone_build:
|
16 |
+
return
|
17 |
+
|
18 |
+
presets = [None, 'anime', 'realistic']
|
19 |
+
|
20 |
+
for preset in presets:
|
21 |
+
win32_cmd_preset = win32_cmd.replace('{cmds}', '' if preset is None else f'--preset {preset}')
|
22 |
+
bat_path = os.path.join(win32_root, 'run.bat' if preset is None else f'run_{preset}.bat')
|
23 |
+
if not os.path.exists(bat_path):
|
24 |
+
with open(bat_path, "w", encoding="utf-8") as f:
|
25 |
+
f.write(win32_cmd_preset)
|
26 |
+
return
|
css/style.css
ADDED
@@ -0,0 +1,198 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
/* based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/v1.6.0/style.css */
|
2 |
+
|
3 |
+
#context-menu{
|
4 |
+
z-index:9999;
|
5 |
+
position:absolute;
|
6 |
+
display:block;
|
7 |
+
padding:0px 0;
|
8 |
+
border:2px solid #a55000;
|
9 |
+
border-radius:8px;
|
10 |
+
box-shadow:1px 1px 2px #CE6400;
|
11 |
+
width: 200px;
|
12 |
+
}
|
13 |
+
|
14 |
+
.context-menu-items{
|
15 |
+
list-style: none;
|
16 |
+
margin: 0;
|
17 |
+
padding: 0;
|
18 |
+
}
|
19 |
+
|
20 |
+
.context-menu-items a{
|
21 |
+
display:block;
|
22 |
+
padding:5px;
|
23 |
+
cursor:pointer;
|
24 |
+
}
|
25 |
+
|
26 |
+
.context-menu-items a:hover{
|
27 |
+
background: #a55000;
|
28 |
+
}
|
29 |
+
|
30 |
+
.canvas-tooltip-info {
|
31 |
+
position: absolute;
|
32 |
+
top: 28px;
|
33 |
+
left: 2px;
|
34 |
+
cursor: help;
|
35 |
+
background-color: rgba(0, 0, 0, 0.3);
|
36 |
+
width: 20px;
|
37 |
+
height: 20px;
|
38 |
+
border-radius: 50%;
|
39 |
+
display: flex;
|
40 |
+
align-items: center;
|
41 |
+
justify-content: center;
|
42 |
+
flex-direction: column;
|
43 |
+
|
44 |
+
z-index: 100;
|
45 |
+
}
|
46 |
+
|
47 |
+
.canvas-tooltip-info::after {
|
48 |
+
content: '';
|
49 |
+
display: block;
|
50 |
+
width: 2px;
|
51 |
+
height: 7px;
|
52 |
+
background-color: white;
|
53 |
+
margin-top: 2px;
|
54 |
+
}
|
55 |
+
|
56 |
+
.canvas-tooltip-info::before {
|
57 |
+
content: '';
|
58 |
+
display: block;
|
59 |
+
width: 2px;
|
60 |
+
height: 2px;
|
61 |
+
background-color: white;
|
62 |
+
}
|
63 |
+
|
64 |
+
.canvas-tooltip-content {
|
65 |
+
display: none;
|
66 |
+
background-color: #f9f9f9;
|
67 |
+
color: #333;
|
68 |
+
border: 1px solid #ddd;
|
69 |
+
padding: 15px;
|
70 |
+
position: absolute;
|
71 |
+
top: 40px;
|
72 |
+
left: 10px;
|
73 |
+
width: 250px;
|
74 |
+
font-size: 16px;
|
75 |
+
opacity: 0;
|
76 |
+
border-radius: 8px;
|
77 |
+
box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2);
|
78 |
+
|
79 |
+
z-index: 100;
|
80 |
+
}
|
81 |
+
|
82 |
+
.canvas-tooltip:hover .canvas-tooltip-content {
|
83 |
+
display: block;
|
84 |
+
animation: fadeIn 0.5s;
|
85 |
+
opacity: 1;
|
86 |
+
}
|
87 |
+
|
88 |
+
@keyframes fadeIn {
|
89 |
+
from {opacity: 0;}
|
90 |
+
to {opacity: 1;}
|
91 |
+
}
|
92 |
+
|
93 |
+
.styler {
|
94 |
+
overflow:inherit !important;
|
95 |
+
}
|
96 |
+
|
97 |
+
.gradio-container{
|
98 |
+
overflow: visible;
|
99 |
+
}
|
100 |
+
|
101 |
+
/* fullpage image viewer */
|
102 |
+
|
103 |
+
#lightboxModal{
|
104 |
+
display: none;
|
105 |
+
position: fixed;
|
106 |
+
z-index: 1001;
|
107 |
+
left: 0;
|
108 |
+
top: 0;
|
109 |
+
width: 100%;
|
110 |
+
height: 100%;
|
111 |
+
overflow: auto;
|
112 |
+
background-color: rgba(20, 20, 20, 0.95);
|
113 |
+
user-select: none;
|
114 |
+
-webkit-user-select: none;
|
115 |
+
flex-direction: column;
|
116 |
+
}
|
117 |
+
|
118 |
+
.modalControls {
|
119 |
+
display: flex;
|
120 |
+
position: absolute;
|
121 |
+
right: 0px;
|
122 |
+
left: 0px;
|
123 |
+
gap: 1em;
|
124 |
+
padding: 1em;
|
125 |
+
background-color:rgba(0,0,0,0);
|
126 |
+
z-index: 1;
|
127 |
+
transition: 0.2s ease background-color;
|
128 |
+
}
|
129 |
+
.modalControls:hover {
|
130 |
+
background-color:rgba(0,0,0,0.9);
|
131 |
+
}
|
132 |
+
.modalClose {
|
133 |
+
margin-left: auto;
|
134 |
+
}
|
135 |
+
.modalControls span{
|
136 |
+
color: white;
|
137 |
+
text-shadow: 0px 0px 0.25em black;
|
138 |
+
font-size: 35px;
|
139 |
+
font-weight: bold;
|
140 |
+
cursor: pointer;
|
141 |
+
width: 1em;
|
142 |
+
}
|
143 |
+
|
144 |
+
.modalControls span:hover, .modalControls span:focus{
|
145 |
+
color: #999;
|
146 |
+
text-decoration: none;
|
147 |
+
}
|
148 |
+
|
149 |
+
#lightboxModal > img {
|
150 |
+
display: block;
|
151 |
+
margin: auto;
|
152 |
+
width: auto;
|
153 |
+
}
|
154 |
+
|
155 |
+
#lightboxModal > img.modalImageFullscreen{
|
156 |
+
object-fit: contain;
|
157 |
+
height: 100%;
|
158 |
+
width: 100%;
|
159 |
+
min-height: 0;
|
160 |
+
}
|
161 |
+
|
162 |
+
.modalPrev,
|
163 |
+
.modalNext {
|
164 |
+
cursor: pointer;
|
165 |
+
position: absolute;
|
166 |
+
top: 50%;
|
167 |
+
width: auto;
|
168 |
+
padding: 16px;
|
169 |
+
margin-top: -50px;
|
170 |
+
color: white;
|
171 |
+
font-weight: bold;
|
172 |
+
font-size: 20px;
|
173 |
+
transition: 0.6s ease;
|
174 |
+
border-radius: 0 3px 3px 0;
|
175 |
+
user-select: none;
|
176 |
+
-webkit-user-select: none;
|
177 |
+
}
|
178 |
+
|
179 |
+
.modalNext {
|
180 |
+
right: 0;
|
181 |
+
border-radius: 3px 0 0 3px;
|
182 |
+
}
|
183 |
+
|
184 |
+
.modalPrev:hover,
|
185 |
+
.modalNext:hover {
|
186 |
+
background-color: rgba(0, 0, 0, 0.8);
|
187 |
+
}
|
188 |
+
|
189 |
+
#imageARPreview {
|
190 |
+
position: absolute;
|
191 |
+
top: 0px;
|
192 |
+
left: 0px;
|
193 |
+
border: 2px solid red;
|
194 |
+
background: rgba(255, 0, 0, 0.3);
|
195 |
+
z-index: 900;
|
196 |
+
pointer-events: none;
|
197 |
+
display: none;
|
198 |
+
}
|
entry_with_update.py
ADDED
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import sys
|
3 |
+
|
4 |
+
|
5 |
+
root = os.path.dirname(os.path.abspath(__file__))
|
6 |
+
sys.path.append(root)
|
7 |
+
os.chdir(root)
|
8 |
+
|
9 |
+
|
10 |
+
try:
|
11 |
+
import pygit2
|
12 |
+
pygit2.option(pygit2.GIT_OPT_SET_OWNER_VALIDATION, 0)
|
13 |
+
|
14 |
+
repo = pygit2.Repository(os.path.abspath(os.path.dirname(__file__)))
|
15 |
+
|
16 |
+
branch_name = repo.head.shorthand
|
17 |
+
|
18 |
+
remote_name = 'origin'
|
19 |
+
remote = repo.remotes[remote_name]
|
20 |
+
|
21 |
+
remote.fetch()
|
22 |
+
|
23 |
+
local_branch_ref = f'refs/heads/{branch_name}'
|
24 |
+
local_branch = repo.lookup_reference(local_branch_ref)
|
25 |
+
|
26 |
+
remote_reference = f'refs/remotes/{remote_name}/{branch_name}'
|
27 |
+
remote_commit = repo.revparse_single(remote_reference)
|
28 |
+
|
29 |
+
merge_result, _ = repo.merge_analysis(remote_commit.id)
|
30 |
+
|
31 |
+
if merge_result & pygit2.GIT_MERGE_ANALYSIS_UP_TO_DATE:
|
32 |
+
print("Already up-to-date")
|
33 |
+
elif merge_result & pygit2.GIT_MERGE_ANALYSIS_FASTFORWARD:
|
34 |
+
local_branch.set_target(remote_commit.id)
|
35 |
+
repo.head.set_target(remote_commit.id)
|
36 |
+
repo.checkout_tree(repo.get(remote_commit.id))
|
37 |
+
repo.reset(local_branch.target, pygit2.GIT_RESET_HARD)
|
38 |
+
print("Fast-forward merge")
|
39 |
+
elif merge_result & pygit2.GIT_MERGE_ANALYSIS_NORMAL:
|
40 |
+
print("Update failed - Did you modify any file?")
|
41 |
+
except Exception as e:
|
42 |
+
print('Update failed.')
|
43 |
+
print(str(e))
|
44 |
+
|
45 |
+
print('Update succeeded.')
|
46 |
+
from launch import *
|
environment.yaml
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
name: fooocus
|
2 |
+
channels:
|
3 |
+
- defaults
|
4 |
+
dependencies:
|
5 |
+
- python=3.10
|
6 |
+
- pip=23.0
|
7 |
+
- packaging
|
experiments_expansion.py
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from modules.expansion import FooocusExpansion
|
2 |
+
|
3 |
+
expansion = FooocusExpansion()
|
4 |
+
|
5 |
+
text = 'a handsome man'
|
6 |
+
|
7 |
+
for i in range(64):
|
8 |
+
print(expansion(text, seed=i))
|
experiments_face.py
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import extras.face_crop as cropper
|
3 |
+
|
4 |
+
|
5 |
+
img = cv2.imread('lena.png')
|
6 |
+
result = cropper.crop_image(img)
|
7 |
+
cv2.imwrite('lena_result.png', result)
|
experiments_interrogate.py
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
from extras.interrogate import default_interrogator as default_interrogator_photo
|
3 |
+
from extras.wd14tagger import default_interrogator as default_interrogator_anime
|
4 |
+
|
5 |
+
img = cv2.imread('./test_imgs/red_box.jpg')[:, :, ::-1].copy()
|
6 |
+
print(default_interrogator_photo(img))
|
7 |
+
img = cv2.imread('./test_imgs/miku.jpg')[:, :, ::-1].copy()
|
8 |
+
print(default_interrogator_anime(img))
|
extras/BLIP/configs/bert_config.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"BertModel"
|
4 |
+
],
|
5 |
+
"attention_probs_dropout_prob": 0.1,
|
6 |
+
"hidden_act": "gelu",
|
7 |
+
"hidden_dropout_prob": 0.1,
|
8 |
+
"hidden_size": 768,
|
9 |
+
"initializer_range": 0.02,
|
10 |
+
"intermediate_size": 3072,
|
11 |
+
"layer_norm_eps": 1e-12,
|
12 |
+
"max_position_embeddings": 512,
|
13 |
+
"model_type": "bert",
|
14 |
+
"num_attention_heads": 12,
|
15 |
+
"num_hidden_layers": 12,
|
16 |
+
"pad_token_id": 0,
|
17 |
+
"type_vocab_size": 2,
|
18 |
+
"vocab_size": 30522,
|
19 |
+
"encoder_width": 768,
|
20 |
+
"add_cross_attention": true
|
21 |
+
}
|
extras/BLIP/configs/caption_coco.yaml
ADDED
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
image_root: '/export/share/datasets/vision/coco/images/'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
coco_gt_root: 'annotation/coco_gt'
|
4 |
+
|
5 |
+
# set pretrained as a file path or an url
|
6 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
|
7 |
+
|
8 |
+
# size of vit model; base or large
|
9 |
+
vit: 'base'
|
10 |
+
vit_grad_ckpt: False
|
11 |
+
vit_ckpt_layer: 0
|
12 |
+
batch_size: 32
|
13 |
+
init_lr: 1e-5
|
14 |
+
|
15 |
+
# vit: 'large'
|
16 |
+
# vit_grad_ckpt: True
|
17 |
+
# vit_ckpt_layer: 5
|
18 |
+
# batch_size: 16
|
19 |
+
# init_lr: 2e-6
|
20 |
+
|
21 |
+
image_size: 384
|
22 |
+
|
23 |
+
# generation configs
|
24 |
+
max_length: 20
|
25 |
+
min_length: 5
|
26 |
+
num_beams: 3
|
27 |
+
prompt: 'a picture of '
|
28 |
+
|
29 |
+
# optimizer
|
30 |
+
weight_decay: 0.05
|
31 |
+
min_lr: 0
|
32 |
+
max_epoch: 5
|
33 |
+
|
extras/BLIP/configs/med_config.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"BertModel"
|
4 |
+
],
|
5 |
+
"attention_probs_dropout_prob": 0.1,
|
6 |
+
"hidden_act": "gelu",
|
7 |
+
"hidden_dropout_prob": 0.1,
|
8 |
+
"hidden_size": 768,
|
9 |
+
"initializer_range": 0.02,
|
10 |
+
"intermediate_size": 3072,
|
11 |
+
"layer_norm_eps": 1e-12,
|
12 |
+
"max_position_embeddings": 512,
|
13 |
+
"model_type": "bert",
|
14 |
+
"num_attention_heads": 12,
|
15 |
+
"num_hidden_layers": 12,
|
16 |
+
"pad_token_id": 0,
|
17 |
+
"type_vocab_size": 2,
|
18 |
+
"vocab_size": 30524,
|
19 |
+
"encoder_width": 768,
|
20 |
+
"add_cross_attention": true
|
21 |
+
}
|
extras/BLIP/configs/nlvr.yaml
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
image_root: '/export/share/datasets/vision/NLVR2/'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
|
4 |
+
# set pretrained as a file path or an url
|
5 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_nlvr.pth'
|
6 |
+
|
7 |
+
#size of vit model; base or large
|
8 |
+
vit: 'base'
|
9 |
+
batch_size_train: 16
|
10 |
+
batch_size_test: 64
|
11 |
+
vit_grad_ckpt: False
|
12 |
+
vit_ckpt_layer: 0
|
13 |
+
max_epoch: 15
|
14 |
+
|
15 |
+
image_size: 384
|
16 |
+
|
17 |
+
# optimizer
|
18 |
+
weight_decay: 0.05
|
19 |
+
init_lr: 3e-5
|
20 |
+
min_lr: 0
|
21 |
+
|
extras/BLIP/configs/nocaps.yaml
ADDED
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
image_root: '/export/share/datasets/vision/nocaps/'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
|
4 |
+
# set pretrained as a file path or an url
|
5 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth'
|
6 |
+
|
7 |
+
vit: 'base'
|
8 |
+
batch_size: 32
|
9 |
+
|
10 |
+
image_size: 384
|
11 |
+
|
12 |
+
max_length: 20
|
13 |
+
min_length: 5
|
14 |
+
num_beams: 3
|
15 |
+
prompt: 'a picture of '
|
extras/BLIP/configs/pretrain.yaml
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
train_file: ['/export/share/junnan-li/VL_pretrain/annotation/coco_karpathy_train.json',
|
2 |
+
'/export/share/junnan-li/VL_pretrain/annotation/vg_caption.json',
|
3 |
+
]
|
4 |
+
laion_path: ''
|
5 |
+
|
6 |
+
# size of vit model; base or large
|
7 |
+
vit: 'base'
|
8 |
+
vit_grad_ckpt: False
|
9 |
+
vit_ckpt_layer: 0
|
10 |
+
|
11 |
+
image_size: 224
|
12 |
+
batch_size: 75
|
13 |
+
|
14 |
+
queue_size: 57600
|
15 |
+
alpha: 0.4
|
16 |
+
|
17 |
+
# optimizer
|
18 |
+
weight_decay: 0.05
|
19 |
+
init_lr: 3e-4
|
20 |
+
min_lr: 1e-6
|
21 |
+
warmup_lr: 1e-6
|
22 |
+
lr_decay_rate: 0.9
|
23 |
+
max_epoch: 20
|
24 |
+
warmup_steps: 3000
|
25 |
+
|
26 |
+
|
27 |
+
|
extras/BLIP/configs/retrieval_coco.yaml
ADDED
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
image_root: '/export/share/datasets/vision/coco/images/'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
dataset: 'coco'
|
4 |
+
|
5 |
+
# set pretrained as a file path or an url
|
6 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'
|
7 |
+
|
8 |
+
# size of vit model; base or large
|
9 |
+
|
10 |
+
vit: 'base'
|
11 |
+
batch_size_train: 32
|
12 |
+
batch_size_test: 64
|
13 |
+
vit_grad_ckpt: True
|
14 |
+
vit_ckpt_layer: 4
|
15 |
+
init_lr: 1e-5
|
16 |
+
|
17 |
+
# vit: 'large'
|
18 |
+
# batch_size_train: 16
|
19 |
+
# batch_size_test: 32
|
20 |
+
# vit_grad_ckpt: True
|
21 |
+
# vit_ckpt_layer: 12
|
22 |
+
# init_lr: 5e-6
|
23 |
+
|
24 |
+
image_size: 384
|
25 |
+
queue_size: 57600
|
26 |
+
alpha: 0.4
|
27 |
+
k_test: 256
|
28 |
+
negative_all_rank: True
|
29 |
+
|
30 |
+
# optimizer
|
31 |
+
weight_decay: 0.05
|
32 |
+
min_lr: 0
|
33 |
+
max_epoch: 6
|
34 |
+
|
extras/BLIP/configs/retrieval_flickr.yaml
ADDED
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
image_root: '/export/share/datasets/vision/flickr30k/'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
dataset: 'flickr'
|
4 |
+
|
5 |
+
# set pretrained as a file path or an url
|
6 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_flickr.pth'
|
7 |
+
|
8 |
+
# size of vit model; base or large
|
9 |
+
|
10 |
+
vit: 'base'
|
11 |
+
batch_size_train: 32
|
12 |
+
batch_size_test: 64
|
13 |
+
vit_grad_ckpt: True
|
14 |
+
vit_ckpt_layer: 4
|
15 |
+
init_lr: 1e-5
|
16 |
+
|
17 |
+
# vit: 'large'
|
18 |
+
# batch_size_train: 16
|
19 |
+
# batch_size_test: 32
|
20 |
+
# vit_grad_ckpt: True
|
21 |
+
# vit_ckpt_layer: 10
|
22 |
+
# init_lr: 5e-6
|
23 |
+
|
24 |
+
image_size: 384
|
25 |
+
queue_size: 57600
|
26 |
+
alpha: 0.4
|
27 |
+
k_test: 128
|
28 |
+
negative_all_rank: False
|
29 |
+
|
30 |
+
# optimizer
|
31 |
+
weight_decay: 0.05
|
32 |
+
min_lr: 0
|
33 |
+
max_epoch: 6
|
34 |
+
|
extras/BLIP/configs/retrieval_msrvtt.yaml
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
video_root: '/export/share/dongxuli/data/msrvtt_retrieval/videos'
|
2 |
+
ann_root: 'annotation'
|
3 |
+
|
4 |
+
# set pretrained as a file path or an url
|
5 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'
|
6 |
+
|
7 |
+
# size of vit model; base or large
|
8 |
+
vit: 'base'
|
9 |
+
batch_size: 64
|
10 |
+
k_test: 128
|
11 |
+
image_size: 384
|
12 |
+
num_frm_test: 8
|
extras/BLIP/configs/vqa.yaml
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
vqa_root: '/export/share/datasets/vision/VQA/Images/mscoco/' #followed by train2014/
|
2 |
+
vg_root: '/export/share/datasets/vision/visual-genome/' #followed by image/
|
3 |
+
train_files: ['vqa_train','vqa_val','vg_qa']
|
4 |
+
ann_root: 'annotation'
|
5 |
+
|
6 |
+
# set pretrained as a file path or an url
|
7 |
+
pretrained: 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'
|
8 |
+
|
9 |
+
# size of vit model; base or large
|
10 |
+
vit: 'base'
|
11 |
+
batch_size_train: 16
|
12 |
+
batch_size_test: 32
|
13 |
+
vit_grad_ckpt: False
|
14 |
+
vit_ckpt_layer: 0
|
15 |
+
init_lr: 2e-5
|
16 |
+
|
17 |
+
image_size: 480
|
18 |
+
|
19 |
+
k_test: 128
|
20 |
+
inference: 'rank'
|
21 |
+
|
22 |
+
# optimizer
|
23 |
+
weight_decay: 0.05
|
24 |
+
min_lr: 0
|
25 |
+
max_epoch: 10
|
extras/BLIP/models/bert_tokenizer/config.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"architectures": [
|
3 |
+
"BertForMaskedLM"
|
4 |
+
],
|
5 |
+
"attention_probs_dropout_prob": 0.1,
|
6 |
+
"gradient_checkpointing": false,
|
7 |
+
"hidden_act": "gelu",
|
8 |
+
"hidden_dropout_prob": 0.1,
|
9 |
+
"hidden_size": 768,
|
10 |
+
"initializer_range": 0.02,
|
11 |
+
"intermediate_size": 3072,
|
12 |
+
"layer_norm_eps": 1e-12,
|
13 |
+
"max_position_embeddings": 512,
|
14 |
+
"model_type": "bert",
|
15 |
+
"num_attention_heads": 12,
|
16 |
+
"num_hidden_layers": 12,
|
17 |
+
"pad_token_id": 0,
|
18 |
+
"position_embedding_type": "absolute",
|
19 |
+
"transformers_version": "4.6.0.dev0",
|
20 |
+
"type_vocab_size": 2,
|
21 |
+
"use_cache": true,
|
22 |
+
"vocab_size": 30522
|
23 |
+
}
|
extras/BLIP/models/bert_tokenizer/tokenizer.json
ADDED
The diff for this file is too large to render.
See raw diff
|
|
extras/BLIP/models/bert_tokenizer/tokenizer_config.json
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"do_lower_case": true
|
3 |
+
}
|
extras/BLIP/models/bert_tokenizer/vocab.txt
ADDED
The diff for this file is too large to render.
See raw diff
|
|
extras/BLIP/models/blip.py
ADDED
@@ -0,0 +1,239 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
* Copyright (c) 2022, salesforce.com, inc.
|
3 |
+
* All rights reserved.
|
4 |
+
* SPDX-License-Identifier: BSD-3-Clause
|
5 |
+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
6 |
+
* By Junnan Li
|
7 |
+
'''
|
8 |
+
import warnings
|
9 |
+
warnings.filterwarnings("ignore")
|
10 |
+
|
11 |
+
from extras.BLIP.models.vit import VisionTransformer, interpolate_pos_embed
|
12 |
+
from extras.BLIP.models.med import BertConfig, BertModel, BertLMHeadModel
|
13 |
+
from transformers import BertTokenizer
|
14 |
+
|
15 |
+
import torch
|
16 |
+
from torch import nn
|
17 |
+
import torch.nn.functional as F
|
18 |
+
|
19 |
+
import os
|
20 |
+
from urllib.parse import urlparse
|
21 |
+
from timm.models.hub import download_cached_file
|
22 |
+
|
23 |
+
class BLIP_Base(nn.Module):
|
24 |
+
def __init__(self,
|
25 |
+
med_config = 'configs/med_config.json',
|
26 |
+
image_size = 224,
|
27 |
+
vit = 'base',
|
28 |
+
vit_grad_ckpt = False,
|
29 |
+
vit_ckpt_layer = 0,
|
30 |
+
):
|
31 |
+
"""
|
32 |
+
Args:
|
33 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
34 |
+
image_size (int): input image size
|
35 |
+
vit (str): model size of vision transformer
|
36 |
+
"""
|
37 |
+
super().__init__()
|
38 |
+
|
39 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
|
40 |
+
self.tokenizer = init_tokenizer()
|
41 |
+
med_config = BertConfig.from_json_file(med_config)
|
42 |
+
med_config.encoder_width = vision_width
|
43 |
+
self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)
|
44 |
+
|
45 |
+
|
46 |
+
def forward(self, image, caption, mode):
|
47 |
+
|
48 |
+
assert mode in ['image', 'text', 'multimodal'], "mode parameter must be image, text, or multimodal"
|
49 |
+
text = self.tokenizer(caption, return_tensors="pt").to(image.device)
|
50 |
+
|
51 |
+
if mode=='image':
|
52 |
+
# return image features
|
53 |
+
image_embeds = self.visual_encoder(image)
|
54 |
+
return image_embeds
|
55 |
+
|
56 |
+
elif mode=='text':
|
57 |
+
# return text features
|
58 |
+
text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,
|
59 |
+
return_dict = True, mode = 'text')
|
60 |
+
return text_output.last_hidden_state
|
61 |
+
|
62 |
+
elif mode=='multimodal':
|
63 |
+
# return multimodel features
|
64 |
+
image_embeds = self.visual_encoder(image)
|
65 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
66 |
+
|
67 |
+
text.input_ids[:,0] = self.tokenizer.enc_token_id
|
68 |
+
output = self.text_encoder(text.input_ids,
|
69 |
+
attention_mask = text.attention_mask,
|
70 |
+
encoder_hidden_states = image_embeds,
|
71 |
+
encoder_attention_mask = image_atts,
|
72 |
+
return_dict = True,
|
73 |
+
)
|
74 |
+
return output.last_hidden_state
|
75 |
+
|
76 |
+
|
77 |
+
|
78 |
+
class BLIP_Decoder(nn.Module):
|
79 |
+
def __init__(self,
|
80 |
+
med_config = 'configs/med_config.json',
|
81 |
+
image_size = 384,
|
82 |
+
vit = 'base',
|
83 |
+
vit_grad_ckpt = False,
|
84 |
+
vit_ckpt_layer = 0,
|
85 |
+
prompt = 'a picture of ',
|
86 |
+
):
|
87 |
+
"""
|
88 |
+
Args:
|
89 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
90 |
+
image_size (int): input image size
|
91 |
+
vit (str): model size of vision transformer
|
92 |
+
"""
|
93 |
+
super().__init__()
|
94 |
+
|
95 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
|
96 |
+
self.tokenizer = init_tokenizer()
|
97 |
+
med_config = BertConfig.from_json_file(med_config)
|
98 |
+
med_config.encoder_width = vision_width
|
99 |
+
self.text_decoder = BertLMHeadModel(config=med_config)
|
100 |
+
|
101 |
+
self.prompt = prompt
|
102 |
+
self.prompt_length = len(self.tokenizer(self.prompt).input_ids)-1
|
103 |
+
|
104 |
+
|
105 |
+
def forward(self, image, caption):
|
106 |
+
|
107 |
+
image_embeds = self.visual_encoder(image)
|
108 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
109 |
+
|
110 |
+
text = self.tokenizer(caption, padding='longest', truncation=True, max_length=40, return_tensors="pt").to(image.device)
|
111 |
+
|
112 |
+
text.input_ids[:,0] = self.tokenizer.bos_token_id
|
113 |
+
|
114 |
+
decoder_targets = text.input_ids.masked_fill(text.input_ids == self.tokenizer.pad_token_id, -100)
|
115 |
+
decoder_targets[:,:self.prompt_length] = -100
|
116 |
+
|
117 |
+
decoder_output = self.text_decoder(text.input_ids,
|
118 |
+
attention_mask = text.attention_mask,
|
119 |
+
encoder_hidden_states = image_embeds,
|
120 |
+
encoder_attention_mask = image_atts,
|
121 |
+
labels = decoder_targets,
|
122 |
+
return_dict = True,
|
123 |
+
)
|
124 |
+
loss_lm = decoder_output.loss
|
125 |
+
|
126 |
+
return loss_lm
|
127 |
+
|
128 |
+
def generate(self, image, sample=False, num_beams=3, max_length=30, min_length=10, top_p=0.9, repetition_penalty=1.0):
|
129 |
+
image_embeds = self.visual_encoder(image)
|
130 |
+
|
131 |
+
if not sample:
|
132 |
+
image_embeds = image_embeds.repeat_interleave(num_beams,dim=0)
|
133 |
+
|
134 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
135 |
+
model_kwargs = {"encoder_hidden_states": image_embeds, "encoder_attention_mask":image_atts}
|
136 |
+
|
137 |
+
prompt = [self.prompt] * image.size(0)
|
138 |
+
input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(image.device)
|
139 |
+
input_ids[:,0] = self.tokenizer.bos_token_id
|
140 |
+
input_ids = input_ids[:, :-1]
|
141 |
+
|
142 |
+
if sample:
|
143 |
+
#nucleus sampling
|
144 |
+
outputs = self.text_decoder.generate(input_ids=input_ids,
|
145 |
+
max_length=max_length,
|
146 |
+
min_length=min_length,
|
147 |
+
do_sample=True,
|
148 |
+
top_p=top_p,
|
149 |
+
num_return_sequences=1,
|
150 |
+
eos_token_id=self.tokenizer.sep_token_id,
|
151 |
+
pad_token_id=self.tokenizer.pad_token_id,
|
152 |
+
repetition_penalty=1.1,
|
153 |
+
**model_kwargs)
|
154 |
+
else:
|
155 |
+
#beam search
|
156 |
+
outputs = self.text_decoder.generate(input_ids=input_ids,
|
157 |
+
max_length=max_length,
|
158 |
+
min_length=min_length,
|
159 |
+
num_beams=num_beams,
|
160 |
+
eos_token_id=self.tokenizer.sep_token_id,
|
161 |
+
pad_token_id=self.tokenizer.pad_token_id,
|
162 |
+
repetition_penalty=repetition_penalty,
|
163 |
+
**model_kwargs)
|
164 |
+
|
165 |
+
captions = []
|
166 |
+
for output in outputs:
|
167 |
+
caption = self.tokenizer.decode(output, skip_special_tokens=True)
|
168 |
+
captions.append(caption[len(self.prompt):])
|
169 |
+
return captions
|
170 |
+
|
171 |
+
|
172 |
+
def blip_decoder(pretrained='',**kwargs):
|
173 |
+
model = BLIP_Decoder(**kwargs)
|
174 |
+
if pretrained:
|
175 |
+
model,msg = load_checkpoint(model,pretrained)
|
176 |
+
assert(len(msg.missing_keys)==0)
|
177 |
+
return model
|
178 |
+
|
179 |
+
def blip_feature_extractor(pretrained='',**kwargs):
|
180 |
+
model = BLIP_Base(**kwargs)
|
181 |
+
if pretrained:
|
182 |
+
model,msg = load_checkpoint(model,pretrained)
|
183 |
+
assert(len(msg.missing_keys)==0)
|
184 |
+
return model
|
185 |
+
|
186 |
+
def init_tokenizer():
|
187 |
+
tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "bert_tokenizer")
|
188 |
+
tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
|
189 |
+
tokenizer.add_special_tokens({'bos_token':'[DEC]'})
|
190 |
+
tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})
|
191 |
+
tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]
|
192 |
+
return tokenizer
|
193 |
+
|
194 |
+
|
195 |
+
def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
|
196 |
+
|
197 |
+
assert vit in ['base', 'large'], "vit parameter must be base or large"
|
198 |
+
if vit=='base':
|
199 |
+
vision_width = 768
|
200 |
+
visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12,
|
201 |
+
num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
|
202 |
+
drop_path_rate=0 or drop_path_rate
|
203 |
+
)
|
204 |
+
elif vit=='large':
|
205 |
+
vision_width = 1024
|
206 |
+
visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24,
|
207 |
+
num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
|
208 |
+
drop_path_rate=0.1 or drop_path_rate
|
209 |
+
)
|
210 |
+
return visual_encoder, vision_width
|
211 |
+
|
212 |
+
def is_url(url_or_filename):
|
213 |
+
parsed = urlparse(url_or_filename)
|
214 |
+
return parsed.scheme in ("http", "https")
|
215 |
+
|
216 |
+
def load_checkpoint(model,url_or_filename):
|
217 |
+
if is_url(url_or_filename):
|
218 |
+
cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
|
219 |
+
checkpoint = torch.load(cached_file, map_location='cpu')
|
220 |
+
elif os.path.isfile(url_or_filename):
|
221 |
+
checkpoint = torch.load(url_or_filename, map_location='cpu')
|
222 |
+
else:
|
223 |
+
raise RuntimeError('checkpoint url or path is invalid')
|
224 |
+
|
225 |
+
state_dict = checkpoint['model']
|
226 |
+
|
227 |
+
state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder)
|
228 |
+
if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
|
229 |
+
state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
|
230 |
+
model.visual_encoder_m)
|
231 |
+
for key in model.state_dict().keys():
|
232 |
+
if key in state_dict.keys():
|
233 |
+
if state_dict[key].shape!=model.state_dict()[key].shape:
|
234 |
+
del state_dict[key]
|
235 |
+
|
236 |
+
msg = model.load_state_dict(state_dict,strict=False)
|
237 |
+
print('load checkpoint from %s'%url_or_filename)
|
238 |
+
return model,msg
|
239 |
+
|
extras/BLIP/models/blip_itm.py
ADDED
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from extras.BLIP.models.med import BertConfig, BertModel
|
2 |
+
from transformers import BertTokenizer
|
3 |
+
|
4 |
+
import torch
|
5 |
+
from torch import nn
|
6 |
+
import torch.nn.functional as F
|
7 |
+
|
8 |
+
from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
|
9 |
+
|
10 |
+
class BLIP_ITM(nn.Module):
|
11 |
+
def __init__(self,
|
12 |
+
med_config = 'configs/med_config.json',
|
13 |
+
image_size = 384,
|
14 |
+
vit = 'base',
|
15 |
+
vit_grad_ckpt = False,
|
16 |
+
vit_ckpt_layer = 0,
|
17 |
+
embed_dim = 256,
|
18 |
+
):
|
19 |
+
"""
|
20 |
+
Args:
|
21 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
22 |
+
image_size (int): input image size
|
23 |
+
vit (str): model size of vision transformer
|
24 |
+
"""
|
25 |
+
super().__init__()
|
26 |
+
|
27 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
|
28 |
+
self.tokenizer = init_tokenizer()
|
29 |
+
med_config = BertConfig.from_json_file(med_config)
|
30 |
+
med_config.encoder_width = vision_width
|
31 |
+
self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)
|
32 |
+
|
33 |
+
text_width = self.text_encoder.config.hidden_size
|
34 |
+
|
35 |
+
self.vision_proj = nn.Linear(vision_width, embed_dim)
|
36 |
+
self.text_proj = nn.Linear(text_width, embed_dim)
|
37 |
+
|
38 |
+
self.itm_head = nn.Linear(text_width, 2)
|
39 |
+
|
40 |
+
|
41 |
+
def forward(self, image, caption, match_head='itm'):
|
42 |
+
|
43 |
+
image_embeds = self.visual_encoder(image)
|
44 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
45 |
+
|
46 |
+
text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=35,
|
47 |
+
return_tensors="pt").to(image.device)
|
48 |
+
|
49 |
+
|
50 |
+
if match_head=='itm':
|
51 |
+
output = self.text_encoder(text.input_ids,
|
52 |
+
attention_mask = text.attention_mask,
|
53 |
+
encoder_hidden_states = image_embeds,
|
54 |
+
encoder_attention_mask = image_atts,
|
55 |
+
return_dict = True,
|
56 |
+
)
|
57 |
+
itm_output = self.itm_head(output.last_hidden_state[:,0,:])
|
58 |
+
return itm_output
|
59 |
+
|
60 |
+
elif match_head=='itc':
|
61 |
+
text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,
|
62 |
+
return_dict = True, mode = 'text')
|
63 |
+
image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)
|
64 |
+
text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)
|
65 |
+
|
66 |
+
sim = image_feat @ text_feat.t()
|
67 |
+
return sim
|
68 |
+
|
69 |
+
|
70 |
+
def blip_itm(pretrained='',**kwargs):
|
71 |
+
model = BLIP_ITM(**kwargs)
|
72 |
+
if pretrained:
|
73 |
+
model,msg = load_checkpoint(model,pretrained)
|
74 |
+
assert(len(msg.missing_keys)==0)
|
75 |
+
return model
|
76 |
+
|
extras/BLIP/models/blip_nlvr.py
ADDED
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from extras.BLIP.models.med import BertConfig
|
2 |
+
from extras.BLIP.models.nlvr_encoder import BertModel
|
3 |
+
from extras.BLIP.models.vit import interpolate_pos_embed
|
4 |
+
from extras.BLIP.models.blip import create_vit, init_tokenizer, is_url
|
5 |
+
|
6 |
+
from timm.models.hub import download_cached_file
|
7 |
+
|
8 |
+
import torch
|
9 |
+
from torch import nn
|
10 |
+
import torch.nn.functional as F
|
11 |
+
from transformers import BertTokenizer
|
12 |
+
import numpy as np
|
13 |
+
import os
|
14 |
+
|
15 |
+
|
16 |
+
class BLIP_NLVR(nn.Module):
|
17 |
+
def __init__(self,
|
18 |
+
med_config = 'configs/med_config.json',
|
19 |
+
image_size = 480,
|
20 |
+
vit = 'base',
|
21 |
+
vit_grad_ckpt = False,
|
22 |
+
vit_ckpt_layer = 0,
|
23 |
+
):
|
24 |
+
"""
|
25 |
+
Args:
|
26 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
27 |
+
image_size (int): input image size
|
28 |
+
vit (str): model size of vision transformer
|
29 |
+
"""
|
30 |
+
super().__init__()
|
31 |
+
|
32 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, drop_path_rate=0.1)
|
33 |
+
self.tokenizer = init_tokenizer()
|
34 |
+
med_config = BertConfig.from_json_file(med_config)
|
35 |
+
med_config.encoder_width = vision_width
|
36 |
+
self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)
|
37 |
+
|
38 |
+
self.cls_head = nn.Sequential(
|
39 |
+
nn.Linear(self.text_encoder.config.hidden_size, self.text_encoder.config.hidden_size),
|
40 |
+
nn.ReLU(),
|
41 |
+
nn.Linear(self.text_encoder.config.hidden_size, 2)
|
42 |
+
)
|
43 |
+
|
44 |
+
def forward(self, image, text, targets, train=True):
|
45 |
+
|
46 |
+
image_embeds = self.visual_encoder(image)
|
47 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
48 |
+
image0_embeds, image1_embeds = torch.split(image_embeds,targets.size(0))
|
49 |
+
|
50 |
+
text = self.tokenizer(text, padding='longest', return_tensors="pt").to(image.device)
|
51 |
+
text.input_ids[:,0] = self.tokenizer.enc_token_id
|
52 |
+
|
53 |
+
output = self.text_encoder(text.input_ids,
|
54 |
+
attention_mask = text.attention_mask,
|
55 |
+
encoder_hidden_states = [image0_embeds,image1_embeds],
|
56 |
+
encoder_attention_mask = [image_atts[:image0_embeds.size(0)],
|
57 |
+
image_atts[image0_embeds.size(0):]],
|
58 |
+
return_dict = True,
|
59 |
+
)
|
60 |
+
hidden_state = output.last_hidden_state[:,0,:]
|
61 |
+
prediction = self.cls_head(hidden_state)
|
62 |
+
|
63 |
+
if train:
|
64 |
+
loss = F.cross_entropy(prediction, targets)
|
65 |
+
return loss
|
66 |
+
else:
|
67 |
+
return prediction
|
68 |
+
|
69 |
+
def blip_nlvr(pretrained='',**kwargs):
|
70 |
+
model = BLIP_NLVR(**kwargs)
|
71 |
+
if pretrained:
|
72 |
+
model,msg = load_checkpoint(model,pretrained)
|
73 |
+
print("missing keys:")
|
74 |
+
print(msg.missing_keys)
|
75 |
+
return model
|
76 |
+
|
77 |
+
|
78 |
+
def load_checkpoint(model,url_or_filename):
|
79 |
+
if is_url(url_or_filename):
|
80 |
+
cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
|
81 |
+
checkpoint = torch.load(cached_file, map_location='cpu')
|
82 |
+
elif os.path.isfile(url_or_filename):
|
83 |
+
checkpoint = torch.load(url_or_filename, map_location='cpu')
|
84 |
+
else:
|
85 |
+
raise RuntimeError('checkpoint url or path is invalid')
|
86 |
+
state_dict = checkpoint['model']
|
87 |
+
|
88 |
+
state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder)
|
89 |
+
|
90 |
+
for key in list(state_dict.keys()):
|
91 |
+
if 'crossattention.self.' in key:
|
92 |
+
new_key0 = key.replace('self','self0')
|
93 |
+
new_key1 = key.replace('self','self1')
|
94 |
+
state_dict[new_key0] = state_dict[key]
|
95 |
+
state_dict[new_key1] = state_dict[key]
|
96 |
+
elif 'crossattention.output.dense.' in key:
|
97 |
+
new_key0 = key.replace('dense','dense0')
|
98 |
+
new_key1 = key.replace('dense','dense1')
|
99 |
+
state_dict[new_key0] = state_dict[key]
|
100 |
+
state_dict[new_key1] = state_dict[key]
|
101 |
+
|
102 |
+
msg = model.load_state_dict(state_dict,strict=False)
|
103 |
+
print('load checkpoint from %s'%url_or_filename)
|
104 |
+
return model,msg
|
105 |
+
|
extras/BLIP/models/blip_pretrain.py
ADDED
@@ -0,0 +1,339 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
* Copyright (c) 2022, salesforce.com, inc.
|
3 |
+
* All rights reserved.
|
4 |
+
* SPDX-License-Identifier: BSD-3-Clause
|
5 |
+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
6 |
+
* By Junnan Li
|
7 |
+
'''
|
8 |
+
from extras.BLIP.models.med import BertConfig, BertModel, BertLMHeadModel
|
9 |
+
from transformers import BertTokenizer
|
10 |
+
import transformers
|
11 |
+
transformers.logging.set_verbosity_error()
|
12 |
+
|
13 |
+
import torch
|
14 |
+
from torch import nn
|
15 |
+
import torch.nn.functional as F
|
16 |
+
|
17 |
+
from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
|
18 |
+
|
19 |
+
class BLIP_Pretrain(nn.Module):
|
20 |
+
def __init__(self,
|
21 |
+
med_config = 'configs/bert_config.json',
|
22 |
+
image_size = 224,
|
23 |
+
vit = 'base',
|
24 |
+
vit_grad_ckpt = False,
|
25 |
+
vit_ckpt_layer = 0,
|
26 |
+
embed_dim = 256,
|
27 |
+
queue_size = 57600,
|
28 |
+
momentum = 0.995,
|
29 |
+
):
|
30 |
+
"""
|
31 |
+
Args:
|
32 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
33 |
+
image_size (int): input image size
|
34 |
+
vit (str): model size of vision transformer
|
35 |
+
"""
|
36 |
+
super().__init__()
|
37 |
+
|
38 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
|
39 |
+
|
40 |
+
if vit=='base':
|
41 |
+
checkpoint = torch.hub.load_state_dict_from_url(
|
42 |
+
url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
|
43 |
+
map_location="cpu", check_hash=True)
|
44 |
+
state_dict = checkpoint["model"]
|
45 |
+
msg = self.visual_encoder.load_state_dict(state_dict,strict=False)
|
46 |
+
elif vit=='large':
|
47 |
+
from timm.models.helpers import load_custom_pretrained
|
48 |
+
from timm.models.vision_transformer import default_cfgs
|
49 |
+
load_custom_pretrained(self.visual_encoder,default_cfgs['vit_large_patch16_224_in21k'])
|
50 |
+
|
51 |
+
self.tokenizer = init_tokenizer()
|
52 |
+
encoder_config = BertConfig.from_json_file(med_config)
|
53 |
+
encoder_config.encoder_width = vision_width
|
54 |
+
self.text_encoder = BertModel.from_pretrained('bert-base-uncased',config=encoder_config, add_pooling_layer=False)
|
55 |
+
self.text_encoder.resize_token_embeddings(len(self.tokenizer))
|
56 |
+
|
57 |
+
text_width = self.text_encoder.config.hidden_size
|
58 |
+
|
59 |
+
self.vision_proj = nn.Linear(vision_width, embed_dim)
|
60 |
+
self.text_proj = nn.Linear(text_width, embed_dim)
|
61 |
+
|
62 |
+
self.itm_head = nn.Linear(text_width, 2)
|
63 |
+
|
64 |
+
# create momentum encoders
|
65 |
+
self.visual_encoder_m, vision_width = create_vit(vit,image_size)
|
66 |
+
self.vision_proj_m = nn.Linear(vision_width, embed_dim)
|
67 |
+
self.text_encoder_m = BertModel(config=encoder_config, add_pooling_layer=False)
|
68 |
+
self.text_proj_m = nn.Linear(text_width, embed_dim)
|
69 |
+
|
70 |
+
self.model_pairs = [[self.visual_encoder,self.visual_encoder_m],
|
71 |
+
[self.vision_proj,self.vision_proj_m],
|
72 |
+
[self.text_encoder,self.text_encoder_m],
|
73 |
+
[self.text_proj,self.text_proj_m],
|
74 |
+
]
|
75 |
+
self.copy_params()
|
76 |
+
|
77 |
+
# create the queue
|
78 |
+
self.register_buffer("image_queue", torch.randn(embed_dim, queue_size))
|
79 |
+
self.register_buffer("text_queue", torch.randn(embed_dim, queue_size))
|
80 |
+
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
|
81 |
+
|
82 |
+
self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
|
83 |
+
self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
|
84 |
+
|
85 |
+
self.queue_size = queue_size
|
86 |
+
self.momentum = momentum
|
87 |
+
self.temp = nn.Parameter(0.07*torch.ones([]))
|
88 |
+
|
89 |
+
# create the decoder
|
90 |
+
decoder_config = BertConfig.from_json_file(med_config)
|
91 |
+
decoder_config.encoder_width = vision_width
|
92 |
+
self.text_decoder = BertLMHeadModel.from_pretrained('bert-base-uncased',config=decoder_config)
|
93 |
+
self.text_decoder.resize_token_embeddings(len(self.tokenizer))
|
94 |
+
tie_encoder_decoder_weights(self.text_encoder,self.text_decoder.bert,'','/attention')
|
95 |
+
|
96 |
+
|
97 |
+
def forward(self, image, caption, alpha):
|
98 |
+
with torch.no_grad():
|
99 |
+
self.temp.clamp_(0.001,0.5)
|
100 |
+
|
101 |
+
image_embeds = self.visual_encoder(image)
|
102 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
103 |
+
image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)
|
104 |
+
|
105 |
+
text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=30,
|
106 |
+
return_tensors="pt").to(image.device)
|
107 |
+
text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,
|
108 |
+
return_dict = True, mode = 'text')
|
109 |
+
text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)
|
110 |
+
|
111 |
+
# get momentum features
|
112 |
+
with torch.no_grad():
|
113 |
+
self._momentum_update()
|
114 |
+
image_embeds_m = self.visual_encoder_m(image)
|
115 |
+
image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)
|
116 |
+
image_feat_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)
|
117 |
+
|
118 |
+
text_output_m = self.text_encoder_m(text.input_ids, attention_mask = text.attention_mask,
|
119 |
+
return_dict = True, mode = 'text')
|
120 |
+
text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1)
|
121 |
+
text_feat_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
|
122 |
+
|
123 |
+
sim_i2t_m = image_feat_m @ text_feat_all / self.temp
|
124 |
+
sim_t2i_m = text_feat_m @ image_feat_all / self.temp
|
125 |
+
|
126 |
+
sim_targets = torch.zeros(sim_i2t_m.size()).to(image.device)
|
127 |
+
sim_targets.fill_diagonal_(1)
|
128 |
+
|
129 |
+
sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
|
130 |
+
sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets
|
131 |
+
|
132 |
+
sim_i2t = image_feat @ text_feat_all / self.temp
|
133 |
+
sim_t2i = text_feat @ image_feat_all / self.temp
|
134 |
+
|
135 |
+
loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
|
136 |
+
loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean()
|
137 |
+
|
138 |
+
loss_ita = (loss_i2t+loss_t2i)/2
|
139 |
+
|
140 |
+
self._dequeue_and_enqueue(image_feat_m, text_feat_m)
|
141 |
+
|
142 |
+
###============== Image-text Matching ===================###
|
143 |
+
encoder_input_ids = text.input_ids.clone()
|
144 |
+
encoder_input_ids[:,0] = self.tokenizer.enc_token_id
|
145 |
+
|
146 |
+
# forward the positve image-text pair
|
147 |
+
bs = image.size(0)
|
148 |
+
output_pos = self.text_encoder(encoder_input_ids,
|
149 |
+
attention_mask = text.attention_mask,
|
150 |
+
encoder_hidden_states = image_embeds,
|
151 |
+
encoder_attention_mask = image_atts,
|
152 |
+
return_dict = True,
|
153 |
+
)
|
154 |
+
with torch.no_grad():
|
155 |
+
weights_t2i = F.softmax(sim_t2i[:,:bs],dim=1)+1e-4
|
156 |
+
weights_t2i.fill_diagonal_(0)
|
157 |
+
weights_i2t = F.softmax(sim_i2t[:,:bs],dim=1)+1e-4
|
158 |
+
weights_i2t.fill_diagonal_(0)
|
159 |
+
|
160 |
+
# select a negative image for each text
|
161 |
+
image_embeds_neg = []
|
162 |
+
for b in range(bs):
|
163 |
+
neg_idx = torch.multinomial(weights_t2i[b], 1).item()
|
164 |
+
image_embeds_neg.append(image_embeds[neg_idx])
|
165 |
+
image_embeds_neg = torch.stack(image_embeds_neg,dim=0)
|
166 |
+
|
167 |
+
# select a negative text for each image
|
168 |
+
text_ids_neg = []
|
169 |
+
text_atts_neg = []
|
170 |
+
for b in range(bs):
|
171 |
+
neg_idx = torch.multinomial(weights_i2t[b], 1).item()
|
172 |
+
text_ids_neg.append(encoder_input_ids[neg_idx])
|
173 |
+
text_atts_neg.append(text.attention_mask[neg_idx])
|
174 |
+
|
175 |
+
text_ids_neg = torch.stack(text_ids_neg,dim=0)
|
176 |
+
text_atts_neg = torch.stack(text_atts_neg,dim=0)
|
177 |
+
|
178 |
+
text_ids_all = torch.cat([encoder_input_ids, text_ids_neg],dim=0)
|
179 |
+
text_atts_all = torch.cat([text.attention_mask, text_atts_neg],dim=0)
|
180 |
+
|
181 |
+
image_embeds_all = torch.cat([image_embeds_neg,image_embeds],dim=0)
|
182 |
+
image_atts_all = torch.cat([image_atts,image_atts],dim=0)
|
183 |
+
|
184 |
+
output_neg = self.text_encoder(text_ids_all,
|
185 |
+
attention_mask = text_atts_all,
|
186 |
+
encoder_hidden_states = image_embeds_all,
|
187 |
+
encoder_attention_mask = image_atts_all,
|
188 |
+
return_dict = True,
|
189 |
+
)
|
190 |
+
|
191 |
+
vl_embeddings = torch.cat([output_pos.last_hidden_state[:,0,:], output_neg.last_hidden_state[:,0,:]],dim=0)
|
192 |
+
vl_output = self.itm_head(vl_embeddings)
|
193 |
+
|
194 |
+
itm_labels = torch.cat([torch.ones(bs,dtype=torch.long),torch.zeros(2*bs,dtype=torch.long)],
|
195 |
+
dim=0).to(image.device)
|
196 |
+
loss_itm = F.cross_entropy(vl_output, itm_labels)
|
197 |
+
|
198 |
+
##================= LM ========================##
|
199 |
+
decoder_input_ids = text.input_ids.clone()
|
200 |
+
decoder_input_ids[:,0] = self.tokenizer.bos_token_id
|
201 |
+
decoder_targets = decoder_input_ids.masked_fill(decoder_input_ids == self.tokenizer.pad_token_id, -100)
|
202 |
+
|
203 |
+
decoder_output = self.text_decoder(decoder_input_ids,
|
204 |
+
attention_mask = text.attention_mask,
|
205 |
+
encoder_hidden_states = image_embeds,
|
206 |
+
encoder_attention_mask = image_atts,
|
207 |
+
labels = decoder_targets,
|
208 |
+
return_dict = True,
|
209 |
+
)
|
210 |
+
|
211 |
+
loss_lm = decoder_output.loss
|
212 |
+
return loss_ita, loss_itm, loss_lm
|
213 |
+
|
214 |
+
|
215 |
+
|
216 |
+
@torch.no_grad()
|
217 |
+
def copy_params(self):
|
218 |
+
for model_pair in self.model_pairs:
|
219 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
220 |
+
param_m.data.copy_(param.data) # initialize
|
221 |
+
param_m.requires_grad = False # not update by gradient
|
222 |
+
|
223 |
+
|
224 |
+
@torch.no_grad()
|
225 |
+
def _momentum_update(self):
|
226 |
+
for model_pair in self.model_pairs:
|
227 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
228 |
+
param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
|
229 |
+
|
230 |
+
|
231 |
+
@torch.no_grad()
|
232 |
+
def _dequeue_and_enqueue(self, image_feat, text_feat):
|
233 |
+
# gather keys before updating queue
|
234 |
+
image_feats = concat_all_gather(image_feat)
|
235 |
+
text_feats = concat_all_gather(text_feat)
|
236 |
+
|
237 |
+
batch_size = image_feats.shape[0]
|
238 |
+
|
239 |
+
ptr = int(self.queue_ptr)
|
240 |
+
assert self.queue_size % batch_size == 0 # for simplicity
|
241 |
+
|
242 |
+
# replace the keys at ptr (dequeue and enqueue)
|
243 |
+
self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
|
244 |
+
self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
|
245 |
+
ptr = (ptr + batch_size) % self.queue_size # move pointer
|
246 |
+
|
247 |
+
self.queue_ptr[0] = ptr
|
248 |
+
|
249 |
+
|
250 |
+
def blip_pretrain(**kwargs):
|
251 |
+
model = BLIP_Pretrain(**kwargs)
|
252 |
+
return model
|
253 |
+
|
254 |
+
|
255 |
+
@torch.no_grad()
|
256 |
+
def concat_all_gather(tensor):
|
257 |
+
"""
|
258 |
+
Performs all_gather operation on the provided tensors.
|
259 |
+
*** Warning ***: torch.distributed.all_gather has no gradient.
|
260 |
+
"""
|
261 |
+
tensors_gather = [torch.ones_like(tensor)
|
262 |
+
for _ in range(torch.distributed.get_world_size())]
|
263 |
+
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
|
264 |
+
|
265 |
+
output = torch.cat(tensors_gather, dim=0)
|
266 |
+
return output
|
267 |
+
|
268 |
+
|
269 |
+
from typing import List
|
270 |
+
def tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str, skip_key:str):
|
271 |
+
uninitialized_encoder_weights: List[str] = []
|
272 |
+
if decoder.__class__ != encoder.__class__:
|
273 |
+
print(
|
274 |
+
f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder weights are correctly initialized."
|
275 |
+
)
|
276 |
+
|
277 |
+
def tie_encoder_to_decoder_recursively(
|
278 |
+
decoder_pointer: nn.Module,
|
279 |
+
encoder_pointer: nn.Module,
|
280 |
+
module_name: str,
|
281 |
+
uninitialized_encoder_weights: List[str],
|
282 |
+
skip_key: str,
|
283 |
+
depth=0,
|
284 |
+
):
|
285 |
+
assert isinstance(decoder_pointer, nn.Module) and isinstance(
|
286 |
+
encoder_pointer, nn.Module
|
287 |
+
), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module"
|
288 |
+
if hasattr(decoder_pointer, "weight") and skip_key not in module_name:
|
289 |
+
assert hasattr(encoder_pointer, "weight")
|
290 |
+
encoder_pointer.weight = decoder_pointer.weight
|
291 |
+
if hasattr(decoder_pointer, "bias"):
|
292 |
+
assert hasattr(encoder_pointer, "bias")
|
293 |
+
encoder_pointer.bias = decoder_pointer.bias
|
294 |
+
print(module_name+' is tied')
|
295 |
+
return
|
296 |
+
|
297 |
+
encoder_modules = encoder_pointer._modules
|
298 |
+
decoder_modules = decoder_pointer._modules
|
299 |
+
if len(decoder_modules) > 0:
|
300 |
+
assert (
|
301 |
+
len(encoder_modules) > 0
|
302 |
+
), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"
|
303 |
+
|
304 |
+
all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
|
305 |
+
encoder_layer_pos = 0
|
306 |
+
for name, module in decoder_modules.items():
|
307 |
+
if name.isdigit():
|
308 |
+
encoder_name = str(int(name) + encoder_layer_pos)
|
309 |
+
decoder_name = name
|
310 |
+
if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
|
311 |
+
encoder_modules
|
312 |
+
) != len(decoder_modules):
|
313 |
+
# this can happen if the name corresponds to the position in a list module list of layers
|
314 |
+
# in this case the decoder has added a cross-attention that the encoder does not have
|
315 |
+
# thus skip this step and subtract one layer pos from encoder
|
316 |
+
encoder_layer_pos -= 1
|
317 |
+
continue
|
318 |
+
elif name not in encoder_modules:
|
319 |
+
continue
|
320 |
+
elif depth > 500:
|
321 |
+
raise ValueError(
|
322 |
+
"Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model."
|
323 |
+
)
|
324 |
+
else:
|
325 |
+
decoder_name = encoder_name = name
|
326 |
+
tie_encoder_to_decoder_recursively(
|
327 |
+
decoder_modules[decoder_name],
|
328 |
+
encoder_modules[encoder_name],
|
329 |
+
module_name + "/" + name,
|
330 |
+
uninitialized_encoder_weights,
|
331 |
+
skip_key,
|
332 |
+
depth=depth + 1,
|
333 |
+
)
|
334 |
+
all_encoder_weights.remove(module_name + "/" + encoder_name)
|
335 |
+
|
336 |
+
uninitialized_encoder_weights += list(all_encoder_weights)
|
337 |
+
|
338 |
+
# tie weights recursively
|
339 |
+
tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights, skip_key)
|
extras/BLIP/models/blip_retrieval.py
ADDED
@@ -0,0 +1,319 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from extras.BLIP.models.med import BertConfig, BertModel
|
2 |
+
from transformers import BertTokenizer
|
3 |
+
|
4 |
+
import torch
|
5 |
+
from torch import nn
|
6 |
+
import torch.nn.functional as F
|
7 |
+
|
8 |
+
from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
|
9 |
+
|
10 |
+
class BLIP_Retrieval(nn.Module):
|
11 |
+
def __init__(self,
|
12 |
+
med_config = 'configs/med_config.json',
|
13 |
+
image_size = 384,
|
14 |
+
vit = 'base',
|
15 |
+
vit_grad_ckpt = False,
|
16 |
+
vit_ckpt_layer = 0,
|
17 |
+
embed_dim = 256,
|
18 |
+
queue_size = 57600,
|
19 |
+
momentum = 0.995,
|
20 |
+
negative_all_rank = False,
|
21 |
+
):
|
22 |
+
"""
|
23 |
+
Args:
|
24 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
25 |
+
image_size (int): input image size
|
26 |
+
vit (str): model size of vision transformer
|
27 |
+
"""
|
28 |
+
super().__init__()
|
29 |
+
|
30 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer)
|
31 |
+
self.tokenizer = init_tokenizer()
|
32 |
+
med_config = BertConfig.from_json_file(med_config)
|
33 |
+
med_config.encoder_width = vision_width
|
34 |
+
self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)
|
35 |
+
|
36 |
+
text_width = self.text_encoder.config.hidden_size
|
37 |
+
|
38 |
+
self.vision_proj = nn.Linear(vision_width, embed_dim)
|
39 |
+
self.text_proj = nn.Linear(text_width, embed_dim)
|
40 |
+
|
41 |
+
self.itm_head = nn.Linear(text_width, 2)
|
42 |
+
|
43 |
+
# create momentum encoders
|
44 |
+
self.visual_encoder_m, vision_width = create_vit(vit,image_size)
|
45 |
+
self.vision_proj_m = nn.Linear(vision_width, embed_dim)
|
46 |
+
self.text_encoder_m = BertModel(config=med_config, add_pooling_layer=False)
|
47 |
+
self.text_proj_m = nn.Linear(text_width, embed_dim)
|
48 |
+
|
49 |
+
self.model_pairs = [[self.visual_encoder,self.visual_encoder_m],
|
50 |
+
[self.vision_proj,self.vision_proj_m],
|
51 |
+
[self.text_encoder,self.text_encoder_m],
|
52 |
+
[self.text_proj,self.text_proj_m],
|
53 |
+
]
|
54 |
+
self.copy_params()
|
55 |
+
|
56 |
+
# create the queue
|
57 |
+
self.register_buffer("image_queue", torch.randn(embed_dim, queue_size))
|
58 |
+
self.register_buffer("text_queue", torch.randn(embed_dim, queue_size))
|
59 |
+
self.register_buffer("idx_queue", torch.full((1,queue_size),-100))
|
60 |
+
self.register_buffer("ptr_queue", torch.zeros(1, dtype=torch.long))
|
61 |
+
|
62 |
+
self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
|
63 |
+
self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
|
64 |
+
|
65 |
+
self.queue_size = queue_size
|
66 |
+
self.momentum = momentum
|
67 |
+
self.temp = nn.Parameter(0.07*torch.ones([]))
|
68 |
+
|
69 |
+
self.negative_all_rank = negative_all_rank
|
70 |
+
|
71 |
+
|
72 |
+
def forward(self, image, caption, alpha, idx):
|
73 |
+
with torch.no_grad():
|
74 |
+
self.temp.clamp_(0.001,0.5)
|
75 |
+
|
76 |
+
image_embeds = self.visual_encoder(image)
|
77 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
78 |
+
image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)
|
79 |
+
|
80 |
+
text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=35,
|
81 |
+
return_tensors="pt").to(image.device)
|
82 |
+
|
83 |
+
text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,
|
84 |
+
return_dict = True, mode = 'text')
|
85 |
+
text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)
|
86 |
+
|
87 |
+
###============== Image-text Contrastive Learning ===================###
|
88 |
+
idx = idx.view(-1,1)
|
89 |
+
idx_all = torch.cat([idx.t(), self.idx_queue.clone().detach()],dim=1)
|
90 |
+
pos_idx = torch.eq(idx, idx_all).float()
|
91 |
+
sim_targets = pos_idx / pos_idx.sum(1,keepdim=True)
|
92 |
+
|
93 |
+
# get momentum features
|
94 |
+
with torch.no_grad():
|
95 |
+
self._momentum_update()
|
96 |
+
image_embeds_m = self.visual_encoder_m(image)
|
97 |
+
image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)
|
98 |
+
image_feat_m_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)
|
99 |
+
|
100 |
+
text_output_m = self.text_encoder_m(text.input_ids, attention_mask = text.attention_mask,
|
101 |
+
return_dict = True, mode = 'text')
|
102 |
+
text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1)
|
103 |
+
text_feat_m_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
|
104 |
+
|
105 |
+
sim_i2t_m = image_feat_m @ text_feat_m_all / self.temp
|
106 |
+
sim_t2i_m = text_feat_m @ image_feat_m_all / self.temp
|
107 |
+
|
108 |
+
sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
|
109 |
+
sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets
|
110 |
+
|
111 |
+
sim_i2t = image_feat @ text_feat_m_all / self.temp
|
112 |
+
sim_t2i = text_feat @ image_feat_m_all / self.temp
|
113 |
+
|
114 |
+
loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
|
115 |
+
loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean()
|
116 |
+
|
117 |
+
loss_ita = (loss_i2t+loss_t2i)/2
|
118 |
+
|
119 |
+
idxs = concat_all_gather(idx)
|
120 |
+
self._dequeue_and_enqueue(image_feat_m, text_feat_m, idxs)
|
121 |
+
|
122 |
+
###============== Image-text Matching ===================###
|
123 |
+
encoder_input_ids = text.input_ids.clone()
|
124 |
+
encoder_input_ids[:,0] = self.tokenizer.enc_token_id
|
125 |
+
|
126 |
+
# forward the positve image-text pair
|
127 |
+
bs = image.size(0)
|
128 |
+
output_pos = self.text_encoder(encoder_input_ids,
|
129 |
+
attention_mask = text.attention_mask,
|
130 |
+
encoder_hidden_states = image_embeds,
|
131 |
+
encoder_attention_mask = image_atts,
|
132 |
+
return_dict = True,
|
133 |
+
)
|
134 |
+
|
135 |
+
|
136 |
+
if self.negative_all_rank:
|
137 |
+
# compute sample similarity
|
138 |
+
with torch.no_grad():
|
139 |
+
mask = torch.eq(idx, idxs.t())
|
140 |
+
|
141 |
+
image_feat_world = concat_all_gather(image_feat)
|
142 |
+
text_feat_world = concat_all_gather(text_feat)
|
143 |
+
|
144 |
+
sim_i2t = image_feat @ text_feat_world.t() / self.temp
|
145 |
+
sim_t2i = text_feat @ image_feat_world.t() / self.temp
|
146 |
+
|
147 |
+
weights_i2t = F.softmax(sim_i2t,dim=1)
|
148 |
+
weights_i2t.masked_fill_(mask, 0)
|
149 |
+
|
150 |
+
weights_t2i = F.softmax(sim_t2i,dim=1)
|
151 |
+
weights_t2i.masked_fill_(mask, 0)
|
152 |
+
|
153 |
+
image_embeds_world = all_gather_with_grad(image_embeds)
|
154 |
+
|
155 |
+
# select a negative image (from all ranks) for each text
|
156 |
+
image_embeds_neg = []
|
157 |
+
for b in range(bs):
|
158 |
+
neg_idx = torch.multinomial(weights_t2i[b], 1).item()
|
159 |
+
image_embeds_neg.append(image_embeds_world[neg_idx])
|
160 |
+
image_embeds_neg = torch.stack(image_embeds_neg,dim=0)
|
161 |
+
|
162 |
+
# select a negative text (from all ranks) for each image
|
163 |
+
input_ids_world = concat_all_gather(encoder_input_ids)
|
164 |
+
att_mask_world = concat_all_gather(text.attention_mask)
|
165 |
+
|
166 |
+
text_ids_neg = []
|
167 |
+
text_atts_neg = []
|
168 |
+
for b in range(bs):
|
169 |
+
neg_idx = torch.multinomial(weights_i2t[b], 1).item()
|
170 |
+
text_ids_neg.append(input_ids_world[neg_idx])
|
171 |
+
text_atts_neg.append(att_mask_world[neg_idx])
|
172 |
+
|
173 |
+
else:
|
174 |
+
with torch.no_grad():
|
175 |
+
mask = torch.eq(idx, idx.t())
|
176 |
+
|
177 |
+
sim_i2t = image_feat @ text_feat.t() / self.temp
|
178 |
+
sim_t2i = text_feat @ image_feat.t() / self.temp
|
179 |
+
|
180 |
+
weights_i2t = F.softmax(sim_i2t,dim=1)
|
181 |
+
weights_i2t.masked_fill_(mask, 0)
|
182 |
+
|
183 |
+
weights_t2i = F.softmax(sim_t2i,dim=1)
|
184 |
+
weights_t2i.masked_fill_(mask, 0)
|
185 |
+
|
186 |
+
# select a negative image (from same rank) for each text
|
187 |
+
image_embeds_neg = []
|
188 |
+
for b in range(bs):
|
189 |
+
neg_idx = torch.multinomial(weights_t2i[b], 1).item()
|
190 |
+
image_embeds_neg.append(image_embeds[neg_idx])
|
191 |
+
image_embeds_neg = torch.stack(image_embeds_neg,dim=0)
|
192 |
+
|
193 |
+
# select a negative text (from same rank) for each image
|
194 |
+
text_ids_neg = []
|
195 |
+
text_atts_neg = []
|
196 |
+
for b in range(bs):
|
197 |
+
neg_idx = torch.multinomial(weights_i2t[b], 1).item()
|
198 |
+
text_ids_neg.append(encoder_input_ids[neg_idx])
|
199 |
+
text_atts_neg.append(text.attention_mask[neg_idx])
|
200 |
+
|
201 |
+
text_ids_neg = torch.stack(text_ids_neg,dim=0)
|
202 |
+
text_atts_neg = torch.stack(text_atts_neg,dim=0)
|
203 |
+
|
204 |
+
text_ids_all = torch.cat([encoder_input_ids, text_ids_neg],dim=0)
|
205 |
+
text_atts_all = torch.cat([text.attention_mask, text_atts_neg],dim=0)
|
206 |
+
|
207 |
+
image_embeds_all = torch.cat([image_embeds_neg,image_embeds],dim=0)
|
208 |
+
image_atts_all = torch.cat([image_atts,image_atts],dim=0)
|
209 |
+
|
210 |
+
output_neg = self.text_encoder(text_ids_all,
|
211 |
+
attention_mask = text_atts_all,
|
212 |
+
encoder_hidden_states = image_embeds_all,
|
213 |
+
encoder_attention_mask = image_atts_all,
|
214 |
+
return_dict = True,
|
215 |
+
)
|
216 |
+
|
217 |
+
|
218 |
+
vl_embeddings = torch.cat([output_pos.last_hidden_state[:,0,:], output_neg.last_hidden_state[:,0,:]],dim=0)
|
219 |
+
vl_output = self.itm_head(vl_embeddings)
|
220 |
+
|
221 |
+
itm_labels = torch.cat([torch.ones(bs,dtype=torch.long),torch.zeros(2*bs,dtype=torch.long)],
|
222 |
+
dim=0).to(image.device)
|
223 |
+
loss_itm = F.cross_entropy(vl_output, itm_labels)
|
224 |
+
|
225 |
+
return loss_ita, loss_itm
|
226 |
+
|
227 |
+
|
228 |
+
@torch.no_grad()
|
229 |
+
def copy_params(self):
|
230 |
+
for model_pair in self.model_pairs:
|
231 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
232 |
+
param_m.data.copy_(param.data) # initialize
|
233 |
+
param_m.requires_grad = False # not update by gradient
|
234 |
+
|
235 |
+
|
236 |
+
@torch.no_grad()
|
237 |
+
def _momentum_update(self):
|
238 |
+
for model_pair in self.model_pairs:
|
239 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
240 |
+
param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
|
241 |
+
|
242 |
+
|
243 |
+
@torch.no_grad()
|
244 |
+
def _dequeue_and_enqueue(self, image_feat, text_feat, idxs):
|
245 |
+
# gather keys before updating queue
|
246 |
+
image_feats = concat_all_gather(image_feat)
|
247 |
+
text_feats = concat_all_gather(text_feat)
|
248 |
+
|
249 |
+
|
250 |
+
batch_size = image_feats.shape[0]
|
251 |
+
|
252 |
+
ptr = int(self.ptr_queue)
|
253 |
+
assert self.queue_size % batch_size == 0 # for simplicity
|
254 |
+
|
255 |
+
# replace the keys at ptr (dequeue and enqueue)
|
256 |
+
self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
|
257 |
+
self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
|
258 |
+
self.idx_queue[:, ptr:ptr + batch_size] = idxs.T
|
259 |
+
ptr = (ptr + batch_size) % self.queue_size # move pointer
|
260 |
+
|
261 |
+
self.ptr_queue[0] = ptr
|
262 |
+
|
263 |
+
|
264 |
+
def blip_retrieval(pretrained='',**kwargs):
|
265 |
+
model = BLIP_Retrieval(**kwargs)
|
266 |
+
if pretrained:
|
267 |
+
model,msg = load_checkpoint(model,pretrained)
|
268 |
+
print("missing keys:")
|
269 |
+
print(msg.missing_keys)
|
270 |
+
return model
|
271 |
+
|
272 |
+
|
273 |
+
@torch.no_grad()
|
274 |
+
def concat_all_gather(tensor):
|
275 |
+
"""
|
276 |
+
Performs all_gather operation on the provided tensors.
|
277 |
+
*** Warning ***: torch.distributed.all_gather has no gradient.
|
278 |
+
"""
|
279 |
+
tensors_gather = [torch.ones_like(tensor)
|
280 |
+
for _ in range(torch.distributed.get_world_size())]
|
281 |
+
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
|
282 |
+
|
283 |
+
output = torch.cat(tensors_gather, dim=0)
|
284 |
+
return output
|
285 |
+
|
286 |
+
|
287 |
+
class GatherLayer(torch.autograd.Function):
|
288 |
+
"""
|
289 |
+
Gather tensors from all workers with support for backward propagation:
|
290 |
+
This implementation does not cut the gradients as torch.distributed.all_gather does.
|
291 |
+
"""
|
292 |
+
|
293 |
+
@staticmethod
|
294 |
+
def forward(ctx, x):
|
295 |
+
output = [torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())]
|
296 |
+
torch.distributed.all_gather(output, x)
|
297 |
+
return tuple(output)
|
298 |
+
|
299 |
+
@staticmethod
|
300 |
+
def backward(ctx, *grads):
|
301 |
+
all_gradients = torch.stack(grads)
|
302 |
+
torch.distributed.all_reduce(all_gradients)
|
303 |
+
return all_gradients[torch.distributed.get_rank()]
|
304 |
+
|
305 |
+
|
306 |
+
def all_gather_with_grad(tensors):
|
307 |
+
"""
|
308 |
+
Performs all_gather operation on the provided tensors.
|
309 |
+
Graph remains connected for backward grad computation.
|
310 |
+
"""
|
311 |
+
# Queue the gathered tensors
|
312 |
+
world_size = torch.distributed.get_world_size()
|
313 |
+
# There is no need for reduction in the single-proc case
|
314 |
+
if world_size == 1:
|
315 |
+
return tensors
|
316 |
+
|
317 |
+
tensor_all = GatherLayer.apply(tensors)
|
318 |
+
|
319 |
+
return torch.cat(tensor_all, dim=0)
|
extras/BLIP/models/blip_vqa.py
ADDED
@@ -0,0 +1,186 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from extras.BLIP.models.med import BertConfig, BertModel, BertLMHeadModel
|
2 |
+
from extras.BLIP.models.blip import create_vit, init_tokenizer, load_checkpoint
|
3 |
+
|
4 |
+
import torch
|
5 |
+
from torch import nn
|
6 |
+
import torch.nn.functional as F
|
7 |
+
from transformers import BertTokenizer
|
8 |
+
import numpy as np
|
9 |
+
|
10 |
+
class BLIP_VQA(nn.Module):
|
11 |
+
def __init__(self,
|
12 |
+
med_config = 'configs/med_config.json',
|
13 |
+
image_size = 480,
|
14 |
+
vit = 'base',
|
15 |
+
vit_grad_ckpt = False,
|
16 |
+
vit_ckpt_layer = 0,
|
17 |
+
):
|
18 |
+
"""
|
19 |
+
Args:
|
20 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
21 |
+
image_size (int): input image size
|
22 |
+
vit (str): model size of vision transformer
|
23 |
+
"""
|
24 |
+
super().__init__()
|
25 |
+
|
26 |
+
self.visual_encoder, vision_width = create_vit(vit, image_size, vit_grad_ckpt, vit_ckpt_layer, drop_path_rate=0.1)
|
27 |
+
self.tokenizer = init_tokenizer()
|
28 |
+
|
29 |
+
encoder_config = BertConfig.from_json_file(med_config)
|
30 |
+
encoder_config.encoder_width = vision_width
|
31 |
+
self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False)
|
32 |
+
|
33 |
+
decoder_config = BertConfig.from_json_file(med_config)
|
34 |
+
self.text_decoder = BertLMHeadModel(config=decoder_config)
|
35 |
+
|
36 |
+
|
37 |
+
def forward(self, image, question, answer=None, n=None, weights=None, train=True, inference='rank', k_test=128):
|
38 |
+
|
39 |
+
image_embeds = self.visual_encoder(image)
|
40 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
41 |
+
|
42 |
+
question = self.tokenizer(question, padding='longest', truncation=True, max_length=35,
|
43 |
+
return_tensors="pt").to(image.device)
|
44 |
+
question.input_ids[:,0] = self.tokenizer.enc_token_id
|
45 |
+
|
46 |
+
if train:
|
47 |
+
'''
|
48 |
+
n: number of answers for each question
|
49 |
+
weights: weight for each answer
|
50 |
+
'''
|
51 |
+
answer = self.tokenizer(answer, padding='longest', return_tensors="pt").to(image.device)
|
52 |
+
answer.input_ids[:,0] = self.tokenizer.bos_token_id
|
53 |
+
answer_targets = answer.input_ids.masked_fill(answer.input_ids == self.tokenizer.pad_token_id, -100)
|
54 |
+
|
55 |
+
question_output = self.text_encoder(question.input_ids,
|
56 |
+
attention_mask = question.attention_mask,
|
57 |
+
encoder_hidden_states = image_embeds,
|
58 |
+
encoder_attention_mask = image_atts,
|
59 |
+
return_dict = True)
|
60 |
+
|
61 |
+
question_states = []
|
62 |
+
question_atts = []
|
63 |
+
for b, n in enumerate(n):
|
64 |
+
question_states += [question_output.last_hidden_state[b]]*n
|
65 |
+
question_atts += [question.attention_mask[b]]*n
|
66 |
+
question_states = torch.stack(question_states,0)
|
67 |
+
question_atts = torch.stack(question_atts,0)
|
68 |
+
|
69 |
+
answer_output = self.text_decoder(answer.input_ids,
|
70 |
+
attention_mask = answer.attention_mask,
|
71 |
+
encoder_hidden_states = question_states,
|
72 |
+
encoder_attention_mask = question_atts,
|
73 |
+
labels = answer_targets,
|
74 |
+
return_dict = True,
|
75 |
+
reduction = 'none',
|
76 |
+
)
|
77 |
+
|
78 |
+
loss = weights * answer_output.loss
|
79 |
+
loss = loss.sum()/image.size(0)
|
80 |
+
|
81 |
+
return loss
|
82 |
+
|
83 |
+
|
84 |
+
else:
|
85 |
+
question_output = self.text_encoder(question.input_ids,
|
86 |
+
attention_mask = question.attention_mask,
|
87 |
+
encoder_hidden_states = image_embeds,
|
88 |
+
encoder_attention_mask = image_atts,
|
89 |
+
return_dict = True)
|
90 |
+
|
91 |
+
if inference=='generate':
|
92 |
+
num_beams = 3
|
93 |
+
question_states = question_output.last_hidden_state.repeat_interleave(num_beams,dim=0)
|
94 |
+
question_atts = torch.ones(question_states.size()[:-1],dtype=torch.long).to(question_states.device)
|
95 |
+
model_kwargs = {"encoder_hidden_states": question_states, "encoder_attention_mask":question_atts}
|
96 |
+
|
97 |
+
bos_ids = torch.full((image.size(0),1),fill_value=self.tokenizer.bos_token_id,device=image.device)
|
98 |
+
|
99 |
+
outputs = self.text_decoder.generate(input_ids=bos_ids,
|
100 |
+
max_length=10,
|
101 |
+
min_length=1,
|
102 |
+
num_beams=num_beams,
|
103 |
+
eos_token_id=self.tokenizer.sep_token_id,
|
104 |
+
pad_token_id=self.tokenizer.pad_token_id,
|
105 |
+
**model_kwargs)
|
106 |
+
|
107 |
+
answers = []
|
108 |
+
for output in outputs:
|
109 |
+
answer = self.tokenizer.decode(output, skip_special_tokens=True)
|
110 |
+
answers.append(answer)
|
111 |
+
return answers
|
112 |
+
|
113 |
+
elif inference=='rank':
|
114 |
+
max_ids = self.rank_answer(question_output.last_hidden_state, question.attention_mask,
|
115 |
+
answer.input_ids, answer.attention_mask, k_test)
|
116 |
+
return max_ids
|
117 |
+
|
118 |
+
|
119 |
+
|
120 |
+
def rank_answer(self, question_states, question_atts, answer_ids, answer_atts, k):
|
121 |
+
|
122 |
+
num_ques = question_states.size(0)
|
123 |
+
start_ids = answer_ids[0,0].repeat(num_ques,1) # bos token
|
124 |
+
|
125 |
+
start_output = self.text_decoder(start_ids,
|
126 |
+
encoder_hidden_states = question_states,
|
127 |
+
encoder_attention_mask = question_atts,
|
128 |
+
return_dict = True,
|
129 |
+
reduction = 'none')
|
130 |
+
logits = start_output.logits[:,0,:] # first token's logit
|
131 |
+
|
132 |
+
# topk_probs: top-k probability
|
133 |
+
# topk_ids: [num_question, k]
|
134 |
+
answer_first_token = answer_ids[:,1]
|
135 |
+
prob_first_token = F.softmax(logits,dim=1).index_select(dim=1, index=answer_first_token)
|
136 |
+
topk_probs, topk_ids = prob_first_token.topk(k,dim=1)
|
137 |
+
|
138 |
+
# answer input: [num_question*k, answer_len]
|
139 |
+
input_ids = []
|
140 |
+
input_atts = []
|
141 |
+
for b, topk_id in enumerate(topk_ids):
|
142 |
+
input_ids.append(answer_ids.index_select(dim=0, index=topk_id))
|
143 |
+
input_atts.append(answer_atts.index_select(dim=0, index=topk_id))
|
144 |
+
input_ids = torch.cat(input_ids,dim=0)
|
145 |
+
input_atts = torch.cat(input_atts,dim=0)
|
146 |
+
|
147 |
+
targets_ids = input_ids.masked_fill(input_ids == self.tokenizer.pad_token_id, -100)
|
148 |
+
|
149 |
+
# repeat encoder's output for top-k answers
|
150 |
+
question_states = tile(question_states, 0, k)
|
151 |
+
question_atts = tile(question_atts, 0, k)
|
152 |
+
|
153 |
+
output = self.text_decoder(input_ids,
|
154 |
+
attention_mask = input_atts,
|
155 |
+
encoder_hidden_states = question_states,
|
156 |
+
encoder_attention_mask = question_atts,
|
157 |
+
labels = targets_ids,
|
158 |
+
return_dict = True,
|
159 |
+
reduction = 'none')
|
160 |
+
|
161 |
+
log_probs_sum = -output.loss
|
162 |
+
log_probs_sum = log_probs_sum.view(num_ques,k)
|
163 |
+
|
164 |
+
max_topk_ids = log_probs_sum.argmax(dim=1)
|
165 |
+
max_ids = topk_ids[max_topk_ids>=0,max_topk_ids]
|
166 |
+
|
167 |
+
return max_ids
|
168 |
+
|
169 |
+
|
170 |
+
def blip_vqa(pretrained='',**kwargs):
|
171 |
+
model = BLIP_VQA(**kwargs)
|
172 |
+
if pretrained:
|
173 |
+
model,msg = load_checkpoint(model,pretrained)
|
174 |
+
# assert(len(msg.missing_keys)==0)
|
175 |
+
return model
|
176 |
+
|
177 |
+
|
178 |
+
def tile(x, dim, n_tile):
|
179 |
+
init_dim = x.size(dim)
|
180 |
+
repeat_idx = [1] * x.dim()
|
181 |
+
repeat_idx[dim] = n_tile
|
182 |
+
x = x.repeat(*(repeat_idx))
|
183 |
+
order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))
|
184 |
+
return torch.index_select(x, dim, order_index.to(x.device))
|
185 |
+
|
186 |
+
|
extras/BLIP/models/med.py
ADDED
@@ -0,0 +1,955 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
* Copyright (c) 2022, salesforce.com, inc.
|
3 |
+
* All rights reserved.
|
4 |
+
* SPDX-License-Identifier: BSD-3-Clause
|
5 |
+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
6 |
+
* By Junnan Li
|
7 |
+
* Based on huggingface code base
|
8 |
+
* https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
|
9 |
+
'''
|
10 |
+
|
11 |
+
import math
|
12 |
+
import os
|
13 |
+
import warnings
|
14 |
+
from dataclasses import dataclass
|
15 |
+
from typing import Optional, Tuple
|
16 |
+
|
17 |
+
import torch
|
18 |
+
from torch import Tensor, device, dtype, nn
|
19 |
+
import torch.utils.checkpoint
|
20 |
+
from torch import nn
|
21 |
+
from torch.nn import CrossEntropyLoss
|
22 |
+
import torch.nn.functional as F
|
23 |
+
|
24 |
+
from transformers.activations import ACT2FN
|
25 |
+
from transformers.file_utils import (
|
26 |
+
ModelOutput,
|
27 |
+
)
|
28 |
+
from transformers.modeling_outputs import (
|
29 |
+
BaseModelOutputWithPastAndCrossAttentions,
|
30 |
+
BaseModelOutputWithPoolingAndCrossAttentions,
|
31 |
+
CausalLMOutputWithCrossAttentions,
|
32 |
+
MaskedLMOutput,
|
33 |
+
MultipleChoiceModelOutput,
|
34 |
+
NextSentencePredictorOutput,
|
35 |
+
QuestionAnsweringModelOutput,
|
36 |
+
SequenceClassifierOutput,
|
37 |
+
TokenClassifierOutput,
|
38 |
+
)
|
39 |
+
from transformers.modeling_utils import (
|
40 |
+
PreTrainedModel,
|
41 |
+
apply_chunking_to_forward,
|
42 |
+
find_pruneable_heads_and_indices,
|
43 |
+
prune_linear_layer,
|
44 |
+
)
|
45 |
+
from transformers.utils import logging
|
46 |
+
from transformers.models.bert.configuration_bert import BertConfig
|
47 |
+
|
48 |
+
|
49 |
+
logger = logging.get_logger(__name__)
|
50 |
+
|
51 |
+
|
52 |
+
class BertEmbeddings(nn.Module):
|
53 |
+
"""Construct the embeddings from word and position embeddings."""
|
54 |
+
|
55 |
+
def __init__(self, config):
|
56 |
+
super().__init__()
|
57 |
+
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
|
58 |
+
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
|
59 |
+
|
60 |
+
# self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
|
61 |
+
# any TensorFlow checkpoint file
|
62 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
63 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
64 |
+
|
65 |
+
# position_ids (1, len position emb) is contiguous in memory and exported when serialized
|
66 |
+
self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
|
67 |
+
self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
|
68 |
+
|
69 |
+
self.config = config
|
70 |
+
|
71 |
+
def forward(
|
72 |
+
self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
|
73 |
+
):
|
74 |
+
if input_ids is not None:
|
75 |
+
input_shape = input_ids.size()
|
76 |
+
else:
|
77 |
+
input_shape = inputs_embeds.size()[:-1]
|
78 |
+
|
79 |
+
seq_length = input_shape[1]
|
80 |
+
|
81 |
+
if position_ids is None:
|
82 |
+
position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
|
83 |
+
|
84 |
+
if inputs_embeds is None:
|
85 |
+
inputs_embeds = self.word_embeddings(input_ids)
|
86 |
+
|
87 |
+
embeddings = inputs_embeds
|
88 |
+
|
89 |
+
if self.position_embedding_type == "absolute":
|
90 |
+
position_embeddings = self.position_embeddings(position_ids)
|
91 |
+
embeddings += position_embeddings
|
92 |
+
embeddings = self.LayerNorm(embeddings)
|
93 |
+
embeddings = self.dropout(embeddings)
|
94 |
+
return embeddings
|
95 |
+
|
96 |
+
|
97 |
+
class BertSelfAttention(nn.Module):
|
98 |
+
def __init__(self, config, is_cross_attention):
|
99 |
+
super().__init__()
|
100 |
+
self.config = config
|
101 |
+
if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
|
102 |
+
raise ValueError(
|
103 |
+
"The hidden size (%d) is not a multiple of the number of attention "
|
104 |
+
"heads (%d)" % (config.hidden_size, config.num_attention_heads)
|
105 |
+
)
|
106 |
+
|
107 |
+
self.num_attention_heads = config.num_attention_heads
|
108 |
+
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
|
109 |
+
self.all_head_size = self.num_attention_heads * self.attention_head_size
|
110 |
+
|
111 |
+
self.query = nn.Linear(config.hidden_size, self.all_head_size)
|
112 |
+
if is_cross_attention:
|
113 |
+
self.key = nn.Linear(config.encoder_width, self.all_head_size)
|
114 |
+
self.value = nn.Linear(config.encoder_width, self.all_head_size)
|
115 |
+
else:
|
116 |
+
self.key = nn.Linear(config.hidden_size, self.all_head_size)
|
117 |
+
self.value = nn.Linear(config.hidden_size, self.all_head_size)
|
118 |
+
|
119 |
+
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
120 |
+
self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
|
121 |
+
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
122 |
+
self.max_position_embeddings = config.max_position_embeddings
|
123 |
+
self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
|
124 |
+
self.save_attention = False
|
125 |
+
|
126 |
+
def save_attn_gradients(self, attn_gradients):
|
127 |
+
self.attn_gradients = attn_gradients
|
128 |
+
|
129 |
+
def get_attn_gradients(self):
|
130 |
+
return self.attn_gradients
|
131 |
+
|
132 |
+
def save_attention_map(self, attention_map):
|
133 |
+
self.attention_map = attention_map
|
134 |
+
|
135 |
+
def get_attention_map(self):
|
136 |
+
return self.attention_map
|
137 |
+
|
138 |
+
def transpose_for_scores(self, x):
|
139 |
+
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
|
140 |
+
x = x.view(*new_x_shape)
|
141 |
+
return x.permute(0, 2, 1, 3)
|
142 |
+
|
143 |
+
def forward(
|
144 |
+
self,
|
145 |
+
hidden_states,
|
146 |
+
attention_mask=None,
|
147 |
+
head_mask=None,
|
148 |
+
encoder_hidden_states=None,
|
149 |
+
encoder_attention_mask=None,
|
150 |
+
past_key_value=None,
|
151 |
+
output_attentions=False,
|
152 |
+
):
|
153 |
+
mixed_query_layer = self.query(hidden_states)
|
154 |
+
|
155 |
+
# If this is instantiated as a cross-attention module, the keys
|
156 |
+
# and values come from an encoder; the attention mask needs to be
|
157 |
+
# such that the encoder's padding tokens are not attended to.
|
158 |
+
is_cross_attention = encoder_hidden_states is not None
|
159 |
+
|
160 |
+
if is_cross_attention:
|
161 |
+
key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
|
162 |
+
value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
|
163 |
+
attention_mask = encoder_attention_mask
|
164 |
+
elif past_key_value is not None:
|
165 |
+
key_layer = self.transpose_for_scores(self.key(hidden_states))
|
166 |
+
value_layer = self.transpose_for_scores(self.value(hidden_states))
|
167 |
+
key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
|
168 |
+
value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
|
169 |
+
else:
|
170 |
+
key_layer = self.transpose_for_scores(self.key(hidden_states))
|
171 |
+
value_layer = self.transpose_for_scores(self.value(hidden_states))
|
172 |
+
|
173 |
+
query_layer = self.transpose_for_scores(mixed_query_layer)
|
174 |
+
|
175 |
+
past_key_value = (key_layer, value_layer)
|
176 |
+
|
177 |
+
# Take the dot product between "query" and "key" to get the raw attention scores.
|
178 |
+
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
|
179 |
+
|
180 |
+
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
181 |
+
seq_length = hidden_states.size()[1]
|
182 |
+
position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
|
183 |
+
position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
|
184 |
+
distance = position_ids_l - position_ids_r
|
185 |
+
positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
|
186 |
+
positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
|
187 |
+
|
188 |
+
if self.position_embedding_type == "relative_key":
|
189 |
+
relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
190 |
+
attention_scores = attention_scores + relative_position_scores
|
191 |
+
elif self.position_embedding_type == "relative_key_query":
|
192 |
+
relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
193 |
+
relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
|
194 |
+
attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
|
195 |
+
|
196 |
+
attention_scores = attention_scores / math.sqrt(self.attention_head_size)
|
197 |
+
if attention_mask is not None:
|
198 |
+
# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
|
199 |
+
attention_scores = attention_scores + attention_mask
|
200 |
+
|
201 |
+
# Normalize the attention scores to probabilities.
|
202 |
+
attention_probs = nn.Softmax(dim=-1)(attention_scores)
|
203 |
+
|
204 |
+
if is_cross_attention and self.save_attention:
|
205 |
+
self.save_attention_map(attention_probs)
|
206 |
+
attention_probs.register_hook(self.save_attn_gradients)
|
207 |
+
|
208 |
+
# This is actually dropping out entire tokens to attend to, which might
|
209 |
+
# seem a bit unusual, but is taken from the original Transformer paper.
|
210 |
+
attention_probs_dropped = self.dropout(attention_probs)
|
211 |
+
|
212 |
+
# Mask heads if we want to
|
213 |
+
if head_mask is not None:
|
214 |
+
attention_probs_dropped = attention_probs_dropped * head_mask
|
215 |
+
|
216 |
+
context_layer = torch.matmul(attention_probs_dropped, value_layer)
|
217 |
+
|
218 |
+
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
|
219 |
+
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
|
220 |
+
context_layer = context_layer.view(*new_context_layer_shape)
|
221 |
+
|
222 |
+
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
|
223 |
+
|
224 |
+
outputs = outputs + (past_key_value,)
|
225 |
+
return outputs
|
226 |
+
|
227 |
+
|
228 |
+
class BertSelfOutput(nn.Module):
|
229 |
+
def __init__(self, config):
|
230 |
+
super().__init__()
|
231 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
232 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
233 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
234 |
+
|
235 |
+
def forward(self, hidden_states, input_tensor):
|
236 |
+
hidden_states = self.dense(hidden_states)
|
237 |
+
hidden_states = self.dropout(hidden_states)
|
238 |
+
hidden_states = self.LayerNorm(hidden_states + input_tensor)
|
239 |
+
return hidden_states
|
240 |
+
|
241 |
+
|
242 |
+
class BertAttention(nn.Module):
|
243 |
+
def __init__(self, config, is_cross_attention=False):
|
244 |
+
super().__init__()
|
245 |
+
self.self = BertSelfAttention(config, is_cross_attention)
|
246 |
+
self.output = BertSelfOutput(config)
|
247 |
+
self.pruned_heads = set()
|
248 |
+
|
249 |
+
def prune_heads(self, heads):
|
250 |
+
if len(heads) == 0:
|
251 |
+
return
|
252 |
+
heads, index = find_pruneable_heads_and_indices(
|
253 |
+
heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
|
254 |
+
)
|
255 |
+
|
256 |
+
# Prune linear layers
|
257 |
+
self.self.query = prune_linear_layer(self.self.query, index)
|
258 |
+
self.self.key = prune_linear_layer(self.self.key, index)
|
259 |
+
self.self.value = prune_linear_layer(self.self.value, index)
|
260 |
+
self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
|
261 |
+
|
262 |
+
# Update hyper params and store pruned heads
|
263 |
+
self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
|
264 |
+
self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
|
265 |
+
self.pruned_heads = self.pruned_heads.union(heads)
|
266 |
+
|
267 |
+
def forward(
|
268 |
+
self,
|
269 |
+
hidden_states,
|
270 |
+
attention_mask=None,
|
271 |
+
head_mask=None,
|
272 |
+
encoder_hidden_states=None,
|
273 |
+
encoder_attention_mask=None,
|
274 |
+
past_key_value=None,
|
275 |
+
output_attentions=False,
|
276 |
+
):
|
277 |
+
self_outputs = self.self(
|
278 |
+
hidden_states,
|
279 |
+
attention_mask,
|
280 |
+
head_mask,
|
281 |
+
encoder_hidden_states,
|
282 |
+
encoder_attention_mask,
|
283 |
+
past_key_value,
|
284 |
+
output_attentions,
|
285 |
+
)
|
286 |
+
attention_output = self.output(self_outputs[0], hidden_states)
|
287 |
+
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
|
288 |
+
return outputs
|
289 |
+
|
290 |
+
|
291 |
+
class BertIntermediate(nn.Module):
|
292 |
+
def __init__(self, config):
|
293 |
+
super().__init__()
|
294 |
+
self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
|
295 |
+
if isinstance(config.hidden_act, str):
|
296 |
+
self.intermediate_act_fn = ACT2FN[config.hidden_act]
|
297 |
+
else:
|
298 |
+
self.intermediate_act_fn = config.hidden_act
|
299 |
+
|
300 |
+
def forward(self, hidden_states):
|
301 |
+
hidden_states = self.dense(hidden_states)
|
302 |
+
hidden_states = self.intermediate_act_fn(hidden_states)
|
303 |
+
return hidden_states
|
304 |
+
|
305 |
+
|
306 |
+
class BertOutput(nn.Module):
|
307 |
+
def __init__(self, config):
|
308 |
+
super().__init__()
|
309 |
+
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
|
310 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
311 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
312 |
+
|
313 |
+
def forward(self, hidden_states, input_tensor):
|
314 |
+
hidden_states = self.dense(hidden_states)
|
315 |
+
hidden_states = self.dropout(hidden_states)
|
316 |
+
hidden_states = self.LayerNorm(hidden_states + input_tensor)
|
317 |
+
return hidden_states
|
318 |
+
|
319 |
+
|
320 |
+
class BertLayer(nn.Module):
|
321 |
+
def __init__(self, config, layer_num):
|
322 |
+
super().__init__()
|
323 |
+
self.config = config
|
324 |
+
self.chunk_size_feed_forward = config.chunk_size_feed_forward
|
325 |
+
self.seq_len_dim = 1
|
326 |
+
self.attention = BertAttention(config)
|
327 |
+
self.layer_num = layer_num
|
328 |
+
if self.config.add_cross_attention:
|
329 |
+
self.crossattention = BertAttention(config, is_cross_attention=self.config.add_cross_attention)
|
330 |
+
self.intermediate = BertIntermediate(config)
|
331 |
+
self.output = BertOutput(config)
|
332 |
+
|
333 |
+
def forward(
|
334 |
+
self,
|
335 |
+
hidden_states,
|
336 |
+
attention_mask=None,
|
337 |
+
head_mask=None,
|
338 |
+
encoder_hidden_states=None,
|
339 |
+
encoder_attention_mask=None,
|
340 |
+
past_key_value=None,
|
341 |
+
output_attentions=False,
|
342 |
+
mode=None,
|
343 |
+
):
|
344 |
+
# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
|
345 |
+
self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
|
346 |
+
self_attention_outputs = self.attention(
|
347 |
+
hidden_states,
|
348 |
+
attention_mask,
|
349 |
+
head_mask,
|
350 |
+
output_attentions=output_attentions,
|
351 |
+
past_key_value=self_attn_past_key_value,
|
352 |
+
)
|
353 |
+
attention_output = self_attention_outputs[0]
|
354 |
+
|
355 |
+
outputs = self_attention_outputs[1:-1]
|
356 |
+
present_key_value = self_attention_outputs[-1]
|
357 |
+
|
358 |
+
if mode=='multimodal':
|
359 |
+
assert encoder_hidden_states is not None, "encoder_hidden_states must be given for cross-attention layers"
|
360 |
+
|
361 |
+
cross_attention_outputs = self.crossattention(
|
362 |
+
attention_output,
|
363 |
+
attention_mask,
|
364 |
+
head_mask,
|
365 |
+
encoder_hidden_states,
|
366 |
+
encoder_attention_mask,
|
367 |
+
output_attentions=output_attentions,
|
368 |
+
)
|
369 |
+
attention_output = cross_attention_outputs[0]
|
370 |
+
outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
|
371 |
+
layer_output = apply_chunking_to_forward(
|
372 |
+
self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
|
373 |
+
)
|
374 |
+
outputs = (layer_output,) + outputs
|
375 |
+
|
376 |
+
outputs = outputs + (present_key_value,)
|
377 |
+
|
378 |
+
return outputs
|
379 |
+
|
380 |
+
def feed_forward_chunk(self, attention_output):
|
381 |
+
intermediate_output = self.intermediate(attention_output)
|
382 |
+
layer_output = self.output(intermediate_output, attention_output)
|
383 |
+
return layer_output
|
384 |
+
|
385 |
+
|
386 |
+
class BertEncoder(nn.Module):
|
387 |
+
def __init__(self, config):
|
388 |
+
super().__init__()
|
389 |
+
self.config = config
|
390 |
+
self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)])
|
391 |
+
self.gradient_checkpointing = False
|
392 |
+
|
393 |
+
def forward(
|
394 |
+
self,
|
395 |
+
hidden_states,
|
396 |
+
attention_mask=None,
|
397 |
+
head_mask=None,
|
398 |
+
encoder_hidden_states=None,
|
399 |
+
encoder_attention_mask=None,
|
400 |
+
past_key_values=None,
|
401 |
+
use_cache=None,
|
402 |
+
output_attentions=False,
|
403 |
+
output_hidden_states=False,
|
404 |
+
return_dict=True,
|
405 |
+
mode='multimodal',
|
406 |
+
):
|
407 |
+
all_hidden_states = () if output_hidden_states else None
|
408 |
+
all_self_attentions = () if output_attentions else None
|
409 |
+
all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
|
410 |
+
|
411 |
+
next_decoder_cache = () if use_cache else None
|
412 |
+
|
413 |
+
for i in range(self.config.num_hidden_layers):
|
414 |
+
layer_module = self.layer[i]
|
415 |
+
if output_hidden_states:
|
416 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
417 |
+
|
418 |
+
layer_head_mask = head_mask[i] if head_mask is not None else None
|
419 |
+
past_key_value = past_key_values[i] if past_key_values is not None else None
|
420 |
+
|
421 |
+
if self.gradient_checkpointing and self.training:
|
422 |
+
|
423 |
+
if use_cache:
|
424 |
+
logger.warn(
|
425 |
+
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
426 |
+
)
|
427 |
+
use_cache = False
|
428 |
+
|
429 |
+
def create_custom_forward(module):
|
430 |
+
def custom_forward(*inputs):
|
431 |
+
return module(*inputs, past_key_value, output_attentions)
|
432 |
+
|
433 |
+
return custom_forward
|
434 |
+
|
435 |
+
layer_outputs = torch.utils.checkpoint.checkpoint(
|
436 |
+
create_custom_forward(layer_module),
|
437 |
+
hidden_states,
|
438 |
+
attention_mask,
|
439 |
+
layer_head_mask,
|
440 |
+
encoder_hidden_states,
|
441 |
+
encoder_attention_mask,
|
442 |
+
mode=mode,
|
443 |
+
)
|
444 |
+
else:
|
445 |
+
layer_outputs = layer_module(
|
446 |
+
hidden_states,
|
447 |
+
attention_mask,
|
448 |
+
layer_head_mask,
|
449 |
+
encoder_hidden_states,
|
450 |
+
encoder_attention_mask,
|
451 |
+
past_key_value,
|
452 |
+
output_attentions,
|
453 |
+
mode=mode,
|
454 |
+
)
|
455 |
+
|
456 |
+
hidden_states = layer_outputs[0]
|
457 |
+
if use_cache:
|
458 |
+
next_decoder_cache += (layer_outputs[-1],)
|
459 |
+
if output_attentions:
|
460 |
+
all_self_attentions = all_self_attentions + (layer_outputs[1],)
|
461 |
+
|
462 |
+
if output_hidden_states:
|
463 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
464 |
+
|
465 |
+
if not return_dict:
|
466 |
+
return tuple(
|
467 |
+
v
|
468 |
+
for v in [
|
469 |
+
hidden_states,
|
470 |
+
next_decoder_cache,
|
471 |
+
all_hidden_states,
|
472 |
+
all_self_attentions,
|
473 |
+
all_cross_attentions,
|
474 |
+
]
|
475 |
+
if v is not None
|
476 |
+
)
|
477 |
+
return BaseModelOutputWithPastAndCrossAttentions(
|
478 |
+
last_hidden_state=hidden_states,
|
479 |
+
past_key_values=next_decoder_cache,
|
480 |
+
hidden_states=all_hidden_states,
|
481 |
+
attentions=all_self_attentions,
|
482 |
+
cross_attentions=all_cross_attentions,
|
483 |
+
)
|
484 |
+
|
485 |
+
|
486 |
+
class BertPooler(nn.Module):
|
487 |
+
def __init__(self, config):
|
488 |
+
super().__init__()
|
489 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
490 |
+
self.activation = nn.Tanh()
|
491 |
+
|
492 |
+
def forward(self, hidden_states):
|
493 |
+
# We "pool" the model by simply taking the hidden state corresponding
|
494 |
+
# to the first token.
|
495 |
+
first_token_tensor = hidden_states[:, 0]
|
496 |
+
pooled_output = self.dense(first_token_tensor)
|
497 |
+
pooled_output = self.activation(pooled_output)
|
498 |
+
return pooled_output
|
499 |
+
|
500 |
+
|
501 |
+
class BertPredictionHeadTransform(nn.Module):
|
502 |
+
def __init__(self, config):
|
503 |
+
super().__init__()
|
504 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
505 |
+
if isinstance(config.hidden_act, str):
|
506 |
+
self.transform_act_fn = ACT2FN[config.hidden_act]
|
507 |
+
else:
|
508 |
+
self.transform_act_fn = config.hidden_act
|
509 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
510 |
+
|
511 |
+
def forward(self, hidden_states):
|
512 |
+
hidden_states = self.dense(hidden_states)
|
513 |
+
hidden_states = self.transform_act_fn(hidden_states)
|
514 |
+
hidden_states = self.LayerNorm(hidden_states)
|
515 |
+
return hidden_states
|
516 |
+
|
517 |
+
|
518 |
+
class BertLMPredictionHead(nn.Module):
|
519 |
+
def __init__(self, config):
|
520 |
+
super().__init__()
|
521 |
+
self.transform = BertPredictionHeadTransform(config)
|
522 |
+
|
523 |
+
# The output weights are the same as the input embeddings, but there is
|
524 |
+
# an output-only bias for each token.
|
525 |
+
self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
526 |
+
|
527 |
+
self.bias = nn.Parameter(torch.zeros(config.vocab_size))
|
528 |
+
|
529 |
+
# Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
|
530 |
+
self.decoder.bias = self.bias
|
531 |
+
|
532 |
+
def forward(self, hidden_states):
|
533 |
+
hidden_states = self.transform(hidden_states)
|
534 |
+
hidden_states = self.decoder(hidden_states)
|
535 |
+
return hidden_states
|
536 |
+
|
537 |
+
|
538 |
+
class BertOnlyMLMHead(nn.Module):
|
539 |
+
def __init__(self, config):
|
540 |
+
super().__init__()
|
541 |
+
self.predictions = BertLMPredictionHead(config)
|
542 |
+
|
543 |
+
def forward(self, sequence_output):
|
544 |
+
prediction_scores = self.predictions(sequence_output)
|
545 |
+
return prediction_scores
|
546 |
+
|
547 |
+
|
548 |
+
class BertPreTrainedModel(PreTrainedModel):
|
549 |
+
"""
|
550 |
+
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
551 |
+
models.
|
552 |
+
"""
|
553 |
+
|
554 |
+
config_class = BertConfig
|
555 |
+
base_model_prefix = "bert"
|
556 |
+
_keys_to_ignore_on_load_missing = [r"position_ids"]
|
557 |
+
|
558 |
+
def _init_weights(self, module):
|
559 |
+
""" Initialize the weights """
|
560 |
+
if isinstance(module, (nn.Linear, nn.Embedding)):
|
561 |
+
# Slightly different from the TF version which uses truncated_normal for initialization
|
562 |
+
# cf https://github.com/pytorch/pytorch/pull/5617
|
563 |
+
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
564 |
+
elif isinstance(module, nn.LayerNorm):
|
565 |
+
module.bias.data.zero_()
|
566 |
+
module.weight.data.fill_(1.0)
|
567 |
+
if isinstance(module, nn.Linear) and module.bias is not None:
|
568 |
+
module.bias.data.zero_()
|
569 |
+
|
570 |
+
|
571 |
+
class BertModel(BertPreTrainedModel):
|
572 |
+
"""
|
573 |
+
The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
|
574 |
+
cross-attention is added between the self-attention layers, following the architecture described in `Attention is
|
575 |
+
all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
|
576 |
+
Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
|
577 |
+
argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
|
578 |
+
input to the forward pass.
|
579 |
+
"""
|
580 |
+
|
581 |
+
def __init__(self, config, add_pooling_layer=True):
|
582 |
+
super().__init__(config)
|
583 |
+
self.config = config
|
584 |
+
|
585 |
+
self.embeddings = BertEmbeddings(config)
|
586 |
+
|
587 |
+
self.encoder = BertEncoder(config)
|
588 |
+
|
589 |
+
self.pooler = BertPooler(config) if add_pooling_layer else None
|
590 |
+
|
591 |
+
self.init_weights()
|
592 |
+
|
593 |
+
|
594 |
+
def get_input_embeddings(self):
|
595 |
+
return self.embeddings.word_embeddings
|
596 |
+
|
597 |
+
def set_input_embeddings(self, value):
|
598 |
+
self.embeddings.word_embeddings = value
|
599 |
+
|
600 |
+
def _prune_heads(self, heads_to_prune):
|
601 |
+
"""
|
602 |
+
Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
|
603 |
+
class PreTrainedModel
|
604 |
+
"""
|
605 |
+
for layer, heads in heads_to_prune.items():
|
606 |
+
self.encoder.layer[layer].attention.prune_heads(heads)
|
607 |
+
|
608 |
+
|
609 |
+
def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:
|
610 |
+
"""
|
611 |
+
Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
|
612 |
+
|
613 |
+
Arguments:
|
614 |
+
attention_mask (:obj:`torch.Tensor`):
|
615 |
+
Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
|
616 |
+
input_shape (:obj:`Tuple[int]`):
|
617 |
+
The shape of the input to the model.
|
618 |
+
device: (:obj:`torch.device`):
|
619 |
+
The device of the input to the model.
|
620 |
+
|
621 |
+
Returns:
|
622 |
+
:obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
|
623 |
+
"""
|
624 |
+
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
|
625 |
+
# ourselves in which case we just need to make it broadcastable to all heads.
|
626 |
+
if attention_mask.dim() == 3:
|
627 |
+
extended_attention_mask = attention_mask[:, None, :, :]
|
628 |
+
elif attention_mask.dim() == 2:
|
629 |
+
# Provided a padding mask of dimensions [batch_size, seq_length]
|
630 |
+
# - if the model is a decoder, apply a causal mask in addition to the padding mask
|
631 |
+
# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
|
632 |
+
if is_decoder:
|
633 |
+
batch_size, seq_length = input_shape
|
634 |
+
|
635 |
+
seq_ids = torch.arange(seq_length, device=device)
|
636 |
+
causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
|
637 |
+
# in case past_key_values are used we need to add a prefix ones mask to the causal mask
|
638 |
+
# causal and attention masks must have same type with pytorch version < 1.3
|
639 |
+
causal_mask = causal_mask.to(attention_mask.dtype)
|
640 |
+
|
641 |
+
if causal_mask.shape[1] < attention_mask.shape[1]:
|
642 |
+
prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
|
643 |
+
causal_mask = torch.cat(
|
644 |
+
[
|
645 |
+
torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
|
646 |
+
causal_mask,
|
647 |
+
],
|
648 |
+
axis=-1,
|
649 |
+
)
|
650 |
+
|
651 |
+
extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
|
652 |
+
else:
|
653 |
+
extended_attention_mask = attention_mask[:, None, None, :]
|
654 |
+
else:
|
655 |
+
raise ValueError(
|
656 |
+
"Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
|
657 |
+
input_shape, attention_mask.shape
|
658 |
+
)
|
659 |
+
)
|
660 |
+
|
661 |
+
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
|
662 |
+
# masked positions, this operation will create a tensor which is 0.0 for
|
663 |
+
# positions we want to attend and -10000.0 for masked positions.
|
664 |
+
# Since we are adding it to the raw scores before the softmax, this is
|
665 |
+
# effectively the same as removing these entirely.
|
666 |
+
extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
|
667 |
+
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
|
668 |
+
return extended_attention_mask
|
669 |
+
|
670 |
+
def forward(
|
671 |
+
self,
|
672 |
+
input_ids=None,
|
673 |
+
attention_mask=None,
|
674 |
+
position_ids=None,
|
675 |
+
head_mask=None,
|
676 |
+
inputs_embeds=None,
|
677 |
+
encoder_embeds=None,
|
678 |
+
encoder_hidden_states=None,
|
679 |
+
encoder_attention_mask=None,
|
680 |
+
past_key_values=None,
|
681 |
+
use_cache=None,
|
682 |
+
output_attentions=None,
|
683 |
+
output_hidden_states=None,
|
684 |
+
return_dict=None,
|
685 |
+
is_decoder=False,
|
686 |
+
mode='multimodal',
|
687 |
+
):
|
688 |
+
r"""
|
689 |
+
encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
|
690 |
+
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
|
691 |
+
the model is configured as a decoder.
|
692 |
+
encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
|
693 |
+
Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
|
694 |
+
the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
|
695 |
+
- 1 for tokens that are **not masked**,
|
696 |
+
- 0 for tokens that are **masked**.
|
697 |
+
past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
|
698 |
+
Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
|
699 |
+
If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
|
700 |
+
(those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
|
701 |
+
instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
|
702 |
+
use_cache (:obj:`bool`, `optional`):
|
703 |
+
If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
|
704 |
+
decoding (see :obj:`past_key_values`).
|
705 |
+
"""
|
706 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
707 |
+
output_hidden_states = (
|
708 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
709 |
+
)
|
710 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
711 |
+
|
712 |
+
if is_decoder:
|
713 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
714 |
+
else:
|
715 |
+
use_cache = False
|
716 |
+
|
717 |
+
if input_ids is not None and inputs_embeds is not None:
|
718 |
+
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
|
719 |
+
elif input_ids is not None:
|
720 |
+
input_shape = input_ids.size()
|
721 |
+
batch_size, seq_length = input_shape
|
722 |
+
device = input_ids.device
|
723 |
+
elif inputs_embeds is not None:
|
724 |
+
input_shape = inputs_embeds.size()[:-1]
|
725 |
+
batch_size, seq_length = input_shape
|
726 |
+
device = inputs_embeds.device
|
727 |
+
elif encoder_embeds is not None:
|
728 |
+
input_shape = encoder_embeds.size()[:-1]
|
729 |
+
batch_size, seq_length = input_shape
|
730 |
+
device = encoder_embeds.device
|
731 |
+
else:
|
732 |
+
raise ValueError("You have to specify either input_ids or inputs_embeds or encoder_embeds")
|
733 |
+
|
734 |
+
# past_key_values_length
|
735 |
+
past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
|
736 |
+
|
737 |
+
if attention_mask is None:
|
738 |
+
attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
|
739 |
+
|
740 |
+
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
|
741 |
+
# ourselves in which case we just need to make it broadcastable to all heads.
|
742 |
+
extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape,
|
743 |
+
device, is_decoder)
|
744 |
+
|
745 |
+
# If a 2D or 3D attention mask is provided for the cross-attention
|
746 |
+
# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
|
747 |
+
if encoder_hidden_states is not None:
|
748 |
+
if type(encoder_hidden_states) == list:
|
749 |
+
encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
|
750 |
+
else:
|
751 |
+
encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
|
752 |
+
encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
|
753 |
+
|
754 |
+
if type(encoder_attention_mask) == list:
|
755 |
+
encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
|
756 |
+
elif encoder_attention_mask is None:
|
757 |
+
encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
|
758 |
+
encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
|
759 |
+
else:
|
760 |
+
encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
|
761 |
+
else:
|
762 |
+
encoder_extended_attention_mask = None
|
763 |
+
|
764 |
+
# Prepare head mask if needed
|
765 |
+
# 1.0 in head_mask indicate we keep the head
|
766 |
+
# attention_probs has shape bsz x n_heads x N x N
|
767 |
+
# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
|
768 |
+
# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
|
769 |
+
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
|
770 |
+
|
771 |
+
if encoder_embeds is None:
|
772 |
+
embedding_output = self.embeddings(
|
773 |
+
input_ids=input_ids,
|
774 |
+
position_ids=position_ids,
|
775 |
+
inputs_embeds=inputs_embeds,
|
776 |
+
past_key_values_length=past_key_values_length,
|
777 |
+
)
|
778 |
+
else:
|
779 |
+
embedding_output = encoder_embeds
|
780 |
+
|
781 |
+
encoder_outputs = self.encoder(
|
782 |
+
embedding_output,
|
783 |
+
attention_mask=extended_attention_mask,
|
784 |
+
head_mask=head_mask,
|
785 |
+
encoder_hidden_states=encoder_hidden_states,
|
786 |
+
encoder_attention_mask=encoder_extended_attention_mask,
|
787 |
+
past_key_values=past_key_values,
|
788 |
+
use_cache=use_cache,
|
789 |
+
output_attentions=output_attentions,
|
790 |
+
output_hidden_states=output_hidden_states,
|
791 |
+
return_dict=return_dict,
|
792 |
+
mode=mode,
|
793 |
+
)
|
794 |
+
sequence_output = encoder_outputs[0]
|
795 |
+
pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
|
796 |
+
|
797 |
+
if not return_dict:
|
798 |
+
return (sequence_output, pooled_output) + encoder_outputs[1:]
|
799 |
+
|
800 |
+
return BaseModelOutputWithPoolingAndCrossAttentions(
|
801 |
+
last_hidden_state=sequence_output,
|
802 |
+
pooler_output=pooled_output,
|
803 |
+
past_key_values=encoder_outputs.past_key_values,
|
804 |
+
hidden_states=encoder_outputs.hidden_states,
|
805 |
+
attentions=encoder_outputs.attentions,
|
806 |
+
cross_attentions=encoder_outputs.cross_attentions,
|
807 |
+
)
|
808 |
+
|
809 |
+
|
810 |
+
|
811 |
+
class BertLMHeadModel(BertPreTrainedModel):
|
812 |
+
|
813 |
+
_keys_to_ignore_on_load_unexpected = [r"pooler"]
|
814 |
+
_keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
|
815 |
+
|
816 |
+
def __init__(self, config):
|
817 |
+
super().__init__(config)
|
818 |
+
|
819 |
+
self.bert = BertModel(config, add_pooling_layer=False)
|
820 |
+
self.cls = BertOnlyMLMHead(config)
|
821 |
+
|
822 |
+
self.init_weights()
|
823 |
+
|
824 |
+
def get_output_embeddings(self):
|
825 |
+
return self.cls.predictions.decoder
|
826 |
+
|
827 |
+
def set_output_embeddings(self, new_embeddings):
|
828 |
+
self.cls.predictions.decoder = new_embeddings
|
829 |
+
|
830 |
+
def forward(
|
831 |
+
self,
|
832 |
+
input_ids=None,
|
833 |
+
attention_mask=None,
|
834 |
+
position_ids=None,
|
835 |
+
head_mask=None,
|
836 |
+
inputs_embeds=None,
|
837 |
+
encoder_hidden_states=None,
|
838 |
+
encoder_attention_mask=None,
|
839 |
+
labels=None,
|
840 |
+
past_key_values=None,
|
841 |
+
use_cache=None,
|
842 |
+
output_attentions=None,
|
843 |
+
output_hidden_states=None,
|
844 |
+
return_dict=None,
|
845 |
+
return_logits=False,
|
846 |
+
is_decoder=True,
|
847 |
+
reduction='mean',
|
848 |
+
mode='multimodal',
|
849 |
+
):
|
850 |
+
r"""
|
851 |
+
encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
|
852 |
+
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
|
853 |
+
the model is configured as a decoder.
|
854 |
+
encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
|
855 |
+
Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
|
856 |
+
the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
|
857 |
+
- 1 for tokens that are **not masked**,
|
858 |
+
- 0 for tokens that are **masked**.
|
859 |
+
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
|
860 |
+
Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
|
861 |
+
``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
|
862 |
+
ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
|
863 |
+
past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
|
864 |
+
Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
|
865 |
+
If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
|
866 |
+
(those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
|
867 |
+
instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
|
868 |
+
use_cache (:obj:`bool`, `optional`):
|
869 |
+
If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
|
870 |
+
decoding (see :obj:`past_key_values`).
|
871 |
+
Returns:
|
872 |
+
Example::
|
873 |
+
>>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
|
874 |
+
>>> import torch
|
875 |
+
>>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
|
876 |
+
>>> config = BertConfig.from_pretrained("bert-base-cased")
|
877 |
+
>>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
|
878 |
+
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
|
879 |
+
>>> outputs = model(**inputs)
|
880 |
+
>>> prediction_logits = outputs.logits
|
881 |
+
"""
|
882 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
883 |
+
if labels is not None:
|
884 |
+
use_cache = False
|
885 |
+
|
886 |
+
outputs = self.bert(
|
887 |
+
input_ids,
|
888 |
+
attention_mask=attention_mask,
|
889 |
+
position_ids=position_ids,
|
890 |
+
head_mask=head_mask,
|
891 |
+
inputs_embeds=inputs_embeds,
|
892 |
+
encoder_hidden_states=encoder_hidden_states,
|
893 |
+
encoder_attention_mask=encoder_attention_mask,
|
894 |
+
past_key_values=past_key_values,
|
895 |
+
use_cache=use_cache,
|
896 |
+
output_attentions=output_attentions,
|
897 |
+
output_hidden_states=output_hidden_states,
|
898 |
+
return_dict=return_dict,
|
899 |
+
is_decoder=is_decoder,
|
900 |
+
mode=mode,
|
901 |
+
)
|
902 |
+
|
903 |
+
sequence_output = outputs[0]
|
904 |
+
prediction_scores = self.cls(sequence_output)
|
905 |
+
|
906 |
+
if return_logits:
|
907 |
+
return prediction_scores[:, :-1, :].contiguous()
|
908 |
+
|
909 |
+
lm_loss = None
|
910 |
+
if labels is not None:
|
911 |
+
# we are doing next-token prediction; shift prediction scores and input ids by one
|
912 |
+
shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
|
913 |
+
labels = labels[:, 1:].contiguous()
|
914 |
+
loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)
|
915 |
+
lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
|
916 |
+
if reduction=='none':
|
917 |
+
lm_loss = lm_loss.view(prediction_scores.size(0),-1).sum(1)
|
918 |
+
|
919 |
+
if not return_dict:
|
920 |
+
output = (prediction_scores,) + outputs[2:]
|
921 |
+
return ((lm_loss,) + output) if lm_loss is not None else output
|
922 |
+
|
923 |
+
return CausalLMOutputWithCrossAttentions(
|
924 |
+
loss=lm_loss,
|
925 |
+
logits=prediction_scores,
|
926 |
+
past_key_values=outputs.past_key_values,
|
927 |
+
hidden_states=outputs.hidden_states,
|
928 |
+
attentions=outputs.attentions,
|
929 |
+
cross_attentions=outputs.cross_attentions,
|
930 |
+
)
|
931 |
+
|
932 |
+
def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **model_kwargs):
|
933 |
+
input_shape = input_ids.shape
|
934 |
+
# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
|
935 |
+
if attention_mask is None:
|
936 |
+
attention_mask = input_ids.new_ones(input_shape)
|
937 |
+
|
938 |
+
# cut decoder_input_ids if past is used
|
939 |
+
if past is not None:
|
940 |
+
input_ids = input_ids[:, -1:]
|
941 |
+
|
942 |
+
return {
|
943 |
+
"input_ids": input_ids,
|
944 |
+
"attention_mask": attention_mask,
|
945 |
+
"past_key_values": past,
|
946 |
+
"encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
|
947 |
+
"encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
|
948 |
+
"is_decoder": True,
|
949 |
+
}
|
950 |
+
|
951 |
+
def _reorder_cache(self, past, beam_idx):
|
952 |
+
reordered_past = ()
|
953 |
+
for layer_past in past:
|
954 |
+
reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
|
955 |
+
return reordered_past
|
extras/BLIP/models/nlvr_encoder.py
ADDED
@@ -0,0 +1,843 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import math
|
2 |
+
import os
|
3 |
+
import warnings
|
4 |
+
from dataclasses import dataclass
|
5 |
+
from typing import Optional, Tuple
|
6 |
+
|
7 |
+
import torch
|
8 |
+
from torch import Tensor, device, dtype, nn
|
9 |
+
import torch.utils.checkpoint
|
10 |
+
from torch import nn
|
11 |
+
from torch.nn import CrossEntropyLoss
|
12 |
+
import torch.nn.functional as F
|
13 |
+
|
14 |
+
from transformers.activations import ACT2FN
|
15 |
+
from transformers.file_utils import (
|
16 |
+
ModelOutput,
|
17 |
+
)
|
18 |
+
from transformers.modeling_outputs import (
|
19 |
+
BaseModelOutputWithPastAndCrossAttentions,
|
20 |
+
BaseModelOutputWithPoolingAndCrossAttentions,
|
21 |
+
CausalLMOutputWithCrossAttentions,
|
22 |
+
MaskedLMOutput,
|
23 |
+
MultipleChoiceModelOutput,
|
24 |
+
NextSentencePredictorOutput,
|
25 |
+
QuestionAnsweringModelOutput,
|
26 |
+
SequenceClassifierOutput,
|
27 |
+
TokenClassifierOutput,
|
28 |
+
)
|
29 |
+
from transformers.modeling_utils import (
|
30 |
+
PreTrainedModel,
|
31 |
+
apply_chunking_to_forward,
|
32 |
+
find_pruneable_heads_and_indices,
|
33 |
+
prune_linear_layer,
|
34 |
+
)
|
35 |
+
from transformers.utils import logging
|
36 |
+
from transformers.models.bert.configuration_bert import BertConfig
|
37 |
+
|
38 |
+
|
39 |
+
logger = logging.get_logger(__name__)
|
40 |
+
|
41 |
+
|
42 |
+
class BertEmbeddings(nn.Module):
|
43 |
+
"""Construct the embeddings from word and position embeddings."""
|
44 |
+
|
45 |
+
def __init__(self, config):
|
46 |
+
super().__init__()
|
47 |
+
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
|
48 |
+
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
|
49 |
+
|
50 |
+
# self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
|
51 |
+
# any TensorFlow checkpoint file
|
52 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
53 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
54 |
+
|
55 |
+
# position_ids (1, len position emb) is contiguous in memory and exported when serialized
|
56 |
+
self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
|
57 |
+
self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
|
58 |
+
|
59 |
+
self.config = config
|
60 |
+
|
61 |
+
def forward(
|
62 |
+
self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
|
63 |
+
):
|
64 |
+
if input_ids is not None:
|
65 |
+
input_shape = input_ids.size()
|
66 |
+
else:
|
67 |
+
input_shape = inputs_embeds.size()[:-1]
|
68 |
+
|
69 |
+
seq_length = input_shape[1]
|
70 |
+
|
71 |
+
if position_ids is None:
|
72 |
+
position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
|
73 |
+
|
74 |
+
if inputs_embeds is None:
|
75 |
+
inputs_embeds = self.word_embeddings(input_ids)
|
76 |
+
|
77 |
+
embeddings = inputs_embeds
|
78 |
+
|
79 |
+
if self.position_embedding_type == "absolute":
|
80 |
+
position_embeddings = self.position_embeddings(position_ids)
|
81 |
+
embeddings += position_embeddings
|
82 |
+
embeddings = self.LayerNorm(embeddings)
|
83 |
+
embeddings = self.dropout(embeddings)
|
84 |
+
return embeddings
|
85 |
+
|
86 |
+
|
87 |
+
class BertSelfAttention(nn.Module):
|
88 |
+
def __init__(self, config, is_cross_attention):
|
89 |
+
super().__init__()
|
90 |
+
self.config = config
|
91 |
+
if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
|
92 |
+
raise ValueError(
|
93 |
+
"The hidden size (%d) is not a multiple of the number of attention "
|
94 |
+
"heads (%d)" % (config.hidden_size, config.num_attention_heads)
|
95 |
+
)
|
96 |
+
|
97 |
+
self.num_attention_heads = config.num_attention_heads
|
98 |
+
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
|
99 |
+
self.all_head_size = self.num_attention_heads * self.attention_head_size
|
100 |
+
|
101 |
+
self.query = nn.Linear(config.hidden_size, self.all_head_size)
|
102 |
+
if is_cross_attention:
|
103 |
+
self.key = nn.Linear(config.encoder_width, self.all_head_size)
|
104 |
+
self.value = nn.Linear(config.encoder_width, self.all_head_size)
|
105 |
+
else:
|
106 |
+
self.key = nn.Linear(config.hidden_size, self.all_head_size)
|
107 |
+
self.value = nn.Linear(config.hidden_size, self.all_head_size)
|
108 |
+
|
109 |
+
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
110 |
+
self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
|
111 |
+
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
112 |
+
self.max_position_embeddings = config.max_position_embeddings
|
113 |
+
self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
|
114 |
+
self.save_attention = False
|
115 |
+
|
116 |
+
def save_attn_gradients(self, attn_gradients):
|
117 |
+
self.attn_gradients = attn_gradients
|
118 |
+
|
119 |
+
def get_attn_gradients(self):
|
120 |
+
return self.attn_gradients
|
121 |
+
|
122 |
+
def save_attention_map(self, attention_map):
|
123 |
+
self.attention_map = attention_map
|
124 |
+
|
125 |
+
def get_attention_map(self):
|
126 |
+
return self.attention_map
|
127 |
+
|
128 |
+
def transpose_for_scores(self, x):
|
129 |
+
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
|
130 |
+
x = x.view(*new_x_shape)
|
131 |
+
return x.permute(0, 2, 1, 3)
|
132 |
+
|
133 |
+
def forward(
|
134 |
+
self,
|
135 |
+
hidden_states,
|
136 |
+
attention_mask=None,
|
137 |
+
head_mask=None,
|
138 |
+
encoder_hidden_states=None,
|
139 |
+
encoder_attention_mask=None,
|
140 |
+
past_key_value=None,
|
141 |
+
output_attentions=False,
|
142 |
+
):
|
143 |
+
mixed_query_layer = self.query(hidden_states)
|
144 |
+
|
145 |
+
# If this is instantiated as a cross-attention module, the keys
|
146 |
+
# and values come from an encoder; the attention mask needs to be
|
147 |
+
# such that the encoder's padding tokens are not attended to.
|
148 |
+
is_cross_attention = encoder_hidden_states is not None
|
149 |
+
|
150 |
+
if is_cross_attention:
|
151 |
+
key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
|
152 |
+
value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
|
153 |
+
attention_mask = encoder_attention_mask
|
154 |
+
elif past_key_value is not None:
|
155 |
+
key_layer = self.transpose_for_scores(self.key(hidden_states))
|
156 |
+
value_layer = self.transpose_for_scores(self.value(hidden_states))
|
157 |
+
key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
|
158 |
+
value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
|
159 |
+
else:
|
160 |
+
key_layer = self.transpose_for_scores(self.key(hidden_states))
|
161 |
+
value_layer = self.transpose_for_scores(self.value(hidden_states))
|
162 |
+
|
163 |
+
query_layer = self.transpose_for_scores(mixed_query_layer)
|
164 |
+
|
165 |
+
past_key_value = (key_layer, value_layer)
|
166 |
+
|
167 |
+
# Take the dot product between "query" and "key" to get the raw attention scores.
|
168 |
+
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
|
169 |
+
|
170 |
+
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
171 |
+
seq_length = hidden_states.size()[1]
|
172 |
+
position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
|
173 |
+
position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
|
174 |
+
distance = position_ids_l - position_ids_r
|
175 |
+
positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
|
176 |
+
positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
|
177 |
+
|
178 |
+
if self.position_embedding_type == "relative_key":
|
179 |
+
relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
180 |
+
attention_scores = attention_scores + relative_position_scores
|
181 |
+
elif self.position_embedding_type == "relative_key_query":
|
182 |
+
relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
183 |
+
relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
|
184 |
+
attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
|
185 |
+
|
186 |
+
attention_scores = attention_scores / math.sqrt(self.attention_head_size)
|
187 |
+
if attention_mask is not None:
|
188 |
+
# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
|
189 |
+
attention_scores = attention_scores + attention_mask
|
190 |
+
|
191 |
+
# Normalize the attention scores to probabilities.
|
192 |
+
attention_probs = nn.Softmax(dim=-1)(attention_scores)
|
193 |
+
|
194 |
+
if is_cross_attention and self.save_attention:
|
195 |
+
self.save_attention_map(attention_probs)
|
196 |
+
attention_probs.register_hook(self.save_attn_gradients)
|
197 |
+
|
198 |
+
# This is actually dropping out entire tokens to attend to, which might
|
199 |
+
# seem a bit unusual, but is taken from the original Transformer paper.
|
200 |
+
attention_probs_dropped = self.dropout(attention_probs)
|
201 |
+
|
202 |
+
# Mask heads if we want to
|
203 |
+
if head_mask is not None:
|
204 |
+
attention_probs_dropped = attention_probs_dropped * head_mask
|
205 |
+
|
206 |
+
context_layer = torch.matmul(attention_probs_dropped, value_layer)
|
207 |
+
|
208 |
+
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
|
209 |
+
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
|
210 |
+
context_layer = context_layer.view(*new_context_layer_shape)
|
211 |
+
|
212 |
+
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
|
213 |
+
|
214 |
+
outputs = outputs + (past_key_value,)
|
215 |
+
return outputs
|
216 |
+
|
217 |
+
|
218 |
+
class BertSelfOutput(nn.Module):
|
219 |
+
def __init__(self, config, twin=False, merge=False):
|
220 |
+
super().__init__()
|
221 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
222 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
223 |
+
if twin:
|
224 |
+
self.dense0 = nn.Linear(config.hidden_size, config.hidden_size)
|
225 |
+
self.dense1 = nn.Linear(config.hidden_size, config.hidden_size)
|
226 |
+
else:
|
227 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
228 |
+
if merge:
|
229 |
+
self.act = ACT2FN[config.hidden_act]
|
230 |
+
self.merge_layer = nn.Linear(config.hidden_size * 2, config.hidden_size)
|
231 |
+
self.merge = True
|
232 |
+
else:
|
233 |
+
self.merge = False
|
234 |
+
|
235 |
+
def forward(self, hidden_states, input_tensor):
|
236 |
+
if type(hidden_states) == list:
|
237 |
+
hidden_states0 = self.dense0(hidden_states[0])
|
238 |
+
hidden_states1 = self.dense1(hidden_states[1])
|
239 |
+
if self.merge:
|
240 |
+
#hidden_states = self.merge_layer(self.act(torch.cat([hidden_states0,hidden_states1],dim=-1)))
|
241 |
+
hidden_states = self.merge_layer(torch.cat([hidden_states0,hidden_states1],dim=-1))
|
242 |
+
else:
|
243 |
+
hidden_states = (hidden_states0+hidden_states1)/2
|
244 |
+
else:
|
245 |
+
hidden_states = self.dense(hidden_states)
|
246 |
+
hidden_states = self.dropout(hidden_states)
|
247 |
+
hidden_states = self.LayerNorm(hidden_states + input_tensor)
|
248 |
+
return hidden_states
|
249 |
+
|
250 |
+
|
251 |
+
class BertAttention(nn.Module):
|
252 |
+
def __init__(self, config, is_cross_attention=False, layer_num=-1):
|
253 |
+
super().__init__()
|
254 |
+
if is_cross_attention:
|
255 |
+
self.self0 = BertSelfAttention(config, is_cross_attention)
|
256 |
+
self.self1 = BertSelfAttention(config, is_cross_attention)
|
257 |
+
else:
|
258 |
+
self.self = BertSelfAttention(config, is_cross_attention)
|
259 |
+
self.output = BertSelfOutput(config, twin=is_cross_attention, merge=(is_cross_attention and layer_num>=6))
|
260 |
+
self.pruned_heads = set()
|
261 |
+
|
262 |
+
def prune_heads(self, heads):
|
263 |
+
if len(heads) == 0:
|
264 |
+
return
|
265 |
+
heads, index = find_pruneable_heads_and_indices(
|
266 |
+
heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
|
267 |
+
)
|
268 |
+
|
269 |
+
# Prune linear layers
|
270 |
+
self.self.query = prune_linear_layer(self.self.query, index)
|
271 |
+
self.self.key = prune_linear_layer(self.self.key, index)
|
272 |
+
self.self.value = prune_linear_layer(self.self.value, index)
|
273 |
+
self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
|
274 |
+
|
275 |
+
# Update hyper params and store pruned heads
|
276 |
+
self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
|
277 |
+
self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
|
278 |
+
self.pruned_heads = self.pruned_heads.union(heads)
|
279 |
+
|
280 |
+
def forward(
|
281 |
+
self,
|
282 |
+
hidden_states,
|
283 |
+
attention_mask=None,
|
284 |
+
head_mask=None,
|
285 |
+
encoder_hidden_states=None,
|
286 |
+
encoder_attention_mask=None,
|
287 |
+
past_key_value=None,
|
288 |
+
output_attentions=False,
|
289 |
+
):
|
290 |
+
if type(encoder_hidden_states)==list:
|
291 |
+
self_outputs0 = self.self0(
|
292 |
+
hidden_states,
|
293 |
+
attention_mask,
|
294 |
+
head_mask,
|
295 |
+
encoder_hidden_states[0],
|
296 |
+
encoder_attention_mask[0],
|
297 |
+
past_key_value,
|
298 |
+
output_attentions,
|
299 |
+
)
|
300 |
+
self_outputs1 = self.self1(
|
301 |
+
hidden_states,
|
302 |
+
attention_mask,
|
303 |
+
head_mask,
|
304 |
+
encoder_hidden_states[1],
|
305 |
+
encoder_attention_mask[1],
|
306 |
+
past_key_value,
|
307 |
+
output_attentions,
|
308 |
+
)
|
309 |
+
attention_output = self.output([self_outputs0[0],self_outputs1[0]], hidden_states)
|
310 |
+
|
311 |
+
outputs = (attention_output,) + self_outputs0[1:] # add attentions if we output them
|
312 |
+
else:
|
313 |
+
self_outputs = self.self(
|
314 |
+
hidden_states,
|
315 |
+
attention_mask,
|
316 |
+
head_mask,
|
317 |
+
encoder_hidden_states,
|
318 |
+
encoder_attention_mask,
|
319 |
+
past_key_value,
|
320 |
+
output_attentions,
|
321 |
+
)
|
322 |
+
attention_output = self.output(self_outputs[0], hidden_states)
|
323 |
+
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
|
324 |
+
return outputs
|
325 |
+
|
326 |
+
|
327 |
+
class BertIntermediate(nn.Module):
|
328 |
+
def __init__(self, config):
|
329 |
+
super().__init__()
|
330 |
+
self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
|
331 |
+
if isinstance(config.hidden_act, str):
|
332 |
+
self.intermediate_act_fn = ACT2FN[config.hidden_act]
|
333 |
+
else:
|
334 |
+
self.intermediate_act_fn = config.hidden_act
|
335 |
+
|
336 |
+
def forward(self, hidden_states):
|
337 |
+
hidden_states = self.dense(hidden_states)
|
338 |
+
hidden_states = self.intermediate_act_fn(hidden_states)
|
339 |
+
return hidden_states
|
340 |
+
|
341 |
+
|
342 |
+
class BertOutput(nn.Module):
|
343 |
+
def __init__(self, config):
|
344 |
+
super().__init__()
|
345 |
+
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
|
346 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
347 |
+
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
348 |
+
|
349 |
+
def forward(self, hidden_states, input_tensor):
|
350 |
+
hidden_states = self.dense(hidden_states)
|
351 |
+
hidden_states = self.dropout(hidden_states)
|
352 |
+
hidden_states = self.LayerNorm(hidden_states + input_tensor)
|
353 |
+
return hidden_states
|
354 |
+
|
355 |
+
|
356 |
+
class BertLayer(nn.Module):
|
357 |
+
def __init__(self, config, layer_num):
|
358 |
+
super().__init__()
|
359 |
+
self.config = config
|
360 |
+
self.chunk_size_feed_forward = config.chunk_size_feed_forward
|
361 |
+
self.seq_len_dim = 1
|
362 |
+
self.attention = BertAttention(config)
|
363 |
+
self.layer_num = layer_num
|
364 |
+
if self.config.add_cross_attention:
|
365 |
+
self.crossattention = BertAttention(config, is_cross_attention=self.config.add_cross_attention, layer_num=layer_num)
|
366 |
+
self.intermediate = BertIntermediate(config)
|
367 |
+
self.output = BertOutput(config)
|
368 |
+
|
369 |
+
def forward(
|
370 |
+
self,
|
371 |
+
hidden_states,
|
372 |
+
attention_mask=None,
|
373 |
+
head_mask=None,
|
374 |
+
encoder_hidden_states=None,
|
375 |
+
encoder_attention_mask=None,
|
376 |
+
past_key_value=None,
|
377 |
+
output_attentions=False,
|
378 |
+
mode=None,
|
379 |
+
):
|
380 |
+
# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
|
381 |
+
self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
|
382 |
+
self_attention_outputs = self.attention(
|
383 |
+
hidden_states,
|
384 |
+
attention_mask,
|
385 |
+
head_mask,
|
386 |
+
output_attentions=output_attentions,
|
387 |
+
past_key_value=self_attn_past_key_value,
|
388 |
+
)
|
389 |
+
attention_output = self_attention_outputs[0]
|
390 |
+
|
391 |
+
outputs = self_attention_outputs[1:-1]
|
392 |
+
present_key_value = self_attention_outputs[-1]
|
393 |
+
|
394 |
+
if mode=='multimodal':
|
395 |
+
assert encoder_hidden_states is not None, "encoder_hidden_states must be given for cross-attention layers"
|
396 |
+
cross_attention_outputs = self.crossattention(
|
397 |
+
attention_output,
|
398 |
+
attention_mask,
|
399 |
+
head_mask,
|
400 |
+
encoder_hidden_states,
|
401 |
+
encoder_attention_mask,
|
402 |
+
output_attentions=output_attentions,
|
403 |
+
)
|
404 |
+
attention_output = cross_attention_outputs[0]
|
405 |
+
outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
|
406 |
+
layer_output = apply_chunking_to_forward(
|
407 |
+
self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
|
408 |
+
)
|
409 |
+
outputs = (layer_output,) + outputs
|
410 |
+
|
411 |
+
outputs = outputs + (present_key_value,)
|
412 |
+
|
413 |
+
return outputs
|
414 |
+
|
415 |
+
def feed_forward_chunk(self, attention_output):
|
416 |
+
intermediate_output = self.intermediate(attention_output)
|
417 |
+
layer_output = self.output(intermediate_output, attention_output)
|
418 |
+
return layer_output
|
419 |
+
|
420 |
+
|
421 |
+
class BertEncoder(nn.Module):
|
422 |
+
def __init__(self, config):
|
423 |
+
super().__init__()
|
424 |
+
self.config = config
|
425 |
+
self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)])
|
426 |
+
self.gradient_checkpointing = False
|
427 |
+
|
428 |
+
def forward(
|
429 |
+
self,
|
430 |
+
hidden_states,
|
431 |
+
attention_mask=None,
|
432 |
+
head_mask=None,
|
433 |
+
encoder_hidden_states=None,
|
434 |
+
encoder_attention_mask=None,
|
435 |
+
past_key_values=None,
|
436 |
+
use_cache=None,
|
437 |
+
output_attentions=False,
|
438 |
+
output_hidden_states=False,
|
439 |
+
return_dict=True,
|
440 |
+
mode='multimodal',
|
441 |
+
):
|
442 |
+
all_hidden_states = () if output_hidden_states else None
|
443 |
+
all_self_attentions = () if output_attentions else None
|
444 |
+
all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
|
445 |
+
|
446 |
+
next_decoder_cache = () if use_cache else None
|
447 |
+
|
448 |
+
for i in range(self.config.num_hidden_layers):
|
449 |
+
layer_module = self.layer[i]
|
450 |
+
if output_hidden_states:
|
451 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
452 |
+
|
453 |
+
layer_head_mask = head_mask[i] if head_mask is not None else None
|
454 |
+
past_key_value = past_key_values[i] if past_key_values is not None else None
|
455 |
+
|
456 |
+
if self.gradient_checkpointing and self.training:
|
457 |
+
|
458 |
+
if use_cache:
|
459 |
+
logger.warn(
|
460 |
+
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
461 |
+
)
|
462 |
+
use_cache = False
|
463 |
+
|
464 |
+
def create_custom_forward(module):
|
465 |
+
def custom_forward(*inputs):
|
466 |
+
return module(*inputs, past_key_value, output_attentions)
|
467 |
+
|
468 |
+
return custom_forward
|
469 |
+
|
470 |
+
layer_outputs = torch.utils.checkpoint.checkpoint(
|
471 |
+
create_custom_forward(layer_module),
|
472 |
+
hidden_states,
|
473 |
+
attention_mask,
|
474 |
+
layer_head_mask,
|
475 |
+
encoder_hidden_states,
|
476 |
+
encoder_attention_mask,
|
477 |
+
mode=mode,
|
478 |
+
)
|
479 |
+
else:
|
480 |
+
layer_outputs = layer_module(
|
481 |
+
hidden_states,
|
482 |
+
attention_mask,
|
483 |
+
layer_head_mask,
|
484 |
+
encoder_hidden_states,
|
485 |
+
encoder_attention_mask,
|
486 |
+
past_key_value,
|
487 |
+
output_attentions,
|
488 |
+
mode=mode,
|
489 |
+
)
|
490 |
+
|
491 |
+
hidden_states = layer_outputs[0]
|
492 |
+
if use_cache:
|
493 |
+
next_decoder_cache += (layer_outputs[-1],)
|
494 |
+
if output_attentions:
|
495 |
+
all_self_attentions = all_self_attentions + (layer_outputs[1],)
|
496 |
+
|
497 |
+
if output_hidden_states:
|
498 |
+
all_hidden_states = all_hidden_states + (hidden_states,)
|
499 |
+
|
500 |
+
if not return_dict:
|
501 |
+
return tuple(
|
502 |
+
v
|
503 |
+
for v in [
|
504 |
+
hidden_states,
|
505 |
+
next_decoder_cache,
|
506 |
+
all_hidden_states,
|
507 |
+
all_self_attentions,
|
508 |
+
all_cross_attentions,
|
509 |
+
]
|
510 |
+
if v is not None
|
511 |
+
)
|
512 |
+
return BaseModelOutputWithPastAndCrossAttentions(
|
513 |
+
last_hidden_state=hidden_states,
|
514 |
+
past_key_values=next_decoder_cache,
|
515 |
+
hidden_states=all_hidden_states,
|
516 |
+
attentions=all_self_attentions,
|
517 |
+
cross_attentions=all_cross_attentions,
|
518 |
+
)
|
519 |
+
|
520 |
+
|
521 |
+
class BertPooler(nn.Module):
|
522 |
+
def __init__(self, config):
|
523 |
+
super().__init__()
|
524 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
525 |
+
self.activation = nn.Tanh()
|
526 |
+
|
527 |
+
def forward(self, hidden_states):
|
528 |
+
# We "pool" the model by simply taking the hidden state corresponding
|
529 |
+
# to the first token.
|
530 |
+
first_token_tensor = hidden_states[:, 0]
|
531 |
+
pooled_output = self.dense(first_token_tensor)
|
532 |
+
pooled_output = self.activation(pooled_output)
|
533 |
+
return pooled_output
|
534 |
+
|
535 |
+
|
536 |
+
class BertPredictionHeadTransform(nn.Module):
|
537 |
+
def __init__(self, config):
|
538 |
+
super().__init__()
|
539 |
+
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
540 |
+
if isinstance(config.hidden_act, str):
|
541 |
+
self.transform_act_fn = ACT2FN[config.hidden_act]
|
542 |
+
else:
|
543 |
+
self.transform_act_fn = config.hidden_act
|
544 |
+
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
545 |
+
|
546 |
+
def forward(self, hidden_states):
|
547 |
+
hidden_states = self.dense(hidden_states)
|
548 |
+
hidden_states = self.transform_act_fn(hidden_states)
|
549 |
+
hidden_states = self.LayerNorm(hidden_states)
|
550 |
+
return hidden_states
|
551 |
+
|
552 |
+
|
553 |
+
class BertLMPredictionHead(nn.Module):
|
554 |
+
def __init__(self, config):
|
555 |
+
super().__init__()
|
556 |
+
self.transform = BertPredictionHeadTransform(config)
|
557 |
+
|
558 |
+
# The output weights are the same as the input embeddings, but there is
|
559 |
+
# an output-only bias for each token.
|
560 |
+
self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
561 |
+
|
562 |
+
self.bias = nn.Parameter(torch.zeros(config.vocab_size))
|
563 |
+
|
564 |
+
# Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
|
565 |
+
self.decoder.bias = self.bias
|
566 |
+
|
567 |
+
def forward(self, hidden_states):
|
568 |
+
hidden_states = self.transform(hidden_states)
|
569 |
+
hidden_states = self.decoder(hidden_states)
|
570 |
+
return hidden_states
|
571 |
+
|
572 |
+
|
573 |
+
class BertOnlyMLMHead(nn.Module):
|
574 |
+
def __init__(self, config):
|
575 |
+
super().__init__()
|
576 |
+
self.predictions = BertLMPredictionHead(config)
|
577 |
+
|
578 |
+
def forward(self, sequence_output):
|
579 |
+
prediction_scores = self.predictions(sequence_output)
|
580 |
+
return prediction_scores
|
581 |
+
|
582 |
+
|
583 |
+
class BertPreTrainedModel(PreTrainedModel):
|
584 |
+
"""
|
585 |
+
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
586 |
+
models.
|
587 |
+
"""
|
588 |
+
|
589 |
+
config_class = BertConfig
|
590 |
+
base_model_prefix = "bert"
|
591 |
+
_keys_to_ignore_on_load_missing = [r"position_ids"]
|
592 |
+
|
593 |
+
def _init_weights(self, module):
|
594 |
+
""" Initialize the weights """
|
595 |
+
if isinstance(module, (nn.Linear, nn.Embedding)):
|
596 |
+
# Slightly different from the TF version which uses truncated_normal for initialization
|
597 |
+
# cf https://github.com/pytorch/pytorch/pull/5617
|
598 |
+
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
599 |
+
elif isinstance(module, nn.LayerNorm):
|
600 |
+
module.bias.data.zero_()
|
601 |
+
module.weight.data.fill_(1.0)
|
602 |
+
if isinstance(module, nn.Linear) and module.bias is not None:
|
603 |
+
module.bias.data.zero_()
|
604 |
+
|
605 |
+
|
606 |
+
class BertModel(BertPreTrainedModel):
|
607 |
+
"""
|
608 |
+
The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
|
609 |
+
cross-attention is added between the self-attention layers, following the architecture described in `Attention is
|
610 |
+
all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
|
611 |
+
Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
|
612 |
+
argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
|
613 |
+
input to the forward pass.
|
614 |
+
"""
|
615 |
+
|
616 |
+
def __init__(self, config, add_pooling_layer=True):
|
617 |
+
super().__init__(config)
|
618 |
+
self.config = config
|
619 |
+
|
620 |
+
self.embeddings = BertEmbeddings(config)
|
621 |
+
|
622 |
+
self.encoder = BertEncoder(config)
|
623 |
+
|
624 |
+
self.pooler = BertPooler(config) if add_pooling_layer else None
|
625 |
+
|
626 |
+
self.init_weights()
|
627 |
+
|
628 |
+
|
629 |
+
def get_input_embeddings(self):
|
630 |
+
return self.embeddings.word_embeddings
|
631 |
+
|
632 |
+
def set_input_embeddings(self, value):
|
633 |
+
self.embeddings.word_embeddings = value
|
634 |
+
|
635 |
+
def _prune_heads(self, heads_to_prune):
|
636 |
+
"""
|
637 |
+
Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
|
638 |
+
class PreTrainedModel
|
639 |
+
"""
|
640 |
+
for layer, heads in heads_to_prune.items():
|
641 |
+
self.encoder.layer[layer].attention.prune_heads(heads)
|
642 |
+
|
643 |
+
|
644 |
+
def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:
|
645 |
+
"""
|
646 |
+
Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
|
647 |
+
|
648 |
+
Arguments:
|
649 |
+
attention_mask (:obj:`torch.Tensor`):
|
650 |
+
Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
|
651 |
+
input_shape (:obj:`Tuple[int]`):
|
652 |
+
The shape of the input to the model.
|
653 |
+
device: (:obj:`torch.device`):
|
654 |
+
The device of the input to the model.
|
655 |
+
|
656 |
+
Returns:
|
657 |
+
:obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
|
658 |
+
"""
|
659 |
+
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
|
660 |
+
# ourselves in which case we just need to make it broadcastable to all heads.
|
661 |
+
if attention_mask.dim() == 3:
|
662 |
+
extended_attention_mask = attention_mask[:, None, :, :]
|
663 |
+
elif attention_mask.dim() == 2:
|
664 |
+
# Provided a padding mask of dimensions [batch_size, seq_length]
|
665 |
+
# - if the model is a decoder, apply a causal mask in addition to the padding mask
|
666 |
+
# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
|
667 |
+
if is_decoder:
|
668 |
+
batch_size, seq_length = input_shape
|
669 |
+
|
670 |
+
seq_ids = torch.arange(seq_length, device=device)
|
671 |
+
causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
|
672 |
+
# in case past_key_values are used we need to add a prefix ones mask to the causal mask
|
673 |
+
# causal and attention masks must have same type with pytorch version < 1.3
|
674 |
+
causal_mask = causal_mask.to(attention_mask.dtype)
|
675 |
+
|
676 |
+
if causal_mask.shape[1] < attention_mask.shape[1]:
|
677 |
+
prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
|
678 |
+
causal_mask = torch.cat(
|
679 |
+
[
|
680 |
+
torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
|
681 |
+
causal_mask,
|
682 |
+
],
|
683 |
+
axis=-1,
|
684 |
+
)
|
685 |
+
|
686 |
+
extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
|
687 |
+
else:
|
688 |
+
extended_attention_mask = attention_mask[:, None, None, :]
|
689 |
+
else:
|
690 |
+
raise ValueError(
|
691 |
+
"Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
|
692 |
+
input_shape, attention_mask.shape
|
693 |
+
)
|
694 |
+
)
|
695 |
+
|
696 |
+
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
|
697 |
+
# masked positions, this operation will create a tensor which is 0.0 for
|
698 |
+
# positions we want to attend and -10000.0 for masked positions.
|
699 |
+
# Since we are adding it to the raw scores before the softmax, this is
|
700 |
+
# effectively the same as removing these entirely.
|
701 |
+
extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
|
702 |
+
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
|
703 |
+
return extended_attention_mask
|
704 |
+
|
705 |
+
def forward(
|
706 |
+
self,
|
707 |
+
input_ids=None,
|
708 |
+
attention_mask=None,
|
709 |
+
position_ids=None,
|
710 |
+
head_mask=None,
|
711 |
+
inputs_embeds=None,
|
712 |
+
encoder_embeds=None,
|
713 |
+
encoder_hidden_states=None,
|
714 |
+
encoder_attention_mask=None,
|
715 |
+
past_key_values=None,
|
716 |
+
use_cache=None,
|
717 |
+
output_attentions=None,
|
718 |
+
output_hidden_states=None,
|
719 |
+
return_dict=None,
|
720 |
+
is_decoder=False,
|
721 |
+
mode='multimodal',
|
722 |
+
):
|
723 |
+
r"""
|
724 |
+
encoder_hidden_states (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
|
725 |
+
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
|
726 |
+
the model is configured as a decoder.
|
727 |
+
encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
|
728 |
+
Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
|
729 |
+
the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
|
730 |
+
- 1 for tokens that are **not masked**,
|
731 |
+
- 0 for tokens that are **masked**.
|
732 |
+
past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
|
733 |
+
Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
|
734 |
+
If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
|
735 |
+
(those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
|
736 |
+
instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
|
737 |
+
use_cache (:obj:`bool`, `optional`):
|
738 |
+
If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
|
739 |
+
decoding (see :obj:`past_key_values`).
|
740 |
+
"""
|
741 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
742 |
+
output_hidden_states = (
|
743 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
744 |
+
)
|
745 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
746 |
+
|
747 |
+
if is_decoder:
|
748 |
+
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
749 |
+
else:
|
750 |
+
use_cache = False
|
751 |
+
|
752 |
+
if input_ids is not None and inputs_embeds is not None:
|
753 |
+
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
|
754 |
+
elif input_ids is not None:
|
755 |
+
input_shape = input_ids.size()
|
756 |
+
batch_size, seq_length = input_shape
|
757 |
+
device = input_ids.device
|
758 |
+
elif inputs_embeds is not None:
|
759 |
+
input_shape = inputs_embeds.size()[:-1]
|
760 |
+
batch_size, seq_length = input_shape
|
761 |
+
device = inputs_embeds.device
|
762 |
+
elif encoder_embeds is not None:
|
763 |
+
input_shape = encoder_embeds.size()[:-1]
|
764 |
+
batch_size, seq_length = input_shape
|
765 |
+
device = encoder_embeds.device
|
766 |
+
else:
|
767 |
+
raise ValueError("You have to specify either input_ids or inputs_embeds or encoder_embeds")
|
768 |
+
|
769 |
+
# past_key_values_length
|
770 |
+
past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
|
771 |
+
|
772 |
+
if attention_mask is None:
|
773 |
+
attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
|
774 |
+
|
775 |
+
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
|
776 |
+
# ourselves in which case we just need to make it broadcastable to all heads.
|
777 |
+
extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape,
|
778 |
+
device, is_decoder)
|
779 |
+
|
780 |
+
# If a 2D or 3D attention mask is provided for the cross-attention
|
781 |
+
# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
|
782 |
+
if encoder_hidden_states is not None:
|
783 |
+
if type(encoder_hidden_states) == list:
|
784 |
+
encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
|
785 |
+
else:
|
786 |
+
encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
|
787 |
+
encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
|
788 |
+
|
789 |
+
if type(encoder_attention_mask) == list:
|
790 |
+
encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
|
791 |
+
elif encoder_attention_mask is None:
|
792 |
+
encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
|
793 |
+
encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
|
794 |
+
else:
|
795 |
+
encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
|
796 |
+
else:
|
797 |
+
encoder_extended_attention_mask = None
|
798 |
+
|
799 |
+
# Prepare head mask if needed
|
800 |
+
# 1.0 in head_mask indicate we keep the head
|
801 |
+
# attention_probs has shape bsz x n_heads x N x N
|
802 |
+
# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
|
803 |
+
# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
|
804 |
+
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
|
805 |
+
|
806 |
+
if encoder_embeds is None:
|
807 |
+
embedding_output = self.embeddings(
|
808 |
+
input_ids=input_ids,
|
809 |
+
position_ids=position_ids,
|
810 |
+
inputs_embeds=inputs_embeds,
|
811 |
+
past_key_values_length=past_key_values_length,
|
812 |
+
)
|
813 |
+
else:
|
814 |
+
embedding_output = encoder_embeds
|
815 |
+
|
816 |
+
encoder_outputs = self.encoder(
|
817 |
+
embedding_output,
|
818 |
+
attention_mask=extended_attention_mask,
|
819 |
+
head_mask=head_mask,
|
820 |
+
encoder_hidden_states=encoder_hidden_states,
|
821 |
+
encoder_attention_mask=encoder_extended_attention_mask,
|
822 |
+
past_key_values=past_key_values,
|
823 |
+
use_cache=use_cache,
|
824 |
+
output_attentions=output_attentions,
|
825 |
+
output_hidden_states=output_hidden_states,
|
826 |
+
return_dict=return_dict,
|
827 |
+
mode=mode,
|
828 |
+
)
|
829 |
+
sequence_output = encoder_outputs[0]
|
830 |
+
pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
|
831 |
+
|
832 |
+
if not return_dict:
|
833 |
+
return (sequence_output, pooled_output) + encoder_outputs[1:]
|
834 |
+
|
835 |
+
return BaseModelOutputWithPoolingAndCrossAttentions(
|
836 |
+
last_hidden_state=sequence_output,
|
837 |
+
pooler_output=pooled_output,
|
838 |
+
past_key_values=encoder_outputs.past_key_values,
|
839 |
+
hidden_states=encoder_outputs.hidden_states,
|
840 |
+
attentions=encoder_outputs.attentions,
|
841 |
+
cross_attentions=encoder_outputs.cross_attentions,
|
842 |
+
)
|
843 |
+
|
extras/BLIP/models/vit.py
ADDED
@@ -0,0 +1,308 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
* Copyright (c) 2022, salesforce.com, inc.
|
3 |
+
* All rights reserved.
|
4 |
+
* SPDX-License-Identifier: BSD-3-Clause
|
5 |
+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
6 |
+
* By Junnan Li
|
7 |
+
* Based on timm code base
|
8 |
+
* https://github.com/rwightman/pytorch-image-models/tree/master/timm
|
9 |
+
'''
|
10 |
+
|
11 |
+
import torch
|
12 |
+
import torch.nn as nn
|
13 |
+
import torch.nn.functional as F
|
14 |
+
from functools import partial
|
15 |
+
|
16 |
+
from timm.models.vision_transformer import _cfg, PatchEmbed
|
17 |
+
from timm.models.registry import register_model
|
18 |
+
from timm.models.layers import trunc_normal_, DropPath
|
19 |
+
from timm.models.helpers import named_apply, adapt_input_conv
|
20 |
+
|
21 |
+
|
22 |
+
def checkpoint_wrapper(x):
|
23 |
+
return x
|
24 |
+
|
25 |
+
|
26 |
+
class Mlp(nn.Module):
|
27 |
+
""" MLP as used in Vision Transformer, MLP-Mixer and related networks
|
28 |
+
"""
|
29 |
+
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
|
30 |
+
super().__init__()
|
31 |
+
out_features = out_features or in_features
|
32 |
+
hidden_features = hidden_features or in_features
|
33 |
+
self.fc1 = nn.Linear(in_features, hidden_features)
|
34 |
+
self.act = act_layer()
|
35 |
+
self.fc2 = nn.Linear(hidden_features, out_features)
|
36 |
+
self.drop = nn.Dropout(drop)
|
37 |
+
|
38 |
+
def forward(self, x):
|
39 |
+
x = self.fc1(x)
|
40 |
+
x = self.act(x)
|
41 |
+
x = self.drop(x)
|
42 |
+
x = self.fc2(x)
|
43 |
+
x = self.drop(x)
|
44 |
+
return x
|
45 |
+
|
46 |
+
|
47 |
+
class Attention(nn.Module):
|
48 |
+
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
|
49 |
+
super().__init__()
|
50 |
+
self.num_heads = num_heads
|
51 |
+
head_dim = dim // num_heads
|
52 |
+
# NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
|
53 |
+
self.scale = qk_scale or head_dim ** -0.5
|
54 |
+
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
55 |
+
self.attn_drop = nn.Dropout(attn_drop)
|
56 |
+
self.proj = nn.Linear(dim, dim)
|
57 |
+
self.proj_drop = nn.Dropout(proj_drop)
|
58 |
+
self.attn_gradients = None
|
59 |
+
self.attention_map = None
|
60 |
+
|
61 |
+
def save_attn_gradients(self, attn_gradients):
|
62 |
+
self.attn_gradients = attn_gradients
|
63 |
+
|
64 |
+
def get_attn_gradients(self):
|
65 |
+
return self.attn_gradients
|
66 |
+
|
67 |
+
def save_attention_map(self, attention_map):
|
68 |
+
self.attention_map = attention_map
|
69 |
+
|
70 |
+
def get_attention_map(self):
|
71 |
+
return self.attention_map
|
72 |
+
|
73 |
+
def forward(self, x, register_hook=False):
|
74 |
+
B, N, C = x.shape
|
75 |
+
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
76 |
+
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
|
77 |
+
|
78 |
+
attn = (q @ k.transpose(-2, -1)) * self.scale
|
79 |
+
attn = attn.softmax(dim=-1)
|
80 |
+
attn = self.attn_drop(attn)
|
81 |
+
|
82 |
+
if register_hook:
|
83 |
+
self.save_attention_map(attn)
|
84 |
+
attn.register_hook(self.save_attn_gradients)
|
85 |
+
|
86 |
+
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
87 |
+
x = self.proj(x)
|
88 |
+
x = self.proj_drop(x)
|
89 |
+
return x
|
90 |
+
|
91 |
+
|
92 |
+
class Block(nn.Module):
|
93 |
+
|
94 |
+
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
|
95 |
+
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_grad_checkpointing=False):
|
96 |
+
super().__init__()
|
97 |
+
self.norm1 = norm_layer(dim)
|
98 |
+
self.attn = Attention(
|
99 |
+
dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
|
100 |
+
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
|
101 |
+
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
102 |
+
self.norm2 = norm_layer(dim)
|
103 |
+
mlp_hidden_dim = int(dim * mlp_ratio)
|
104 |
+
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
|
105 |
+
|
106 |
+
if use_grad_checkpointing:
|
107 |
+
self.attn = checkpoint_wrapper(self.attn)
|
108 |
+
self.mlp = checkpoint_wrapper(self.mlp)
|
109 |
+
|
110 |
+
def forward(self, x, register_hook=False):
|
111 |
+
x = x + self.drop_path(self.attn(self.norm1(x), register_hook=register_hook))
|
112 |
+
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
113 |
+
return x
|
114 |
+
|
115 |
+
|
116 |
+
class VisionTransformer(nn.Module):
|
117 |
+
""" Vision Transformer
|
118 |
+
A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` -
|
119 |
+
https://arxiv.org/abs/2010.11929
|
120 |
+
"""
|
121 |
+
def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
|
122 |
+
num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
|
123 |
+
drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=None,
|
124 |
+
use_grad_checkpointing=False, ckpt_layer=0):
|
125 |
+
"""
|
126 |
+
Args:
|
127 |
+
img_size (int, tuple): input image size
|
128 |
+
patch_size (int, tuple): patch size
|
129 |
+
in_chans (int): number of input channels
|
130 |
+
num_classes (int): number of classes for classification head
|
131 |
+
embed_dim (int): embedding dimension
|
132 |
+
depth (int): depth of transformer
|
133 |
+
num_heads (int): number of attention heads
|
134 |
+
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
|
135 |
+
qkv_bias (bool): enable bias for qkv if True
|
136 |
+
qk_scale (float): override default qk scale of head_dim ** -0.5 if set
|
137 |
+
representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
|
138 |
+
drop_rate (float): dropout rate
|
139 |
+
attn_drop_rate (float): attention dropout rate
|
140 |
+
drop_path_rate (float): stochastic depth rate
|
141 |
+
norm_layer: (nn.Module): normalization layer
|
142 |
+
"""
|
143 |
+
super().__init__()
|
144 |
+
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
145 |
+
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
|
146 |
+
|
147 |
+
self.patch_embed = PatchEmbed(
|
148 |
+
img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
149 |
+
|
150 |
+
num_patches = self.patch_embed.num_patches
|
151 |
+
|
152 |
+
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
153 |
+
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
|
154 |
+
self.pos_drop = nn.Dropout(p=drop_rate)
|
155 |
+
|
156 |
+
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
157 |
+
self.blocks = nn.ModuleList([
|
158 |
+
Block(
|
159 |
+
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
|
160 |
+
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
|
161 |
+
use_grad_checkpointing=(use_grad_checkpointing and i>=depth-ckpt_layer)
|
162 |
+
)
|
163 |
+
for i in range(depth)])
|
164 |
+
self.norm = norm_layer(embed_dim)
|
165 |
+
|
166 |
+
trunc_normal_(self.pos_embed, std=.02)
|
167 |
+
trunc_normal_(self.cls_token, std=.02)
|
168 |
+
self.apply(self._init_weights)
|
169 |
+
|
170 |
+
def _init_weights(self, m):
|
171 |
+
if isinstance(m, nn.Linear):
|
172 |
+
trunc_normal_(m.weight, std=.02)
|
173 |
+
if isinstance(m, nn.Linear) and m.bias is not None:
|
174 |
+
nn.init.constant_(m.bias, 0)
|
175 |
+
elif isinstance(m, nn.LayerNorm):
|
176 |
+
nn.init.constant_(m.bias, 0)
|
177 |
+
nn.init.constant_(m.weight, 1.0)
|
178 |
+
|
179 |
+
@torch.jit.ignore
|
180 |
+
def no_weight_decay(self):
|
181 |
+
return {'pos_embed', 'cls_token'}
|
182 |
+
|
183 |
+
def forward(self, x, register_blk=-1):
|
184 |
+
B = x.shape[0]
|
185 |
+
x = self.patch_embed(x)
|
186 |
+
|
187 |
+
cls_tokens = self.cls_token.expand(B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks
|
188 |
+
x = torch.cat((cls_tokens, x), dim=1)
|
189 |
+
|
190 |
+
x = x + self.pos_embed[:,:x.size(1),:]
|
191 |
+
x = self.pos_drop(x)
|
192 |
+
|
193 |
+
for i,blk in enumerate(self.blocks):
|
194 |
+
x = blk(x, register_blk==i)
|
195 |
+
x = self.norm(x)
|
196 |
+
|
197 |
+
return x
|
198 |
+
|
199 |
+
@torch.jit.ignore()
|
200 |
+
def load_pretrained(self, checkpoint_path, prefix=''):
|
201 |
+
_load_weights(self, checkpoint_path, prefix)
|
202 |
+
|
203 |
+
|
204 |
+
@torch.no_grad()
|
205 |
+
def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str = ''):
|
206 |
+
""" Load weights from .npz checkpoints for official Google Brain Flax implementation
|
207 |
+
"""
|
208 |
+
import numpy as np
|
209 |
+
|
210 |
+
def _n2p(w, t=True):
|
211 |
+
if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:
|
212 |
+
w = w.flatten()
|
213 |
+
if t:
|
214 |
+
if w.ndim == 4:
|
215 |
+
w = w.transpose([3, 2, 0, 1])
|
216 |
+
elif w.ndim == 3:
|
217 |
+
w = w.transpose([2, 0, 1])
|
218 |
+
elif w.ndim == 2:
|
219 |
+
w = w.transpose([1, 0])
|
220 |
+
return torch.from_numpy(w)
|
221 |
+
|
222 |
+
w = np.load(checkpoint_path)
|
223 |
+
if not prefix and 'opt/target/embedding/kernel' in w:
|
224 |
+
prefix = 'opt/target/'
|
225 |
+
|
226 |
+
if hasattr(model.patch_embed, 'backbone'):
|
227 |
+
# hybrid
|
228 |
+
backbone = model.patch_embed.backbone
|
229 |
+
stem_only = not hasattr(backbone, 'stem')
|
230 |
+
stem = backbone if stem_only else backbone.stem
|
231 |
+
stem.conv.weight.copy_(adapt_input_conv(stem.conv.weight.shape[1], _n2p(w[f'{prefix}conv_root/kernel'])))
|
232 |
+
stem.norm.weight.copy_(_n2p(w[f'{prefix}gn_root/scale']))
|
233 |
+
stem.norm.bias.copy_(_n2p(w[f'{prefix}gn_root/bias']))
|
234 |
+
if not stem_only:
|
235 |
+
for i, stage in enumerate(backbone.stages):
|
236 |
+
for j, block in enumerate(stage.blocks):
|
237 |
+
bp = f'{prefix}block{i + 1}/unit{j + 1}/'
|
238 |
+
for r in range(3):
|
239 |
+
getattr(block, f'conv{r + 1}').weight.copy_(_n2p(w[f'{bp}conv{r + 1}/kernel']))
|
240 |
+
getattr(block, f'norm{r + 1}').weight.copy_(_n2p(w[f'{bp}gn{r + 1}/scale']))
|
241 |
+
getattr(block, f'norm{r + 1}').bias.copy_(_n2p(w[f'{bp}gn{r + 1}/bias']))
|
242 |
+
if block.downsample is not None:
|
243 |
+
block.downsample.conv.weight.copy_(_n2p(w[f'{bp}conv_proj/kernel']))
|
244 |
+
block.downsample.norm.weight.copy_(_n2p(w[f'{bp}gn_proj/scale']))
|
245 |
+
block.downsample.norm.bias.copy_(_n2p(w[f'{bp}gn_proj/bias']))
|
246 |
+
embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])
|
247 |
+
else:
|
248 |
+
embed_conv_w = adapt_input_conv(
|
249 |
+
model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))
|
250 |
+
model.patch_embed.proj.weight.copy_(embed_conv_w)
|
251 |
+
model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
|
252 |
+
model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
|
253 |
+
pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)
|
254 |
+
if pos_embed_w.shape != model.pos_embed.shape:
|
255 |
+
pos_embed_w = resize_pos_embed( # resize pos embedding when different size from pretrained weights
|
256 |
+
pos_embed_w, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)
|
257 |
+
model.pos_embed.copy_(pos_embed_w)
|
258 |
+
model.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))
|
259 |
+
model.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))
|
260 |
+
# if isinstance(model.head, nn.Linear) and model.head.bias.shape[0] == w[f'{prefix}head/bias'].shape[-1]:
|
261 |
+
# model.head.weight.copy_(_n2p(w[f'{prefix}head/kernel']))
|
262 |
+
# model.head.bias.copy_(_n2p(w[f'{prefix}head/bias']))
|
263 |
+
# if isinstance(getattr(model.pre_logits, 'fc', None), nn.Linear) and f'{prefix}pre_logits/bias' in w:
|
264 |
+
# model.pre_logits.fc.weight.copy_(_n2p(w[f'{prefix}pre_logits/kernel']))
|
265 |
+
# model.pre_logits.fc.bias.copy_(_n2p(w[f'{prefix}pre_logits/bias']))
|
266 |
+
for i, block in enumerate(model.blocks.children()):
|
267 |
+
block_prefix = f'{prefix}Transformer/encoderblock_{i}/'
|
268 |
+
mha_prefix = block_prefix + 'MultiHeadDotProductAttention_1/'
|
269 |
+
block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
|
270 |
+
block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
|
271 |
+
block.attn.qkv.weight.copy_(torch.cat([
|
272 |
+
_n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))
|
273 |
+
block.attn.qkv.bias.copy_(torch.cat([
|
274 |
+
_n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))
|
275 |
+
block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
|
276 |
+
block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
|
277 |
+
for r in range(2):
|
278 |
+
getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/kernel']))
|
279 |
+
getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/bias']))
|
280 |
+
block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/scale']))
|
281 |
+
block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/bias']))
|
282 |
+
|
283 |
+
|
284 |
+
def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):
|
285 |
+
# interpolate position embedding
|
286 |
+
embedding_size = pos_embed_checkpoint.shape[-1]
|
287 |
+
num_patches = visual_encoder.patch_embed.num_patches
|
288 |
+
num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches
|
289 |
+
# height (== width) for the checkpoint position embedding
|
290 |
+
orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
|
291 |
+
# height (== width) for the new position embedding
|
292 |
+
new_size = int(num_patches ** 0.5)
|
293 |
+
|
294 |
+
if orig_size!=new_size:
|
295 |
+
# class_token and dist_token are kept unchanged
|
296 |
+
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
|
297 |
+
# only the position tokens are interpolated
|
298 |
+
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
|
299 |
+
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
|
300 |
+
pos_tokens = torch.nn.functional.interpolate(
|
301 |
+
pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
|
302 |
+
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
|
303 |
+
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
|
304 |
+
print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))
|
305 |
+
|
306 |
+
return new_pos_embed
|
307 |
+
else:
|
308 |
+
return pos_embed_checkpoint
|
extras/expansion.py
ADDED
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Fooocus GPT2 Expansion
|
2 |
+
# Algorithm created by Lvmin Zhang at 2023, Stanford
|
3 |
+
# If used inside Fooocus, any use is permitted.
|
4 |
+
# If used outside Fooocus, only non-commercial use is permitted (CC-By NC 4.0).
|
5 |
+
# This applies to the word list, vocab, model, and algorithm.
|
6 |
+
|
7 |
+
|
8 |
+
import os
|
9 |
+
import torch
|
10 |
+
import math
|
11 |
+
import ldm_patched.modules.model_management as model_management
|
12 |
+
|
13 |
+
from transformers.generation.logits_process import LogitsProcessorList
|
14 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
|
15 |
+
from modules.config import path_fooocus_expansion
|
16 |
+
from ldm_patched.modules.model_patcher import ModelPatcher
|
17 |
+
|
18 |
+
|
19 |
+
# limitation of np.random.seed(), called from transformers.set_seed()
|
20 |
+
SEED_LIMIT_NUMPY = 2**32
|
21 |
+
neg_inf = - 8192.0
|
22 |
+
|
23 |
+
|
24 |
+
def safe_str(x):
|
25 |
+
x = str(x)
|
26 |
+
for _ in range(16):
|
27 |
+
x = x.replace(' ', ' ')
|
28 |
+
return x.strip(",. \r\n")
|
29 |
+
|
30 |
+
|
31 |
+
def remove_pattern(x, pattern):
|
32 |
+
for p in pattern:
|
33 |
+
x = x.replace(p, '')
|
34 |
+
return x
|
35 |
+
|
36 |
+
|
37 |
+
class FooocusExpansion:
|
38 |
+
def __init__(self):
|
39 |
+
self.tokenizer = AutoTokenizer.from_pretrained(path_fooocus_expansion)
|
40 |
+
|
41 |
+
positive_words = open(os.path.join(path_fooocus_expansion, 'positive.txt'),
|
42 |
+
encoding='utf-8').read().splitlines()
|
43 |
+
positive_words = ['Ġ' + x.lower() for x in positive_words if x != '']
|
44 |
+
|
45 |
+
self.logits_bias = torch.zeros((1, len(self.tokenizer.vocab)), dtype=torch.float32) + neg_inf
|
46 |
+
|
47 |
+
debug_list = []
|
48 |
+
for k, v in self.tokenizer.vocab.items():
|
49 |
+
if k in positive_words:
|
50 |
+
self.logits_bias[0, v] = 0
|
51 |
+
debug_list.append(k[1:])
|
52 |
+
|
53 |
+
print(f'Fooocus V2 Expansion: Vocab with {len(debug_list)} words.')
|
54 |
+
|
55 |
+
# debug_list = '\n'.join(sorted(debug_list))
|
56 |
+
# print(debug_list)
|
57 |
+
|
58 |
+
# t11 = self.tokenizer(',', return_tensors="np")
|
59 |
+
# t198 = self.tokenizer('\n', return_tensors="np")
|
60 |
+
# eos = self.tokenizer.eos_token_id
|
61 |
+
|
62 |
+
self.model = AutoModelForCausalLM.from_pretrained(path_fooocus_expansion)
|
63 |
+
self.model.eval()
|
64 |
+
|
65 |
+
load_device = model_management.text_encoder_device()
|
66 |
+
offload_device = model_management.text_encoder_offload_device()
|
67 |
+
|
68 |
+
# MPS hack
|
69 |
+
if model_management.is_device_mps(load_device):
|
70 |
+
load_device = torch.device('cpu')
|
71 |
+
offload_device = torch.device('cpu')
|
72 |
+
|
73 |
+
use_fp16 = model_management.should_use_fp16(device=load_device)
|
74 |
+
|
75 |
+
if use_fp16:
|
76 |
+
self.model.half()
|
77 |
+
|
78 |
+
self.patcher = ModelPatcher(self.model, load_device=load_device, offload_device=offload_device)
|
79 |
+
print(f'Fooocus Expansion engine loaded for {load_device}, use_fp16 = {use_fp16}.')
|
80 |
+
|
81 |
+
@torch.no_grad()
|
82 |
+
@torch.inference_mode()
|
83 |
+
def logits_processor(self, input_ids, scores):
|
84 |
+
assert scores.ndim == 2 and scores.shape[0] == 1
|
85 |
+
self.logits_bias = self.logits_bias.to(scores)
|
86 |
+
|
87 |
+
bias = self.logits_bias.clone()
|
88 |
+
bias[0, input_ids[0].to(bias.device).long()] = neg_inf
|
89 |
+
bias[0, 11] = 0
|
90 |
+
|
91 |
+
return scores + bias
|
92 |
+
|
93 |
+
@torch.no_grad()
|
94 |
+
@torch.inference_mode()
|
95 |
+
def __call__(self, prompt, seed):
|
96 |
+
if prompt == '':
|
97 |
+
return ''
|
98 |
+
|
99 |
+
if self.patcher.current_device != self.patcher.load_device:
|
100 |
+
print('Fooocus Expansion loaded by itself.')
|
101 |
+
model_management.load_model_gpu(self.patcher)
|
102 |
+
|
103 |
+
seed = int(seed) % SEED_LIMIT_NUMPY
|
104 |
+
set_seed(seed)
|
105 |
+
prompt = safe_str(prompt) + ','
|
106 |
+
|
107 |
+
tokenized_kwargs = self.tokenizer(prompt, return_tensors="pt")
|
108 |
+
tokenized_kwargs.data['input_ids'] = tokenized_kwargs.data['input_ids'].to(self.patcher.load_device)
|
109 |
+
tokenized_kwargs.data['attention_mask'] = tokenized_kwargs.data['attention_mask'].to(self.patcher.load_device)
|
110 |
+
|
111 |
+
current_token_length = int(tokenized_kwargs.data['input_ids'].shape[1])
|
112 |
+
max_token_length = 75 * int(math.ceil(float(current_token_length) / 75.0))
|
113 |
+
max_new_tokens = max_token_length - current_token_length
|
114 |
+
|
115 |
+
# https://huggingface.co/blog/introducing-csearch
|
116 |
+
# https://huggingface.co/docs/transformers/generation_strategies
|
117 |
+
features = self.model.generate(**tokenized_kwargs,
|
118 |
+
top_k=100,
|
119 |
+
max_new_tokens=max_new_tokens,
|
120 |
+
do_sample=True,
|
121 |
+
logits_processor=LogitsProcessorList([self.logits_processor]))
|
122 |
+
|
123 |
+
response = self.tokenizer.batch_decode(features, skip_special_tokens=True)
|
124 |
+
result = safe_str(response[0])
|
125 |
+
|
126 |
+
return result
|
extras/face_crop.py
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import numpy as np
|
3 |
+
import modules.config
|
4 |
+
|
5 |
+
|
6 |
+
faceRestoreHelper = None
|
7 |
+
|
8 |
+
|
9 |
+
def align_warp_face(self, landmark, border_mode='constant'):
|
10 |
+
affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template, method=cv2.LMEDS)[0]
|
11 |
+
self.affine_matrices.append(affine_matrix)
|
12 |
+
if border_mode == 'constant':
|
13 |
+
border_mode = cv2.BORDER_CONSTANT
|
14 |
+
elif border_mode == 'reflect101':
|
15 |
+
border_mode = cv2.BORDER_REFLECT101
|
16 |
+
elif border_mode == 'reflect':
|
17 |
+
border_mode = cv2.BORDER_REFLECT
|
18 |
+
input_img = self.input_img
|
19 |
+
cropped_face = cv2.warpAffine(input_img, affine_matrix, self.face_size,
|
20 |
+
borderMode=border_mode, borderValue=(135, 133, 132))
|
21 |
+
return cropped_face
|
22 |
+
|
23 |
+
|
24 |
+
def crop_image(img_rgb):
|
25 |
+
global faceRestoreHelper
|
26 |
+
|
27 |
+
if faceRestoreHelper is None:
|
28 |
+
from extras.facexlib.utils.face_restoration_helper import FaceRestoreHelper
|
29 |
+
faceRestoreHelper = FaceRestoreHelper(
|
30 |
+
upscale_factor=1,
|
31 |
+
model_rootpath=modules.config.path_controlnet,
|
32 |
+
device='cpu' # use cpu is safer since we are out of memory management
|
33 |
+
)
|
34 |
+
|
35 |
+
faceRestoreHelper.clean_all()
|
36 |
+
faceRestoreHelper.read_image(np.ascontiguousarray(img_rgb[:, :, ::-1].copy()))
|
37 |
+
faceRestoreHelper.get_face_landmarks_5()
|
38 |
+
|
39 |
+
landmarks = faceRestoreHelper.all_landmarks_5
|
40 |
+
# landmarks are already sorted with confidence.
|
41 |
+
|
42 |
+
if len(landmarks) == 0:
|
43 |
+
print('No face detected')
|
44 |
+
return img_rgb
|
45 |
+
else:
|
46 |
+
print(f'Detected {len(landmarks)} faces')
|
47 |
+
|
48 |
+
result = align_warp_face(faceRestoreHelper, landmarks[0])
|
49 |
+
|
50 |
+
return np.ascontiguousarray(result[:, :, ::-1].copy())
|
extras/facexlib/detection/__init__.py
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from copy import deepcopy
|
3 |
+
|
4 |
+
from extras.facexlib.utils import load_file_from_url
|
5 |
+
from .retinaface import RetinaFace
|
6 |
+
|
7 |
+
|
8 |
+
def init_detection_model(model_name, half=False, device='cuda', model_rootpath=None):
|
9 |
+
if model_name == 'retinaface_resnet50':
|
10 |
+
model = RetinaFace(network_name='resnet50', half=half, device=device)
|
11 |
+
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth'
|
12 |
+
elif model_name == 'retinaface_mobile0.25':
|
13 |
+
model = RetinaFace(network_name='mobile0.25', half=half, device=device)
|
14 |
+
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
|
15 |
+
else:
|
16 |
+
raise NotImplementedError(f'{model_name} is not implemented.')
|
17 |
+
|
18 |
+
model_path = load_file_from_url(
|
19 |
+
url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
|
20 |
+
|
21 |
+
# TODO: clean pretrained model
|
22 |
+
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
|
23 |
+
# remove unnecessary 'module.'
|
24 |
+
for k, v in deepcopy(load_net).items():
|
25 |
+
if k.startswith('module.'):
|
26 |
+
load_net[k[7:]] = v
|
27 |
+
load_net.pop(k)
|
28 |
+
model.load_state_dict(load_net, strict=True)
|
29 |
+
model.eval()
|
30 |
+
model = model.to(device)
|
31 |
+
return model
|
extras/facexlib/detection/align_trans.py
ADDED
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import numpy as np
|
3 |
+
|
4 |
+
from .matlab_cp2tform import get_similarity_transform_for_cv2
|
5 |
+
|
6 |
+
# reference facial points, a list of coordinates (x,y)
|
7 |
+
REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051], [65.53179932, 51.50139999], [48.02519989, 71.73660278],
|
8 |
+
[33.54930115, 92.3655014], [62.72990036, 92.20410156]]
|
9 |
+
|
10 |
+
DEFAULT_CROP_SIZE = (96, 112)
|
11 |
+
|
12 |
+
|
13 |
+
class FaceWarpException(Exception):
|
14 |
+
|
15 |
+
def __str__(self):
|
16 |
+
return 'In File {}:{}'.format(__file__, super.__str__(self))
|
17 |
+
|
18 |
+
|
19 |
+
def get_reference_facial_points(output_size=None, inner_padding_factor=0.0, outer_padding=(0, 0), default_square=False):
|
20 |
+
"""
|
21 |
+
Function:
|
22 |
+
----------
|
23 |
+
get reference 5 key points according to crop settings:
|
24 |
+
0. Set default crop_size:
|
25 |
+
if default_square:
|
26 |
+
crop_size = (112, 112)
|
27 |
+
else:
|
28 |
+
crop_size = (96, 112)
|
29 |
+
1. Pad the crop_size by inner_padding_factor in each side;
|
30 |
+
2. Resize crop_size into (output_size - outer_padding*2),
|
31 |
+
pad into output_size with outer_padding;
|
32 |
+
3. Output reference_5point;
|
33 |
+
Parameters:
|
34 |
+
----------
|
35 |
+
@output_size: (w, h) or None
|
36 |
+
size of aligned face image
|
37 |
+
@inner_padding_factor: (w_factor, h_factor)
|
38 |
+
padding factor for inner (w, h)
|
39 |
+
@outer_padding: (w_pad, h_pad)
|
40 |
+
each row is a pair of coordinates (x, y)
|
41 |
+
@default_square: True or False
|
42 |
+
if True:
|
43 |
+
default crop_size = (112, 112)
|
44 |
+
else:
|
45 |
+
default crop_size = (96, 112);
|
46 |
+
!!! make sure, if output_size is not None:
|
47 |
+
(output_size - outer_padding)
|
48 |
+
= some_scale * (default crop_size * (1.0 +
|
49 |
+
inner_padding_factor))
|
50 |
+
Returns:
|
51 |
+
----------
|
52 |
+
@reference_5point: 5x2 np.array
|
53 |
+
each row is a pair of transformed coordinates (x, y)
|
54 |
+
"""
|
55 |
+
|
56 |
+
tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
|
57 |
+
tmp_crop_size = np.array(DEFAULT_CROP_SIZE)
|
58 |
+
|
59 |
+
# 0) make the inner region a square
|
60 |
+
if default_square:
|
61 |
+
size_diff = max(tmp_crop_size) - tmp_crop_size
|
62 |
+
tmp_5pts += size_diff / 2
|
63 |
+
tmp_crop_size += size_diff
|
64 |
+
|
65 |
+
if (output_size and output_size[0] == tmp_crop_size[0] and output_size[1] == tmp_crop_size[1]):
|
66 |
+
|
67 |
+
return tmp_5pts
|
68 |
+
|
69 |
+
if (inner_padding_factor == 0 and outer_padding == (0, 0)):
|
70 |
+
if output_size is None:
|
71 |
+
return tmp_5pts
|
72 |
+
else:
|
73 |
+
raise FaceWarpException('No paddings to do, output_size must be None or {}'.format(tmp_crop_size))
|
74 |
+
|
75 |
+
# check output size
|
76 |
+
if not (0 <= inner_padding_factor <= 1.0):
|
77 |
+
raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')
|
78 |
+
|
79 |
+
if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0) and output_size is None):
|
80 |
+
output_size = tmp_crop_size * \
|
81 |
+
(1 + inner_padding_factor * 2).astype(np.int32)
|
82 |
+
output_size += np.array(outer_padding)
|
83 |
+
if not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1]):
|
84 |
+
raise FaceWarpException('Not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1])')
|
85 |
+
|
86 |
+
# 1) pad the inner region according inner_padding_factor
|
87 |
+
if inner_padding_factor > 0:
|
88 |
+
size_diff = tmp_crop_size * inner_padding_factor * 2
|
89 |
+
tmp_5pts += size_diff / 2
|
90 |
+
tmp_crop_size += np.round(size_diff).astype(np.int32)
|
91 |
+
|
92 |
+
# 2) resize the padded inner region
|
93 |
+
size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2
|
94 |
+
|
95 |
+
if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:
|
96 |
+
raise FaceWarpException('Must have (output_size - outer_padding)'
|
97 |
+
'= some_scale * (crop_size * (1.0 + inner_padding_factor)')
|
98 |
+
|
99 |
+
scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]
|
100 |
+
tmp_5pts = tmp_5pts * scale_factor
|
101 |
+
# size_diff = tmp_crop_size * (scale_factor - min(scale_factor))
|
102 |
+
# tmp_5pts = tmp_5pts + size_diff / 2
|
103 |
+
tmp_crop_size = size_bf_outer_pad
|
104 |
+
|
105 |
+
# 3) add outer_padding to make output_size
|
106 |
+
reference_5point = tmp_5pts + np.array(outer_padding)
|
107 |
+
tmp_crop_size = output_size
|
108 |
+
|
109 |
+
return reference_5point
|
110 |
+
|
111 |
+
|
112 |
+
def get_affine_transform_matrix(src_pts, dst_pts):
|
113 |
+
"""
|
114 |
+
Function:
|
115 |
+
----------
|
116 |
+
get affine transform matrix 'tfm' from src_pts to dst_pts
|
117 |
+
Parameters:
|
118 |
+
----------
|
119 |
+
@src_pts: Kx2 np.array
|
120 |
+
source points matrix, each row is a pair of coordinates (x, y)
|
121 |
+
@dst_pts: Kx2 np.array
|
122 |
+
destination points matrix, each row is a pair of coordinates (x, y)
|
123 |
+
Returns:
|
124 |
+
----------
|
125 |
+
@tfm: 2x3 np.array
|
126 |
+
transform matrix from src_pts to dst_pts
|
127 |
+
"""
|
128 |
+
|
129 |
+
tfm = np.float32([[1, 0, 0], [0, 1, 0]])
|
130 |
+
n_pts = src_pts.shape[0]
|
131 |
+
ones = np.ones((n_pts, 1), src_pts.dtype)
|
132 |
+
src_pts_ = np.hstack([src_pts, ones])
|
133 |
+
dst_pts_ = np.hstack([dst_pts, ones])
|
134 |
+
|
135 |
+
A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)
|
136 |
+
|
137 |
+
if rank == 3:
|
138 |
+
tfm = np.float32([[A[0, 0], A[1, 0], A[2, 0]], [A[0, 1], A[1, 1], A[2, 1]]])
|
139 |
+
elif rank == 2:
|
140 |
+
tfm = np.float32([[A[0, 0], A[1, 0], 0], [A[0, 1], A[1, 1], 0]])
|
141 |
+
|
142 |
+
return tfm
|
143 |
+
|
144 |
+
|
145 |
+
def warp_and_crop_face(src_img, facial_pts, reference_pts=None, crop_size=(96, 112), align_type='smilarity'):
|
146 |
+
"""
|
147 |
+
Function:
|
148 |
+
----------
|
149 |
+
apply affine transform 'trans' to uv
|
150 |
+
Parameters:
|
151 |
+
----------
|
152 |
+
@src_img: 3x3 np.array
|
153 |
+
input image
|
154 |
+
@facial_pts: could be
|
155 |
+
1)a list of K coordinates (x,y)
|
156 |
+
or
|
157 |
+
2) Kx2 or 2xK np.array
|
158 |
+
each row or col is a pair of coordinates (x, y)
|
159 |
+
@reference_pts: could be
|
160 |
+
1) a list of K coordinates (x,y)
|
161 |
+
or
|
162 |
+
2) Kx2 or 2xK np.array
|
163 |
+
each row or col is a pair of coordinates (x, y)
|
164 |
+
or
|
165 |
+
3) None
|
166 |
+
if None, use default reference facial points
|
167 |
+
@crop_size: (w, h)
|
168 |
+
output face image size
|
169 |
+
@align_type: transform type, could be one of
|
170 |
+
1) 'similarity': use similarity transform
|
171 |
+
2) 'cv2_affine': use the first 3 points to do affine transform,
|
172 |
+
by calling cv2.getAffineTransform()
|
173 |
+
3) 'affine': use all points to do affine transform
|
174 |
+
Returns:
|
175 |
+
----------
|
176 |
+
@face_img: output face image with size (w, h) = @crop_size
|
177 |
+
"""
|
178 |
+
|
179 |
+
if reference_pts is None:
|
180 |
+
if crop_size[0] == 96 and crop_size[1] == 112:
|
181 |
+
reference_pts = REFERENCE_FACIAL_POINTS
|
182 |
+
else:
|
183 |
+
default_square = False
|
184 |
+
inner_padding_factor = 0
|
185 |
+
outer_padding = (0, 0)
|
186 |
+
output_size = crop_size
|
187 |
+
|
188 |
+
reference_pts = get_reference_facial_points(output_size, inner_padding_factor, outer_padding,
|
189 |
+
default_square)
|
190 |
+
|
191 |
+
ref_pts = np.float32(reference_pts)
|
192 |
+
ref_pts_shp = ref_pts.shape
|
193 |
+
if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
|
194 |
+
raise FaceWarpException('reference_pts.shape must be (K,2) or (2,K) and K>2')
|
195 |
+
|
196 |
+
if ref_pts_shp[0] == 2:
|
197 |
+
ref_pts = ref_pts.T
|
198 |
+
|
199 |
+
src_pts = np.float32(facial_pts)
|
200 |
+
src_pts_shp = src_pts.shape
|
201 |
+
if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
|
202 |
+
raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')
|
203 |
+
|
204 |
+
if src_pts_shp[0] == 2:
|
205 |
+
src_pts = src_pts.T
|
206 |
+
|
207 |
+
if src_pts.shape != ref_pts.shape:
|
208 |
+
raise FaceWarpException('facial_pts and reference_pts must have the same shape')
|
209 |
+
|
210 |
+
if align_type == 'cv2_affine':
|
211 |
+
tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
|
212 |
+
elif align_type == 'affine':
|
213 |
+
tfm = get_affine_transform_matrix(src_pts, ref_pts)
|
214 |
+
else:
|
215 |
+
tfm = get_similarity_transform_for_cv2(src_pts, ref_pts)
|
216 |
+
|
217 |
+
face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]))
|
218 |
+
|
219 |
+
return face_img
|
extras/facexlib/detection/matlab_cp2tform.py
ADDED
@@ -0,0 +1,317 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
from numpy.linalg import inv, lstsq
|
3 |
+
from numpy.linalg import matrix_rank as rank
|
4 |
+
from numpy.linalg import norm
|
5 |
+
|
6 |
+
|
7 |
+
class MatlabCp2tormException(Exception):
|
8 |
+
|
9 |
+
def __str__(self):
|
10 |
+
return 'In File {}:{}'.format(__file__, super.__str__(self))
|
11 |
+
|
12 |
+
|
13 |
+
def tformfwd(trans, uv):
|
14 |
+
"""
|
15 |
+
Function:
|
16 |
+
----------
|
17 |
+
apply affine transform 'trans' to uv
|
18 |
+
|
19 |
+
Parameters:
|
20 |
+
----------
|
21 |
+
@trans: 3x3 np.array
|
22 |
+
transform matrix
|
23 |
+
@uv: Kx2 np.array
|
24 |
+
each row is a pair of coordinates (x, y)
|
25 |
+
|
26 |
+
Returns:
|
27 |
+
----------
|
28 |
+
@xy: Kx2 np.array
|
29 |
+
each row is a pair of transformed coordinates (x, y)
|
30 |
+
"""
|
31 |
+
uv = np.hstack((uv, np.ones((uv.shape[0], 1))))
|
32 |
+
xy = np.dot(uv, trans)
|
33 |
+
xy = xy[:, 0:-1]
|
34 |
+
return xy
|
35 |
+
|
36 |
+
|
37 |
+
def tforminv(trans, uv):
|
38 |
+
"""
|
39 |
+
Function:
|
40 |
+
----------
|
41 |
+
apply the inverse of affine transform 'trans' to uv
|
42 |
+
|
43 |
+
Parameters:
|
44 |
+
----------
|
45 |
+
@trans: 3x3 np.array
|
46 |
+
transform matrix
|
47 |
+
@uv: Kx2 np.array
|
48 |
+
each row is a pair of coordinates (x, y)
|
49 |
+
|
50 |
+
Returns:
|
51 |
+
----------
|
52 |
+
@xy: Kx2 np.array
|
53 |
+
each row is a pair of inverse-transformed coordinates (x, y)
|
54 |
+
"""
|
55 |
+
Tinv = inv(trans)
|
56 |
+
xy = tformfwd(Tinv, uv)
|
57 |
+
return xy
|
58 |
+
|
59 |
+
|
60 |
+
def findNonreflectiveSimilarity(uv, xy, options=None):
|
61 |
+
options = {'K': 2}
|
62 |
+
|
63 |
+
K = options['K']
|
64 |
+
M = xy.shape[0]
|
65 |
+
x = xy[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
|
66 |
+
y = xy[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
|
67 |
+
|
68 |
+
tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
|
69 |
+
tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
|
70 |
+
X = np.vstack((tmp1, tmp2))
|
71 |
+
|
72 |
+
u = uv[:, 0].reshape((-1, 1)) # use reshape to keep a column vector
|
73 |
+
v = uv[:, 1].reshape((-1, 1)) # use reshape to keep a column vector
|
74 |
+
U = np.vstack((u, v))
|
75 |
+
|
76 |
+
# We know that X * r = U
|
77 |
+
if rank(X) >= 2 * K:
|
78 |
+
r, _, _, _ = lstsq(X, U, rcond=-1)
|
79 |
+
r = np.squeeze(r)
|
80 |
+
else:
|
81 |
+
raise Exception('cp2tform:twoUniquePointsReq')
|
82 |
+
sc = r[0]
|
83 |
+
ss = r[1]
|
84 |
+
tx = r[2]
|
85 |
+
ty = r[3]
|
86 |
+
|
87 |
+
Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
|
88 |
+
T = inv(Tinv)
|
89 |
+
T[:, 2] = np.array([0, 0, 1])
|
90 |
+
|
91 |
+
return T, Tinv
|
92 |
+
|
93 |
+
|
94 |
+
def findSimilarity(uv, xy, options=None):
|
95 |
+
options = {'K': 2}
|
96 |
+
|
97 |
+
# uv = np.array(uv)
|
98 |
+
# xy = np.array(xy)
|
99 |
+
|
100 |
+
# Solve for trans1
|
101 |
+
trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
|
102 |
+
|
103 |
+
# Solve for trans2
|
104 |
+
|
105 |
+
# manually reflect the xy data across the Y-axis
|
106 |
+
xyR = xy
|
107 |
+
xyR[:, 0] = -1 * xyR[:, 0]
|
108 |
+
|
109 |
+
trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
|
110 |
+
|
111 |
+
# manually reflect the tform to undo the reflection done on xyR
|
112 |
+
TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
|
113 |
+
|
114 |
+
trans2 = np.dot(trans2r, TreflectY)
|
115 |
+
|
116 |
+
# Figure out if trans1 or trans2 is better
|
117 |
+
xy1 = tformfwd(trans1, uv)
|
118 |
+
norm1 = norm(xy1 - xy)
|
119 |
+
|
120 |
+
xy2 = tformfwd(trans2, uv)
|
121 |
+
norm2 = norm(xy2 - xy)
|
122 |
+
|
123 |
+
if norm1 <= norm2:
|
124 |
+
return trans1, trans1_inv
|
125 |
+
else:
|
126 |
+
trans2_inv = inv(trans2)
|
127 |
+
return trans2, trans2_inv
|
128 |
+
|
129 |
+
|
130 |
+
def get_similarity_transform(src_pts, dst_pts, reflective=True):
|
131 |
+
"""
|
132 |
+
Function:
|
133 |
+
----------
|
134 |
+
Find Similarity Transform Matrix 'trans':
|
135 |
+
u = src_pts[:, 0]
|
136 |
+
v = src_pts[:, 1]
|
137 |
+
x = dst_pts[:, 0]
|
138 |
+
y = dst_pts[:, 1]
|
139 |
+
[x, y, 1] = [u, v, 1] * trans
|
140 |
+
|
141 |
+
Parameters:
|
142 |
+
----------
|
143 |
+
@src_pts: Kx2 np.array
|
144 |
+
source points, each row is a pair of coordinates (x, y)
|
145 |
+
@dst_pts: Kx2 np.array
|
146 |
+
destination points, each row is a pair of transformed
|
147 |
+
coordinates (x, y)
|
148 |
+
@reflective: True or False
|
149 |
+
if True:
|
150 |
+
use reflective similarity transform
|
151 |
+
else:
|
152 |
+
use non-reflective similarity transform
|
153 |
+
|
154 |
+
Returns:
|
155 |
+
----------
|
156 |
+
@trans: 3x3 np.array
|
157 |
+
transform matrix from uv to xy
|
158 |
+
trans_inv: 3x3 np.array
|
159 |
+
inverse of trans, transform matrix from xy to uv
|
160 |
+
"""
|
161 |
+
|
162 |
+
if reflective:
|
163 |
+
trans, trans_inv = findSimilarity(src_pts, dst_pts)
|
164 |
+
else:
|
165 |
+
trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
|
166 |
+
|
167 |
+
return trans, trans_inv
|
168 |
+
|
169 |
+
|
170 |
+
def cvt_tform_mat_for_cv2(trans):
|
171 |
+
"""
|
172 |
+
Function:
|
173 |
+
----------
|
174 |
+
Convert Transform Matrix 'trans' into 'cv2_trans' which could be
|
175 |
+
directly used by cv2.warpAffine():
|
176 |
+
u = src_pts[:, 0]
|
177 |
+
v = src_pts[:, 1]
|
178 |
+
x = dst_pts[:, 0]
|
179 |
+
y = dst_pts[:, 1]
|
180 |
+
[x, y].T = cv_trans * [u, v, 1].T
|
181 |
+
|
182 |
+
Parameters:
|
183 |
+
----------
|
184 |
+
@trans: 3x3 np.array
|
185 |
+
transform matrix from uv to xy
|
186 |
+
|
187 |
+
Returns:
|
188 |
+
----------
|
189 |
+
@cv2_trans: 2x3 np.array
|
190 |
+
transform matrix from src_pts to dst_pts, could be directly used
|
191 |
+
for cv2.warpAffine()
|
192 |
+
"""
|
193 |
+
cv2_trans = trans[:, 0:2].T
|
194 |
+
|
195 |
+
return cv2_trans
|
196 |
+
|
197 |
+
|
198 |
+
def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
|
199 |
+
"""
|
200 |
+
Function:
|
201 |
+
----------
|
202 |
+
Find Similarity Transform Matrix 'cv2_trans' which could be
|
203 |
+
directly used by cv2.warpAffine():
|
204 |
+
u = src_pts[:, 0]
|
205 |
+
v = src_pts[:, 1]
|
206 |
+
x = dst_pts[:, 0]
|
207 |
+
y = dst_pts[:, 1]
|
208 |
+
[x, y].T = cv_trans * [u, v, 1].T
|
209 |
+
|
210 |
+
Parameters:
|
211 |
+
----------
|
212 |
+
@src_pts: Kx2 np.array
|
213 |
+
source points, each row is a pair of coordinates (x, y)
|
214 |
+
@dst_pts: Kx2 np.array
|
215 |
+
destination points, each row is a pair of transformed
|
216 |
+
coordinates (x, y)
|
217 |
+
reflective: True or False
|
218 |
+
if True:
|
219 |
+
use reflective similarity transform
|
220 |
+
else:
|
221 |
+
use non-reflective similarity transform
|
222 |
+
|
223 |
+
Returns:
|
224 |
+
----------
|
225 |
+
@cv2_trans: 2x3 np.array
|
226 |
+
transform matrix from src_pts to dst_pts, could be directly used
|
227 |
+
for cv2.warpAffine()
|
228 |
+
"""
|
229 |
+
trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
|
230 |
+
cv2_trans = cvt_tform_mat_for_cv2(trans)
|
231 |
+
|
232 |
+
return cv2_trans
|
233 |
+
|
234 |
+
|
235 |
+
if __name__ == '__main__':
|
236 |
+
"""
|
237 |
+
u = [0, 6, -2]
|
238 |
+
v = [0, 3, 5]
|
239 |
+
x = [-1, 0, 4]
|
240 |
+
y = [-1, -10, 4]
|
241 |
+
|
242 |
+
# In Matlab, run:
|
243 |
+
#
|
244 |
+
# uv = [u'; v'];
|
245 |
+
# xy = [x'; y'];
|
246 |
+
# tform_sim=cp2tform(uv,xy,'similarity');
|
247 |
+
#
|
248 |
+
# trans = tform_sim.tdata.T
|
249 |
+
# ans =
|
250 |
+
# -0.0764 -1.6190 0
|
251 |
+
# 1.6190 -0.0764 0
|
252 |
+
# -3.2156 0.0290 1.0000
|
253 |
+
# trans_inv = tform_sim.tdata.Tinv
|
254 |
+
# ans =
|
255 |
+
#
|
256 |
+
# -0.0291 0.6163 0
|
257 |
+
# -0.6163 -0.0291 0
|
258 |
+
# -0.0756 1.9826 1.0000
|
259 |
+
# xy_m=tformfwd(tform_sim, u,v)
|
260 |
+
#
|
261 |
+
# xy_m =
|
262 |
+
#
|
263 |
+
# -3.2156 0.0290
|
264 |
+
# 1.1833 -9.9143
|
265 |
+
# 5.0323 2.8853
|
266 |
+
# uv_m=tforminv(tform_sim, x,y)
|
267 |
+
#
|
268 |
+
# uv_m =
|
269 |
+
#
|
270 |
+
# 0.5698 1.3953
|
271 |
+
# 6.0872 2.2733
|
272 |
+
# -2.6570 4.3314
|
273 |
+
"""
|
274 |
+
u = [0, 6, -2]
|
275 |
+
v = [0, 3, 5]
|
276 |
+
x = [-1, 0, 4]
|
277 |
+
y = [-1, -10, 4]
|
278 |
+
|
279 |
+
uv = np.array((u, v)).T
|
280 |
+
xy = np.array((x, y)).T
|
281 |
+
|
282 |
+
print('\n--->uv:')
|
283 |
+
print(uv)
|
284 |
+
print('\n--->xy:')
|
285 |
+
print(xy)
|
286 |
+
|
287 |
+
trans, trans_inv = get_similarity_transform(uv, xy)
|
288 |
+
|
289 |
+
print('\n--->trans matrix:')
|
290 |
+
print(trans)
|
291 |
+
|
292 |
+
print('\n--->trans_inv matrix:')
|
293 |
+
print(trans_inv)
|
294 |
+
|
295 |
+
print('\n---> apply transform to uv')
|
296 |
+
print('\nxy_m = uv_augmented * trans')
|
297 |
+
uv_aug = np.hstack((uv, np.ones((uv.shape[0], 1))))
|
298 |
+
xy_m = np.dot(uv_aug, trans)
|
299 |
+
print(xy_m)
|
300 |
+
|
301 |
+
print('\nxy_m = tformfwd(trans, uv)')
|
302 |
+
xy_m = tformfwd(trans, uv)
|
303 |
+
print(xy_m)
|
304 |
+
|
305 |
+
print('\n---> apply inverse transform to xy')
|
306 |
+
print('\nuv_m = xy_augmented * trans_inv')
|
307 |
+
xy_aug = np.hstack((xy, np.ones((xy.shape[0], 1))))
|
308 |
+
uv_m = np.dot(xy_aug, trans_inv)
|
309 |
+
print(uv_m)
|
310 |
+
|
311 |
+
print('\nuv_m = tformfwd(trans_inv, xy)')
|
312 |
+
uv_m = tformfwd(trans_inv, xy)
|
313 |
+
print(uv_m)
|
314 |
+
|
315 |
+
uv_m = tforminv(trans, xy)
|
316 |
+
print('\nuv_m = tforminv(trans, xy)')
|
317 |
+
print(uv_m)
|
extras/facexlib/detection/retinaface.py
ADDED
@@ -0,0 +1,366 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import numpy as np
|
3 |
+
import torch
|
4 |
+
import torch.nn as nn
|
5 |
+
import torch.nn.functional as F
|
6 |
+
from PIL import Image
|
7 |
+
from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter
|
8 |
+
|
9 |
+
from extras.facexlib.detection.align_trans import get_reference_facial_points, warp_and_crop_face
|
10 |
+
from extras.facexlib.detection.retinaface_net import FPN, SSH, MobileNetV1, make_bbox_head, make_class_head, make_landmark_head
|
11 |
+
from extras.facexlib.detection.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
|
12 |
+
py_cpu_nms)
|
13 |
+
|
14 |
+
|
15 |
+
def generate_config(network_name):
|
16 |
+
|
17 |
+
cfg_mnet = {
|
18 |
+
'name': 'mobilenet0.25',
|
19 |
+
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
20 |
+
'steps': [8, 16, 32],
|
21 |
+
'variance': [0.1, 0.2],
|
22 |
+
'clip': False,
|
23 |
+
'loc_weight': 2.0,
|
24 |
+
'gpu_train': True,
|
25 |
+
'batch_size': 32,
|
26 |
+
'ngpu': 1,
|
27 |
+
'epoch': 250,
|
28 |
+
'decay1': 190,
|
29 |
+
'decay2': 220,
|
30 |
+
'image_size': 640,
|
31 |
+
'return_layers': {
|
32 |
+
'stage1': 1,
|
33 |
+
'stage2': 2,
|
34 |
+
'stage3': 3
|
35 |
+
},
|
36 |
+
'in_channel': 32,
|
37 |
+
'out_channel': 64
|
38 |
+
}
|
39 |
+
|
40 |
+
cfg_re50 = {
|
41 |
+
'name': 'Resnet50',
|
42 |
+
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
43 |
+
'steps': [8, 16, 32],
|
44 |
+
'variance': [0.1, 0.2],
|
45 |
+
'clip': False,
|
46 |
+
'loc_weight': 2.0,
|
47 |
+
'gpu_train': True,
|
48 |
+
'batch_size': 24,
|
49 |
+
'ngpu': 4,
|
50 |
+
'epoch': 100,
|
51 |
+
'decay1': 70,
|
52 |
+
'decay2': 90,
|
53 |
+
'image_size': 840,
|
54 |
+
'return_layers': {
|
55 |
+
'layer2': 1,
|
56 |
+
'layer3': 2,
|
57 |
+
'layer4': 3
|
58 |
+
},
|
59 |
+
'in_channel': 256,
|
60 |
+
'out_channel': 256
|
61 |
+
}
|
62 |
+
|
63 |
+
if network_name == 'mobile0.25':
|
64 |
+
return cfg_mnet
|
65 |
+
elif network_name == 'resnet50':
|
66 |
+
return cfg_re50
|
67 |
+
else:
|
68 |
+
raise NotImplementedError(f'network_name={network_name}')
|
69 |
+
|
70 |
+
|
71 |
+
class RetinaFace(nn.Module):
|
72 |
+
|
73 |
+
def __init__(self, network_name='resnet50', half=False, phase='test', device=None):
|
74 |
+
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
|
75 |
+
|
76 |
+
super(RetinaFace, self).__init__()
|
77 |
+
self.half_inference = half
|
78 |
+
cfg = generate_config(network_name)
|
79 |
+
self.backbone = cfg['name']
|
80 |
+
|
81 |
+
self.model_name = f'retinaface_{network_name}'
|
82 |
+
self.cfg = cfg
|
83 |
+
self.phase = phase
|
84 |
+
self.target_size, self.max_size = 1600, 2150
|
85 |
+
self.resize, self.scale, self.scale1 = 1., None, None
|
86 |
+
self.mean_tensor = torch.tensor([[[[104.]], [[117.]], [[123.]]]], device=self.device)
|
87 |
+
self.reference = get_reference_facial_points(default_square=True)
|
88 |
+
# Build network.
|
89 |
+
backbone = None
|
90 |
+
if cfg['name'] == 'mobilenet0.25':
|
91 |
+
backbone = MobileNetV1()
|
92 |
+
self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
|
93 |
+
elif cfg['name'] == 'Resnet50':
|
94 |
+
import torchvision.models as models
|
95 |
+
backbone = models.resnet50(weights=None)
|
96 |
+
self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
|
97 |
+
|
98 |
+
in_channels_stage2 = cfg['in_channel']
|
99 |
+
in_channels_list = [
|
100 |
+
in_channels_stage2 * 2,
|
101 |
+
in_channels_stage2 * 4,
|
102 |
+
in_channels_stage2 * 8,
|
103 |
+
]
|
104 |
+
|
105 |
+
out_channels = cfg['out_channel']
|
106 |
+
self.fpn = FPN(in_channels_list, out_channels)
|
107 |
+
self.ssh1 = SSH(out_channels, out_channels)
|
108 |
+
self.ssh2 = SSH(out_channels, out_channels)
|
109 |
+
self.ssh3 = SSH(out_channels, out_channels)
|
110 |
+
|
111 |
+
self.ClassHead = make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
|
112 |
+
self.BboxHead = make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
|
113 |
+
self.LandmarkHead = make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
|
114 |
+
|
115 |
+
self.to(self.device)
|
116 |
+
self.eval()
|
117 |
+
if self.half_inference:
|
118 |
+
self.half()
|
119 |
+
|
120 |
+
def forward(self, inputs):
|
121 |
+
out = self.body(inputs)
|
122 |
+
|
123 |
+
if self.backbone == 'mobilenet0.25' or self.backbone == 'Resnet50':
|
124 |
+
out = list(out.values())
|
125 |
+
# FPN
|
126 |
+
fpn = self.fpn(out)
|
127 |
+
|
128 |
+
# SSH
|
129 |
+
feature1 = self.ssh1(fpn[0])
|
130 |
+
feature2 = self.ssh2(fpn[1])
|
131 |
+
feature3 = self.ssh3(fpn[2])
|
132 |
+
features = [feature1, feature2, feature3]
|
133 |
+
|
134 |
+
bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
135 |
+
classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
136 |
+
tmp = [self.LandmarkHead[i](feature) for i, feature in enumerate(features)]
|
137 |
+
ldm_regressions = (torch.cat(tmp, dim=1))
|
138 |
+
|
139 |
+
if self.phase == 'train':
|
140 |
+
output = (bbox_regressions, classifications, ldm_regressions)
|
141 |
+
else:
|
142 |
+
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
143 |
+
return output
|
144 |
+
|
145 |
+
def __detect_faces(self, inputs):
|
146 |
+
# get scale
|
147 |
+
height, width = inputs.shape[2:]
|
148 |
+
self.scale = torch.tensor([width, height, width, height], dtype=torch.float32, device=self.device)
|
149 |
+
tmp = [width, height, width, height, width, height, width, height, width, height]
|
150 |
+
self.scale1 = torch.tensor(tmp, dtype=torch.float32, device=self.device)
|
151 |
+
|
152 |
+
# forawrd
|
153 |
+
inputs = inputs.to(self.device)
|
154 |
+
if self.half_inference:
|
155 |
+
inputs = inputs.half()
|
156 |
+
loc, conf, landmarks = self(inputs)
|
157 |
+
|
158 |
+
# get priorbox
|
159 |
+
priorbox = PriorBox(self.cfg, image_size=inputs.shape[2:])
|
160 |
+
priors = priorbox.forward().to(self.device)
|
161 |
+
|
162 |
+
return loc, conf, landmarks, priors
|
163 |
+
|
164 |
+
# single image detection
|
165 |
+
def transform(self, image, use_origin_size):
|
166 |
+
# convert to opencv format
|
167 |
+
if isinstance(image, Image.Image):
|
168 |
+
image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
|
169 |
+
image = image.astype(np.float32)
|
170 |
+
|
171 |
+
# testing scale
|
172 |
+
im_size_min = np.min(image.shape[0:2])
|
173 |
+
im_size_max = np.max(image.shape[0:2])
|
174 |
+
resize = float(self.target_size) / float(im_size_min)
|
175 |
+
|
176 |
+
# prevent bigger axis from being more than max_size
|
177 |
+
if np.round(resize * im_size_max) > self.max_size:
|
178 |
+
resize = float(self.max_size) / float(im_size_max)
|
179 |
+
resize = 1 if use_origin_size else resize
|
180 |
+
|
181 |
+
# resize
|
182 |
+
if resize != 1:
|
183 |
+
image = cv2.resize(image, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
|
184 |
+
|
185 |
+
# convert to torch.tensor format
|
186 |
+
# image -= (104, 117, 123)
|
187 |
+
image = image.transpose(2, 0, 1)
|
188 |
+
image = torch.from_numpy(image).unsqueeze(0)
|
189 |
+
|
190 |
+
return image, resize
|
191 |
+
|
192 |
+
def detect_faces(
|
193 |
+
self,
|
194 |
+
image,
|
195 |
+
conf_threshold=0.8,
|
196 |
+
nms_threshold=0.4,
|
197 |
+
use_origin_size=True,
|
198 |
+
):
|
199 |
+
image, self.resize = self.transform(image, use_origin_size)
|
200 |
+
image = image.to(self.device)
|
201 |
+
if self.half_inference:
|
202 |
+
image = image.half()
|
203 |
+
image = image - self.mean_tensor
|
204 |
+
|
205 |
+
loc, conf, landmarks, priors = self.__detect_faces(image)
|
206 |
+
|
207 |
+
boxes = decode(loc.data.squeeze(0), priors.data, self.cfg['variance'])
|
208 |
+
boxes = boxes * self.scale / self.resize
|
209 |
+
boxes = boxes.cpu().numpy()
|
210 |
+
|
211 |
+
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
212 |
+
|
213 |
+
landmarks = decode_landm(landmarks.squeeze(0), priors, self.cfg['variance'])
|
214 |
+
landmarks = landmarks * self.scale1 / self.resize
|
215 |
+
landmarks = landmarks.cpu().numpy()
|
216 |
+
|
217 |
+
# ignore low scores
|
218 |
+
inds = np.where(scores > conf_threshold)[0]
|
219 |
+
boxes, landmarks, scores = boxes[inds], landmarks[inds], scores[inds]
|
220 |
+
|
221 |
+
# sort
|
222 |
+
order = scores.argsort()[::-1]
|
223 |
+
boxes, landmarks, scores = boxes[order], landmarks[order], scores[order]
|
224 |
+
|
225 |
+
# do NMS
|
226 |
+
bounding_boxes = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
227 |
+
keep = py_cpu_nms(bounding_boxes, nms_threshold)
|
228 |
+
bounding_boxes, landmarks = bounding_boxes[keep, :], landmarks[keep]
|
229 |
+
# self.t['forward_pass'].toc()
|
230 |
+
# print(self.t['forward_pass'].average_time)
|
231 |
+
# import sys
|
232 |
+
# sys.stdout.flush()
|
233 |
+
return np.concatenate((bounding_boxes, landmarks), axis=1)
|
234 |
+
|
235 |
+
def __align_multi(self, image, boxes, landmarks, limit=None):
|
236 |
+
|
237 |
+
if len(boxes) < 1:
|
238 |
+
return [], []
|
239 |
+
|
240 |
+
if limit:
|
241 |
+
boxes = boxes[:limit]
|
242 |
+
landmarks = landmarks[:limit]
|
243 |
+
|
244 |
+
faces = []
|
245 |
+
for landmark in landmarks:
|
246 |
+
facial5points = [[landmark[2 * j], landmark[2 * j + 1]] for j in range(5)]
|
247 |
+
|
248 |
+
warped_face = warp_and_crop_face(np.array(image), facial5points, self.reference, crop_size=(112, 112))
|
249 |
+
faces.append(warped_face)
|
250 |
+
|
251 |
+
return np.concatenate((boxes, landmarks), axis=1), faces
|
252 |
+
|
253 |
+
def align_multi(self, img, conf_threshold=0.8, limit=None):
|
254 |
+
|
255 |
+
rlt = self.detect_faces(img, conf_threshold=conf_threshold)
|
256 |
+
boxes, landmarks = rlt[:, 0:5], rlt[:, 5:]
|
257 |
+
|
258 |
+
return self.__align_multi(img, boxes, landmarks, limit)
|
259 |
+
|
260 |
+
# batched detection
|
261 |
+
def batched_transform(self, frames, use_origin_size):
|
262 |
+
"""
|
263 |
+
Arguments:
|
264 |
+
frames: a list of PIL.Image, or torch.Tensor(shape=[n, h, w, c],
|
265 |
+
type=np.float32, BGR format).
|
266 |
+
use_origin_size: whether to use origin size.
|
267 |
+
"""
|
268 |
+
from_PIL = True if isinstance(frames[0], Image.Image) else False
|
269 |
+
|
270 |
+
# convert to opencv format
|
271 |
+
if from_PIL:
|
272 |
+
frames = [cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR) for frame in frames]
|
273 |
+
frames = np.asarray(frames, dtype=np.float32)
|
274 |
+
|
275 |
+
# testing scale
|
276 |
+
im_size_min = np.min(frames[0].shape[0:2])
|
277 |
+
im_size_max = np.max(frames[0].shape[0:2])
|
278 |
+
resize = float(self.target_size) / float(im_size_min)
|
279 |
+
|
280 |
+
# prevent bigger axis from being more than max_size
|
281 |
+
if np.round(resize * im_size_max) > self.max_size:
|
282 |
+
resize = float(self.max_size) / float(im_size_max)
|
283 |
+
resize = 1 if use_origin_size else resize
|
284 |
+
|
285 |
+
# resize
|
286 |
+
if resize != 1:
|
287 |
+
if not from_PIL:
|
288 |
+
frames = F.interpolate(frames, scale_factor=resize)
|
289 |
+
else:
|
290 |
+
frames = [
|
291 |
+
cv2.resize(frame, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
|
292 |
+
for frame in frames
|
293 |
+
]
|
294 |
+
|
295 |
+
# convert to torch.tensor format
|
296 |
+
if not from_PIL:
|
297 |
+
frames = frames.transpose(1, 2).transpose(1, 3).contiguous()
|
298 |
+
else:
|
299 |
+
frames = frames.transpose((0, 3, 1, 2))
|
300 |
+
frames = torch.from_numpy(frames)
|
301 |
+
|
302 |
+
return frames, resize
|
303 |
+
|
304 |
+
def batched_detect_faces(self, frames, conf_threshold=0.8, nms_threshold=0.4, use_origin_size=True):
|
305 |
+
"""
|
306 |
+
Arguments:
|
307 |
+
frames: a list of PIL.Image, or np.array(shape=[n, h, w, c],
|
308 |
+
type=np.uint8, BGR format).
|
309 |
+
conf_threshold: confidence threshold.
|
310 |
+
nms_threshold: nms threshold.
|
311 |
+
use_origin_size: whether to use origin size.
|
312 |
+
Returns:
|
313 |
+
final_bounding_boxes: list of np.array ([n_boxes, 5],
|
314 |
+
type=np.float32).
|
315 |
+
final_landmarks: list of np.array ([n_boxes, 10], type=np.float32).
|
316 |
+
"""
|
317 |
+
# self.t['forward_pass'].tic()
|
318 |
+
frames, self.resize = self.batched_transform(frames, use_origin_size)
|
319 |
+
frames = frames.to(self.device)
|
320 |
+
frames = frames - self.mean_tensor
|
321 |
+
|
322 |
+
b_loc, b_conf, b_landmarks, priors = self.__detect_faces(frames)
|
323 |
+
|
324 |
+
final_bounding_boxes, final_landmarks = [], []
|
325 |
+
|
326 |
+
# decode
|
327 |
+
priors = priors.unsqueeze(0)
|
328 |
+
b_loc = batched_decode(b_loc, priors, self.cfg['variance']) * self.scale / self.resize
|
329 |
+
b_landmarks = batched_decode_landm(b_landmarks, priors, self.cfg['variance']) * self.scale1 / self.resize
|
330 |
+
b_conf = b_conf[:, :, 1]
|
331 |
+
|
332 |
+
# index for selection
|
333 |
+
b_indice = b_conf > conf_threshold
|
334 |
+
|
335 |
+
# concat
|
336 |
+
b_loc_and_conf = torch.cat((b_loc, b_conf.unsqueeze(-1)), dim=2).float()
|
337 |
+
|
338 |
+
for pred, landm, inds in zip(b_loc_and_conf, b_landmarks, b_indice):
|
339 |
+
|
340 |
+
# ignore low scores
|
341 |
+
pred, landm = pred[inds, :], landm[inds, :]
|
342 |
+
if pred.shape[0] == 0:
|
343 |
+
final_bounding_boxes.append(np.array([], dtype=np.float32))
|
344 |
+
final_landmarks.append(np.array([], dtype=np.float32))
|
345 |
+
continue
|
346 |
+
|
347 |
+
# sort
|
348 |
+
# order = score.argsort(descending=True)
|
349 |
+
# box, landm, score = box[order], landm[order], score[order]
|
350 |
+
|
351 |
+
# to CPU
|
352 |
+
bounding_boxes, landm = pred.cpu().numpy(), landm.cpu().numpy()
|
353 |
+
|
354 |
+
# NMS
|
355 |
+
keep = py_cpu_nms(bounding_boxes, nms_threshold)
|
356 |
+
bounding_boxes, landmarks = bounding_boxes[keep, :], landm[keep]
|
357 |
+
|
358 |
+
# append
|
359 |
+
final_bounding_boxes.append(bounding_boxes)
|
360 |
+
final_landmarks.append(landmarks)
|
361 |
+
# self.t['forward_pass'].toc(average=True)
|
362 |
+
# self.batch_time += self.t['forward_pass'].diff
|
363 |
+
# self.total_frame += len(frames)
|
364 |
+
# print(self.batch_time / self.total_frame)
|
365 |
+
|
366 |
+
return final_bounding_boxes, final_landmarks
|
extras/facexlib/detection/retinaface_net.py
ADDED
@@ -0,0 +1,196 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import torch.nn.functional as F
|
4 |
+
|
5 |
+
|
6 |
+
def conv_bn(inp, oup, stride=1, leaky=0):
|
7 |
+
return nn.Sequential(
|
8 |
+
nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),
|
9 |
+
nn.LeakyReLU(negative_slope=leaky, inplace=True))
|
10 |
+
|
11 |
+
|
12 |
+
def conv_bn_no_relu(inp, oup, stride):
|
13 |
+
return nn.Sequential(
|
14 |
+
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
|
15 |
+
nn.BatchNorm2d(oup),
|
16 |
+
)
|
17 |
+
|
18 |
+
|
19 |
+
def conv_bn1X1(inp, oup, stride, leaky=0):
|
20 |
+
return nn.Sequential(
|
21 |
+
nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False), nn.BatchNorm2d(oup),
|
22 |
+
nn.LeakyReLU(negative_slope=leaky, inplace=True))
|
23 |
+
|
24 |
+
|
25 |
+
def conv_dw(inp, oup, stride, leaky=0.1):
|
26 |
+
return nn.Sequential(
|
27 |
+
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
|
28 |
+
nn.BatchNorm2d(inp),
|
29 |
+
nn.LeakyReLU(negative_slope=leaky, inplace=True),
|
30 |
+
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
|
31 |
+
nn.BatchNorm2d(oup),
|
32 |
+
nn.LeakyReLU(negative_slope=leaky, inplace=True),
|
33 |
+
)
|
34 |
+
|
35 |
+
|
36 |
+
class SSH(nn.Module):
|
37 |
+
|
38 |
+
def __init__(self, in_channel, out_channel):
|
39 |
+
super(SSH, self).__init__()
|
40 |
+
assert out_channel % 4 == 0
|
41 |
+
leaky = 0
|
42 |
+
if (out_channel <= 64):
|
43 |
+
leaky = 0.1
|
44 |
+
self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
|
45 |
+
|
46 |
+
self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
|
47 |
+
self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
|
48 |
+
|
49 |
+
self.conv7X7_2 = conv_bn(out_channel // 4, out_channel // 4, stride=1, leaky=leaky)
|
50 |
+
self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
|
51 |
+
|
52 |
+
def forward(self, input):
|
53 |
+
conv3X3 = self.conv3X3(input)
|
54 |
+
|
55 |
+
conv5X5_1 = self.conv5X5_1(input)
|
56 |
+
conv5X5 = self.conv5X5_2(conv5X5_1)
|
57 |
+
|
58 |
+
conv7X7_2 = self.conv7X7_2(conv5X5_1)
|
59 |
+
conv7X7 = self.conv7x7_3(conv7X7_2)
|
60 |
+
|
61 |
+
out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
|
62 |
+
out = F.relu(out)
|
63 |
+
return out
|
64 |
+
|
65 |
+
|
66 |
+
class FPN(nn.Module):
|
67 |
+
|
68 |
+
def __init__(self, in_channels_list, out_channels):
|
69 |
+
super(FPN, self).__init__()
|
70 |
+
leaky = 0
|
71 |
+
if (out_channels <= 64):
|
72 |
+
leaky = 0.1
|
73 |
+
self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride=1, leaky=leaky)
|
74 |
+
self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride=1, leaky=leaky)
|
75 |
+
self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride=1, leaky=leaky)
|
76 |
+
|
77 |
+
self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
|
78 |
+
self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
|
79 |
+
|
80 |
+
def forward(self, input):
|
81 |
+
# names = list(input.keys())
|
82 |
+
# input = list(input.values())
|
83 |
+
|
84 |
+
output1 = self.output1(input[0])
|
85 |
+
output2 = self.output2(input[1])
|
86 |
+
output3 = self.output3(input[2])
|
87 |
+
|
88 |
+
up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode='nearest')
|
89 |
+
output2 = output2 + up3
|
90 |
+
output2 = self.merge2(output2)
|
91 |
+
|
92 |
+
up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode='nearest')
|
93 |
+
output1 = output1 + up2
|
94 |
+
output1 = self.merge1(output1)
|
95 |
+
|
96 |
+
out = [output1, output2, output3]
|
97 |
+
return out
|
98 |
+
|
99 |
+
|
100 |
+
class MobileNetV1(nn.Module):
|
101 |
+
|
102 |
+
def __init__(self):
|
103 |
+
super(MobileNetV1, self).__init__()
|
104 |
+
self.stage1 = nn.Sequential(
|
105 |
+
conv_bn(3, 8, 2, leaky=0.1), # 3
|
106 |
+
conv_dw(8, 16, 1), # 7
|
107 |
+
conv_dw(16, 32, 2), # 11
|
108 |
+
conv_dw(32, 32, 1), # 19
|
109 |
+
conv_dw(32, 64, 2), # 27
|
110 |
+
conv_dw(64, 64, 1), # 43
|
111 |
+
)
|
112 |
+
self.stage2 = nn.Sequential(
|
113 |
+
conv_dw(64, 128, 2), # 43 + 16 = 59
|
114 |
+
conv_dw(128, 128, 1), # 59 + 32 = 91
|
115 |
+
conv_dw(128, 128, 1), # 91 + 32 = 123
|
116 |
+
conv_dw(128, 128, 1), # 123 + 32 = 155
|
117 |
+
conv_dw(128, 128, 1), # 155 + 32 = 187
|
118 |
+
conv_dw(128, 128, 1), # 187 + 32 = 219
|
119 |
+
)
|
120 |
+
self.stage3 = nn.Sequential(
|
121 |
+
conv_dw(128, 256, 2), # 219 +3 2 = 241
|
122 |
+
conv_dw(256, 256, 1), # 241 + 64 = 301
|
123 |
+
)
|
124 |
+
self.avg = nn.AdaptiveAvgPool2d((1, 1))
|
125 |
+
self.fc = nn.Linear(256, 1000)
|
126 |
+
|
127 |
+
def forward(self, x):
|
128 |
+
x = self.stage1(x)
|
129 |
+
x = self.stage2(x)
|
130 |
+
x = self.stage3(x)
|
131 |
+
x = self.avg(x)
|
132 |
+
# x = self.model(x)
|
133 |
+
x = x.view(-1, 256)
|
134 |
+
x = self.fc(x)
|
135 |
+
return x
|
136 |
+
|
137 |
+
|
138 |
+
class ClassHead(nn.Module):
|
139 |
+
|
140 |
+
def __init__(self, inchannels=512, num_anchors=3):
|
141 |
+
super(ClassHead, self).__init__()
|
142 |
+
self.num_anchors = num_anchors
|
143 |
+
self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
|
144 |
+
|
145 |
+
def forward(self, x):
|
146 |
+
out = self.conv1x1(x)
|
147 |
+
out = out.permute(0, 2, 3, 1).contiguous()
|
148 |
+
|
149 |
+
return out.view(out.shape[0], -1, 2)
|
150 |
+
|
151 |
+
|
152 |
+
class BboxHead(nn.Module):
|
153 |
+
|
154 |
+
def __init__(self, inchannels=512, num_anchors=3):
|
155 |
+
super(BboxHead, self).__init__()
|
156 |
+
self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0)
|
157 |
+
|
158 |
+
def forward(self, x):
|
159 |
+
out = self.conv1x1(x)
|
160 |
+
out = out.permute(0, 2, 3, 1).contiguous()
|
161 |
+
|
162 |
+
return out.view(out.shape[0], -1, 4)
|
163 |
+
|
164 |
+
|
165 |
+
class LandmarkHead(nn.Module):
|
166 |
+
|
167 |
+
def __init__(self, inchannels=512, num_anchors=3):
|
168 |
+
super(LandmarkHead, self).__init__()
|
169 |
+
self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0)
|
170 |
+
|
171 |
+
def forward(self, x):
|
172 |
+
out = self.conv1x1(x)
|
173 |
+
out = out.permute(0, 2, 3, 1).contiguous()
|
174 |
+
|
175 |
+
return out.view(out.shape[0], -1, 10)
|
176 |
+
|
177 |
+
|
178 |
+
def make_class_head(fpn_num=3, inchannels=64, anchor_num=2):
|
179 |
+
classhead = nn.ModuleList()
|
180 |
+
for i in range(fpn_num):
|
181 |
+
classhead.append(ClassHead(inchannels, anchor_num))
|
182 |
+
return classhead
|
183 |
+
|
184 |
+
|
185 |
+
def make_bbox_head(fpn_num=3, inchannels=64, anchor_num=2):
|
186 |
+
bboxhead = nn.ModuleList()
|
187 |
+
for i in range(fpn_num):
|
188 |
+
bboxhead.append(BboxHead(inchannels, anchor_num))
|
189 |
+
return bboxhead
|
190 |
+
|
191 |
+
|
192 |
+
def make_landmark_head(fpn_num=3, inchannels=64, anchor_num=2):
|
193 |
+
landmarkhead = nn.ModuleList()
|
194 |
+
for i in range(fpn_num):
|
195 |
+
landmarkhead.append(LandmarkHead(inchannels, anchor_num))
|
196 |
+
return landmarkhead
|
extras/facexlib/detection/retinaface_utils.py
ADDED
@@ -0,0 +1,421 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import torch
|
3 |
+
import torchvision
|
4 |
+
from itertools import product as product
|
5 |
+
from math import ceil
|
6 |
+
|
7 |
+
|
8 |
+
class PriorBox(object):
|
9 |
+
|
10 |
+
def __init__(self, cfg, image_size=None, phase='train'):
|
11 |
+
super(PriorBox, self).__init__()
|
12 |
+
self.min_sizes = cfg['min_sizes']
|
13 |
+
self.steps = cfg['steps']
|
14 |
+
self.clip = cfg['clip']
|
15 |
+
self.image_size = image_size
|
16 |
+
self.feature_maps = [[ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)] for step in self.steps]
|
17 |
+
self.name = 's'
|
18 |
+
|
19 |
+
def forward(self):
|
20 |
+
anchors = []
|
21 |
+
for k, f in enumerate(self.feature_maps):
|
22 |
+
min_sizes = self.min_sizes[k]
|
23 |
+
for i, j in product(range(f[0]), range(f[1])):
|
24 |
+
for min_size in min_sizes:
|
25 |
+
s_kx = min_size / self.image_size[1]
|
26 |
+
s_ky = min_size / self.image_size[0]
|
27 |
+
dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
|
28 |
+
dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
|
29 |
+
for cy, cx in product(dense_cy, dense_cx):
|
30 |
+
anchors += [cx, cy, s_kx, s_ky]
|
31 |
+
|
32 |
+
# back to torch land
|
33 |
+
output = torch.Tensor(anchors).view(-1, 4)
|
34 |
+
if self.clip:
|
35 |
+
output.clamp_(max=1, min=0)
|
36 |
+
return output
|
37 |
+
|
38 |
+
|
39 |
+
def py_cpu_nms(dets, thresh):
|
40 |
+
"""Pure Python NMS baseline."""
|
41 |
+
keep = torchvision.ops.nms(
|
42 |
+
boxes=torch.Tensor(dets[:, :4]),
|
43 |
+
scores=torch.Tensor(dets[:, 4]),
|
44 |
+
iou_threshold=thresh,
|
45 |
+
)
|
46 |
+
|
47 |
+
return list(keep)
|
48 |
+
|
49 |
+
|
50 |
+
def point_form(boxes):
|
51 |
+
""" Convert prior_boxes to (xmin, ymin, xmax, ymax)
|
52 |
+
representation for comparison to point form ground truth data.
|
53 |
+
Args:
|
54 |
+
boxes: (tensor) center-size default boxes from priorbox layers.
|
55 |
+
Return:
|
56 |
+
boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
|
57 |
+
"""
|
58 |
+
return torch.cat(
|
59 |
+
(
|
60 |
+
boxes[:, :2] - boxes[:, 2:] / 2, # xmin, ymin
|
61 |
+
boxes[:, :2] + boxes[:, 2:] / 2),
|
62 |
+
1) # xmax, ymax
|
63 |
+
|
64 |
+
|
65 |
+
def center_size(boxes):
|
66 |
+
""" Convert prior_boxes to (cx, cy, w, h)
|
67 |
+
representation for comparison to center-size form ground truth data.
|
68 |
+
Args:
|
69 |
+
boxes: (tensor) point_form boxes
|
70 |
+
Return:
|
71 |
+
boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
|
72 |
+
"""
|
73 |
+
return torch.cat(
|
74 |
+
(boxes[:, 2:] + boxes[:, :2]) / 2, # cx, cy
|
75 |
+
boxes[:, 2:] - boxes[:, :2],
|
76 |
+
1) # w, h
|
77 |
+
|
78 |
+
|
79 |
+
def intersect(box_a, box_b):
|
80 |
+
""" We resize both tensors to [A,B,2] without new malloc:
|
81 |
+
[A,2] -> [A,1,2] -> [A,B,2]
|
82 |
+
[B,2] -> [1,B,2] -> [A,B,2]
|
83 |
+
Then we compute the area of intersect between box_a and box_b.
|
84 |
+
Args:
|
85 |
+
box_a: (tensor) bounding boxes, Shape: [A,4].
|
86 |
+
box_b: (tensor) bounding boxes, Shape: [B,4].
|
87 |
+
Return:
|
88 |
+
(tensor) intersection area, Shape: [A,B].
|
89 |
+
"""
|
90 |
+
A = box_a.size(0)
|
91 |
+
B = box_b.size(0)
|
92 |
+
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
|
93 |
+
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
|
94 |
+
inter = torch.clamp((max_xy - min_xy), min=0)
|
95 |
+
return inter[:, :, 0] * inter[:, :, 1]
|
96 |
+
|
97 |
+
|
98 |
+
def jaccard(box_a, box_b):
|
99 |
+
"""Compute the jaccard overlap of two sets of boxes. The jaccard overlap
|
100 |
+
is simply the intersection over union of two boxes. Here we operate on
|
101 |
+
ground truth boxes and default boxes.
|
102 |
+
E.g.:
|
103 |
+
A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
|
104 |
+
Args:
|
105 |
+
box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
|
106 |
+
box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
|
107 |
+
Return:
|
108 |
+
jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
|
109 |
+
"""
|
110 |
+
inter = intersect(box_a, box_b)
|
111 |
+
area_a = ((box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B]
|
112 |
+
area_b = ((box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B]
|
113 |
+
union = area_a + area_b - inter
|
114 |
+
return inter / union # [A,B]
|
115 |
+
|
116 |
+
|
117 |
+
def matrix_iou(a, b):
|
118 |
+
"""
|
119 |
+
return iou of a and b, numpy version for data augenmentation
|
120 |
+
"""
|
121 |
+
lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
|
122 |
+
rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
|
123 |
+
|
124 |
+
area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
|
125 |
+
area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
|
126 |
+
area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
|
127 |
+
return area_i / (area_a[:, np.newaxis] + area_b - area_i)
|
128 |
+
|
129 |
+
|
130 |
+
def matrix_iof(a, b):
|
131 |
+
"""
|
132 |
+
return iof of a and b, numpy version for data augenmentation
|
133 |
+
"""
|
134 |
+
lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
|
135 |
+
rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
|
136 |
+
|
137 |
+
area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
|
138 |
+
area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
|
139 |
+
return area_i / np.maximum(area_a[:, np.newaxis], 1)
|
140 |
+
|
141 |
+
|
142 |
+
def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):
|
143 |
+
"""Match each prior box with the ground truth box of the highest jaccard
|
144 |
+
overlap, encode the bounding boxes, then return the matched indices
|
145 |
+
corresponding to both confidence and location preds.
|
146 |
+
Args:
|
147 |
+
threshold: (float) The overlap threshold used when matching boxes.
|
148 |
+
truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
|
149 |
+
priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
|
150 |
+
variances: (tensor) Variances corresponding to each prior coord,
|
151 |
+
Shape: [num_priors, 4].
|
152 |
+
labels: (tensor) All the class labels for the image, Shape: [num_obj].
|
153 |
+
landms: (tensor) Ground truth landms, Shape [num_obj, 10].
|
154 |
+
loc_t: (tensor) Tensor to be filled w/ encoded location targets.
|
155 |
+
conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
|
156 |
+
landm_t: (tensor) Tensor to be filled w/ encoded landm targets.
|
157 |
+
idx: (int) current batch index
|
158 |
+
Return:
|
159 |
+
The matched indices corresponding to 1)location 2)confidence
|
160 |
+
3)landm preds.
|
161 |
+
"""
|
162 |
+
# jaccard index
|
163 |
+
overlaps = jaccard(truths, point_form(priors))
|
164 |
+
# (Bipartite Matching)
|
165 |
+
# [1,num_objects] best prior for each ground truth
|
166 |
+
best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
|
167 |
+
|
168 |
+
# ignore hard gt
|
169 |
+
valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
|
170 |
+
best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
|
171 |
+
if best_prior_idx_filter.shape[0] <= 0:
|
172 |
+
loc_t[idx] = 0
|
173 |
+
conf_t[idx] = 0
|
174 |
+
return
|
175 |
+
|
176 |
+
# [1,num_priors] best ground truth for each prior
|
177 |
+
best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
|
178 |
+
best_truth_idx.squeeze_(0)
|
179 |
+
best_truth_overlap.squeeze_(0)
|
180 |
+
best_prior_idx.squeeze_(1)
|
181 |
+
best_prior_idx_filter.squeeze_(1)
|
182 |
+
best_prior_overlap.squeeze_(1)
|
183 |
+
best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2) # ensure best prior
|
184 |
+
# TODO refactor: index best_prior_idx with long tensor
|
185 |
+
# ensure every gt matches with its prior of max overlap
|
186 |
+
for j in range(best_prior_idx.size(0)): # 判别此anchor是预测哪一个boxes
|
187 |
+
best_truth_idx[best_prior_idx[j]] = j
|
188 |
+
matches = truths[best_truth_idx] # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
|
189 |
+
conf = labels[best_truth_idx] # Shape: [num_priors] 此处为每一个anchor对应的label取出来
|
190 |
+
conf[best_truth_overlap < threshold] = 0 # label as background overlap<0.35的全部作为负样本
|
191 |
+
loc = encode(matches, priors, variances)
|
192 |
+
|
193 |
+
matches_landm = landms[best_truth_idx]
|
194 |
+
landm = encode_landm(matches_landm, priors, variances)
|
195 |
+
loc_t[idx] = loc # [num_priors,4] encoded offsets to learn
|
196 |
+
conf_t[idx] = conf # [num_priors] top class label for each prior
|
197 |
+
landm_t[idx] = landm
|
198 |
+
|
199 |
+
|
200 |
+
def encode(matched, priors, variances):
|
201 |
+
"""Encode the variances from the priorbox layers into the ground truth boxes
|
202 |
+
we have matched (based on jaccard overlap) with the prior boxes.
|
203 |
+
Args:
|
204 |
+
matched: (tensor) Coords of ground truth for each prior in point-form
|
205 |
+
Shape: [num_priors, 4].
|
206 |
+
priors: (tensor) Prior boxes in center-offset form
|
207 |
+
Shape: [num_priors,4].
|
208 |
+
variances: (list[float]) Variances of priorboxes
|
209 |
+
Return:
|
210 |
+
encoded boxes (tensor), Shape: [num_priors, 4]
|
211 |
+
"""
|
212 |
+
|
213 |
+
# dist b/t match center and prior's center
|
214 |
+
g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
|
215 |
+
# encode variance
|
216 |
+
g_cxcy /= (variances[0] * priors[:, 2:])
|
217 |
+
# match wh / prior wh
|
218 |
+
g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
|
219 |
+
g_wh = torch.log(g_wh) / variances[1]
|
220 |
+
# return target for smooth_l1_loss
|
221 |
+
return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4]
|
222 |
+
|
223 |
+
|
224 |
+
def encode_landm(matched, priors, variances):
|
225 |
+
"""Encode the variances from the priorbox layers into the ground truth boxes
|
226 |
+
we have matched (based on jaccard overlap) with the prior boxes.
|
227 |
+
Args:
|
228 |
+
matched: (tensor) Coords of ground truth for each prior in point-form
|
229 |
+
Shape: [num_priors, 10].
|
230 |
+
priors: (tensor) Prior boxes in center-offset form
|
231 |
+
Shape: [num_priors,4].
|
232 |
+
variances: (list[float]) Variances of priorboxes
|
233 |
+
Return:
|
234 |
+
encoded landm (tensor), Shape: [num_priors, 10]
|
235 |
+
"""
|
236 |
+
|
237 |
+
# dist b/t match center and prior's center
|
238 |
+
matched = torch.reshape(matched, (matched.size(0), 5, 2))
|
239 |
+
priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
|
240 |
+
priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
|
241 |
+
priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
|
242 |
+
priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
|
243 |
+
priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2)
|
244 |
+
g_cxcy = matched[:, :, :2] - priors[:, :, :2]
|
245 |
+
# encode variance
|
246 |
+
g_cxcy /= (variances[0] * priors[:, :, 2:])
|
247 |
+
# g_cxcy /= priors[:, :, 2:]
|
248 |
+
g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)
|
249 |
+
# return target for smooth_l1_loss
|
250 |
+
return g_cxcy
|
251 |
+
|
252 |
+
|
253 |
+
# Adapted from https://github.com/Hakuyume/chainer-ssd
|
254 |
+
def decode(loc, priors, variances):
|
255 |
+
"""Decode locations from predictions using priors to undo
|
256 |
+
the encoding we did for offset regression at train time.
|
257 |
+
Args:
|
258 |
+
loc (tensor): location predictions for loc layers,
|
259 |
+
Shape: [num_priors,4]
|
260 |
+
priors (tensor): Prior boxes in center-offset form.
|
261 |
+
Shape: [num_priors,4].
|
262 |
+
variances: (list[float]) Variances of priorboxes
|
263 |
+
Return:
|
264 |
+
decoded bounding box predictions
|
265 |
+
"""
|
266 |
+
|
267 |
+
boxes = torch.cat((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
268 |
+
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
|
269 |
+
boxes[:, :2] -= boxes[:, 2:] / 2
|
270 |
+
boxes[:, 2:] += boxes[:, :2]
|
271 |
+
return boxes
|
272 |
+
|
273 |
+
|
274 |
+
def decode_landm(pre, priors, variances):
|
275 |
+
"""Decode landm from predictions using priors to undo
|
276 |
+
the encoding we did for offset regression at train time.
|
277 |
+
Args:
|
278 |
+
pre (tensor): landm predictions for loc layers,
|
279 |
+
Shape: [num_priors,10]
|
280 |
+
priors (tensor): Prior boxes in center-offset form.
|
281 |
+
Shape: [num_priors,4].
|
282 |
+
variances: (list[float]) Variances of priorboxes
|
283 |
+
Return:
|
284 |
+
decoded landm predictions
|
285 |
+
"""
|
286 |
+
tmp = (
|
287 |
+
priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
|
288 |
+
priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
|
289 |
+
priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
|
290 |
+
priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
|
291 |
+
priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
|
292 |
+
)
|
293 |
+
landms = torch.cat(tmp, dim=1)
|
294 |
+
return landms
|
295 |
+
|
296 |
+
|
297 |
+
def batched_decode(b_loc, priors, variances):
|
298 |
+
"""Decode locations from predictions using priors to undo
|
299 |
+
the encoding we did for offset regression at train time.
|
300 |
+
Args:
|
301 |
+
b_loc (tensor): location predictions for loc layers,
|
302 |
+
Shape: [num_batches,num_priors,4]
|
303 |
+
priors (tensor): Prior boxes in center-offset form.
|
304 |
+
Shape: [1,num_priors,4].
|
305 |
+
variances: (list[float]) Variances of priorboxes
|
306 |
+
Return:
|
307 |
+
decoded bounding box predictions
|
308 |
+
"""
|
309 |
+
boxes = (
|
310 |
+
priors[:, :, :2] + b_loc[:, :, :2] * variances[0] * priors[:, :, 2:],
|
311 |
+
priors[:, :, 2:] * torch.exp(b_loc[:, :, 2:] * variances[1]),
|
312 |
+
)
|
313 |
+
boxes = torch.cat(boxes, dim=2)
|
314 |
+
|
315 |
+
boxes[:, :, :2] -= boxes[:, :, 2:] / 2
|
316 |
+
boxes[:, :, 2:] += boxes[:, :, :2]
|
317 |
+
return boxes
|
318 |
+
|
319 |
+
|
320 |
+
def batched_decode_landm(pre, priors, variances):
|
321 |
+
"""Decode landm from predictions using priors to undo
|
322 |
+
the encoding we did for offset regression at train time.
|
323 |
+
Args:
|
324 |
+
pre (tensor): landm predictions for loc layers,
|
325 |
+
Shape: [num_batches,num_priors,10]
|
326 |
+
priors (tensor): Prior boxes in center-offset form.
|
327 |
+
Shape: [1,num_priors,4].
|
328 |
+
variances: (list[float]) Variances of priorboxes
|
329 |
+
Return:
|
330 |
+
decoded landm predictions
|
331 |
+
"""
|
332 |
+
landms = (
|
333 |
+
priors[:, :, :2] + pre[:, :, :2] * variances[0] * priors[:, :, 2:],
|
334 |
+
priors[:, :, :2] + pre[:, :, 2:4] * variances[0] * priors[:, :, 2:],
|
335 |
+
priors[:, :, :2] + pre[:, :, 4:6] * variances[0] * priors[:, :, 2:],
|
336 |
+
priors[:, :, :2] + pre[:, :, 6:8] * variances[0] * priors[:, :, 2:],
|
337 |
+
priors[:, :, :2] + pre[:, :, 8:10] * variances[0] * priors[:, :, 2:],
|
338 |
+
)
|
339 |
+
landms = torch.cat(landms, dim=2)
|
340 |
+
return landms
|
341 |
+
|
342 |
+
|
343 |
+
def log_sum_exp(x):
|
344 |
+
"""Utility function for computing log_sum_exp while determining
|
345 |
+
This will be used to determine unaveraged confidence loss across
|
346 |
+
all examples in a batch.
|
347 |
+
Args:
|
348 |
+
x (Variable(tensor)): conf_preds from conf layers
|
349 |
+
"""
|
350 |
+
x_max = x.data.max()
|
351 |
+
return torch.log(torch.sum(torch.exp(x - x_max), 1, keepdim=True)) + x_max
|
352 |
+
|
353 |
+
|
354 |
+
# Original author: Francisco Massa:
|
355 |
+
# https://github.com/fmassa/object-detection.torch
|
356 |
+
# Ported to PyTorch by Max deGroot (02/01/2017)
|
357 |
+
def nms(boxes, scores, overlap=0.5, top_k=200):
|
358 |
+
"""Apply non-maximum suppression at test time to avoid detecting too many
|
359 |
+
overlapping bounding boxes for a given object.
|
360 |
+
Args:
|
361 |
+
boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
|
362 |
+
scores: (tensor) The class predscores for the img, Shape:[num_priors].
|
363 |
+
overlap: (float) The overlap thresh for suppressing unnecessary boxes.
|
364 |
+
top_k: (int) The Maximum number of box preds to consider.
|
365 |
+
Return:
|
366 |
+
The indices of the kept boxes with respect to num_priors.
|
367 |
+
"""
|
368 |
+
|
369 |
+
keep = torch.Tensor(scores.size(0)).fill_(0).long()
|
370 |
+
if boxes.numel() == 0:
|
371 |
+
return keep
|
372 |
+
x1 = boxes[:, 0]
|
373 |
+
y1 = boxes[:, 1]
|
374 |
+
x2 = boxes[:, 2]
|
375 |
+
y2 = boxes[:, 3]
|
376 |
+
area = torch.mul(x2 - x1, y2 - y1)
|
377 |
+
v, idx = scores.sort(0) # sort in ascending order
|
378 |
+
# I = I[v >= 0.01]
|
379 |
+
idx = idx[-top_k:] # indices of the top-k largest vals
|
380 |
+
xx1 = boxes.new()
|
381 |
+
yy1 = boxes.new()
|
382 |
+
xx2 = boxes.new()
|
383 |
+
yy2 = boxes.new()
|
384 |
+
w = boxes.new()
|
385 |
+
h = boxes.new()
|
386 |
+
|
387 |
+
# keep = torch.Tensor()
|
388 |
+
count = 0
|
389 |
+
while idx.numel() > 0:
|
390 |
+
i = idx[-1] # index of current largest val
|
391 |
+
# keep.append(i)
|
392 |
+
keep[count] = i
|
393 |
+
count += 1
|
394 |
+
if idx.size(0) == 1:
|
395 |
+
break
|
396 |
+
idx = idx[:-1] # remove kept element from view
|
397 |
+
# load bboxes of next highest vals
|
398 |
+
torch.index_select(x1, 0, idx, out=xx1)
|
399 |
+
torch.index_select(y1, 0, idx, out=yy1)
|
400 |
+
torch.index_select(x2, 0, idx, out=xx2)
|
401 |
+
torch.index_select(y2, 0, idx, out=yy2)
|
402 |
+
# store element-wise max with next highest score
|
403 |
+
xx1 = torch.clamp(xx1, min=x1[i])
|
404 |
+
yy1 = torch.clamp(yy1, min=y1[i])
|
405 |
+
xx2 = torch.clamp(xx2, max=x2[i])
|
406 |
+
yy2 = torch.clamp(yy2, max=y2[i])
|
407 |
+
w.resize_as_(xx2)
|
408 |
+
h.resize_as_(yy2)
|
409 |
+
w = xx2 - xx1
|
410 |
+
h = yy2 - yy1
|
411 |
+
# check sizes of xx1 and xx2.. after each iteration
|
412 |
+
w = torch.clamp(w, min=0.0)
|
413 |
+
h = torch.clamp(h, min=0.0)
|
414 |
+
inter = w * h
|
415 |
+
# IoU = i / (area(a) + area(b) - i)
|
416 |
+
rem_areas = torch.index_select(area, 0, idx) # load remaining areas)
|
417 |
+
union = (rem_areas - inter) + area[i]
|
418 |
+
IoU = inter / union # store result in iou
|
419 |
+
# keep only elements with an IoU <= overlap
|
420 |
+
idx = idx[IoU.le(overlap)]
|
421 |
+
return keep, count
|
extras/facexlib/parsing/__init__.py
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from extras.facexlib.utils import load_file_from_url
|
4 |
+
from .bisenet import BiSeNet
|
5 |
+
from .parsenet import ParseNet
|
6 |
+
|
7 |
+
|
8 |
+
def init_parsing_model(model_name='bisenet', half=False, device='cuda', model_rootpath=None):
|
9 |
+
if model_name == 'bisenet':
|
10 |
+
model = BiSeNet(num_class=19)
|
11 |
+
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.0/parsing_bisenet.pth'
|
12 |
+
elif model_name == 'parsenet':
|
13 |
+
model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
|
14 |
+
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth'
|
15 |
+
else:
|
16 |
+
raise NotImplementedError(f'{model_name} is not implemented.')
|
17 |
+
|
18 |
+
model_path = load_file_from_url(
|
19 |
+
url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
|
20 |
+
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
|
21 |
+
model.load_state_dict(load_net, strict=True)
|
22 |
+
model.eval()
|
23 |
+
model = model.to(device)
|
24 |
+
return model
|
extras/facexlib/parsing/bisenet.py
ADDED
@@ -0,0 +1,140 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import torch.nn.functional as F
|
4 |
+
|
5 |
+
from .resnet import ResNet18
|
6 |
+
|
7 |
+
|
8 |
+
class ConvBNReLU(nn.Module):
|
9 |
+
|
10 |
+
def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1):
|
11 |
+
super(ConvBNReLU, self).__init__()
|
12 |
+
self.conv = nn.Conv2d(in_chan, out_chan, kernel_size=ks, stride=stride, padding=padding, bias=False)
|
13 |
+
self.bn = nn.BatchNorm2d(out_chan)
|
14 |
+
|
15 |
+
def forward(self, x):
|
16 |
+
x = self.conv(x)
|
17 |
+
x = F.relu(self.bn(x))
|
18 |
+
return x
|
19 |
+
|
20 |
+
|
21 |
+
class BiSeNetOutput(nn.Module):
|
22 |
+
|
23 |
+
def __init__(self, in_chan, mid_chan, num_class):
|
24 |
+
super(BiSeNetOutput, self).__init__()
|
25 |
+
self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
|
26 |
+
self.conv_out = nn.Conv2d(mid_chan, num_class, kernel_size=1, bias=False)
|
27 |
+
|
28 |
+
def forward(self, x):
|
29 |
+
feat = self.conv(x)
|
30 |
+
out = self.conv_out(feat)
|
31 |
+
return out, feat
|
32 |
+
|
33 |
+
|
34 |
+
class AttentionRefinementModule(nn.Module):
|
35 |
+
|
36 |
+
def __init__(self, in_chan, out_chan):
|
37 |
+
super(AttentionRefinementModule, self).__init__()
|
38 |
+
self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
|
39 |
+
self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False)
|
40 |
+
self.bn_atten = nn.BatchNorm2d(out_chan)
|
41 |
+
self.sigmoid_atten = nn.Sigmoid()
|
42 |
+
|
43 |
+
def forward(self, x):
|
44 |
+
feat = self.conv(x)
|
45 |
+
atten = F.avg_pool2d(feat, feat.size()[2:])
|
46 |
+
atten = self.conv_atten(atten)
|
47 |
+
atten = self.bn_atten(atten)
|
48 |
+
atten = self.sigmoid_atten(atten)
|
49 |
+
out = torch.mul(feat, atten)
|
50 |
+
return out
|
51 |
+
|
52 |
+
|
53 |
+
class ContextPath(nn.Module):
|
54 |
+
|
55 |
+
def __init__(self):
|
56 |
+
super(ContextPath, self).__init__()
|
57 |
+
self.resnet = ResNet18()
|
58 |
+
self.arm16 = AttentionRefinementModule(256, 128)
|
59 |
+
self.arm32 = AttentionRefinementModule(512, 128)
|
60 |
+
self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
|
61 |
+
self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
|
62 |
+
self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
|
63 |
+
|
64 |
+
def forward(self, x):
|
65 |
+
feat8, feat16, feat32 = self.resnet(x)
|
66 |
+
h8, w8 = feat8.size()[2:]
|
67 |
+
h16, w16 = feat16.size()[2:]
|
68 |
+
h32, w32 = feat32.size()[2:]
|
69 |
+
|
70 |
+
avg = F.avg_pool2d(feat32, feat32.size()[2:])
|
71 |
+
avg = self.conv_avg(avg)
|
72 |
+
avg_up = F.interpolate(avg, (h32, w32), mode='nearest')
|
73 |
+
|
74 |
+
feat32_arm = self.arm32(feat32)
|
75 |
+
feat32_sum = feat32_arm + avg_up
|
76 |
+
feat32_up = F.interpolate(feat32_sum, (h16, w16), mode='nearest')
|
77 |
+
feat32_up = self.conv_head32(feat32_up)
|
78 |
+
|
79 |
+
feat16_arm = self.arm16(feat16)
|
80 |
+
feat16_sum = feat16_arm + feat32_up
|
81 |
+
feat16_up = F.interpolate(feat16_sum, (h8, w8), mode='nearest')
|
82 |
+
feat16_up = self.conv_head16(feat16_up)
|
83 |
+
|
84 |
+
return feat8, feat16_up, feat32_up # x8, x8, x16
|
85 |
+
|
86 |
+
|
87 |
+
class FeatureFusionModule(nn.Module):
|
88 |
+
|
89 |
+
def __init__(self, in_chan, out_chan):
|
90 |
+
super(FeatureFusionModule, self).__init__()
|
91 |
+
self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
|
92 |
+
self.conv1 = nn.Conv2d(out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False)
|
93 |
+
self.conv2 = nn.Conv2d(out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False)
|
94 |
+
self.relu = nn.ReLU(inplace=True)
|
95 |
+
self.sigmoid = nn.Sigmoid()
|
96 |
+
|
97 |
+
def forward(self, fsp, fcp):
|
98 |
+
fcat = torch.cat([fsp, fcp], dim=1)
|
99 |
+
feat = self.convblk(fcat)
|
100 |
+
atten = F.avg_pool2d(feat, feat.size()[2:])
|
101 |
+
atten = self.conv1(atten)
|
102 |
+
atten = self.relu(atten)
|
103 |
+
atten = self.conv2(atten)
|
104 |
+
atten = self.sigmoid(atten)
|
105 |
+
feat_atten = torch.mul(feat, atten)
|
106 |
+
feat_out = feat_atten + feat
|
107 |
+
return feat_out
|
108 |
+
|
109 |
+
|
110 |
+
class BiSeNet(nn.Module):
|
111 |
+
|
112 |
+
def __init__(self, num_class):
|
113 |
+
super(BiSeNet, self).__init__()
|
114 |
+
self.cp = ContextPath()
|
115 |
+
self.ffm = FeatureFusionModule(256, 256)
|
116 |
+
self.conv_out = BiSeNetOutput(256, 256, num_class)
|
117 |
+
self.conv_out16 = BiSeNetOutput(128, 64, num_class)
|
118 |
+
self.conv_out32 = BiSeNetOutput(128, 64, num_class)
|
119 |
+
|
120 |
+
def forward(self, x, return_feat=False):
|
121 |
+
h, w = x.size()[2:]
|
122 |
+
feat_res8, feat_cp8, feat_cp16 = self.cp(x) # return res3b1 feature
|
123 |
+
feat_sp = feat_res8 # replace spatial path feature with res3b1 feature
|
124 |
+
feat_fuse = self.ffm(feat_sp, feat_cp8)
|
125 |
+
|
126 |
+
out, feat = self.conv_out(feat_fuse)
|
127 |
+
out16, feat16 = self.conv_out16(feat_cp8)
|
128 |
+
out32, feat32 = self.conv_out32(feat_cp16)
|
129 |
+
|
130 |
+
out = F.interpolate(out, (h, w), mode='bilinear', align_corners=True)
|
131 |
+
out16 = F.interpolate(out16, (h, w), mode='bilinear', align_corners=True)
|
132 |
+
out32 = F.interpolate(out32, (h, w), mode='bilinear', align_corners=True)
|
133 |
+
|
134 |
+
if return_feat:
|
135 |
+
feat = F.interpolate(feat, (h, w), mode='bilinear', align_corners=True)
|
136 |
+
feat16 = F.interpolate(feat16, (h, w), mode='bilinear', align_corners=True)
|
137 |
+
feat32 = F.interpolate(feat32, (h, w), mode='bilinear', align_corners=True)
|
138 |
+
return out, out16, out32, feat, feat16, feat32
|
139 |
+
else:
|
140 |
+
return out, out16, out32
|
extras/facexlib/parsing/parsenet.py
ADDED
@@ -0,0 +1,194 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""Modified from https://github.com/chaofengc/PSFRGAN
|
2 |
+
"""
|
3 |
+
import numpy as np
|
4 |
+
import torch.nn as nn
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
|
8 |
+
class NormLayer(nn.Module):
|
9 |
+
"""Normalization Layers.
|
10 |
+
|
11 |
+
Args:
|
12 |
+
channels: input channels, for batch norm and instance norm.
|
13 |
+
input_size: input shape without batch size, for layer norm.
|
14 |
+
"""
|
15 |
+
|
16 |
+
def __init__(self, channels, normalize_shape=None, norm_type='bn'):
|
17 |
+
super(NormLayer, self).__init__()
|
18 |
+
norm_type = norm_type.lower()
|
19 |
+
self.norm_type = norm_type
|
20 |
+
if norm_type == 'bn':
|
21 |
+
self.norm = nn.BatchNorm2d(channels, affine=True)
|
22 |
+
elif norm_type == 'in':
|
23 |
+
self.norm = nn.InstanceNorm2d(channels, affine=False)
|
24 |
+
elif norm_type == 'gn':
|
25 |
+
self.norm = nn.GroupNorm(32, channels, affine=True)
|
26 |
+
elif norm_type == 'pixel':
|
27 |
+
self.norm = lambda x: F.normalize(x, p=2, dim=1)
|
28 |
+
elif norm_type == 'layer':
|
29 |
+
self.norm = nn.LayerNorm(normalize_shape)
|
30 |
+
elif norm_type == 'none':
|
31 |
+
self.norm = lambda x: x * 1.0
|
32 |
+
else:
|
33 |
+
assert 1 == 0, f'Norm type {norm_type} not support.'
|
34 |
+
|
35 |
+
def forward(self, x, ref=None):
|
36 |
+
if self.norm_type == 'spade':
|
37 |
+
return self.norm(x, ref)
|
38 |
+
else:
|
39 |
+
return self.norm(x)
|
40 |
+
|
41 |
+
|
42 |
+
class ReluLayer(nn.Module):
|
43 |
+
"""Relu Layer.
|
44 |
+
|
45 |
+
Args:
|
46 |
+
relu type: type of relu layer, candidates are
|
47 |
+
- ReLU
|
48 |
+
- LeakyReLU: default relu slope 0.2
|
49 |
+
- PRelu
|
50 |
+
- SELU
|
51 |
+
- none: direct pass
|
52 |
+
"""
|
53 |
+
|
54 |
+
def __init__(self, channels, relu_type='relu'):
|
55 |
+
super(ReluLayer, self).__init__()
|
56 |
+
relu_type = relu_type.lower()
|
57 |
+
if relu_type == 'relu':
|
58 |
+
self.func = nn.ReLU(True)
|
59 |
+
elif relu_type == 'leakyrelu':
|
60 |
+
self.func = nn.LeakyReLU(0.2, inplace=True)
|
61 |
+
elif relu_type == 'prelu':
|
62 |
+
self.func = nn.PReLU(channels)
|
63 |
+
elif relu_type == 'selu':
|
64 |
+
self.func = nn.SELU(True)
|
65 |
+
elif relu_type == 'none':
|
66 |
+
self.func = lambda x: x * 1.0
|
67 |
+
else:
|
68 |
+
assert 1 == 0, f'Relu type {relu_type} not support.'
|
69 |
+
|
70 |
+
def forward(self, x):
|
71 |
+
return self.func(x)
|
72 |
+
|
73 |
+
|
74 |
+
class ConvLayer(nn.Module):
|
75 |
+
|
76 |
+
def __init__(self,
|
77 |
+
in_channels,
|
78 |
+
out_channels,
|
79 |
+
kernel_size=3,
|
80 |
+
scale='none',
|
81 |
+
norm_type='none',
|
82 |
+
relu_type='none',
|
83 |
+
use_pad=True,
|
84 |
+
bias=True):
|
85 |
+
super(ConvLayer, self).__init__()
|
86 |
+
self.use_pad = use_pad
|
87 |
+
self.norm_type = norm_type
|
88 |
+
if norm_type in ['bn']:
|
89 |
+
bias = False
|
90 |
+
|
91 |
+
stride = 2 if scale == 'down' else 1
|
92 |
+
|
93 |
+
self.scale_func = lambda x: x
|
94 |
+
if scale == 'up':
|
95 |
+
self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
|
96 |
+
|
97 |
+
self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
|
98 |
+
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
|
99 |
+
|
100 |
+
self.relu = ReluLayer(out_channels, relu_type)
|
101 |
+
self.norm = NormLayer(out_channels, norm_type=norm_type)
|
102 |
+
|
103 |
+
def forward(self, x):
|
104 |
+
out = self.scale_func(x)
|
105 |
+
if self.use_pad:
|
106 |
+
out = self.reflection_pad(out)
|
107 |
+
out = self.conv2d(out)
|
108 |
+
out = self.norm(out)
|
109 |
+
out = self.relu(out)
|
110 |
+
return out
|
111 |
+
|
112 |
+
|
113 |
+
class ResidualBlock(nn.Module):
|
114 |
+
"""
|
115 |
+
Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
|
116 |
+
"""
|
117 |
+
|
118 |
+
def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
|
119 |
+
super(ResidualBlock, self).__init__()
|
120 |
+
|
121 |
+
if scale == 'none' and c_in == c_out:
|
122 |
+
self.shortcut_func = lambda x: x
|
123 |
+
else:
|
124 |
+
self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
|
125 |
+
|
126 |
+
scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
|
127 |
+
scale_conf = scale_config_dict[scale]
|
128 |
+
|
129 |
+
self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
|
130 |
+
self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
|
131 |
+
|
132 |
+
def forward(self, x):
|
133 |
+
identity = self.shortcut_func(x)
|
134 |
+
|
135 |
+
res = self.conv1(x)
|
136 |
+
res = self.conv2(res)
|
137 |
+
return identity + res
|
138 |
+
|
139 |
+
|
140 |
+
class ParseNet(nn.Module):
|
141 |
+
|
142 |
+
def __init__(self,
|
143 |
+
in_size=128,
|
144 |
+
out_size=128,
|
145 |
+
min_feat_size=32,
|
146 |
+
base_ch=64,
|
147 |
+
parsing_ch=19,
|
148 |
+
res_depth=10,
|
149 |
+
relu_type='LeakyReLU',
|
150 |
+
norm_type='bn',
|
151 |
+
ch_range=[32, 256]):
|
152 |
+
super().__init__()
|
153 |
+
self.res_depth = res_depth
|
154 |
+
act_args = {'norm_type': norm_type, 'relu_type': relu_type}
|
155 |
+
min_ch, max_ch = ch_range
|
156 |
+
|
157 |
+
ch_clip = lambda x: max(min_ch, min(x, max_ch)) # noqa: E731
|
158 |
+
min_feat_size = min(in_size, min_feat_size)
|
159 |
+
|
160 |
+
down_steps = int(np.log2(in_size // min_feat_size))
|
161 |
+
up_steps = int(np.log2(out_size // min_feat_size))
|
162 |
+
|
163 |
+
# =============== define encoder-body-decoder ====================
|
164 |
+
self.encoder = []
|
165 |
+
self.encoder.append(ConvLayer(3, base_ch, 3, 1))
|
166 |
+
head_ch = base_ch
|
167 |
+
for i in range(down_steps):
|
168 |
+
cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
|
169 |
+
self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
|
170 |
+
head_ch = head_ch * 2
|
171 |
+
|
172 |
+
self.body = []
|
173 |
+
for i in range(res_depth):
|
174 |
+
self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
|
175 |
+
|
176 |
+
self.decoder = []
|
177 |
+
for i in range(up_steps):
|
178 |
+
cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
|
179 |
+
self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
|
180 |
+
head_ch = head_ch // 2
|
181 |
+
|
182 |
+
self.encoder = nn.Sequential(*self.encoder)
|
183 |
+
self.body = nn.Sequential(*self.body)
|
184 |
+
self.decoder = nn.Sequential(*self.decoder)
|
185 |
+
self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
|
186 |
+
self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
|
187 |
+
|
188 |
+
def forward(self, x):
|
189 |
+
feat = self.encoder(x)
|
190 |
+
x = feat + self.body(feat)
|
191 |
+
x = self.decoder(x)
|
192 |
+
out_img = self.out_img_conv(x)
|
193 |
+
out_mask = self.out_mask_conv(x)
|
194 |
+
return out_mask, out_img
|
extras/facexlib/parsing/resnet.py
ADDED
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch.nn as nn
|
2 |
+
import torch.nn.functional as F
|
3 |
+
|
4 |
+
|
5 |
+
def conv3x3(in_planes, out_planes, stride=1):
|
6 |
+
"""3x3 convolution with padding"""
|
7 |
+
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
|
8 |
+
|
9 |
+
|
10 |
+
class BasicBlock(nn.Module):
|
11 |
+
|
12 |
+
def __init__(self, in_chan, out_chan, stride=1):
|
13 |
+
super(BasicBlock, self).__init__()
|
14 |
+
self.conv1 = conv3x3(in_chan, out_chan, stride)
|
15 |
+
self.bn1 = nn.BatchNorm2d(out_chan)
|
16 |
+
self.conv2 = conv3x3(out_chan, out_chan)
|
17 |
+
self.bn2 = nn.BatchNorm2d(out_chan)
|
18 |
+
self.relu = nn.ReLU(inplace=True)
|
19 |
+
self.downsample = None
|
20 |
+
if in_chan != out_chan or stride != 1:
|
21 |
+
self.downsample = nn.Sequential(
|
22 |
+
nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False),
|
23 |
+
nn.BatchNorm2d(out_chan),
|
24 |
+
)
|
25 |
+
|
26 |
+
def forward(self, x):
|
27 |
+
residual = self.conv1(x)
|
28 |
+
residual = F.relu(self.bn1(residual))
|
29 |
+
residual = self.conv2(residual)
|
30 |
+
residual = self.bn2(residual)
|
31 |
+
|
32 |
+
shortcut = x
|
33 |
+
if self.downsample is not None:
|
34 |
+
shortcut = self.downsample(x)
|
35 |
+
|
36 |
+
out = shortcut + residual
|
37 |
+
out = self.relu(out)
|
38 |
+
return out
|
39 |
+
|
40 |
+
|
41 |
+
def create_layer_basic(in_chan, out_chan, bnum, stride=1):
|
42 |
+
layers = [BasicBlock(in_chan, out_chan, stride=stride)]
|
43 |
+
for i in range(bnum - 1):
|
44 |
+
layers.append(BasicBlock(out_chan, out_chan, stride=1))
|
45 |
+
return nn.Sequential(*layers)
|
46 |
+
|
47 |
+
|
48 |
+
class ResNet18(nn.Module):
|
49 |
+
|
50 |
+
def __init__(self):
|
51 |
+
super(ResNet18, self).__init__()
|
52 |
+
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
53 |
+
self.bn1 = nn.BatchNorm2d(64)
|
54 |
+
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
55 |
+
self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
|
56 |
+
self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
|
57 |
+
self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
|
58 |
+
self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
|
59 |
+
|
60 |
+
def forward(self, x):
|
61 |
+
x = self.conv1(x)
|
62 |
+
x = F.relu(self.bn1(x))
|
63 |
+
x = self.maxpool(x)
|
64 |
+
|
65 |
+
x = self.layer1(x)
|
66 |
+
feat8 = self.layer2(x) # 1/8
|
67 |
+
feat16 = self.layer3(feat8) # 1/16
|
68 |
+
feat32 = self.layer4(feat16) # 1/32
|
69 |
+
return feat8, feat16, feat32
|
extras/facexlib/utils/__init__.py
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .face_utils import align_crop_face_landmarks, compute_increased_bbox, get_valid_bboxes, paste_face_back
|
2 |
+
from .misc import img2tensor, load_file_from_url, scandir
|
3 |
+
|
4 |
+
__all__ = [
|
5 |
+
'align_crop_face_landmarks', 'compute_increased_bbox', 'get_valid_bboxes', 'load_file_from_url', 'paste_face_back',
|
6 |
+
'img2tensor', 'scandir'
|
7 |
+
]
|