init
This commit is contained in:
colomi
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
Binary file not shown.
@@ -0,0 +1,252 @@
|
||||
# ComfyUI's ControlNet Auxiliary Preprocessors
|
||||
Plug-and-play [ComfyUI](https://github.com/comfyanonymous/ComfyUI) node sets for making [ControlNet](https://github.com/lllyasviel/ControlNet/) hint images
|
||||
|
||||
"anime style, a protest in the street, cyberpunk city, a woman with pink hair and golden eyes (looking at the viewer) is holding a sign with the text "ComfyUI ControlNet Aux" in bold, neon pink" on Flux.1 Dev
|
||||
|
||||
|
||||
|
||||
The code is copy-pasted from the respective folders in https://github.com/lllyasviel/ControlNet/tree/main/annotator and connected to [the 🤗 Hub](https://huggingface.co/lllyasviel/Annotators).
|
||||
|
||||
All credit & copyright goes to https://github.com/lllyasviel.
|
||||
|
||||
# Updates
|
||||
Go to [Update page](./UPDATES.md) to follow updates
|
||||
|
||||
# Installation:
|
||||
## Using ComfyUI Manager (recommended):
|
||||
Install [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager) and do steps introduced there to install this repo.
|
||||
|
||||
## Alternative:
|
||||
If you're running on Linux, or non-admin account on windows you'll want to ensure `/ComfyUI/custom_nodes` and `comfyui_controlnet_aux` has write permissions.
|
||||
|
||||
There is now a **install.bat** you can run to install to portable if detected. Otherwise it will default to system and assume you followed ConfyUI's manual installation steps.
|
||||
|
||||
If you can't run **install.bat** (e.g. you are a Linux user). Open the CMD/Shell and do the following:
|
||||
- Navigate to your `/ComfyUI/custom_nodes/` folder
|
||||
- Run `git clone https://github.com/Fannovel16/comfyui_controlnet_aux/`
|
||||
- Navigate to your `comfyui_controlnet_aux` folder
|
||||
- Portable/venv:
|
||||
- Run `path/to/ComfUI/python_embeded/python.exe -s -m pip install -r requirements.txt`
|
||||
- With system python
|
||||
- Run `pip install -r requirements.txt`
|
||||
- Start ComfyUI
|
||||
|
||||
# Nodes
|
||||
Please note that this repo only supports preprocessors making hint images (e.g. stickman, canny edge, etc).
|
||||
All preprocessors except Inpaint are intergrated into `AIO Aux Preprocessor` node.
|
||||
This node allow you to quickly get the preprocessor but a preprocessor's own threshold parameters won't be able to set.
|
||||
You need to use its node directly to set thresholds.
|
||||
|
||||
# Nodes (sections are categories in Comfy menu)
|
||||
## Line Extractors
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| Binary Lines | binary | control_scribble |
|
||||
| Canny Edge | canny | control_v11p_sd15_canny <br> control_canny <br> t2iadapter_canny |
|
||||
| HED Soft-Edge Lines | hed | control_v11p_sd15_softedge <br> control_hed |
|
||||
| Standard Lineart | standard_lineart | control_v11p_sd15_lineart |
|
||||
| Realistic Lineart | lineart (or `lineart_coarse` if `coarse` is enabled) | control_v11p_sd15_lineart |
|
||||
| Anime Lineart | lineart_anime | control_v11p_sd15s2_lineart_anime |
|
||||
| Manga Lineart | lineart_anime_denoise | control_v11p_sd15s2_lineart_anime |
|
||||
| M-LSD Lines | mlsd | control_v11p_sd15_mlsd <br> control_mlsd |
|
||||
| PiDiNet Soft-Edge Lines | pidinet | control_v11p_sd15_softedge <br> control_scribble |
|
||||
| Scribble Lines | scribble | control_v11p_sd15_scribble <br> control_scribble |
|
||||
| Scribble XDoG Lines | scribble_xdog | control_v11p_sd15_scribble <br> control_scribble |
|
||||
| Fake Scribble Lines | scribble_hed | control_v11p_sd15_scribble <br> control_scribble |
|
||||
| TEED Soft-Edge Lines | teed | [controlnet-sd-xl-1.0-softedge-dexined](https://huggingface.co/SargeZT/controlnet-sd-xl-1.0-softedge-dexined/blob/main/controlnet-sd-xl-1.0-softedge-dexined.safetensors) <br> control_v11p_sd15_softedge (Theoretically)
|
||||
| Scribble PiDiNet Lines | scribble_pidinet | control_v11p_sd15_scribble <br> control_scribble |
|
||||
| AnyLine Lineart | | mistoLine_fp16.safetensors <br> mistoLine_rank256 <br> control_v11p_sd15s2_lineart_anime <br> control_v11p_sd15_lineart |
|
||||
|
||||
## Normal and Depth Estimators
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| MiDaS Depth Map | (normal) depth | control_v11f1p_sd15_depth <br> control_depth <br> t2iadapter_depth |
|
||||
| LeReS Depth Map | depth_leres | control_v11f1p_sd15_depth <br> control_depth <br> t2iadapter_depth |
|
||||
| Zoe Depth Map | depth_zoe | control_v11f1p_sd15_depth <br> control_depth <br> t2iadapter_depth |
|
||||
| MiDaS Normal Map | normal_map | control_normal |
|
||||
| BAE Normal Map | normal_bae | control_v11p_sd15_normalbae |
|
||||
| MeshGraphormer Hand Refiner ([HandRefinder](https://github.com/wenquanlu/HandRefiner)) | depth_hand_refiner | [control_sd15_inpaint_depth_hand_fp16](https://huggingface.co/hr16/ControlNet-HandRefiner-pruned/blob/main/control_sd15_inpaint_depth_hand_fp16.safetensors) |
|
||||
| Depth Anything | depth_anything | [Depth-Anything](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints_controlnet/diffusion_pytorch_model.safetensors) |
|
||||
| Zoe Depth Anything <br> (Basically Zoe but the encoder is replaced with DepthAnything) | depth_anything | [Depth-Anything](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints_controlnet/diffusion_pytorch_model.safetensors) |
|
||||
| Normal DSINE | | control_normal/control_v11p_sd15_normalbae |
|
||||
| Metric3D Depth | | control_v11f1p_sd15_depth <br> control_depth <br> t2iadapter_depth |
|
||||
| Metric3D Normal | | control_v11p_sd15_normalbae |
|
||||
| Depth Anything V2 | | [Depth-Anything](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints_controlnet/diffusion_pytorch_model.safetensors) |
|
||||
|
||||
## Faces and Poses Estimators
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| DWPose Estimator | dw_openpose_full | control_v11p_sd15_openpose <br> control_openpose <br> t2iadapter_openpose |
|
||||
| OpenPose Estimator | openpose (detect_body) <br> openpose_hand (detect_body + detect_hand) <br> openpose_faceonly (detect_face) <br> openpose_full (detect_hand + detect_body + detect_face) | control_v11p_sd15_openpose <br> control_openpose <br> t2iadapter_openpose |
|
||||
| MediaPipe Face Mesh | mediapipe_face | controlnet_sd21_laion_face_v2 |
|
||||
| Animal Estimator | animal_openpose | [control_sd15_animal_openpose_fp16](https://huggingface.co/huchenlei/animal_openpose/blob/main/control_sd15_animal_openpose_fp16.pth) |
|
||||
|
||||
## Optical Flow Estimators
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| Unimatch Optical Flow | | [DragNUWA](https://github.com/ProjectNUWA/DragNUWA) |
|
||||
|
||||
### How to get OpenPose-format JSON?
|
||||
#### User-side
|
||||
This workflow will save images to ComfyUI's output folder (the same location as output images). If you haven't found `Save Pose Keypoints` node, update this extension
|
||||
|
||||
|
||||
#### Dev-side
|
||||
An array of [OpenPose-format JSON](https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md#json-output-format) corresponsding to each frame in an IMAGE batch can be gotten from DWPose and OpenPose using `app.nodeOutputs` on the UI or `/history` API endpoint. JSON output from AnimalPose uses a kinda similar format to OpenPose JSON:
|
||||
```
|
||||
[
|
||||
{
|
||||
"version": "ap10k",
|
||||
"animals": [
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
...
|
||||
],
|
||||
"canvas_height": 512,
|
||||
"canvas_width": 768
|
||||
},
|
||||
...
|
||||
]
|
||||
```
|
||||
|
||||
For extension developers (e.g. Openpose editor):
|
||||
```js
|
||||
const poseNodes = app.graph._nodes.filter(node => ["OpenposePreprocessor", "DWPreprocessor", "AnimalPosePreprocessor"].includes(node.type))
|
||||
for (const poseNode of poseNodes) {
|
||||
const openposeResults = JSON.parse(app.nodeOutputs[poseNode.id].openpose_json[0])
|
||||
console.log(openposeResults) //An array containing Openpose JSON for each frame
|
||||
}
|
||||
```
|
||||
|
||||
For API users:
|
||||
Javascript
|
||||
```js
|
||||
import fetch from "node-fetch" //Remember to add "type": "module" to "package.json"
|
||||
async function main() {
|
||||
const promptId = '792c1905-ecfe-41f4-8114-83e6a4a09a9f' //Too lazy to POST /queue
|
||||
let history = await fetch(`http://127.0.0.1:8188/history/${promptId}`).then(re => re.json())
|
||||
history = history[promptId]
|
||||
const nodeOutputs = Object.values(history.outputs).filter(output => output.openpose_json)
|
||||
for (const nodeOutput of nodeOutputs) {
|
||||
const openposeResults = JSON.parse(nodeOutput.openpose_json[0])
|
||||
console.log(openposeResults) //An array containing Openpose JSON for each frame
|
||||
}
|
||||
}
|
||||
main()
|
||||
```
|
||||
|
||||
Python
|
||||
```py
|
||||
import json, urllib.request
|
||||
|
||||
server_address = "127.0.0.1:8188"
|
||||
prompt_id = '' #Too lazy to POST /queue
|
||||
|
||||
def get_history(prompt_id):
|
||||
with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
|
||||
return json.loads(response.read())
|
||||
|
||||
history = get_history(prompt_id)[prompt_id]
|
||||
for o in history['outputs']:
|
||||
for node_id in history['outputs']:
|
||||
node_output = history['outputs'][node_id]
|
||||
if 'openpose_json' in node_output:
|
||||
print(json.loads(node_output['openpose_json'][0])) #An list containing Openpose JSON for each frame
|
||||
```
|
||||
## Semantic Segmentation
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| OneFormer ADE20K Segmentor | oneformer_ade20k | control_v11p_sd15_seg |
|
||||
| OneFormer COCO Segmentor | oneformer_coco | control_v11p_sd15_seg |
|
||||
| UniFormer Segmentor | segmentation |control_sd15_seg <br> control_v11p_sd15_seg|
|
||||
|
||||
## T2IAdapter-only
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| Color Pallete | color | t2iadapter_color |
|
||||
| Content Shuffle | shuffle | t2iadapter_style |
|
||||
|
||||
## Recolor
|
||||
| Preprocessor Node | sd-webui-controlnet/other | ControlNet/T2I-Adapter |
|
||||
|-----------------------------|---------------------------|-------------------------------------------|
|
||||
| Image Luminance | recolor_luminance | [ioclab_sd15_recolor](https://huggingface.co/lllyasviel/sd_control_collection/resolve/main/ioclab_sd15_recolor.safetensors) <br> [sai_xl_recolor_256lora](https://huggingface.co/lllyasviel/sd_control_collection/resolve/main/sai_xl_recolor_256lora.safetensors) <br> [bdsqlsz_controlllite_xl_recolor_luminance](https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite/resolve/main/bdsqlsz_controlllite_xl_recolor_luminance.safetensors) |
|
||||
| Image Intensity | recolor_intensity | Idk. Maybe same as above? |
|
||||
|
||||
# Examples
|
||||
> A picture is worth a thousand words
|
||||
|
||||
|
||||
|
||||
|
||||
# Testing workflow
|
||||
https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/examples/ExecuteAll.png
|
||||
Input image: https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/examples/comfyui-controlnet-aux-logo.png
|
||||
|
||||
# Q&A:
|
||||
## Why some nodes doesn't appear after I installed this repo?
|
||||
|
||||
This repo has a new mechanism which will skip any custom node can't be imported. If you meet this case, please create a issue on [Issues tab](https://github.com/Fannovel16/comfyui_controlnet_aux/issues) with the log from the command line.
|
||||
|
||||
## DWPose/AnimalPose only uses CPU so it's so slow. How can I make it use GPU?
|
||||
There are two ways to speed-up DWPose: using TorchScript checkpoints (.torchscript.pt) checkpoints or ONNXRuntime (.onnx). TorchScript way is little bit slower than ONNXRuntime but doesn't require any additional library and still way way faster than CPU.
|
||||
|
||||
A torchscript bbox detector is compatiable with an onnx pose estimator and vice versa.
|
||||
### TorchScript
|
||||
Set `bbox_detector` and `pose_estimator` according to this picture. You can try other bbox detector endings with `.torchscript.pt` to reduce bbox detection time if input images are ideal.
|
||||
|
||||
### ONNXRuntime
|
||||
If onnxruntime is installed successfully and the checkpoint used endings with `.onnx`, it will replace default cv2 backend to take advantage of GPU. Note that if you are using NVidia card, this method currently can only works on CUDA 11.8 (ComfyUI_windows_portable_nvidia_cu118_or_cpu.7z) unless you compile onnxruntime yourself.
|
||||
|
||||
1. Know your onnxruntime build:
|
||||
* * NVidia CUDA 11.x or bellow/AMD GPU: `onnxruntime-gpu`
|
||||
* * NVidia CUDA 12.x: `onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`
|
||||
* * DirectML: `onnxruntime-directml`
|
||||
* * OpenVINO: `onnxruntime-openvino`
|
||||
|
||||
Note that if this is your first time using ComfyUI, please test if it can run on your device before doing next steps.
|
||||
|
||||
2. Add it into `requirements.txt`
|
||||
|
||||
3. Run `install.bat` or pip command mentioned in Installation
|
||||
|
||||
|
||||
|
||||
# Assets files of preprocessors
|
||||
* anime_face_segment: [bdsqlsz/qinglong_controlnet-lllite/Annotators/UNet.pth](https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite/blob/main/Annotators/UNet.pth), [anime-seg/isnetis.ckpt](https://huggingface.co/skytnt/anime-seg/blob/main/isnetis.ckpt)
|
||||
* densepose: [LayerNorm/DensePose-TorchScript-with-hint-image/densepose_r50_fpn_dl.torchscript](https://huggingface.co/LayerNorm/DensePose-TorchScript-with-hint-image/blob/main/densepose_r50_fpn_dl.torchscript)
|
||||
* dwpose:
|
||||
* * bbox_detector: Either [yzd-v/DWPose/yolox_l.onnx](https://huggingface.co/yzd-v/DWPose/blob/main/yolox_l.onnx), [hr16/yolox-onnx/yolox_l.torchscript.pt](https://huggingface.co/hr16/yolox-onnx/blob/main/yolox_l.torchscript.pt), [hr16/yolo-nas-fp16/yolo_nas_l_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_l_fp16.onnx), [hr16/yolo-nas-fp16/yolo_nas_m_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_m_fp16.onnx), [hr16/yolo-nas-fp16/yolo_nas_s_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_s_fp16.onnx)
|
||||
* * pose_estimator: Either [hr16/DWPose-TorchScript-BatchSize5/dw-ll_ucoco_384_bs5.torchscript.pt](https://huggingface.co/hr16/DWPose-TorchScript-BatchSize5/blob/main/dw-ll_ucoco_384_bs5.torchscript.pt), [yzd-v/DWPose/dw-ll_ucoco_384.onnx](https://huggingface.co/yzd-v/DWPose/blob/main/dw-ll_ucoco_384.onnx)
|
||||
* animal_pose (ap10k):
|
||||
* * bbox_detector: Either [yzd-v/DWPose/yolox_l.onnx](https://huggingface.co/yzd-v/DWPose/blob/main/yolox_l.onnx), [hr16/yolox-onnx/yolox_l.torchscript.pt](https://huggingface.co/hr16/yolox-onnx/blob/main/yolox_l.torchscript.pt), [hr16/yolo-nas-fp16/yolo_nas_l_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_l_fp16.onnx), [hr16/yolo-nas-fp16/yolo_nas_m_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_m_fp16.onnx), [hr16/yolo-nas-fp16/yolo_nas_s_fp16.onnx](https://huggingface.co/hr16/yolo-nas-fp16/blob/main/yolo_nas_s_fp16.onnx)
|
||||
* * pose_estimator: Either [hr16/DWPose-TorchScript-BatchSize5/rtmpose-m_ap10k_256_bs5.torchscript.pt](https://huggingface.co/hr16/DWPose-TorchScript-BatchSize5/blob/main/rtmpose-m_ap10k_256_bs5.torchscript.pt), [hr16/UnJIT-DWPose/rtmpose-m_ap10k_256.onnx](https://huggingface.co/hr16/UnJIT-DWPose/blob/main/rtmpose-m_ap10k_256.onnx)
|
||||
* hed: [lllyasviel/Annotators/ControlNetHED.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/ControlNetHED.pth)
|
||||
* leres: [lllyasviel/Annotators/res101.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/res101.pth), [lllyasviel/Annotators/latest_net_G.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/latest_net_G.pth)
|
||||
* lineart: [lllyasviel/Annotators/sk_model.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/sk_model.pth), [lllyasviel/Annotators/sk_model2.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/sk_model2.pth)
|
||||
* lineart_anime: [lllyasviel/Annotators/netG.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/netG.pth)
|
||||
* manga_line: [lllyasviel/Annotators/erika.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/erika.pth)
|
||||
* mesh_graphormer: [hr16/ControlNet-HandRefiner-pruned/graphormer_hand_state_dict.bin](https://huggingface.co/hr16/ControlNet-HandRefiner-pruned/blob/main/graphormer_hand_state_dict.bin), [hr16/ControlNet-HandRefiner-pruned/hrnetv2_w64_imagenet_pretrained.pth](https://huggingface.co/hr16/ControlNet-HandRefiner-pruned/blob/main/hrnetv2_w64_imagenet_pretrained.pth)
|
||||
* midas: [lllyasviel/Annotators/dpt_hybrid-midas-501f0c75.pt](https://huggingface.co/lllyasviel/Annotators/blob/main/dpt_hybrid-midas-501f0c75.pt)
|
||||
* mlsd: [lllyasviel/Annotators/mlsd_large_512_fp32.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/mlsd_large_512_fp32.pth)
|
||||
* normalbae: [lllyasviel/Annotators/scannet.pt](https://huggingface.co/lllyasviel/Annotators/blob/main/scannet.pt)
|
||||
* oneformer: [lllyasviel/Annotators/250_16_swin_l_oneformer_ade20k_160k.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/250_16_swin_l_oneformer_ade20k_160k.pth)
|
||||
* open_pose: [lllyasviel/Annotators/body_pose_model.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/body_pose_model.pth), [lllyasviel/Annotators/hand_pose_model.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/hand_pose_model.pth), [lllyasviel/Annotators/facenet.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/facenet.pth)
|
||||
* pidi: [lllyasviel/Annotators/table5_pidinet.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/table5_pidinet.pth)
|
||||
* sam: [dhkim2810/MobileSAM/mobile_sam.pt](https://huggingface.co/dhkim2810/MobileSAM/blob/main/mobile_sam.pt)
|
||||
* uniformer: [lllyasviel/Annotators/upernet_global_small.pth](https://huggingface.co/lllyasviel/Annotators/blob/main/upernet_global_small.pth)
|
||||
* zoe: [lllyasviel/Annotators/ZoeD_M12_N.pt](https://huggingface.co/lllyasviel/Annotators/blob/main/ZoeD_M12_N.pt)
|
||||
* teed: [bdsqlsz/qinglong_controlnet-lllite/7_model.pth](https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite/blob/main/Annotators/7_model.pth)
|
||||
* depth_anything: Either [LiheYoung/Depth-Anything/checkpoints/depth_anything_vitl14.pth](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints/depth_anything_vitl14.pth), [LiheYoung/Depth-Anything/checkpoints/depth_anything_vitb14.pth](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints/depth_anything_vitb14.pth) or [LiheYoung/Depth-Anything/checkpoints/depth_anything_vits14.pth](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints/depth_anything_vits14.pth)
|
||||
* diffusion_edge: Either [hr16/Diffusion-Edge/diffusion_edge_indoor.pt](https://huggingface.co/hr16/Diffusion-Edge/blob/main/diffusion_edge_indoor.pt), [hr16/Diffusion-Edge/diffusion_edge_urban.pt](https://huggingface.co/hr16/Diffusion-Edge/blob/main/diffusion_edge_urban.pt) or [hr16/Diffusion-Edge/diffusion_edge_natrual.pt](https://huggingface.co/hr16/Diffusion-Edge/blob/main/diffusion_edge_natrual.pt)
|
||||
* unimatch: Either [hr16/Unimatch/gmflow-scale2-regrefine6-mixdata.pth](https://huggingface.co/hr16/Unimatch/blob/main/gmflow-scale2-regrefine6-mixdata.pth), [hr16/Unimatch/gmflow-scale2-mixdata.pth](https://huggingface.co/hr16/Unimatch/blob/main/gmflow-scale2-mixdata.pth) or [hr16/Unimatch/gmflow-scale1-mixdata.pth](https://huggingface.co/hr16/Unimatch/blob/main/gmflow-scale1-mixdata.pth)
|
||||
* zoe_depth_anything: Either [LiheYoung/Depth-Anything/checkpoints_metric_depth/depth_anything_metric_depth_indoor.pt](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints_metric_depth/depth_anything_metric_depth_indoor.pt) or [LiheYoung/Depth-Anything/checkpoints_metric_depth/depth_anything_metric_depth_outdoor.pt](https://huggingface.co/spaces/LiheYoung/Depth-Anything/blob/main/checkpoints_metric_depth/depth_anything_metric_depth_outdoor.pt)
|
||||
# 2000 Stars 😄
|
||||
<a href="https://star-history.com/#Fannovel16/comfyui_controlnet_aux&Date">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=Fannovel16/comfyui_controlnet_aux&type=Date&theme=dark" />
|
||||
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=Fannovel16/comfyui_controlnet_aux&type=Date" />
|
||||
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=Fannovel16/comfyui_controlnet_aux&type=Date" />
|
||||
</picture>
|
||||
</a>
|
||||
|
||||
Thanks for yalls supports. I never thought the graph for stars would be linear lol.
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
* `AIO Aux Preprocessor` intergrating all loadable aux preprocessors as dropdown options. Easy to copy, paste and get the preprocessor faster.
|
||||
* Added OpenPose-format JSON output from OpenPose Preprocessor and DWPose Preprocessor. Checks [here](#faces-and-poses).
|
||||
* Fixed wrong model path when downloading DWPose.
|
||||
* Make hint images less blurry.
|
||||
* Added `resolution` option, `PixelPerfectResolution` and `HintImageEnchance` nodes (TODO: Documentation).
|
||||
* Added `RAFT Optical Flow Embedder` for TemporalNet2 (TODO: Workflow example).
|
||||
* Fixed opencv's conflicts between this extension, [ReActor](https://github.com/Gourieff/comfyui-reactor-node) and Roop. Thanks `Gourieff` for [the solution](https://github.com/Fannovel16/comfyui_controlnet_aux/issues/7#issuecomment-1734319075)!
|
||||
* RAFT is removed as the code behind it doesn't match what what the original code does
|
||||
* Changed `lineart`'s display name from `Normal Lineart` to `Realistic Lineart`. This change won't affect old workflows
|
||||
* Added support for `onnxruntime` to speed-up DWPose (see the Q&A)
|
||||
* Fixed TypeError: expected size to be one of int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], but got size with types [<class 'numpy.int64'>, <class 'numpy.int64'>]: [Issue](https://github.com/Fannovel16/comfyui_controlnet_aux/issues/2), [PR](https://github.com/Fannovel16/comfyui_controlnet_aux/pull/71))
|
||||
* Fixed ImageGenResolutionFromImage mishape (https://github.com/Fannovel16/comfyui_controlnet_aux/pull/74)
|
||||
* Fixed LeRes and MiDaS's incomatipility with MPS device
|
||||
* Fixed checking DWPose onnxruntime session multiple times: https://github.com/Fannovel16/comfyui_controlnet_aux/issues/89)
|
||||
* Added `Anime Face Segmentor` (in `ControlNet Preprocessors/Semantic Segmentation`) for [ControlNet AnimeFaceSegmentV2](https://huggingface.co/bdsqlsz/qinglong_controlnet-lllite#animefacesegmentv2). Checks [here](#anime-face-segmentor)
|
||||
* Change download functions and fix [download error](https://github.com/Fannovel16/comfyui_controlnet_aux/issues/39): [PR](https://github.com/Fannovel16/comfyui_controlnet_aux/pull/96)
|
||||
* Caching DWPose Onnxruntime during the first use of DWPose node instead of ComfyUI startup
|
||||
* Added alternative YOLOX models for faster speed when using DWPose
|
||||
* Added alternative DWPose models
|
||||
* Implemented the preprocessor for [AnimalPose ControlNet](https://github.com/abehonest/ControlNet_AnimalPose/tree/main). Check [Animal Pose AP-10K](#animal-pose-ap-10k)
|
||||
* Added YOLO-NAS models which are drop-in replacements of YOLOX
|
||||
* Fixed Openpose Face/Hands no longer detecting: https://github.com/Fannovel16/comfyui_controlnet_aux/issues/54
|
||||
* Added TorchScript implementation of DWPose and AnimalPose
|
||||
* Added TorchScript implementation of DensePose from [Colab notebook](https://colab.research.google.com/drive/16hcaaKs210ivpxjoyGNuvEXZD4eqOOSQ) which doesn't require detectron2. [Example](#densepose). Thanks [@LayerNome](https://github.com/Layer-norm) for fixing bugs related.
|
||||
* Added Standard Lineart Preprocessor
|
||||
* Fixed OpenPose misplacements in some cases
|
||||
* Added Mesh Graphormer - Hand Depth Map & Mask
|
||||
* Misaligned hands bug from MeshGraphormer was fixed
|
||||
* Added more mask options for MeshGraphormer
|
||||
* Added Save Pose Keypoint node for editing
|
||||
* Added Unimatch Optical Flow
|
||||
* Added Depth Anything & Zoe Depth Anything
|
||||
* Removed resolution field from Unimatch Optical Flow as that interpolating optical flow seems unstable
|
||||
* Added TEED Soft-Edge Preprocessor
|
||||
* Added DiffusionEdge
|
||||
* Added Image Luminance and Image Intensity
|
||||
* Added Normal DSINE
|
||||
* Added TTPlanet Tile (09/05/2024, DD/MM/YYYY)
|
||||
* Added AnyLine, Metric3D (18/05/2024)
|
||||
* Added Depth Anything V2 (16/06/2024)
|
||||
* Added Union model of ControlNet and preprocessors
|
||||
|
||||
* Refactor INPUT_TYPES and add Execute All node during the process of learning [Execution Model Inversion](https://github.com/comfyanonymous/ComfyUI/pull/2666)
|
||||
* Added scale_stick_for_xinsr_cn (https://github.com/Fannovel16/comfyui_controlnet_aux/issues/447) (09/04/2024)
|
||||
* PyTorch 2.7 compatibility fixes - eliminated custom_timm, custom_detectron2, and custom_midas_repo dependencies causing hanging issues. Refactored 7 major preprocessors including OneFormer (now using HuggingFace transformers), ZOE, DSINE, MiDaS, BAE, Metric3D, and Uniformer. Resolved ~59 GitHub issues related to import failures, hanging, and extension conflicts. Full modernization to actively maintained packages.
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
import sys, os
|
||||
|
||||
# Disable NPU device initialization and problematic MMCV ops to prevent RuntimeError
|
||||
# Must be set BEFORE any MMCV imports happen anywhere in ComfyUI
|
||||
os.environ['NPU_DEVICE_COUNT'] = '0'
|
||||
os.environ['MMCV_WITH_OPS'] = '0'
|
||||
from .utils import here, define_preprocessor_inputs, INPUT
|
||||
from pathlib import Path
|
||||
import traceback
|
||||
import importlib
|
||||
from .log import log, blue_text, cyan_text, get_summary, get_label
|
||||
from .hint_image_enchance import NODE_CLASS_MAPPINGS as HIE_NODE_CLASS_MAPPINGS
|
||||
from .hint_image_enchance import NODE_DISPLAY_NAME_MAPPINGS as HIE_NODE_DISPLAY_NAME_MAPPINGS
|
||||
#Ref: https://github.com/comfyanonymous/ComfyUI/blob/76d53c4622fc06372975ed2a43ad345935b8a551/nodes.py#L17
|
||||
sys.path.insert(0, str(Path(here, "src").resolve()))
|
||||
for pkg_name in ["custom_controlnet_aux", "custom_mmpkg"]:
|
||||
sys.path.append(str(Path(here, "src", pkg_name).resolve()))
|
||||
|
||||
#Enable CPU fallback for ops not being supported by MPS like upsample_bicubic2d.out
|
||||
#https://github.com/pytorch/pytorch/issues/77764
|
||||
#https://github.com/Fannovel16/comfyui_controlnet_aux/issues/2#issuecomment-1763579485
|
||||
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = os.getenv("PYTORCH_ENABLE_MPS_FALLBACK", '1')
|
||||
|
||||
|
||||
def load_nodes():
|
||||
shorted_errors = []
|
||||
full_error_messages = []
|
||||
node_class_mappings = {}
|
||||
node_display_name_mappings = {}
|
||||
|
||||
for filename in (here / "node_wrappers").iterdir():
|
||||
module_name = filename.stem
|
||||
if module_name.startswith('.'): continue #Skip hidden files created by the OS (e.g. [.DS_Store](https://en.wikipedia.org/wiki/.DS_Store))
|
||||
try:
|
||||
module = importlib.import_module(
|
||||
f".node_wrappers.{module_name}", package=__package__
|
||||
)
|
||||
node_class_mappings.update(getattr(module, "NODE_CLASS_MAPPINGS"))
|
||||
if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
|
||||
node_display_name_mappings.update(getattr(module, "NODE_DISPLAY_NAME_MAPPINGS"))
|
||||
|
||||
log.debug(f"Imported {module_name} nodes")
|
||||
|
||||
except AttributeError:
|
||||
pass # wip nodes
|
||||
except Exception:
|
||||
error_message = traceback.format_exc()
|
||||
full_error_messages.append(error_message)
|
||||
error_message = error_message.splitlines()[-1]
|
||||
shorted_errors.append(
|
||||
f"Failed to import module {module_name} because {error_message}"
|
||||
)
|
||||
|
||||
if len(shorted_errors) > 0:
|
||||
full_err_log = '\n\n'.join(full_error_messages)
|
||||
print(f"\n\nFull error log from comfyui_controlnet_aux: \n{full_err_log}\n\n")
|
||||
log.info(
|
||||
f"Some nodes failed to load:\n\t"
|
||||
+ "\n\t".join(shorted_errors)
|
||||
+ "\n\n"
|
||||
+ "Check that you properly installed the dependencies.\n"
|
||||
+ "If you think this is a bug, please report it on the github page (https://github.com/Fannovel16/comfyui_controlnet_aux/issues)"
|
||||
)
|
||||
return node_class_mappings, node_display_name_mappings
|
||||
|
||||
AUX_NODE_MAPPINGS, AUX_DISPLAY_NAME_MAPPINGS = load_nodes()
|
||||
|
||||
#For nodes not mapping image to image or has special requirements
|
||||
AIO_NOT_SUPPORTED = ["InpaintPreprocessor", "MeshGraphormer+ImpactDetector-DepthMapPreprocessor", "DiffusionEdge_Preprocessor"]
|
||||
AIO_NOT_SUPPORTED += ["SavePoseKpsAsJsonFile", "FacialPartColoringFromPoseKps", "UpperBodyTrackingFromPoseKps", "RenderPeopleKps", "RenderAnimalKps"]
|
||||
AIO_NOT_SUPPORTED += ["Unimatch_OptFlowPreprocessor", "MaskOptFlow"]
|
||||
|
||||
def preprocessor_options():
|
||||
auxs = list(AUX_NODE_MAPPINGS.keys())
|
||||
auxs.insert(0, "none")
|
||||
for name in AIO_NOT_SUPPORTED:
|
||||
if name in auxs:
|
||||
auxs.remove(name)
|
||||
return auxs
|
||||
|
||||
|
||||
PREPROCESSOR_OPTIONS = preprocessor_options()
|
||||
|
||||
class AIO_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
preprocessor=INPUT.COMBO(PREPROCESSOR_OPTIONS, default="none"),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def execute(self, preprocessor, image, resolution=512):
|
||||
if preprocessor == "none":
|
||||
return (image, )
|
||||
else:
|
||||
aux_class = AUX_NODE_MAPPINGS[preprocessor]
|
||||
input_types = aux_class.INPUT_TYPES()
|
||||
input_types = {
|
||||
**input_types["required"],
|
||||
**(input_types["optional"] if "optional" in input_types else {})
|
||||
}
|
||||
params = {}
|
||||
for name, input_type in input_types.items():
|
||||
if name == "image":
|
||||
params[name] = image
|
||||
continue
|
||||
|
||||
if name == "resolution":
|
||||
params[name] = resolution
|
||||
continue
|
||||
|
||||
if len(input_type) == 2 and ("default" in input_type[1]):
|
||||
params[name] = input_type[1]["default"]
|
||||
continue
|
||||
|
||||
default_values = { "INT": 0, "FLOAT": 0.0 }
|
||||
if type(input_type[0]) is list:
|
||||
for input_type_value in input_type[0]:
|
||||
if input_type_value in default_values:
|
||||
params[name] = default_values[input_type[0]]
|
||||
else:
|
||||
if input_type[0] in default_values:
|
||||
params[name] = default_values[input_type[0]]
|
||||
|
||||
return getattr(aux_class(), aux_class.FUNCTION)(**params)
|
||||
|
||||
class ControlNetAuxSimpleAddText:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return dict(
|
||||
required=dict(image=INPUT.IMAGE(), text=INPUT.STRING())
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
def execute(self, image, text):
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
font = ImageFont.truetype(str((here / "NotoSans-Regular.ttf").resolve()), 40)
|
||||
img = Image.fromarray(image[0].cpu().numpy().__mul__(255.).astype(np.uint8))
|
||||
ImageDraw.Draw(img).text((0,0), text, fill=(0,255,0), font=font)
|
||||
return (torch.from_numpy(np.array(img)).unsqueeze(0) / 255.,)
|
||||
|
||||
class ExecuteAllControlNetPreprocessors:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def execute(self, image, resolution=512):
|
||||
try:
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
except:
|
||||
raise RuntimeError("ExecuteAllControlNetPreprocessor requries [Execution Model Inversion](https://github.com/comfyanonymous/ComfyUI/commit/5cfe38). Update ComfyUI/SwarmUI to get this feature")
|
||||
|
||||
graph = GraphBuilder()
|
||||
curr_outputs = []
|
||||
for preprocc in PREPROCESSOR_OPTIONS:
|
||||
preprocc_node = graph.node("AIO_Preprocessor", preprocessor=preprocc, image=image, resolution=resolution)
|
||||
hint_img = preprocc_node.out(0)
|
||||
add_text_node = graph.node("ControlNetAuxSimpleAddText", image=hint_img, text=preprocc)
|
||||
curr_outputs.append(add_text_node.out(0))
|
||||
|
||||
while len(curr_outputs) > 1:
|
||||
_outputs = []
|
||||
for i in range(0, len(curr_outputs), 2):
|
||||
if i+1 < len(curr_outputs):
|
||||
image_batch = graph.node("ImageBatch", image1=curr_outputs[i], image2=curr_outputs[i+1])
|
||||
_outputs.append(image_batch.out(0))
|
||||
else:
|
||||
_outputs.append(curr_outputs[i])
|
||||
curr_outputs = _outputs
|
||||
|
||||
return {
|
||||
"result": (curr_outputs[0],),
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
|
||||
class ControlNetPreprocessorSelector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"preprocessor": (PREPROCESSOR_OPTIONS,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (PREPROCESSOR_OPTIONS,)
|
||||
RETURN_NAMES = ("preprocessor",)
|
||||
FUNCTION = "get_preprocessor"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def get_preprocessor(self, preprocessor: str):
|
||||
return (preprocessor,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
**AUX_NODE_MAPPINGS,
|
||||
"AIO_Preprocessor": AIO_Preprocessor,
|
||||
"ControlNetPreprocessorSelector": ControlNetPreprocessorSelector,
|
||||
**HIE_NODE_CLASS_MAPPINGS,
|
||||
"ExecuteAllControlNetPreprocessors": ExecuteAllControlNetPreprocessors,
|
||||
"ControlNetAuxSimpleAddText": ControlNetAuxSimpleAddText
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
**AUX_DISPLAY_NAME_MAPPINGS,
|
||||
"AIO_Preprocessor": "AIO Aux Preprocessor",
|
||||
"ControlNetPreprocessorSelector": "Preprocessor Selector",
|
||||
**HIE_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
"ExecuteAllControlNetPreprocessors": "Execute All ControlNet Preprocessors"
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
# this is an example for config.yaml file, you can rename it to config.yaml if you want to use it
|
||||
# ###############################################################################################
|
||||
# This path is for custom pressesor models base folder. default is "./ckpts"
|
||||
# you can also use absolute paths like: "/root/ComfyUI/custom_nodes/comfyui_controlnet_aux/ckpts" or "D:\\ComfyUI\\custom_nodes\\comfyui_controlnet_aux\\ckpts"
|
||||
annotator_ckpts_path: "./ckpts"
|
||||
# ###############################################################################################
|
||||
# This path is for downloading temporary files.
|
||||
# You SHOULD use absolute path for this like"D:\\temp", DO NOT use relative paths. Empty for default.
|
||||
custom_temp_path:
|
||||
# ###############################################################################################
|
||||
# if you already have downloaded ckpts via huggingface hub into default cache path like: ~/.cache/huggingface/hub, you can set this True to use symlinks to save space
|
||||
USE_SYMLINKS: False
|
||||
# ###############################################################################################
|
||||
# EP_list is a list of execution providers for onnxruntime, if one of them is not available or not working well, you can delete that provider from here(config.yaml)
|
||||
# you can find all available providers here: https://onnxruntime.ai/docs/execution-providers
|
||||
# for example, if you have CUDA installed, you can set it to: ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
# empty list or only keep ["CPUExecutionProvider"] means you use cv2.dnn.readNetFromONNX to load onnx models
|
||||
# if your onnx models can only run on the CPU or have other issues, we recommend using pt model instead.
|
||||
# default value is ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CPUExecutionProvider"]
|
||||
EP_list: ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CPUExecutionProvider"]
|
||||
@@ -0,0 +1,6 @@
|
||||
from pathlib import Path
|
||||
from utils import here
|
||||
import sys
|
||||
sys.path.append(str(Path(here, "src")))
|
||||
|
||||
from custom_controlnet_aux import *
|
||||
@@ -0,0 +1,233 @@
|
||||
from .log import log
|
||||
from .utils import ResizeMode, safe_numpy
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
from .utils import get_unique_axis0
|
||||
from .lvminthin import nake_nms, lvmin_thin
|
||||
|
||||
MAX_IMAGEGEN_RESOLUTION = 8192 #https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L42
|
||||
RESIZE_MODES = [ResizeMode.RESIZE.value, ResizeMode.INNER_FIT.value, ResizeMode.OUTER_FIT.value]
|
||||
|
||||
#Port from https://github.com/Mikubill/sd-webui-controlnet/blob/e67e017731aad05796b9615dc6eadce911298ea1/internal_controlnet/external_code.py#L89
|
||||
class PixelPerfectResolution:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"original_image": ("IMAGE", ),
|
||||
"image_gen_width": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}),
|
||||
"image_gen_height": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}),
|
||||
#https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L854
|
||||
"resize_mode": (RESIZE_MODES, {"default": ResizeMode.RESIZE.value})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("RESOLUTION (INT)", )
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def execute(self, original_image, image_gen_width, image_gen_height, resize_mode):
|
||||
_, raw_H, raw_W, _ = original_image.shape
|
||||
|
||||
k0 = float(image_gen_height) / float(raw_H)
|
||||
k1 = float(image_gen_width) / float(raw_W)
|
||||
|
||||
if resize_mode == ResizeMode.OUTER_FIT.value:
|
||||
estimation = min(k0, k1) * float(min(raw_H, raw_W))
|
||||
else:
|
||||
estimation = max(k0, k1) * float(min(raw_H, raw_W))
|
||||
|
||||
log.debug(f"Pixel Perfect Computation:")
|
||||
log.debug(f"resize_mode = {resize_mode}")
|
||||
log.debug(f"raw_H = {raw_H}")
|
||||
log.debug(f"raw_W = {raw_W}")
|
||||
log.debug(f"target_H = {image_gen_height}")
|
||||
log.debug(f"target_W = {image_gen_width}")
|
||||
log.debug(f"estimation = {estimation}")
|
||||
|
||||
return (int(np.round(estimation)), )
|
||||
|
||||
class HintImageEnchance:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"hint_image": ("IMAGE", ),
|
||||
"image_gen_width": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}),
|
||||
"image_gen_height": ("INT", {"default": 512, "min": 64, "max": MAX_IMAGEGEN_RESOLUTION, "step": 8}),
|
||||
#https://github.com/comfyanonymous/ComfyUI/blob/c910b4a01ca58b04e5d4ab4c747680b996ada02b/nodes.py#L854
|
||||
"resize_mode": (RESIZE_MODES, {"default": ResizeMode.RESIZE.value})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
def execute(self, hint_image, image_gen_width, image_gen_height, resize_mode):
|
||||
outs = []
|
||||
for single_hint_image in hint_image:
|
||||
np_hint_image = np.asarray(single_hint_image * 255., dtype=np.uint8)
|
||||
|
||||
if resize_mode == ResizeMode.RESIZE.value:
|
||||
np_hint_image = self.execute_resize(np_hint_image, image_gen_width, image_gen_height)
|
||||
elif resize_mode == ResizeMode.OUTER_FIT.value:
|
||||
np_hint_image = self.execute_outer_fit(np_hint_image, image_gen_width, image_gen_height)
|
||||
else:
|
||||
np_hint_image = self.execute_inner_fit(np_hint_image, image_gen_width, image_gen_height)
|
||||
|
||||
outs.append(torch.from_numpy(np_hint_image.astype(np.float32) / 255.0))
|
||||
|
||||
return (torch.stack(outs, dim=0),)
|
||||
|
||||
def execute_resize(self, detected_map, w, h):
|
||||
detected_map = self.high_quality_resize(detected_map, (w, h))
|
||||
detected_map = safe_numpy(detected_map)
|
||||
return detected_map
|
||||
|
||||
def execute_outer_fit(self, detected_map, w, h):
|
||||
old_h, old_w, _ = detected_map.shape
|
||||
old_w = float(old_w)
|
||||
old_h = float(old_h)
|
||||
k0 = float(h) / old_h
|
||||
k1 = float(w) / old_w
|
||||
safeint = lambda x: int(np.round(x))
|
||||
k = min(k0, k1)
|
||||
|
||||
borders = np.concatenate([detected_map[0, :, :], detected_map[-1, :, :], detected_map[:, 0, :], detected_map[:, -1, :]], axis=0)
|
||||
high_quality_border_color = np.median(borders, axis=0).astype(detected_map.dtype)
|
||||
if len(high_quality_border_color) == 4:
|
||||
# Inpaint hijack
|
||||
high_quality_border_color[3] = 255
|
||||
high_quality_background = np.tile(high_quality_border_color[None, None], [h, w, 1])
|
||||
detected_map = self.high_quality_resize(detected_map, (safeint(old_w * k), safeint(old_h * k)))
|
||||
new_h, new_w, _ = detected_map.shape
|
||||
pad_h = max(0, (h - new_h) // 2)
|
||||
pad_w = max(0, (w - new_w) // 2)
|
||||
high_quality_background[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = detected_map
|
||||
detected_map = high_quality_background
|
||||
detected_map = safe_numpy(detected_map)
|
||||
return detected_map
|
||||
|
||||
def execute_inner_fit(self, detected_map, w, h):
|
||||
old_h, old_w, _ = detected_map.shape
|
||||
old_w = float(old_w)
|
||||
old_h = float(old_h)
|
||||
k0 = float(h) / old_h
|
||||
k1 = float(w) / old_w
|
||||
safeint = lambda x: int(np.round(x))
|
||||
k = max(k0, k1)
|
||||
|
||||
detected_map = self.high_quality_resize(detected_map, (safeint(old_w * k), safeint(old_h * k)))
|
||||
new_h, new_w, _ = detected_map.shape
|
||||
pad_h = max(0, (new_h - h) // 2)
|
||||
pad_w = max(0, (new_w - w) // 2)
|
||||
detected_map = detected_map[pad_h:pad_h+h, pad_w:pad_w+w]
|
||||
detected_map = safe_numpy(detected_map)
|
||||
return detected_map
|
||||
|
||||
def high_quality_resize(self, x, size):
|
||||
# Written by lvmin
|
||||
# Super high-quality control map up-scaling, considering binary, seg, and one-pixel edges
|
||||
|
||||
inpaint_mask = None
|
||||
if x.ndim == 3 and x.shape[2] == 4:
|
||||
inpaint_mask = x[:, :, 3]
|
||||
x = x[:, :, 0:3]
|
||||
|
||||
if x.shape[0] != size[1] or x.shape[1] != size[0]:
|
||||
new_size_is_smaller = (size[0] * size[1]) < (x.shape[0] * x.shape[1])
|
||||
new_size_is_bigger = (size[0] * size[1]) > (x.shape[0] * x.shape[1])
|
||||
unique_color_count = len(get_unique_axis0(x.reshape(-1, x.shape[2])))
|
||||
is_one_pixel_edge = False
|
||||
is_binary = False
|
||||
if unique_color_count == 2:
|
||||
is_binary = np.min(x) < 16 and np.max(x) > 240
|
||||
if is_binary:
|
||||
xc = x
|
||||
xc = cv2.erode(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
|
||||
xc = cv2.dilate(xc, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
|
||||
one_pixel_edge_count = np.where(xc < x)[0].shape[0]
|
||||
all_edge_count = np.where(x > 127)[0].shape[0]
|
||||
is_one_pixel_edge = one_pixel_edge_count * 2 > all_edge_count
|
||||
|
||||
if 2 < unique_color_count < 200:
|
||||
interpolation = cv2.INTER_NEAREST
|
||||
elif new_size_is_smaller:
|
||||
interpolation = cv2.INTER_AREA
|
||||
else:
|
||||
interpolation = cv2.INTER_CUBIC # Must be CUBIC because we now use nms. NEVER CHANGE THIS
|
||||
|
||||
y = cv2.resize(x, size, interpolation=interpolation)
|
||||
if inpaint_mask is not None:
|
||||
inpaint_mask = cv2.resize(inpaint_mask, size, interpolation=interpolation)
|
||||
|
||||
if is_binary:
|
||||
y = np.mean(y.astype(np.float32), axis=2).clip(0, 255).astype(np.uint8)
|
||||
if is_one_pixel_edge:
|
||||
y = nake_nms(y)
|
||||
_, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||||
y = lvmin_thin(y, prunings=new_size_is_bigger)
|
||||
else:
|
||||
_, y = cv2.threshold(y, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||||
y = np.stack([y] * 3, axis=2)
|
||||
else:
|
||||
y = x
|
||||
|
||||
if inpaint_mask is not None:
|
||||
inpaint_mask = (inpaint_mask > 127).astype(np.float32) * 255.0
|
||||
inpaint_mask = inpaint_mask[:, :, None].clip(0, 255).astype(np.uint8)
|
||||
y = np.concatenate([y, inpaint_mask], axis=2)
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class ImageGenResolutionFromLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": { "latent": ("LATENT", ) }
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT")
|
||||
RETURN_NAMES = ("IMAGE_GEN_WIDTH (INT)", "IMAGE_GEN_HEIGHT (INT)")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def execute(self, latent):
|
||||
_, _, H, W = latent["samples"].shape
|
||||
return (W * 8, H * 8)
|
||||
|
||||
class ImageGenResolutionFromImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": { "image": ("IMAGE", ) }
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT")
|
||||
RETURN_NAMES = ("IMAGE_GEN_WIDTH (INT)", "IMAGE_GEN_HEIGHT (INT)")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors"
|
||||
|
||||
def execute(self, image):
|
||||
_, H, W, _ = image.shape
|
||||
return (W, H)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PixelPerfectResolution": PixelPerfectResolution,
|
||||
"ImageGenResolutionFromImage": ImageGenResolutionFromImage,
|
||||
"ImageGenResolutionFromLatent": ImageGenResolutionFromLatent,
|
||||
"HintImageEnchance": HintImageEnchance
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PixelPerfectResolution": "Pixel Perfect Resolution",
|
||||
"ImageGenResolutionFromImage": "Generation Resolution From Image",
|
||||
"ImageGenResolutionFromLatent": "Generation Resolution From Latent",
|
||||
"HintImageEnchance": "Enchance And Resize Hint Images"
|
||||
}
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
@echo off
|
||||
|
||||
set "requirements_txt=%~dp0\requirements.txt"
|
||||
set "python_exec=..\..\..\python_embeded\python.exe"
|
||||
|
||||
echo Installing ComfyUI's ControlNet Auxiliary Preprocessors..
|
||||
|
||||
if exist "%python_exec%" (
|
||||
echo Installing with ComfyUI Portable
|
||||
for /f "delims=" %%i in (%requirements_txt%) do (
|
||||
%python_exec% -s -m pip install "%%i"
|
||||
)
|
||||
) else (
|
||||
echo Installing with system Python
|
||||
for /f "delims=" %%i in (%requirements_txt%) do (
|
||||
pip install "%%i"
|
||||
)
|
||||
)
|
||||
|
||||
pause
|
||||
@@ -0,0 +1,80 @@
|
||||
#Cre: https://github.com/melMass/comfy_mtb/blob/main/log.py
|
||||
import logging
|
||||
import re
|
||||
import os
|
||||
|
||||
base_log_level = logging.INFO
|
||||
|
||||
|
||||
# Custom object that discards the output
|
||||
class NullWriter:
|
||||
def write(self, text):
|
||||
pass
|
||||
|
||||
|
||||
class Formatter(logging.Formatter):
|
||||
grey = "\x1b[38;20m"
|
||||
cyan = "\x1b[36;20m"
|
||||
purple = "\x1b[35;20m"
|
||||
yellow = "\x1b[33;20m"
|
||||
red = "\x1b[31;20m"
|
||||
bold_red = "\x1b[31;1m"
|
||||
reset = "\x1b[0m"
|
||||
# format = "%(asctime)s - [%(name)s] - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
|
||||
format = "[%(name)s] | %(levelname)s -> %(message)s"
|
||||
|
||||
FORMATS = {
|
||||
logging.DEBUG: purple + format + reset,
|
||||
logging.INFO: cyan + format + reset,
|
||||
logging.WARNING: yellow + format + reset,
|
||||
logging.ERROR: red + format + reset,
|
||||
logging.CRITICAL: bold_red + format + reset,
|
||||
}
|
||||
|
||||
def format(self, record):
|
||||
log_fmt = self.FORMATS.get(record.levelno)
|
||||
formatter = logging.Formatter(log_fmt)
|
||||
return formatter.format(record)
|
||||
|
||||
|
||||
def mklog(name, level=base_log_level):
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(level)
|
||||
|
||||
for handler in logger.handlers:
|
||||
logger.removeHandler(handler)
|
||||
|
||||
ch = logging.StreamHandler()
|
||||
ch.setLevel(level)
|
||||
ch.setFormatter(Formatter())
|
||||
logger.addHandler(ch)
|
||||
|
||||
# Disable log propagation
|
||||
logger.propagate = False
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
# - The main app logger
|
||||
log = mklog(__package__, base_log_level)
|
||||
|
||||
|
||||
def log_user(arg):
|
||||
print("\033[34mComfyUI ControlNet AUX:\033[0m {arg}")
|
||||
|
||||
|
||||
def get_summary(docstring):
|
||||
return docstring.strip().split("\n\n", 1)[0]
|
||||
|
||||
|
||||
def blue_text(text):
|
||||
return f"\033[94m{text}\033[0m"
|
||||
|
||||
|
||||
def cyan_text(text):
|
||||
return f"\033[96m{text}\033[0m"
|
||||
|
||||
|
||||
def get_label(label):
|
||||
words = re.findall(r"(?:^|[A-Z])[a-z]*", label)
|
||||
return " ".join(words).strip()
|
||||
@@ -0,0 +1,87 @@
|
||||
# High Quality Edge Thinning using Pure Python
|
||||
# Written by Lvmin Zhang
|
||||
# 2023 April
|
||||
# Stanford University
|
||||
# If you use this, please Cite "High Quality Edge Thinning using Pure Python", Lvmin Zhang, In Mikubill/sd-webui-controlnet.
|
||||
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
lvmin_kernels_raw = [
|
||||
np.array([
|
||||
[-1, -1, -1],
|
||||
[0, 1, 0],
|
||||
[1, 1, 1]
|
||||
], dtype=np.int32),
|
||||
np.array([
|
||||
[0, -1, -1],
|
||||
[1, 1, -1],
|
||||
[0, 1, 0]
|
||||
], dtype=np.int32)
|
||||
]
|
||||
|
||||
lvmin_kernels = []
|
||||
lvmin_kernels += [np.rot90(x, k=0, axes=(0, 1)) for x in lvmin_kernels_raw]
|
||||
lvmin_kernels += [np.rot90(x, k=1, axes=(0, 1)) for x in lvmin_kernels_raw]
|
||||
lvmin_kernels += [np.rot90(x, k=2, axes=(0, 1)) for x in lvmin_kernels_raw]
|
||||
lvmin_kernels += [np.rot90(x, k=3, axes=(0, 1)) for x in lvmin_kernels_raw]
|
||||
|
||||
lvmin_prunings_raw = [
|
||||
np.array([
|
||||
[-1, -1, -1],
|
||||
[-1, 1, -1],
|
||||
[0, 0, -1]
|
||||
], dtype=np.int32),
|
||||
np.array([
|
||||
[-1, -1, -1],
|
||||
[-1, 1, -1],
|
||||
[-1, 0, 0]
|
||||
], dtype=np.int32)
|
||||
]
|
||||
|
||||
lvmin_prunings = []
|
||||
lvmin_prunings += [np.rot90(x, k=0, axes=(0, 1)) for x in lvmin_prunings_raw]
|
||||
lvmin_prunings += [np.rot90(x, k=1, axes=(0, 1)) for x in lvmin_prunings_raw]
|
||||
lvmin_prunings += [np.rot90(x, k=2, axes=(0, 1)) for x in lvmin_prunings_raw]
|
||||
lvmin_prunings += [np.rot90(x, k=3, axes=(0, 1)) for x in lvmin_prunings_raw]
|
||||
|
||||
|
||||
def remove_pattern(x, kernel):
|
||||
objects = cv2.morphologyEx(x, cv2.MORPH_HITMISS, kernel)
|
||||
objects = np.where(objects > 127)
|
||||
x[objects] = 0
|
||||
return x, objects[0].shape[0] > 0
|
||||
|
||||
|
||||
def thin_one_time(x, kernels):
|
||||
y = x
|
||||
is_done = True
|
||||
for k in kernels:
|
||||
y, has_update = remove_pattern(y, k)
|
||||
if has_update:
|
||||
is_done = False
|
||||
return y, is_done
|
||||
|
||||
|
||||
def lvmin_thin(x, prunings=True):
|
||||
y = x
|
||||
for i in range(32):
|
||||
y, is_done = thin_one_time(y, lvmin_kernels)
|
||||
if is_done:
|
||||
break
|
||||
if prunings:
|
||||
y, _ = thin_one_time(y, lvmin_prunings)
|
||||
return y
|
||||
|
||||
|
||||
def nake_nms(x):
|
||||
f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
|
||||
f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
|
||||
f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
|
||||
f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
|
||||
y = np.zeros_like(x)
|
||||
for f in [f1, f2, f3, f4]:
|
||||
np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
|
||||
return y
|
||||
@@ -0,0 +1,43 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
class AnimeFace_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#This preprocessor is only trained on 512x resolution
|
||||
#https://github.com/siyeong0/Anime-Face-Segmentation/blob/main/predict.py#L25
|
||||
return define_preprocessor_inputs(
|
||||
remove_background_using_abg=INPUT.BOOLEAN(True),
|
||||
resolution=INPUT.RESOLUTION(default=512, min=512, max=512)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "ABG_CHARACTER_MASK (MASK)")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def execute(self, image, remove_background_using_abg=True, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.anime_face_segment import AnimeFaceSegmentor
|
||||
|
||||
model = AnimeFaceSegmentor.from_pretrained().to(model_management.get_torch_device())
|
||||
if remove_background_using_abg:
|
||||
out_image_with_mask = common_annotator_call(model, image, resolution=resolution, remove_background=True)
|
||||
out_image = out_image_with_mask[..., :3]
|
||||
mask = out_image_with_mask[..., 3:]
|
||||
mask = rearrange(mask, "n h w c -> n c h w")
|
||||
else:
|
||||
out_image = common_annotator_call(model, image, resolution=resolution, remove_background=False)
|
||||
N, H, W, C = out_image.shape
|
||||
mask = torch.ones(N, C, H, W)
|
||||
del model
|
||||
return (out_image, mask)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnimeFace_SemSegPreprocessor": AnimeFace_SemSegPreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnimeFace_SemSegPreprocessor": "Anime Face Segmentor"
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import comfy.model_management as model_management
|
||||
import comfy.utils
|
||||
|
||||
# Requires comfyui_controlnet_aux funcsions and classes
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
|
||||
def get_intensity_mask(image_array, lower_bound, upper_bound):
|
||||
mask = image_array[:, :, 0]
|
||||
mask = np.where((mask >= lower_bound) & (mask <= upper_bound), mask, 0)
|
||||
mask = np.expand_dims(mask, 2).repeat(3, axis=2)
|
||||
return mask
|
||||
|
||||
def combine_layers(base_layer, top_layer):
|
||||
mask = top_layer.astype(bool)
|
||||
temp = 1 - (1 - top_layer) * (1 - base_layer)
|
||||
result = base_layer * (~mask) + temp * mask
|
||||
return result
|
||||
|
||||
class AnyLinePreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
merge_with_lineart=INPUT.COMBO(["lineart_standard", "lineart_realisitic", "lineart_anime", "manga_line"], default="lineart_standard"),
|
||||
resolution=INPUT.RESOLUTION(default=1280, step=8),
|
||||
lineart_lower_bound=INPUT.FLOAT(default=0),
|
||||
lineart_upper_bound=INPUT.FLOAT(default=1),
|
||||
object_min_size=INPUT.INT(default=36, min=1),
|
||||
object_connectivity=INPUT.INT(default=1, min=1)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
|
||||
FUNCTION = "get_anyline"
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def __init__(self):
|
||||
self.device = model_management.get_torch_device()
|
||||
|
||||
def get_anyline(self, image, merge_with_lineart="lineart_standard", resolution=512, lineart_lower_bound=0, lineart_upper_bound=1, object_min_size=36, object_connectivity=1):
|
||||
from custom_controlnet_aux.teed import TEDDetector
|
||||
from skimage import morphology
|
||||
pbar = comfy.utils.ProgressBar(3)
|
||||
|
||||
# Process the image with MTEED model
|
||||
mteed_model = TEDDetector.from_pretrained("TheMistoAI/MistoLine", "MTEED.pth", subfolder="Anyline").to(self.device)
|
||||
mteed_result = common_annotator_call(mteed_model, image, resolution=resolution, show_pbar=False)
|
||||
mteed_result = mteed_result.numpy()
|
||||
del mteed_model
|
||||
pbar.update(1)
|
||||
|
||||
# Process the image with the lineart standard preprocessor
|
||||
if merge_with_lineart == "lineart_standard":
|
||||
from custom_controlnet_aux.lineart_standard import LineartStandardDetector
|
||||
lineart_standard_detector = LineartStandardDetector()
|
||||
lineart_result = common_annotator_call(lineart_standard_detector, image, guassian_sigma=2, intensity_threshold=3, resolution=resolution, show_pbar=False).numpy()
|
||||
del lineart_standard_detector
|
||||
else:
|
||||
from custom_controlnet_aux.lineart import LineartDetector
|
||||
from custom_controlnet_aux.lineart_anime import LineartAnimeDetector
|
||||
from custom_controlnet_aux.manga_line import LineartMangaDetector
|
||||
lineart_detector = dict(lineart_realisitic=LineartDetector, lineart_anime=LineartAnimeDetector, manga_line=LineartMangaDetector)[merge_with_lineart]
|
||||
lineart_detector = lineart_detector.from_pretrained().to(self.device)
|
||||
lineart_result = common_annotator_call(lineart_detector, image, resolution=resolution, show_pbar=False).numpy()
|
||||
del lineart_detector
|
||||
pbar.update(1)
|
||||
|
||||
final_result = []
|
||||
for i in range(len(image)):
|
||||
_lineart_result = get_intensity_mask(lineart_result[i], lower_bound=lineart_lower_bound, upper_bound=lineart_upper_bound)
|
||||
_cleaned = morphology.remove_small_objects(_lineart_result.astype(bool), min_size=object_min_size, connectivity=object_connectivity)
|
||||
_lineart_result = _lineart_result * _cleaned
|
||||
_mteed_result = mteed_result[i]
|
||||
|
||||
# Combine the results
|
||||
final_result.append(torch.from_numpy(combine_layers(_mteed_result, _lineart_result)))
|
||||
pbar.update(1)
|
||||
return (torch.stack(final_result),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnyLineArtPreprocessor_aux": AnyLinePreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnyLineArtPreprocessor_aux": "AnyLine Lineart"
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Binary_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
bin_threshold=INPUT.INT(default=100, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, bin_threshold=100, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.binary import BinaryDetector
|
||||
|
||||
return (common_annotator_call(BinaryDetector(), image, bin_threshold=bin_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BinaryPreprocessor": Binary_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BinaryPreprocessor": "Binary Lines"
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Canny_Edge_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
low_threshold=INPUT.INT(default=100, max=255),
|
||||
high_threshold=INPUT.INT(default=200, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, low_threshold=100, high_threshold=200, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.canny import CannyDetector
|
||||
|
||||
return (common_annotator_call(CannyDetector(), image, low_threshold=low_threshold, high_threshold=high_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CannyEdgePreprocessor": Canny_Edge_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CannyEdgePreprocessor": "Canny Edge"
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Color_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/T2IAdapter-only"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.color import ColorDetector
|
||||
|
||||
return (common_annotator_call(ColorDetector(), image, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ColorPreprocessor": Color_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ColorPreprocessor": "Color Pallete"
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class DensePose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
model=INPUT.COMBO(["densepose_r50_fpn_dl.torchscript", "densepose_r101_fpn_dl.torchscript"]),
|
||||
cmap=INPUT.COMBO(["Viridis (MagicAnimate)", "Parula (CivitAI)"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def execute(self, image, model="densepose_r50_fpn_dl.torchscript", cmap="Viridis (MagicAnimate)", resolution=512):
|
||||
from custom_controlnet_aux.densepose import DenseposeDetector
|
||||
model = DenseposeDetector \
|
||||
.from_pretrained(filename=model) \
|
||||
.to(model_management.get_torch_device())
|
||||
return (common_annotator_call(model, image, cmap="viridis" if "Viridis" in cmap else "parula", resolution=resolution), )
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DensePosePreprocessor": DensePose_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DensePosePreprocessor": "DensePose Estimator"
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Depth_Anything_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
ckpt_name=INPUT.COMBO(
|
||||
["depth_anything_vitl14.pth", "depth_anything_vitb14.pth", "depth_anything_vits14.pth"]
|
||||
),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, ckpt_name="depth_anything_vitl14.pth", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything import DepthAnythingDetector
|
||||
|
||||
model = DepthAnythingDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Zoe_Depth_Anything_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
environment=INPUT.COMBO(["indoor", "outdoor"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, environment="indoor", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.zoe import ZoeDepthAnythingDetector
|
||||
ckpt_name = "depth_anything_metric_depth_indoor.pt" if environment == "indoor" else "depth_anything_metric_depth_outdoor.pt"
|
||||
model = ZoeDepthAnythingDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DepthAnythingPreprocessor": Depth_Anything_Preprocessor,
|
||||
"Zoe_DepthAnythingPreprocessor": Zoe_Depth_Anything_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DepthAnythingPreprocessor": "Depth Anything",
|
||||
"Zoe_DepthAnythingPreprocessor": "Zoe Depth Anything"
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Depth_Anything_V2_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
ckpt_name=INPUT.COMBO(
|
||||
["depth_anything_v2_vitg.pth", "depth_anything_v2_vitl.pth", "depth_anything_v2_vitb.pth", "depth_anything_v2_vits.pth"],
|
||||
default="depth_anything_v2_vitl.pth"
|
||||
),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, ckpt_name="depth_anything_v2_vitl.pth", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything_v2 import DepthAnythingV2Detector
|
||||
|
||||
model = DepthAnythingV2Detector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, max_depth=1)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
""" class Depth_Anything_Metric_V2_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return create_node_input_types(
|
||||
environment=(["indoor", "outdoor"], {"default": "indoor"}),
|
||||
max_depth=("FLOAT", {"min": 0, "max": 100, "default": 20.0, "step": 0.01})
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, environment, resolution=512, max_depth=20.0, **kwargs):
|
||||
from custom_controlnet_aux.depth_anything_v2 import DepthAnythingV2Detector
|
||||
filename = dict(indoor="depth_anything_v2_metric_hypersim_vitl.pth", outdoor="depth_anything_v2_metric_vkitti_vitl.pth")[environment]
|
||||
model = DepthAnythingV2Detector.from_pretrained(filename=filename).to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, max_depth=max_depth)
|
||||
del model
|
||||
return (out, ) """
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DepthAnythingV2Preprocessor": Depth_Anything_V2_Preprocessor,
|
||||
#"Metric_DepthAnythingV2Preprocessor": Depth_Anything_Metric_V2_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DepthAnythingV2Preprocessor": "Depth Anything V2 - Relative",
|
||||
#"Metric_DepthAnythingV2Preprocessor": "Depth Anything V2 - Metric"
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, run_script
|
||||
import comfy.model_management as model_management
|
||||
import sys
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import sklearn
|
||||
except:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'scikit-learn'])
|
||||
|
||||
class DiffusionEdge_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
environment=INPUT.COMBO(["indoor", "urban", "natrual"]),
|
||||
patch_batch_size=INPUT.INT(default=4, min=1, max=16),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, environment="indoor", patch_batch_size=4, resolution=512, **kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.diffusion_edge import DiffusionEdgeDetector
|
||||
|
||||
model = DiffusionEdgeDetector \
|
||||
.from_pretrained(filename = f"diffusion_edge_{environment}.pt") \
|
||||
.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, patch_batch_size=patch_batch_size)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DiffusionEdge_Preprocessor": DiffusionEdge_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DiffusionEdge_Preprocessor": "Diffusion Edge (batch size ↑ => speed ↑, VRAM ↑)",
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class DSINE_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
fov=INPUT.FLOAT(max=365.0, default=60.0),
|
||||
iterations=INPUT.INT(min=1, max=20, default=5),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, fov=60.0, iterations=5, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.dsine import DsineDetector
|
||||
|
||||
model = DsineDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, fov=fov, iterations=iterations, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DSINE-NormalMapPreprocessor": DSINE_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DSINE-NormalMapPreprocessor": "DSINE Normal Map"
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
import warnings
|
||||
from custom_controlnet_aux.dwpose import DwposeDetector, AnimalposeDetector
|
||||
import os
|
||||
import json
|
||||
|
||||
DWPOSE_MODEL_NAME = "yzd-v/DWPose"
|
||||
#Trigger startup caching for onnxruntime
|
||||
GPU_PROVIDERS = ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CoreMLExecutionProvider"]
|
||||
def check_ort_gpu():
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
for provider in GPU_PROVIDERS:
|
||||
if provider in ort.get_available_providers():
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
if not os.environ.get("DWPOSE_ONNXRT_CHECKED"):
|
||||
if check_ort_gpu():
|
||||
print("DWPose: Onnxruntime with acceleration providers detected")
|
||||
else:
|
||||
warnings.warn("DWPose: Onnxruntime not found or doesn't come with acceleration providers, switch to OpenCV with CPU device. DWPose might run very slowly")
|
||||
os.environ['AUX_ORT_PROVIDERS'] = ''
|
||||
os.environ["DWPOSE_ONNXRT_CHECKED"] = '1'
|
||||
|
||||
class DWPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
detect_hand=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_body=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_face=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
bbox_detector=INPUT.COMBO(
|
||||
["None"] + ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
|
||||
default="yolox_l.onnx"
|
||||
),
|
||||
pose_estimator=INPUT.COMBO(
|
||||
["dw-ll_ucoco_384_bs5.torchscript.pt", "dw-ll_ucoco_384.onnx", "dw-ll_ucoco.onnx"],
|
||||
default="dw-ll_ucoco_384_bs5.torchscript.pt"
|
||||
),
|
||||
scale_stick_for_xinsr_cn=INPUT.COMBO(["disable", "enable"])
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, detect_hand="enable", detect_body="enable", detect_face="enable", resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="dw-ll_ucoco_384.onnx", scale_stick_for_xinsr_cn="disable", **kwargs):
|
||||
if bbox_detector == "None":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif bbox_detector == "yolox_l.onnx":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif "yolox" in bbox_detector:
|
||||
yolo_repo = "hr16/yolox-onnx"
|
||||
elif "yolo_nas" in bbox_detector:
|
||||
yolo_repo = "hr16/yolo-nas-fp16"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {bbox_detector}")
|
||||
|
||||
if pose_estimator == "dw-ll_ucoco_384.onnx":
|
||||
pose_repo = DWPOSE_MODEL_NAME
|
||||
elif pose_estimator.endswith(".onnx"):
|
||||
pose_repo = "hr16/UnJIT-DWPose"
|
||||
elif pose_estimator.endswith(".torchscript.pt"):
|
||||
pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {pose_estimator}")
|
||||
|
||||
model = DwposeDetector.from_pretrained(
|
||||
pose_repo,
|
||||
yolo_repo,
|
||||
det_filename=(None if bbox_detector == "None" else bbox_detector), pose_filename=pose_estimator,
|
||||
torchscript_device=model_management.get_torch_device()
|
||||
)
|
||||
detect_hand = detect_hand == "enable"
|
||||
detect_body = detect_body == "enable"
|
||||
detect_face = detect_face == "enable"
|
||||
scale_stick_for_xinsr_cn = scale_stick_for_xinsr_cn == "enable"
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, include_hand=detect_hand, include_face=detect_face, include_body=detect_body, image_and_json=True, resolution=resolution, xinsr_stick_scaling=scale_stick_for_xinsr_cn)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
class AnimalPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
bbox_detector = INPUT.COMBO(
|
||||
["None"] + ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
|
||||
default="yolox_l.torchscript.pt"
|
||||
),
|
||||
pose_estimator = INPUT.COMBO(
|
||||
["rtmpose-m_ap10k_256_bs5.torchscript.pt", "rtmpose-m_ap10k_256.onnx"],
|
||||
default="rtmpose-m_ap10k_256_bs5.torchscript.pt"
|
||||
),
|
||||
resolution = INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="rtmpose-m_ap10k_256.onnx", **kwargs):
|
||||
if bbox_detector == "None":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif bbox_detector == "yolox_l.onnx":
|
||||
yolo_repo = DWPOSE_MODEL_NAME
|
||||
elif "yolox" in bbox_detector:
|
||||
yolo_repo = "hr16/yolox-onnx"
|
||||
elif "yolo_nas" in bbox_detector:
|
||||
yolo_repo = "hr16/yolo-nas-fp16"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {bbox_detector}")
|
||||
|
||||
if pose_estimator == "dw-ll_ucoco_384.onnx":
|
||||
pose_repo = DWPOSE_MODEL_NAME
|
||||
elif pose_estimator.endswith(".onnx"):
|
||||
pose_repo = "hr16/UnJIT-DWPose"
|
||||
elif pose_estimator.endswith(".torchscript.pt"):
|
||||
pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
|
||||
else:
|
||||
raise NotImplementedError(f"Download mechanism for {pose_estimator}")
|
||||
|
||||
model = AnimalposeDetector.from_pretrained(
|
||||
pose_repo,
|
||||
yolo_repo,
|
||||
det_filename=(None if bbox_detector == "None" else bbox_detector), pose_filename=pose_estimator,
|
||||
torchscript_device=model_management.get_torch_device()
|
||||
)
|
||||
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, image_and_json=True, resolution=resolution)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DWPreprocessor": DWPose_Preprocessor,
|
||||
"AnimalPosePreprocessor": AnimalPose_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DWPreprocessor": "DWPose Estimator",
|
||||
"AnimalPosePreprocessor": "AnimalPose Estimator (AP10K)"
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class HED_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.hed import HEDdetector
|
||||
|
||||
model = HEDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe = kwargs["safe"] == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Fake_Scribble_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.hed import HEDdetector
|
||||
|
||||
model = HEDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, scribble=True, safe=kwargs["safe"]=="enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"HEDPreprocessor": HED_Preprocessor,
|
||||
"FakeScribblePreprocessor": Fake_Scribble_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"HEDPreprocessor": "HED Soft-Edge Lines",
|
||||
"FakeScribblePreprocessor": "Fake Scribble Lines (aka scribble_hed)"
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
from ..utils import INPUT
|
||||
|
||||
class InpaintPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return dict(
|
||||
required=dict(image=INPUT.IMAGE(), mask=INPUT.MASK()),
|
||||
optional=dict(black_pixel_for_xinsir_cn=INPUT.BOOLEAN(False))
|
||||
)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "preprocess"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/others"
|
||||
|
||||
def preprocess(self, image, mask, black_pixel_for_xinsir_cn=False):
|
||||
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(image.shape[1], image.shape[2]), mode="bilinear")
|
||||
mask = mask.movedim(1,-1).expand((-1,-1,-1,3))
|
||||
image = image.clone()
|
||||
if black_pixel_for_xinsir_cn:
|
||||
masked_pixel = 0.0
|
||||
else:
|
||||
masked_pixel = -1.0
|
||||
image[mask > 0.5] = masked_pixel
|
||||
return (image,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"InpaintPreprocessor": InpaintPreprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"InpaintPreprocessor": "Inpaint Preprocessor"
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class LERES_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
rm_nearest=INPUT.FLOAT(max=100.0),
|
||||
rm_background=INPUT.FLOAT(max=100.0),
|
||||
boost=INPUT.COMBO(["disable", "enable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, rm_nearest=0, rm_background=0, resolution=512, boost="disable", **kwargs):
|
||||
from custom_controlnet_aux.leres import LeresDetector
|
||||
|
||||
model = LeresDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, thr_a=rm_nearest, thr_b=rm_background, boost=boost == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LeReS-DepthMapPreprocessor": LERES_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LeReS-DepthMapPreprocessor": "LeReS Depth Map (enable boost for leres++)"
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class LineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
coarse=INPUT.COMBO((["disable", "enable"])),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart import LineartDetector
|
||||
|
||||
model = LineartDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, coarse = kwargs["coarse"] == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LineArtPreprocessor": LineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LineArtPreprocessor": "Realistic Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class AnimeLineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart_anime import LineartAnimeDetector
|
||||
|
||||
model = LineartAnimeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AnimeLineArtPreprocessor": AnimeLineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AnimeLineArtPreprocessor": "Anime Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Lineart_Standard_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
guassian_sigma=INPUT.FLOAT(default=6.0, max=100.0),
|
||||
intensity_threshold=INPUT.INT(default=8, max=16),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, guassian_sigma=6, intensity_threshold=8, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.lineart_standard import LineartStandardDetector
|
||||
return (common_annotator_call(LineartStandardDetector(), image, guassian_sigma=guassian_sigma, intensity_threshold=intensity_threshold, resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LineartStandardPreprocessor": Lineart_Standard_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LineartStandardPreprocessor": "Standard Lineart"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Manga2Anime_LineArt_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.manga_line import LineartMangaDetector
|
||||
|
||||
model = LineartMangaDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Manga2Anime_LineArt_Preprocessor": Manga2Anime_LineArt_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Manga2Anime_LineArt_Preprocessor": "Manga Lineart (aka lineart_anime_denoise)"
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, run_script
|
||||
import comfy.model_management as model_management
|
||||
import os, sys
|
||||
import subprocess, threading
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import mediapipe
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'mediapipe'])
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', '--upgrade', 'protobuf'])
|
||||
|
||||
class Media_Pipe_Face_Mesh_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
max_faces=INPUT.INT(default=10, min=1, max=50), #Which image has more than 50 detectable faces?
|
||||
min_confidence=INPUT.FLOAT(default=0.5, min=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "detect"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def detect(self, image, max_faces=10, min_confidence=0.5, resolution=512):
|
||||
#Ref: https://github.com/Fannovel16/comfy_controlnet_preprocessors/issues/70#issuecomment-1677967369
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mediapipe_face import MediapipeFaceDetector
|
||||
return (common_annotator_call(MediapipeFaceDetector(), image, max_faces=max_faces, min_confidence=min_confidence, resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MediaPipe-FaceMeshPreprocessor": Media_Pipe_Face_Mesh_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MediaPipe-FaceMeshPreprocessor": "MediaPipe Face Mesh"
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION, run_script
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
import os, sys
|
||||
import subprocess, threading
|
||||
import scipy.ndimage
|
||||
import cv2
|
||||
import torch.nn.functional as F
|
||||
|
||||
def install_deps():
|
||||
try:
|
||||
import mediapipe
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'mediapipe'])
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', '--upgrade', 'protobuf'])
|
||||
|
||||
try:
|
||||
import trimesh
|
||||
except ImportError:
|
||||
run_script([sys.executable, '-s', '-m', 'pip', 'install', 'trimesh[easy]'])
|
||||
|
||||
#Sauce: https://github.com/comfyanonymous/ComfyUI/blob/8c6493578b3dda233e9b9a953feeaf1e6ca434ad/comfy_extras/nodes_mask.py#L309
|
||||
def expand_mask(mask, expand, tapered_corners):
|
||||
c = 0 if tapered_corners else 1
|
||||
kernel = np.array([[c, 1, c],
|
||||
[1, 1, 1],
|
||||
[c, 1, c]])
|
||||
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
|
||||
out = []
|
||||
for m in mask:
|
||||
output = m.numpy()
|
||||
for _ in range(abs(expand)):
|
||||
if expand < 0:
|
||||
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
||||
else:
|
||||
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
||||
output = torch.from_numpy(output)
|
||||
out.append(output)
|
||||
return torch.stack(out, dim=0)
|
||||
|
||||
class Mesh_Graphormer_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
mask_bbox_padding=("INT", {"default": 30, "min": 0, "max": 100}),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
mask_type=INPUT.COMBO(["based_on_depth", "tight_bboxes", "original"]),
|
||||
mask_expand=INPUT.INT(default=5, min=-MAX_RESOLUTION, max=MAX_RESOLUTION),
|
||||
rand_seed=INPUT.INT(default=88, min=0, max=0xffffffffffffffff),
|
||||
detect_thr=INPUT.FLOAT(default=0.6, min=0.1),
|
||||
presence_thr=INPUT.FLOAT(default=0.6, min=0.1)
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "INPAINTING_MASK")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, mask_bbox_padding=30, mask_type="based_on_depth", mask_expand=5, resolution=512, rand_seed=88, detect_thr=0.6, presence_thr=0.6, **kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mesh_graphormer import MeshGraphormerDetector
|
||||
model = kwargs["model"] if "model" in kwargs \
|
||||
else MeshGraphormerDetector.from_pretrained(detect_thr=detect_thr, presence_thr=presence_thr).to(model_management.get_torch_device())
|
||||
|
||||
depth_map_list = []
|
||||
mask_list = []
|
||||
for single_image in image:
|
||||
np_image = np.asarray(single_image.cpu() * 255., dtype=np.uint8)
|
||||
depth_map, mask, info = model(np_image, output_type="np", detect_resolution=resolution, mask_bbox_padding=mask_bbox_padding, seed=rand_seed)
|
||||
if mask_type == "based_on_depth":
|
||||
H, W = mask.shape[:2]
|
||||
mask = cv2.resize(depth_map.copy(), (W, H))
|
||||
mask[mask > 0] = 255
|
||||
|
||||
elif mask_type == "tight_bboxes":
|
||||
mask = np.zeros_like(mask)
|
||||
hand_bboxes = (info or {}).get("abs_boxes") or []
|
||||
for hand_bbox in hand_bboxes:
|
||||
x_min, x_max, y_min, y_max = hand_bbox
|
||||
mask[y_min:y_max+1, x_min:x_max+1, :] = 255 #HWC
|
||||
|
||||
mask = mask[:, :, :1]
|
||||
depth_map_list.append(torch.from_numpy(depth_map.astype(np.float32) / 255.0))
|
||||
mask_list.append(torch.from_numpy(mask.astype(np.float32) / 255.0))
|
||||
depth_maps, masks = torch.stack(depth_map_list, dim=0), rearrange(torch.stack(mask_list, dim=0), "n h w 1 -> n 1 h w")
|
||||
return depth_maps, expand_mask(masks, mask_expand, tapered_corners=True)
|
||||
|
||||
def normalize_size_base_64(w, h):
|
||||
short_side = min(w, h)
|
||||
remainder = short_side % 64
|
||||
return short_side - remainder + (64 if remainder > 0 else 0)
|
||||
|
||||
class Mesh_Graphormer_With_ImpactDetector_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
types = define_preprocessor_inputs(
|
||||
# Impact pack
|
||||
bbox_threshold=INPUT.FLOAT(default=0.5, min=0.1),
|
||||
bbox_dilation=INPUT.INT(default=10, min=-512, max=512),
|
||||
bbox_crop_factor=INPUT.FLOAT(default=3.0, min=1.0, max=10.0),
|
||||
drop_size=INPUT.INT(default=10, min=1, max=MAX_RESOLUTION),
|
||||
# Mesh Graphormer
|
||||
mask_bbox_padding=INPUT.INT(default=30, min=0, max=100),
|
||||
mask_type=INPUT.COMBO(["based_on_depth", "tight_bboxes", "original"]),
|
||||
mask_expand=INPUT.INT(default=5, min=-MAX_RESOLUTION, max=MAX_RESOLUTION),
|
||||
rand_seed=INPUT.INT(default=88, min=0, max=0xffffffffffffffff),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
types["required"]["bbox_detector"] = ("BBOX_DETECTOR", )
|
||||
return types
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("IMAGE", "INPAINTING_MASK")
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, bbox_detector, bbox_threshold=0.5, bbox_dilation=10, bbox_crop_factor=3.0, drop_size=10, resolution=512, **mesh_graphormer_kwargs):
|
||||
install_deps()
|
||||
from custom_controlnet_aux.mesh_graphormer import MeshGraphormerDetector
|
||||
mesh_graphormer_node = Mesh_Graphormer_Depth_Map_Preprocessor()
|
||||
model = MeshGraphormerDetector.from_pretrained(detect_thr=0.6, presence_thr=0.6).to(model_management.get_torch_device())
|
||||
mesh_graphormer_kwargs["model"] = model
|
||||
|
||||
frames = image
|
||||
depth_maps, masks = [], []
|
||||
for idx in range(len(frames)):
|
||||
frame = frames[idx:idx+1,...] #Impact Pack's BBOX_DETECTOR only supports single batch image
|
||||
bbox_detector.setAux('face') # make default prompt as 'face' if empty prompt for CLIPSeg
|
||||
_, segs = bbox_detector.detect(frame, bbox_threshold, bbox_dilation, bbox_crop_factor, drop_size)
|
||||
bbox_detector.setAux(None)
|
||||
|
||||
n, h, w, _ = frame.shape
|
||||
depth_map, mask = torch.zeros_like(frame), torch.zeros(n, 1, h, w)
|
||||
for i, seg in enumerate(segs):
|
||||
x1, y1, x2, y2 = seg.crop_region
|
||||
cropped_image = frame[:, y1:y2, x1:x2, :] # Never use seg.cropped_image to handle overlapping area
|
||||
mesh_graphormer_kwargs["resolution"] = 0 #Disable resizing
|
||||
sub_depth_map, sub_mask = mesh_graphormer_node.execute(cropped_image, **mesh_graphormer_kwargs)
|
||||
depth_map[:, y1:y2, x1:x2, :] = sub_depth_map
|
||||
mask[:, :, y1:y2, x1:x2] = sub_mask
|
||||
|
||||
depth_maps.append(depth_map)
|
||||
masks.append(mask)
|
||||
|
||||
return (torch.cat(depth_maps), torch.cat(masks))
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MeshGraphormer-DepthMapPreprocessor": Mesh_Graphormer_Depth_Map_Preprocessor,
|
||||
"MeshGraphormer+ImpactDetector-DepthMapPreprocessor": Mesh_Graphormer_With_ImpactDetector_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MeshGraphormer-DepthMapPreprocessor": "MeshGraphormer Hand Refiner",
|
||||
"MeshGraphormer+ImpactDetector-DepthMapPreprocessor": "MeshGraphormer Hand Refiner With External Detector"
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
# Disable NPU device initialization and problematic MMCV ops to prevent RuntimeError
|
||||
os.environ['NPU_DEVICE_COUNT'] = '0'
|
||||
os.environ['MMCV_WITH_OPS'] = '0'
|
||||
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Metric3D_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
backbone=INPUT.COMBO(["vit-small", "vit-large", "vit-giant2"]),
|
||||
fx=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
fy=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, backbone="vit-small", fx=1000, fy=1000, resolution=512):
|
||||
from custom_controlnet_aux.metric3d import Metric3DDetector
|
||||
model = Metric3DDetector.from_pretrained(filename=f"metric_depth_{backbone.replace('-', '_')}_800k.pth").to(model_management.get_torch_device())
|
||||
cb = lambda image, **kwargs: model(image, **kwargs)[0]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, fx=fx, fy=fy, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class Metric3D_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
backbone=INPUT.COMBO(["vit-small", "vit-large", "vit-giant2"]),
|
||||
fx=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
fy=INPUT.INT(default=1000, min=1, max=MAX_RESOLUTION),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, backbone="vit-small", fx=1000, fy=1000, resolution=512):
|
||||
from custom_controlnet_aux.metric3d import Metric3DDetector
|
||||
model = Metric3DDetector.from_pretrained(filename=f"metric_depth_{backbone.replace('-', '_')}_800k.pth").to(model_management.get_torch_device())
|
||||
cb = lambda image, **kwargs: model(image, **kwargs)[1]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, fx=fx, fy=fy, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Metric3D-DepthMapPreprocessor": Metric3D_Depth_Map_Preprocessor,
|
||||
"Metric3D-NormalMapPreprocessor": Metric3D_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Metric3D-DepthMapPreprocessor": "Metric3D Depth Map",
|
||||
"Metric3D-NormalMapPreprocessor": "Metric3D Normal Map"
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
|
||||
class MIDAS_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
a=INPUT.FLOAT(default=np.pi * 2.0, min=0.0, max=np.pi * 5.0),
|
||||
bg_threshold=INPUT.FLOAT(default=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, a=np.pi * 2.0, bg_threshold=0.1, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.midas import MidasDetector
|
||||
|
||||
model = MidasDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
#Dirty hack :))
|
||||
cb = lambda image, **kargs: model(image, **kargs)[1]
|
||||
out = common_annotator_call(cb, image, resolution=resolution, a=a, bg_th=bg_threshold, depth_and_normal=True)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
class MIDAS_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
a=INPUT.FLOAT(default=np.pi * 2.0, min=0.0, max=np.pi * 5.0),
|
||||
bg_threshold=INPUT.FLOAT(default=0.1),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, a=np.pi * 2.0, bg_threshold=0.1, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.midas import MidasDetector
|
||||
|
||||
# Ref: https://github.com/lllyasviel/ControlNet/blob/main/gradio_depth2image.py
|
||||
model = MidasDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, a=a, bg_th=bg_threshold)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MiDaS-NormalMapPreprocessor": MIDAS_Normal_Map_Preprocessor,
|
||||
"MiDaS-DepthMapPreprocessor": MIDAS_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MiDaS-NormalMapPreprocessor": "MiDaS Normal Map",
|
||||
"MiDaS-DepthMapPreprocessor": "MiDaS Depth Map"
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import numpy as np
|
||||
|
||||
class MLSD_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
score_threshold=INPUT.FLOAT(default=0.1, min=0.01, max=2.0),
|
||||
dist_threshold=INPUT.FLOAT(default=0.1, min=0.01, max=20.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, score_threshold, dist_threshold, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.mlsd import MLSDdetector
|
||||
|
||||
model = MLSDdetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, thr_v=score_threshold, thr_d=dist_threshold)
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"M-LSDPreprocessor": MLSD_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"M-LSDPreprocessor": "M-LSD Lines"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class BAE_Normal_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.normalbae import NormalBaeDetector
|
||||
|
||||
model = NormalBaeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BAE-NormalMapPreprocessor": BAE_Normal_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BAE-NormalMapPreprocessor": "BAE Normal Map"
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class OneFormer_COCO_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.oneformer import OneformerSegmentor
|
||||
|
||||
model = OneformerSegmentor.from_pretrained(filename="150_16_swin_l_oneformer_coco_100ep.pth")
|
||||
model = model.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
class OneFormer_ADE20K_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.oneformer import OneformerSegmentor
|
||||
|
||||
model = OneformerSegmentor.from_pretrained(filename="250_16_swin_l_oneformer_ade20k_160k.pth")
|
||||
model = model.to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out,)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OneFormer-COCO-SemSegPreprocessor": OneFormer_COCO_SemSegPreprocessor,
|
||||
"OneFormer-ADE20K-SemSegPreprocessor": OneFormer_ADE20K_SemSegPreprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OneFormer-COCO-SemSegPreprocessor": "OneFormer COCO Segmentor",
|
||||
"OneFormer-ADE20K-SemSegPreprocessor": "OneFormer ADE20K Segmentor"
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
import json
|
||||
|
||||
class OpenPose_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
detect_hand=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_body=INPUT.COMBO(["enable", "disable"]),
|
||||
detect_face=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
scale_stick_for_xinsr_cn=INPUT.COMBO(["disable", "enable"])
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
|
||||
FUNCTION = "estimate_pose"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Faces and Poses Estimators"
|
||||
|
||||
def estimate_pose(self, image, detect_hand="enable", detect_body="enable", detect_face="enable", scale_stick_for_xinsr_cn="disable", resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.open_pose import OpenposeDetector
|
||||
|
||||
detect_hand = detect_hand == "enable"
|
||||
detect_body = detect_body == "enable"
|
||||
detect_face = detect_face == "enable"
|
||||
scale_stick_for_xinsr_cn = scale_stick_for_xinsr_cn == "enable"
|
||||
|
||||
model = OpenposeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
self.openpose_dicts = []
|
||||
def func(image, **kwargs):
|
||||
pose_img, openpose_dict = model(image, **kwargs)
|
||||
self.openpose_dicts.append(openpose_dict)
|
||||
return pose_img
|
||||
|
||||
out = common_annotator_call(func, image, include_hand=detect_hand, include_face=detect_face, include_body=detect_body, image_and_json=True, xinsr_stick_scaling=scale_stick_for_xinsr_cn, resolution=resolution)
|
||||
del model
|
||||
return {
|
||||
'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
|
||||
"result": (out, self.openpose_dicts)
|
||||
}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OpenposePreprocessor": OpenPose_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OpenposePreprocessor": "OpenPose Pose",
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class PIDINET_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=INPUT.COMBO(["enable", "disable"]),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.pidi import PidiNetDetector
|
||||
|
||||
model = PidiNetDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe = safe == "enable")
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PiDiNetPreprocessor": PIDINET_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PiDiNetPreprocessor": "PiDiNet Soft-Edge Lines"
|
||||
}
|
||||
@@ -0,0 +1,398 @@
|
||||
import folder_paths
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import ImageColor
|
||||
from einops import rearrange
|
||||
import torch
|
||||
import itertools
|
||||
|
||||
from ..src.custom_controlnet_aux.dwpose import draw_poses, draw_animalposes, decode_json_as_poses, my_draw_poses, draw_poses_mask
|
||||
|
||||
|
||||
"""
|
||||
Format of POSE_KEYPOINT (AP10K keypoints):
|
||||
[{
|
||||
"version": "ap10k",
|
||||
"animals": [
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
[[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]],
|
||||
...
|
||||
],
|
||||
"canvas_height": 512,
|
||||
"canvas_width": 768
|
||||
},...]
|
||||
Format of POSE_KEYPOINT (OpenPose keypoints):
|
||||
[{
|
||||
"people": [
|
||||
{
|
||||
'pose_keypoints_2d': [[x1, y1, 1], [x2, y2, 1],..., [x17, y17, 1]]
|
||||
"face_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x68, y68, 1]],
|
||||
"hand_left_keypoints_2d": [[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]],
|
||||
"hand_right_keypoints_2d":[[x1, y1, 1], [x2, y2, 1],..., [x21, y21, 1]],
|
||||
}
|
||||
],
|
||||
"canvas_height": canvas_height,
|
||||
"canvas_width": canvas_width,
|
||||
},...]
|
||||
"""
|
||||
|
||||
class SavePoseKpsAsJsonFile:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pose_kps": ("POSE_KEYPOINT",),
|
||||
"filename_prefix": ("STRING", {"default": "PoseKeypoint"})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_pose_kps"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
def save_pose_kps(self, pose_kps, filename_prefix):
|
||||
filename_prefix += self.prefix_append
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
||||
folder_paths.get_save_image_path(filename_prefix, self.output_dir, pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"])
|
||||
file = f"{filename}_{counter:05}.json"
|
||||
with open(os.path.join(full_output_folder, file), 'w') as f:
|
||||
json.dump(pose_kps , f)
|
||||
return {}
|
||||
|
||||
#COCO-Wholebody doesn't have eyebrows as it inherits 68 keypoints format
|
||||
#Perhaps eyebrows can be estimated tho
|
||||
FACIAL_PARTS = ["skin", "left_eye", "right_eye", "nose", "upper_lip", "inner_mouth", "lower_lip"]
|
||||
LAPA_COLORS = dict(
|
||||
skin="rgb(0, 153, 255)",
|
||||
left_eye="rgb(0, 204, 153)",
|
||||
right_eye="rgb(255, 153, 0)",
|
||||
nose="rgb(255, 102, 255)",
|
||||
upper_lip="rgb(102, 0, 51)",
|
||||
inner_mouth="rgb(255, 204, 255)",
|
||||
lower_lip="rgb(255, 0, 102)"
|
||||
)
|
||||
|
||||
#One-based index
|
||||
def kps_idxs(start, end):
|
||||
step = -1 if start > end else 1
|
||||
return list(range(start-1, end+1-1, step))
|
||||
|
||||
#Source: https://www.researchgate.net/profile/Fabrizio-Falchi/publication/338048224/figure/fig1/AS:837860722741255@1576772971540/68-facial-landmarks.jpg
|
||||
FACIAL_PART_RANGES = dict(
|
||||
skin=kps_idxs(1, 17) + kps_idxs(27, 18),
|
||||
nose=kps_idxs(28, 36),
|
||||
left_eye=kps_idxs(37, 42),
|
||||
right_eye=kps_idxs(43, 48),
|
||||
upper_lip=kps_idxs(49, 55) + kps_idxs(65, 61),
|
||||
lower_lip=kps_idxs(61, 68),
|
||||
inner_mouth=kps_idxs(61, 65) + kps_idxs(55, 49)
|
||||
)
|
||||
|
||||
def is_normalized(keypoints) -> bool:
|
||||
point_normalized = [
|
||||
0 <= np.abs(k[0]) <= 1 and 0 <= np.abs(k[1]) <= 1
|
||||
for k in keypoints
|
||||
if k is not None
|
||||
]
|
||||
if not point_normalized:
|
||||
return False
|
||||
return np.all(point_normalized)
|
||||
|
||||
class FacialPartColoringFromPoseKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
input_types = {
|
||||
"required": {"pose_kps": ("POSE_KEYPOINT",), "mode": (["point", "polygon"], {"default": "polygon"})}
|
||||
}
|
||||
for facial_part in FACIAL_PARTS:
|
||||
input_types["required"][facial_part] = ("STRING", {"default": LAPA_COLORS[facial_part], "multiline": False})
|
||||
return input_types
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "colorize"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
def colorize(self, pose_kps, mode, **facial_part_colors):
|
||||
pose_frames = pose_kps
|
||||
np_frames = [self.draw_kps(pose_frame, mode, **facial_part_colors) for pose_frame in pose_frames]
|
||||
np_frames = np.stack(np_frames, axis=0)
|
||||
return (torch.from_numpy(np_frames).float() / 255.,)
|
||||
|
||||
def draw_kps(self, pose_frame, mode, **facial_part_colors):
|
||||
width, height = pose_frame["canvas_width"], pose_frame["canvas_height"]
|
||||
canvas = np.zeros((height, width, 3), dtype=np.uint8)
|
||||
for person, part_name in itertools.product(pose_frame["people"], FACIAL_PARTS):
|
||||
n = len(person["face_keypoints_2d"]) // 3
|
||||
facial_kps = rearrange(np.array(person["face_keypoints_2d"]), "(n c) -> n c", n=n, c=3)[:, :2]
|
||||
if is_normalized(facial_kps):
|
||||
facial_kps *= (width, height)
|
||||
facial_kps = facial_kps.astype(np.int32)
|
||||
part_color = ImageColor.getrgb(facial_part_colors[part_name])[:3]
|
||||
part_contours = facial_kps[FACIAL_PART_RANGES[part_name], :]
|
||||
if mode == "point":
|
||||
for pt in part_contours:
|
||||
cv2.circle(canvas, pt, radius=2, color=part_color, thickness=-1)
|
||||
else:
|
||||
cv2.fillPoly(canvas, pts=[part_contours], color=part_color)
|
||||
return canvas
|
||||
|
||||
# https://raw.githubusercontent.com/CMU-Perceptual-Computing-Lab/openpose/master/.github/media/keypoints_pose_18.png
|
||||
BODY_PART_INDEXES = {
|
||||
"Head": (16, 14, 0, 15, 17),
|
||||
"Neck": (0, 1),
|
||||
"Shoulder": (2, 5),
|
||||
"Torso": (2, 5, 8, 11),
|
||||
"RArm": (2, 3),
|
||||
"RForearm": (3, 4),
|
||||
"LArm": (5, 6),
|
||||
"LForearm": (6, 7),
|
||||
"RThigh": (8, 9),
|
||||
"RLeg": (9, 10),
|
||||
"LThigh": (11, 12),
|
||||
"LLeg": (12, 13)
|
||||
}
|
||||
BODY_PART_DEFAULT_W_H = {
|
||||
"Head": "256, 256",
|
||||
"Neck": "100, 100",
|
||||
"Shoulder": '',
|
||||
"Torso": "350, 450",
|
||||
"RArm": "128, 256",
|
||||
"RForearm": "128, 256",
|
||||
"LArm": "128, 256",
|
||||
"LForearm": "128, 256",
|
||||
"RThigh": "128, 256",
|
||||
"RLeg": "128, 256",
|
||||
"LThigh": "128, 256",
|
||||
"LLeg": "128, 256"
|
||||
}
|
||||
|
||||
class SinglePersonProcess:
|
||||
@classmethod
|
||||
def sort_and_get_max_people(s, pose_kps):
|
||||
for idx in range(len(pose_kps)):
|
||||
pose_kps[idx]["people"] = sorted(pose_kps[idx]["people"], key=lambda person:person["pose_keypoints_2d"][0])
|
||||
return pose_kps, max(len(frame["people"]) for frame in pose_kps)
|
||||
|
||||
def __init__(self, pose_kps, person_idx=0) -> None:
|
||||
self.width, self.height = pose_kps[0]["canvas_width"], pose_kps[0]["canvas_height"]
|
||||
self.poses = [
|
||||
self.normalize(pose_frame["people"][person_idx]["pose_keypoints_2d"])
|
||||
if person_idx < len(pose_frame["people"])
|
||||
else None
|
||||
for pose_frame in pose_kps
|
||||
]
|
||||
|
||||
def normalize(self, pose_kps_2d):
|
||||
n = len(pose_kps_2d) // 3
|
||||
pose_kps_2d = rearrange(np.array(pose_kps_2d), "(n c) -> n c", n=n, c=3)
|
||||
pose_kps_2d[np.argwhere(pose_kps_2d[:,2]==0), :] = np.iinfo(np.int32).max // 2 #Safe large value
|
||||
pose_kps_2d = pose_kps_2d[:, :2]
|
||||
if is_normalized(pose_kps_2d):
|
||||
pose_kps_2d *= (self.width, self.height)
|
||||
return pose_kps_2d
|
||||
|
||||
def get_xyxy_bboxes(self, part_name, bbox_size=(128, 256)):
|
||||
width, height = bbox_size
|
||||
xyxy_bboxes = {}
|
||||
for idx, pose in enumerate(self.poses):
|
||||
if pose is None:
|
||||
xyxy_bboxes[idx] = (np.iinfo(np.int32).max // 2,) * 4
|
||||
continue
|
||||
pts = pose[BODY_PART_INDEXES[part_name], :]
|
||||
|
||||
#top_left = np.min(pts[:,0]), np.min(pts[:,1])
|
||||
#bottom_right = np.max(pts[:,0]), np.max(pts[:,1])
|
||||
#pad_width = np.maximum(width - (bottom_right[0]-top_left[0]), 0) / 2
|
||||
#pad_height = np.maximum(height - (bottom_right[1]-top_left[1]), 0) / 2
|
||||
#xyxy_bboxes.append((
|
||||
# top_left[0] - pad_width, top_left[1] - pad_height,
|
||||
# bottom_right[0] + pad_width, bottom_right[1] + pad_height,
|
||||
#))
|
||||
|
||||
x_mid, y_mid = np.mean(pts[:, 0]), np.mean(pts[:, 1])
|
||||
xyxy_bboxes[idx] = (
|
||||
x_mid - width/2, y_mid - height/2,
|
||||
x_mid + width/2, y_mid + height/2
|
||||
)
|
||||
return xyxy_bboxes
|
||||
|
||||
class UpperBodyTrackingFromPoseKps:
|
||||
PART_NAMES = ["Head", "Neck", "Shoulder", "Torso", "RArm", "RForearm", "LArm", "LForearm"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pose_kps": ("POSE_KEYPOINT",),
|
||||
"id_include": ("STRING", {"default": '', "multiline": False}),
|
||||
**{part_name + "_width_height": ("STRING", {"default": BODY_PART_DEFAULT_W_H[part_name], "multiline": False}) for part_name in s.PART_NAMES}
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("TRACKING", "STRING")
|
||||
RETURN_NAMES = ("tracking", "prompt")
|
||||
FUNCTION = "convert"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def convert(self, pose_kps, id_include, **parts_width_height):
|
||||
parts_width_height = {part_name.replace("_width_height", ''): value for part_name, value in parts_width_height.items()}
|
||||
enabled_part_names = [part_name for part_name in self.PART_NAMES if len(parts_width_height[part_name].strip())]
|
||||
tracked = {part_name: {} for part_name in enabled_part_names}
|
||||
id_include = id_include.strip()
|
||||
id_include = list(map(int, id_include.split(','))) if len(id_include) else []
|
||||
prompt_string = ''
|
||||
pose_kps, max_people = SinglePersonProcess.sort_and_get_max_people(pose_kps)
|
||||
|
||||
for person_idx in range(max_people):
|
||||
if len(id_include) and person_idx not in id_include:
|
||||
continue
|
||||
processor = SinglePersonProcess(pose_kps, person_idx)
|
||||
for part_name in enabled_part_names:
|
||||
bbox_size = tuple(map(int, parts_width_height[part_name].split(',')))
|
||||
part_bboxes = processor.get_xyxy_bboxes(part_name, bbox_size)
|
||||
id_coordinates = {idx: part_bbox+(processor.width, processor.height) for idx, part_bbox in part_bboxes.items()}
|
||||
tracked[part_name][person_idx] = id_coordinates
|
||||
|
||||
for class_name, class_data in tracked.items():
|
||||
for class_id in class_data.keys():
|
||||
class_id_str = str(class_id)
|
||||
# Use the incoming prompt for each class name and ID
|
||||
_class_name = class_name.replace('L', '').replace('R', '').lower()
|
||||
prompt_string += f'"{class_id_str}.{class_name}": "({_class_name})",\n'
|
||||
|
||||
return (tracked, prompt_string)
|
||||
|
||||
|
||||
def numpy2torch(np_image: np.ndarray) -> torch.Tensor:
|
||||
""" [H, W, C] => [B=1, H, W, C]"""
|
||||
return torch.from_numpy(np_image.astype(np.float32) / 255).unsqueeze(0)
|
||||
|
||||
class RenderPeopleKpsMask:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
"changdu": (["yao", "kuan", "da_tui", "xi_gai", "xiao_tui", "jiao_huai"], {"default": "kuan"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps, changdu) -> tuple[torch.Tensor]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
poses, _, height, width = decode_json_as_poses(kps)
|
||||
mask = draw_poses_mask(
|
||||
poses,
|
||||
height,
|
||||
width,
|
||||
changdu,
|
||||
)
|
||||
return (mask,)
|
||||
|
||||
class RenderPeopleKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
"render_body": ("BOOLEAN", {"default": True}),
|
||||
"render_hand": ("BOOLEAN", {"default": True}),
|
||||
"render_face": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps, render_body, render_hand, render_face) -> tuple[np.ndarray]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
poses, _, height, width = decode_json_as_poses(kps)
|
||||
np_image = draw_poses(
|
||||
poses,
|
||||
height,
|
||||
width,
|
||||
render_body,
|
||||
render_hand,
|
||||
render_face,
|
||||
)
|
||||
return (numpy2torch(np_image),)
|
||||
|
||||
class MyRenderPeopleKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps) -> tuple[np.ndarray]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
poses, _, height, width = decode_json_as_poses(kps)
|
||||
np_image = my_draw_poses(
|
||||
poses,
|
||||
height,
|
||||
width,
|
||||
True,
|
||||
False,
|
||||
False,
|
||||
)
|
||||
return (numpy2torch(np_image),)
|
||||
|
||||
class RenderAnimalKps:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"kps": ("POSE_KEYPOINT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "render"
|
||||
CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
|
||||
|
||||
def render(self, kps) -> tuple[np.ndarray]:
|
||||
if isinstance(kps, list):
|
||||
kps = kps[0]
|
||||
|
||||
_, poses, height, width = decode_json_as_poses(kps)
|
||||
np_image = draw_animalposes(poses, height, width)
|
||||
return (numpy2torch(np_image),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SavePoseKpsAsJsonFile": SavePoseKpsAsJsonFile,
|
||||
"FacialPartColoringFromPoseKps": FacialPartColoringFromPoseKps,
|
||||
"UpperBodyTrackingFromPoseKps": UpperBodyTrackingFromPoseKps,
|
||||
"RenderPeopleKps": RenderPeopleKps,
|
||||
"MyRenderPeopleKps": MyRenderPeopleKps,
|
||||
"RenderPeopleKpsMask": RenderPeopleKpsMask,
|
||||
"RenderAnimalKps": RenderAnimalKps,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SavePoseKpsAsJsonFile": "Save Pose Keypoints",
|
||||
"FacialPartColoringFromPoseKps": "Colorize Facial Parts from PoseKPS",
|
||||
"UpperBodyTrackingFromPoseKps": "Upper Body Tracking From PoseKps (InstanceDiffusion)",
|
||||
"RenderPeopleKps": "Render Pose JSON (Human)",
|
||||
"MyRenderPeopleKps": "My RenderPeopleKps",
|
||||
"RenderPeopleKpsMask": "My RenderPeopleKps Mask",
|
||||
"RenderAnimalKps": "Render Pose JSON (Animal)",
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, INPUT, define_preprocessor_inputs
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class PyraCanny_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
low_threshold=INPUT.INT(default=64, max=255),
|
||||
high_threshold=INPUT.INT(default=128, max=255),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, low_threshold=64, high_threshold=128, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.pyracanny import PyraCannyDetector
|
||||
|
||||
return (common_annotator_call(PyraCannyDetector(), image, low_threshold=low_threshold, high_threshold=high_threshold, resolution=resolution), )
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PyraCannyPreprocessor": PyraCanny_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PyraCannyPreprocessor": "PyraCanny"
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
|
||||
class ImageLuminanceDetector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#https://github.com/Mikubill/sd-webui-controlnet/blob/416c345072c9c2066101e225964e3986abe6945e/scripts/processor.py#L1229
|
||||
return define_preprocessor_inputs(
|
||||
gamma_correction=INPUT.FLOAT(default=1.0, min=0.1, max=2.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Recolor"
|
||||
|
||||
def execute(self, image, gamma_correction=1.0, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.recolor import Recolorizer
|
||||
return (common_annotator_call(Recolorizer(), image, mode="luminance", gamma_correction=gamma_correction , resolution=resolution), )
|
||||
|
||||
class ImageIntensityDetector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
#https://github.com/Mikubill/sd-webui-controlnet/blob/416c345072c9c2066101e225964e3986abe6945e/scripts/processor.py#L1229
|
||||
return define_preprocessor_inputs(
|
||||
gamma_correction=INPUT.FLOAT(default=1.0, min=0.1, max=2.0),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Recolor"
|
||||
|
||||
def execute(self, image, gamma_correction=1.0, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.recolor import Recolorizer
|
||||
return (common_annotator_call(Recolorizer(), image, mode="intensity", gamma_correction=gamma_correction , resolution=resolution), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageLuminanceDetector": ImageLuminanceDetector,
|
||||
"ImageIntensityDetector": ImageIntensityDetector
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageLuminanceDetector": "Image Luminance",
|
||||
"ImageIntensityDetector": "Image Intensity"
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, nms
|
||||
import comfy.model_management as model_management
|
||||
import cv2
|
||||
|
||||
class Scribble_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.scribble import ScribbleDetector
|
||||
|
||||
model = ScribbleDetector()
|
||||
return (common_annotator_call(model, image, resolution=resolution), )
|
||||
|
||||
class Scribble_XDoG_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
threshold=INPUT.INT(default=32, min=1, max=64),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, threshold=32, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.scribble import ScribbleXDog_Detector
|
||||
|
||||
model = ScribbleXDog_Detector()
|
||||
return (common_annotator_call(model, image, resolution=resolution, thr_a=threshold), )
|
||||
|
||||
class Scribble_PiDiNet_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe=(["enable", "disable"],),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe="enable", resolution=512):
|
||||
def model(img, **kwargs):
|
||||
from custom_controlnet_aux.pidi import PidiNetDetector
|
||||
pidinet = PidiNetDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
result = pidinet(img, scribble=True, **kwargs)
|
||||
result = nms(result, 127, 3.0)
|
||||
result = cv2.GaussianBlur(result, (0, 0), 3.0)
|
||||
result[result > 4] = 255
|
||||
result[result < 255] = 0
|
||||
return result
|
||||
return (common_annotator_call(model, image, resolution=resolution, safe=safe=="enable"),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ScribblePreprocessor": Scribble_Preprocessor,
|
||||
"Scribble_XDoG_Preprocessor": Scribble_XDoG_Preprocessor,
|
||||
"Scribble_PiDiNet_Preprocessor": Scribble_PiDiNet_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ScribblePreprocessor": "Scribble Lines",
|
||||
"Scribble_XDoG_Preprocessor": "Scribble XDoG Lines",
|
||||
"Scribble_PiDiNet_Preprocessor": "Scribble PiDiNet Lines"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class SAM_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/others"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.sam import SamDetector
|
||||
|
||||
mobile_sam = SamDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(mobile_sam, image, resolution=resolution)
|
||||
del mobile_sam
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SAMPreprocessor": SAM_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SAMPreprocessor": "SAM Segmentor"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT, MAX_RESOLUTION
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Shuffle_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
resolution=INPUT.RESOLUTION(),
|
||||
seed=INPUT.SEED()
|
||||
)
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "preprocess"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/T2IAdapter-only"
|
||||
|
||||
def preprocess(self, image, resolution=512, seed=0):
|
||||
from custom_controlnet_aux.shuffle import ContentShuffleDetector
|
||||
|
||||
return (common_annotator_call(ContentShuffleDetector(), image, resolution=resolution, seed=seed), )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ShufflePreprocessor": Shuffle_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ShufflePreprocessor": "Content Shuffle"
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class TEED_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
safe_steps=INPUT.INT(default=2, max=10),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Line Extractors"
|
||||
|
||||
def execute(self, image, safe_steps=2, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.teed import TEDDetector
|
||||
|
||||
model = TEDDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution, safe_steps=safe_steps)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TEEDPreprocessor": TEED_Preprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TEED_Preprocessor": "TEED Soft-Edge Lines",
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
|
||||
|
||||
class Tile_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
pyrUp_iters=INPUT.INT(default=3, min=1, max=10),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, pyrUp_iters, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.tile import TileDetector
|
||||
|
||||
return (common_annotator_call(TileDetector(), image, pyrUp_iters=pyrUp_iters, resolution=resolution),)
|
||||
|
||||
class TTPlanet_TileGF_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
scale_factor=INPUT.FLOAT(default=1.00, min=1.000, max=8.00),
|
||||
blur_strength=INPUT.FLOAT(default=2.0, min=1.0, max=10.0),
|
||||
radius=INPUT.INT(default=7, min=1, max=20),
|
||||
eps=INPUT.FLOAT(default=0.01, min=0.001, max=0.1, step=0.001),
|
||||
resolution=INPUT.RESOLUTION()
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, scale_factor, blur_strength, radius, eps, **kwargs):
|
||||
from custom_controlnet_aux.tile import TTPlanet_Tile_Detector_GF
|
||||
|
||||
return (common_annotator_call(TTPlanet_Tile_Detector_GF(), image, scale_factor=scale_factor, blur_strength=blur_strength, radius=radius, eps=eps),)
|
||||
|
||||
class TTPlanet_TileSimple_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(
|
||||
scale_factor=INPUT.FLOAT(default=1.00, min=1.000, max=8.00),
|
||||
blur_strength=INPUT.FLOAT(default=2.0, min=1.0, max=10.0),
|
||||
)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/tile"
|
||||
|
||||
def execute(self, image, scale_factor, blur_strength):
|
||||
from custom_controlnet_aux.tile import TTPLanet_Tile_Detector_Simple
|
||||
|
||||
return (common_annotator_call(TTPLanet_Tile_Detector_Simple(), image, scale_factor=scale_factor, blur_strength=blur_strength),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"TilePreprocessor": Tile_Preprocessor,
|
||||
"TTPlanet_TileGF_Preprocessor": TTPlanet_TileGF_Preprocessor,
|
||||
"TTPlanet_TileSimple_Preprocessor": TTPlanet_TileSimple_Preprocessor
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TilePreprocessor": "Tile",
|
||||
"TTPlanet_TileGF_Preprocessor": "TTPlanet Tile GuidedFilter",
|
||||
"TTPlanet_TileSimple_Preprocessor": "TTPlanet Tile Simple"
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
# Disable NPU device initialization and problematic MMCV ops to prevent RuntimeError
|
||||
os.environ['NPU_DEVICE_COUNT'] = '0'
|
||||
os.environ['MMCV_WITH_OPS'] = '0'
|
||||
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Uniformer_SemSegPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "semantic_segmentate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Semantic Segmentation"
|
||||
|
||||
def semantic_segmentate(self, image, resolution=512):
|
||||
from custom_controlnet_aux.uniformer import UniformerSegmentor
|
||||
|
||||
model = UniformerSegmentor.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UniFormer-SemSegPreprocessor": Uniformer_SemSegPreprocessor,
|
||||
"SemSegPreprocessor": Uniformer_SemSegPreprocessor,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UniFormer-SemSegPreprocessor": "UniFormer Segmentor",
|
||||
"SemSegPreprocessor": "Semantic Segmentor (legacy, alias for UniFormer)",
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
from ..utils import common_annotator_call
|
||||
import comfy.model_management as model_management
|
||||
import torch
|
||||
import numpy as np
|
||||
from einops import rearrange
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Unimatch_OptFlowPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": dict(
|
||||
image=("IMAGE",),
|
||||
ckpt_name=(
|
||||
["gmflow-scale1-mixdata.pth", "gmflow-scale2-mixdata.pth", "gmflow-scale2-regrefine6-mixdata.pth"],
|
||||
{"default": "gmflow-scale2-regrefine6-mixdata.pth"}
|
||||
),
|
||||
backward_flow=("BOOLEAN", {"default": False}),
|
||||
bidirectional_flow=("BOOLEAN", {"default": False})
|
||||
)
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTICAL_FLOW", "IMAGE")
|
||||
RETURN_NAMES = ("OPTICAL_FLOW", "PREVIEW_IMAGE")
|
||||
FUNCTION = "estimate"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Optical Flow"
|
||||
|
||||
def estimate(self, image, ckpt_name, backward_flow=False, bidirectional_flow=False):
|
||||
assert len(image) > 1, "[Unimatch] Requiring as least two frames as an optical flow estimator. Only use this node on video input."
|
||||
from custom_controlnet_aux.unimatch import UnimatchDetector
|
||||
tensor_images = image
|
||||
model = UnimatchDetector.from_pretrained(filename=ckpt_name).to(model_management.get_torch_device())
|
||||
flows, vis_flows = [], []
|
||||
for i in range(len(tensor_images) - 1):
|
||||
image0, image1 = np.asarray(image[i:i+2].cpu() * 255., dtype=np.uint8)
|
||||
flow, vis_flow = model(image0, image1, output_type="np", pred_bwd_flow=backward_flow, pred_bidir_flow=bidirectional_flow)
|
||||
flows.append(torch.from_numpy(flow).float())
|
||||
vis_flows.append(torch.from_numpy(vis_flow).float() / 255.)
|
||||
del model
|
||||
return (torch.stack(flows, dim=0), torch.stack(vis_flows, dim=0))
|
||||
|
||||
class MaskOptFlow:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": dict(optical_flow=("OPTICAL_FLOW",), mask=("MASK",))
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OPTICAL_FLOW", "IMAGE")
|
||||
RETURN_NAMES = ("OPTICAL_FLOW", "PREVIEW_IMAGE")
|
||||
FUNCTION = "mask_opt_flow"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Optical Flow"
|
||||
|
||||
def mask_opt_flow(self, optical_flow, mask):
|
||||
from custom_controlnet_aux.unimatch import flow_to_image
|
||||
assert len(mask) >= len(optical_flow), f"Not enough masks to mask optical flow: {len(mask)} vs {len(optical_flow)}"
|
||||
mask = mask[:optical_flow.shape[0]]
|
||||
mask = F.interpolate(mask, optical_flow.shape[1:3])
|
||||
mask = rearrange(mask, "n 1 h w -> n h w 1")
|
||||
vis_flows = torch.stack([torch.from_numpy(flow_to_image(flow)).float() / 255. for flow in optical_flow.numpy()], dim=0)
|
||||
vis_flows *= mask
|
||||
optical_flow *= mask
|
||||
return (optical_flow, vis_flows)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Unimatch_OptFlowPreprocessor": Unimatch_OptFlowPreprocessor,
|
||||
"MaskOptFlow": MaskOptFlow
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Unimatch_OptFlowPreprocessor": "Unimatch Optical Flow",
|
||||
"MaskOptFlow": "Mask Optical Flow (DragNUWA)"
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
from ..utils import common_annotator_call, define_preprocessor_inputs, INPUT
|
||||
import comfy.model_management as model_management
|
||||
|
||||
class Zoe_Depth_Map_Preprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return define_preprocessor_inputs(resolution=INPUT.RESOLUTION())
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
|
||||
|
||||
def execute(self, image, resolution=512, **kwargs):
|
||||
from custom_controlnet_aux.zoe import ZoeDetector
|
||||
|
||||
model = ZoeDetector.from_pretrained().to(model_management.get_torch_device())
|
||||
out = common_annotator_call(model, image, resolution=resolution)
|
||||
del model
|
||||
return (out, )
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Zoe-DepthMapPreprocessor": Zoe_Depth_Map_Preprocessor
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Zoe-DepthMapPreprocessor": "Zoe Depth Map"
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
[project]
|
||||
name = "comfyui_controlnet_aux"
|
||||
description = "Plug-and-play ComfyUI node sets for making ControlNet hint images"
|
||||
|
||||
version = "1.1.0"
|
||||
dependencies = ["torch", "importlib_metadata", "huggingface_hub", "scipy", "opencv-python>=4.7.0.72", "filelock", "numpy", "Pillow", "einops", "torchvision", "pyyaml", "scikit-image", "python-dateutil", "mediapipe", "svglib", "fvcore", "yapf", "omegaconf", "ftfy", "addict", "yacs", "trimesh[easy]", "albumentations", "scikit-learn", "matplotlib"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/Fannovel16/comfyui_controlnet_aux"
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "fannovel16"
|
||||
DisplayName = "comfyui_controlnet_aux"
|
||||
Icon = ""
|
||||
@@ -0,0 +1,26 @@
|
||||
torch
|
||||
importlib_metadata
|
||||
huggingface_hub
|
||||
scipy
|
||||
opencv-python>=4.7.0.72
|
||||
filelock
|
||||
numpy
|
||||
Pillow
|
||||
einops
|
||||
torchvision
|
||||
pyyaml
|
||||
scikit-image
|
||||
python-dateutil
|
||||
mediapipe
|
||||
svglib
|
||||
fvcore
|
||||
yapf
|
||||
omegaconf
|
||||
ftfy
|
||||
addict
|
||||
yacs
|
||||
yapf
|
||||
trimesh[easy]
|
||||
albumentations
|
||||
scikit-learn
|
||||
matplotlib
|
||||
@@ -0,0 +1,56 @@
|
||||
from pathlib import Path
|
||||
import os
|
||||
import re
|
||||
#Thanks ChatGPT
|
||||
pattern = r'\bfrom_pretrained\(.*?pretrained_model_or_path\s*=\s*(.*?)(?:,|\))|filename\s*=\s*(.*?)(?:,|\))|(\w+_filename)\s*=\s*(.*?)(?:,|\))'
|
||||
aux_dir = Path(__file__).parent / 'src' / 'custom_controlnet_aux'
|
||||
VAR_DICT = dict(
|
||||
HF_MODEL_NAME = "lllyasviel/Annotators",
|
||||
DWPOSE_MODEL_NAME = "yzd-v/DWPose",
|
||||
BDS_MODEL_NAME = "bdsqlsz/qinglong_controlnet-lllite",
|
||||
DENSEPOSE_MODEL_NAME = "LayerNorm/DensePose-TorchScript-with-hint-image",
|
||||
MESH_GRAPHORMER_MODEL_NAME = "hr16/ControlNet-HandRefiner-pruned",
|
||||
SAM_MODEL_NAME = "dhkim2810/MobileSAM",
|
||||
UNIMATCH_MODEL_NAME = "hr16/Unimatch",
|
||||
DEPTH_ANYTHING_MODEL_NAME = "LiheYoung/Depth-Anything", #HF Space
|
||||
DIFFUSION_EDGE_MODEL_NAME = "hr16/Diffusion-Edge"
|
||||
)
|
||||
re_result_dict = {}
|
||||
for preprocc in os.listdir(aux_dir):
|
||||
if preprocc in ["__pycache__", 'tests']: continue
|
||||
if '.py' in preprocc: continue
|
||||
f = open(aux_dir / preprocc / '__init__.py', 'r')
|
||||
code = f.read()
|
||||
matches = re.findall(pattern, code)
|
||||
result = [match[0] or match[1] or match[3] for match in matches]
|
||||
if not len(result):
|
||||
print(preprocc)
|
||||
continue
|
||||
result = [el.replace("'", '').replace('"', '') for el in result]
|
||||
result = [VAR_DICT.get(el, el) for el in result]
|
||||
re_result_dict[preprocc] = result
|
||||
f.close()
|
||||
|
||||
for preprocc, re_result in re_result_dict.items():
|
||||
model_name, filenames = re_result[0], re_result[1:]
|
||||
print(f"* {preprocc}: ", end=' ')
|
||||
assests_md = ', '.join([f"[{model_name}/{filename}](https://huggingface.co/{model_name}/blob/main/{filename})" for filename in filenames])
|
||||
print(assests_md)
|
||||
|
||||
preprocc = "dwpose"
|
||||
model_name, filenames = VAR_DICT['DWPOSE_MODEL_NAME'], ["yolox_l.onnx", "dw-ll_ucoco_384.onnx"]
|
||||
print(f"* {preprocc}: ", end=' ')
|
||||
assests_md = ', '.join([f"[{model_name}/{filename}](https://huggingface.co/{model_name}/blob/main/{filename})" for filename in filenames])
|
||||
print(assests_md)
|
||||
|
||||
preprocc = "yolo-nas"
|
||||
model_name, filenames = "hr16/yolo-nas-fp16", ["yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"]
|
||||
print(f"* {preprocc}: ", end=' ')
|
||||
assests_md = ', '.join([f"[{model_name}/{filename}](https://huggingface.co/{model_name}/blob/main/{filename})" for filename in filenames])
|
||||
print(assests_md)
|
||||
|
||||
preprocc = "dwpose-torchscript"
|
||||
model_name, filenames = "hr16/DWPose-TorchScript-BatchSize5", ["dw-ll_ucoco_384_bs5.torchscript.pt", "rtmpose-m_ap10k_256_bs5.torchscript.pt"]
|
||||
print(f"* {preprocc}: ", end=' ')
|
||||
assests_md = ', '.join([f"[{model_name}/{filename}](https://huggingface.co/{model_name}/blob/main/{filename})" for filename in filenames])
|
||||
print(assests_md)
|
||||
@@ -0,0 +1 @@
|
||||
#Dummy file ensuring this package will be recognized
|
||||
@@ -0,0 +1 @@
|
||||
#Dummy file ensuring this package will be recognized
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,66 @@
|
||||
from .network import UNet
|
||||
from .util import seg2img
|
||||
import torch
|
||||
import os
|
||||
import cv2
|
||||
from custom_controlnet_aux.util import HWC3, resize_image_with_pad, common_input_validate, custom_hf_download, BDS_MODEL_NAME
|
||||
from huggingface_hub import hf_hub_download
|
||||
from PIL import Image
|
||||
from einops import rearrange
|
||||
from .anime_segmentation import AnimeSegmentation
|
||||
import numpy as np
|
||||
|
||||
class AnimeFaceSegmentor:
|
||||
def __init__(self, model, seg_model):
|
||||
self.model = model
|
||||
self.seg_model = seg_model
|
||||
self.device = "cpu"
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_or_path=BDS_MODEL_NAME, filename="UNet.pth", seg_filename="isnetis.ckpt"):
|
||||
model_path = custom_hf_download(pretrained_model_or_path, filename, subfolder="Annotators")
|
||||
seg_model_path = custom_hf_download("skytnt/anime-seg", seg_filename)
|
||||
|
||||
model = UNet()
|
||||
ckpt = torch.load(model_path, map_location="cpu")
|
||||
model.load_state_dict(ckpt)
|
||||
model.eval()
|
||||
|
||||
seg_model = AnimeSegmentation(seg_model_path)
|
||||
seg_model.net.eval()
|
||||
return cls(model, seg_model)
|
||||
|
||||
def to(self, device):
|
||||
self.model.to(device)
|
||||
self.seg_model.net.to(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def __call__(self, input_image, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", remove_background=True, **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
|
||||
with torch.no_grad():
|
||||
if remove_background:
|
||||
print(input_image.shape)
|
||||
mask, input_image = self.seg_model(input_image, 0) #Don't resize image as it is resized
|
||||
image_feed = torch.from_numpy(input_image).float().to(self.device)
|
||||
image_feed = rearrange(image_feed, 'h w c -> 1 c h w')
|
||||
image_feed = image_feed / 255
|
||||
seg = self.model(image_feed).squeeze(dim=0)
|
||||
result = seg2img(seg.cpu().detach().numpy())
|
||||
|
||||
detected_map = HWC3(result)
|
||||
detected_map = remove_pad(detected_map)
|
||||
if remove_background:
|
||||
mask = remove_pad(mask)
|
||||
H, W, C = detected_map.shape
|
||||
tmp = np.zeros([H, W, C + 1])
|
||||
tmp[:,:,:C] = detected_map
|
||||
tmp[:,:,3:] = mask
|
||||
detected_map = tmp
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map[..., :3])
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,58 @@
|
||||
#https://github.com/SkyTNT/anime-segmentation/tree/main
|
||||
#Only adapt isnet_is (https://huggingface.co/skytnt/anime-seg/blob/main/isnetis.ckpt)
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
from .isnet import ISNetDIS
|
||||
import numpy as np
|
||||
import cv2
|
||||
from comfy.model_management import get_torch_device
|
||||
DEVICE = get_torch_device()
|
||||
|
||||
class AnimeSegmentation:
|
||||
def __init__(self, ckpt_path):
|
||||
super(AnimeSegmentation).__init__()
|
||||
sd = torch.load(ckpt_path, map_location="cpu")
|
||||
self.net = ISNetDIS()
|
||||
#gt_encoder isn't used during inference
|
||||
self.net.load_state_dict({k.replace("net.", ''):v for k, v in sd.items() if k.startswith("net.")})
|
||||
self.net = self.net.to(DEVICE)
|
||||
self.net.eval()
|
||||
|
||||
def get_mask(self, input_img, s=640):
|
||||
input_img = (input_img / 255).astype(np.float32)
|
||||
if s == 0:
|
||||
img_input = np.transpose(input_img, (2, 0, 1))
|
||||
img_input = img_input[np.newaxis, :]
|
||||
tmpImg = torch.from_numpy(img_input).float().to(DEVICE)
|
||||
with torch.no_grad():
|
||||
pred = self.net(tmpImg)[0][0].sigmoid() #https://github.com/SkyTNT/anime-segmentation/blob/main/train.py#L92C20-L92C47
|
||||
pred = pred.cpu().numpy()[0]
|
||||
pred = np.transpose(pred, (1, 2, 0))
|
||||
#pred = pred[:, :, np.newaxis]
|
||||
return pred
|
||||
|
||||
h, w = h0, w0 = input_img.shape[:-1]
|
||||
h, w = (s, int(s * w / h)) if h > w else (int(s * h / w), s)
|
||||
ph, pw = s - h, s - w
|
||||
img_input = np.zeros([s, s, 3], dtype=np.float32)
|
||||
img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(input_img, (w, h))
|
||||
img_input = np.transpose(img_input, (2, 0, 1))
|
||||
img_input = img_input[np.newaxis, :]
|
||||
tmpImg = torch.from_numpy(img_input).float().to(DEVICE)
|
||||
with torch.no_grad():
|
||||
pred = self.net(tmpImg)[0][0].sigmoid() #https://github.com/SkyTNT/anime-segmentation/blob/main/train.py#L92C20-L92C47
|
||||
pred = pred.cpu().numpy()[0]
|
||||
pred = np.transpose(pred, (1, 2, 0))
|
||||
pred = pred[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
|
||||
#pred = cv2.resize(pred, (w0, h0))[:, :, np.newaxis]
|
||||
pred = cv2.resize(pred, (w0, h0))
|
||||
return pred
|
||||
|
||||
def __call__(self, np_img, img_size):
|
||||
mask = self.get_mask(np_img, int(img_size))
|
||||
np_img = (mask * np_img + 255 * (1 - mask)).astype(np.uint8)
|
||||
mask = (mask * 255).astype(np.uint8)
|
||||
#np_img = np.concatenate([np_img, mask], axis=2, dtype=np.uint8)
|
||||
#mask = mask.repeat(3, axis=2)
|
||||
return mask, np_img
|
||||
|
||||
@@ -0,0 +1,619 @@
|
||||
# Codes are borrowed from
|
||||
# https://github.com/xuebinqin/DIS/blob/main/IS-Net/models/isnet.py
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision import models
|
||||
|
||||
bce_loss = nn.BCEWithLogitsLoss(reduction="mean")
|
||||
|
||||
|
||||
def muti_loss_fusion(preds, target):
|
||||
loss0 = 0.0
|
||||
loss = 0.0
|
||||
|
||||
for i in range(0, len(preds)):
|
||||
if preds[i].shape[2] != target.shape[2] or preds[i].shape[3] != target.shape[3]:
|
||||
tmp_target = F.interpolate(
|
||||
target, size=preds[i].size()[2:], mode="bilinear", align_corners=True
|
||||
)
|
||||
loss = loss + bce_loss(preds[i], tmp_target)
|
||||
else:
|
||||
loss = loss + bce_loss(preds[i], target)
|
||||
if i == 0:
|
||||
loss0 = loss
|
||||
return loss0, loss
|
||||
|
||||
|
||||
fea_loss = nn.MSELoss(reduction="mean")
|
||||
kl_loss = nn.KLDivLoss(reduction="mean")
|
||||
l1_loss = nn.L1Loss(reduction="mean")
|
||||
smooth_l1_loss = nn.SmoothL1Loss(reduction="mean")
|
||||
|
||||
|
||||
def muti_loss_fusion_kl(preds, target, dfs, fs, mode="MSE"):
|
||||
loss0 = 0.0
|
||||
loss = 0.0
|
||||
|
||||
for i in range(0, len(preds)):
|
||||
if preds[i].shape[2] != target.shape[2] or preds[i].shape[3] != target.shape[3]:
|
||||
tmp_target = F.interpolate(
|
||||
target, size=preds[i].size()[2:], mode="bilinear", align_corners=True
|
||||
)
|
||||
loss = loss + bce_loss(preds[i], tmp_target)
|
||||
else:
|
||||
loss = loss + bce_loss(preds[i], target)
|
||||
if i == 0:
|
||||
loss0 = loss
|
||||
|
||||
for i in range(0, len(dfs)):
|
||||
df = dfs[i]
|
||||
fs_i = fs[i]
|
||||
if mode == "MSE":
|
||||
loss = loss + fea_loss(
|
||||
df, fs_i
|
||||
) ### add the mse loss of features as additional constraints
|
||||
elif mode == "KL":
|
||||
loss = loss + kl_loss(F.log_softmax(df, dim=1), F.softmax(fs_i, dim=1))
|
||||
elif mode == "MAE":
|
||||
loss = loss + l1_loss(df, fs_i)
|
||||
elif mode == "SmoothL1":
|
||||
loss = loss + smooth_l1_loss(df, fs_i)
|
||||
|
||||
return loss0, loss
|
||||
|
||||
|
||||
class REBNCONV(nn.Module):
|
||||
def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
|
||||
super(REBNCONV, self).__init__()
|
||||
|
||||
self.conv_s1 = nn.Conv2d(
|
||||
in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate, stride=stride
|
||||
)
|
||||
self.bn_s1 = nn.BatchNorm2d(out_ch)
|
||||
self.relu_s1 = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
|
||||
|
||||
return xout
|
||||
|
||||
|
||||
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
|
||||
def _upsample_like(src, tar):
|
||||
src = F.interpolate(src, size=tar.shape[2:], mode="bilinear", align_corners=False)
|
||||
|
||||
return src
|
||||
|
||||
|
||||
### RSU-7 ###
|
||||
class RSU7(nn.Module):
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
|
||||
super(RSU7, self).__init__()
|
||||
|
||||
self.in_ch = in_ch
|
||||
self.mid_ch = mid_ch
|
||||
self.out_ch = out_ch
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) ## 1 -> 1/2
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
|
||||
self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
|
||||
|
||||
self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, h, w = x.shape
|
||||
|
||||
hx = x
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
hx = self.pool4(hx4)
|
||||
|
||||
hx5 = self.rebnconv5(hx)
|
||||
hx = self.pool5(hx5)
|
||||
|
||||
hx6 = self.rebnconv6(hx)
|
||||
|
||||
hx7 = self.rebnconv7(hx6)
|
||||
|
||||
hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
|
||||
hx6dup = _upsample_like(hx6d, hx5)
|
||||
|
||||
hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
|
||||
hx5dup = _upsample_like(hx5d, hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
|
||||
hx4dup = _upsample_like(hx4d, hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
|
||||
### RSU-6 ###
|
||||
class RSU6(nn.Module):
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU6, self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
|
||||
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
|
||||
|
||||
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
hx = self.pool4(hx4)
|
||||
|
||||
hx5 = self.rebnconv5(hx)
|
||||
|
||||
hx6 = self.rebnconv6(hx5)
|
||||
|
||||
hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
|
||||
hx5dup = _upsample_like(hx5d, hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
|
||||
hx4dup = _upsample_like(hx4d, hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
|
||||
### RSU-5 ###
|
||||
class RSU5(nn.Module):
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU5, self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
|
||||
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
|
||||
|
||||
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
hx = self.pool3(hx3)
|
||||
|
||||
hx4 = self.rebnconv4(hx)
|
||||
|
||||
hx5 = self.rebnconv5(hx4)
|
||||
|
||||
hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
|
||||
hx4dup = _upsample_like(hx4d, hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
|
||||
### RSU-4 ###
|
||||
class RSU4(nn.Module):
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU4, self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
|
||||
self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
|
||||
|
||||
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
|
||||
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx = self.pool1(hx1)
|
||||
|
||||
hx2 = self.rebnconv2(hx)
|
||||
hx = self.pool2(hx2)
|
||||
|
||||
hx3 = self.rebnconv3(hx)
|
||||
|
||||
hx4 = self.rebnconv4(hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
|
||||
### RSU-4F ###
|
||||
class RSU4F(nn.Module):
|
||||
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
|
||||
super(RSU4F, self).__init__()
|
||||
|
||||
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
|
||||
|
||||
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
|
||||
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
|
||||
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
|
||||
|
||||
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
|
||||
|
||||
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
|
||||
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
|
||||
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.rebnconvin(hx)
|
||||
|
||||
hx1 = self.rebnconv1(hxin)
|
||||
hx2 = self.rebnconv2(hx1)
|
||||
hx3 = self.rebnconv3(hx2)
|
||||
|
||||
hx4 = self.rebnconv4(hx3)
|
||||
|
||||
hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
|
||||
hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
|
||||
hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
|
||||
|
||||
return hx1d + hxin
|
||||
|
||||
|
||||
class myrebnconv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_ch=3,
|
||||
out_ch=1,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
super(myrebnconv, self).__init__()
|
||||
|
||||
self.conv = nn.Conv2d(
|
||||
in_ch,
|
||||
out_ch,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
)
|
||||
self.bn = nn.BatchNorm2d(out_ch)
|
||||
self.rl = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.rl(self.bn(self.conv(x)))
|
||||
|
||||
|
||||
class ISNetGTEncoder(nn.Module):
|
||||
def __init__(self, in_ch=1, out_ch=1):
|
||||
super(ISNetGTEncoder, self).__init__()
|
||||
|
||||
self.conv_in = myrebnconv(
|
||||
in_ch, 16, 3, stride=2, padding=1
|
||||
) # nn.Conv2d(in_ch,64,3,stride=2,padding=1)
|
||||
|
||||
self.stage1 = RSU7(16, 16, 64)
|
||||
self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage2 = RSU6(64, 16, 64)
|
||||
self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage3 = RSU5(64, 32, 128)
|
||||
self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage4 = RSU4(128, 32, 256)
|
||||
self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage5 = RSU4F(256, 64, 512)
|
||||
self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage6 = RSU4F(512, 64, 512)
|
||||
|
||||
self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
|
||||
self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
|
||||
self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
|
||||
self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
|
||||
self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
|
||||
self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
|
||||
|
||||
@staticmethod
|
||||
def compute_loss(args):
|
||||
preds, targets = args
|
||||
return muti_loss_fusion(preds, targets)
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.conv_in(hx)
|
||||
# hx = self.pool_in(hxin)
|
||||
|
||||
# stage 1
|
||||
hx1 = self.stage1(hxin)
|
||||
hx = self.pool12(hx1)
|
||||
|
||||
# stage 2
|
||||
hx2 = self.stage2(hx)
|
||||
hx = self.pool23(hx2)
|
||||
|
||||
# stage 3
|
||||
hx3 = self.stage3(hx)
|
||||
hx = self.pool34(hx3)
|
||||
|
||||
# stage 4
|
||||
hx4 = self.stage4(hx)
|
||||
hx = self.pool45(hx4)
|
||||
|
||||
# stage 5
|
||||
hx5 = self.stage5(hx)
|
||||
hx = self.pool56(hx5)
|
||||
|
||||
# stage 6
|
||||
hx6 = self.stage6(hx)
|
||||
|
||||
# side output
|
||||
d1 = self.side1(hx1)
|
||||
d1 = _upsample_like(d1, x)
|
||||
|
||||
d2 = self.side2(hx2)
|
||||
d2 = _upsample_like(d2, x)
|
||||
|
||||
d3 = self.side3(hx3)
|
||||
d3 = _upsample_like(d3, x)
|
||||
|
||||
d4 = self.side4(hx4)
|
||||
d4 = _upsample_like(d4, x)
|
||||
|
||||
d5 = self.side5(hx5)
|
||||
d5 = _upsample_like(d5, x)
|
||||
|
||||
d6 = self.side6(hx6)
|
||||
d6 = _upsample_like(d6, x)
|
||||
|
||||
# d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
|
||||
|
||||
# return [torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(d4), torch.sigmoid(d5), torch.sigmoid(d6)], [hx1, hx2, hx3, hx4, hx5, hx6]
|
||||
return [d1, d2, d3, d4, d5, d6], [hx1, hx2, hx3, hx4, hx5, hx6]
|
||||
|
||||
|
||||
class ISNetDIS(nn.Module):
|
||||
def __init__(self, in_ch=3, out_ch=1):
|
||||
super(ISNetDIS, self).__init__()
|
||||
|
||||
self.conv_in = nn.Conv2d(in_ch, 64, 3, stride=2, padding=1)
|
||||
self.pool_in = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage1 = RSU7(64, 32, 64)
|
||||
self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage2 = RSU6(64, 32, 128)
|
||||
self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage3 = RSU5(128, 64, 256)
|
||||
self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage4 = RSU4(256, 128, 512)
|
||||
self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage5 = RSU4F(512, 256, 512)
|
||||
self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
|
||||
|
||||
self.stage6 = RSU4F(512, 256, 512)
|
||||
|
||||
# decoder
|
||||
self.stage5d = RSU4F(1024, 256, 512)
|
||||
self.stage4d = RSU4(1024, 128, 256)
|
||||
self.stage3d = RSU5(512, 64, 128)
|
||||
self.stage2d = RSU6(256, 32, 64)
|
||||
self.stage1d = RSU7(128, 16, 64)
|
||||
|
||||
self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
|
||||
self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
|
||||
self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
|
||||
self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
|
||||
self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
|
||||
self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
|
||||
|
||||
# self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
|
||||
|
||||
@staticmethod
|
||||
def compute_loss_kl(preds, targets, dfs, fs, mode="MSE"):
|
||||
return muti_loss_fusion_kl(preds, targets, dfs, fs, mode=mode)
|
||||
|
||||
@staticmethod
|
||||
def compute_loss(args):
|
||||
if len(args) == 3:
|
||||
ds, dfs, labels = args
|
||||
return muti_loss_fusion(ds, labels)
|
||||
else:
|
||||
ds, dfs, labels, fs = args
|
||||
return muti_loss_fusion_kl(ds, labels, dfs, fs, mode="MSE")
|
||||
|
||||
def forward(self, x):
|
||||
hx = x
|
||||
|
||||
hxin = self.conv_in(hx)
|
||||
hx = self.pool_in(hxin)
|
||||
|
||||
# stage 1
|
||||
hx1 = self.stage1(hxin)
|
||||
hx = self.pool12(hx1)
|
||||
|
||||
# stage 2
|
||||
hx2 = self.stage2(hx)
|
||||
hx = self.pool23(hx2)
|
||||
|
||||
# stage 3
|
||||
hx3 = self.stage3(hx)
|
||||
hx = self.pool34(hx3)
|
||||
|
||||
# stage 4
|
||||
hx4 = self.stage4(hx)
|
||||
hx = self.pool45(hx4)
|
||||
|
||||
# stage 5
|
||||
hx5 = self.stage5(hx)
|
||||
hx = self.pool56(hx5)
|
||||
|
||||
# stage 6
|
||||
hx6 = self.stage6(hx)
|
||||
hx6up = _upsample_like(hx6, hx5)
|
||||
|
||||
# -------------------- decoder --------------------
|
||||
hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
|
||||
hx5dup = _upsample_like(hx5d, hx4)
|
||||
|
||||
hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
|
||||
hx4dup = _upsample_like(hx4d, hx3)
|
||||
|
||||
hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
|
||||
hx3dup = _upsample_like(hx3d, hx2)
|
||||
|
||||
hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
|
||||
hx2dup = _upsample_like(hx2d, hx1)
|
||||
|
||||
hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
|
||||
|
||||
# side output
|
||||
d1 = self.side1(hx1d)
|
||||
d1 = _upsample_like(d1, x)
|
||||
|
||||
d2 = self.side2(hx2d)
|
||||
d2 = _upsample_like(d2, x)
|
||||
|
||||
d3 = self.side3(hx3d)
|
||||
d3 = _upsample_like(d3, x)
|
||||
|
||||
d4 = self.side4(hx4d)
|
||||
d4 = _upsample_like(d4, x)
|
||||
|
||||
d5 = self.side5(hx5d)
|
||||
d5 = _upsample_like(d5, x)
|
||||
|
||||
d6 = self.side6(hx6)
|
||||
d6 = _upsample_like(d6, x)
|
||||
|
||||
# d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
|
||||
|
||||
# return [torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(d4), torch.sigmoid(d5), torch.sigmoid(d6)], [hx1d, hx2d, hx3d, hx4d, hx5d, hx6]
|
||||
return [d1, d2, d3, d4, d5, d6], [hx1d, hx2d, hx3d, hx4d, hx5d, hx6]
|
||||
@@ -0,0 +1,100 @@
|
||||
#https://github.com/siyeong0/Anime-Face-Segmentation/blob/main/network.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
from custom_controlnet_aux.util import custom_torch_download
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(UNet, self).__init__()
|
||||
self.NUM_SEG_CLASSES = 7 # Background, hair, face, eye, mouth, skin, clothes
|
||||
|
||||
mobilenet_v2 = torchvision.models.mobilenet_v2(pretrained=False)
|
||||
mobilenet_v2.load_state_dict(torch.load(custom_torch_download(filename="mobilenet_v2-b0353104.pth")), strict=True)
|
||||
mob_blocks = mobilenet_v2.features
|
||||
|
||||
# Encoder
|
||||
self.en_block0 = nn.Sequential( # in_ch=3 out_ch=16
|
||||
mob_blocks[0],
|
||||
mob_blocks[1]
|
||||
)
|
||||
self.en_block1 = nn.Sequential( # in_ch=16 out_ch=24
|
||||
mob_blocks[2],
|
||||
mob_blocks[3],
|
||||
)
|
||||
self.en_block2 = nn.Sequential( # in_ch=24 out_ch=32
|
||||
mob_blocks[4],
|
||||
mob_blocks[5],
|
||||
mob_blocks[6],
|
||||
)
|
||||
self.en_block3 = nn.Sequential( # in_ch=32 out_ch=96
|
||||
mob_blocks[7],
|
||||
mob_blocks[8],
|
||||
mob_blocks[9],
|
||||
mob_blocks[10],
|
||||
mob_blocks[11],
|
||||
mob_blocks[12],
|
||||
mob_blocks[13],
|
||||
)
|
||||
self.en_block4 = nn.Sequential( # in_ch=96 out_ch=160
|
||||
mob_blocks[14],
|
||||
mob_blocks[15],
|
||||
mob_blocks[16],
|
||||
)
|
||||
|
||||
# Decoder
|
||||
self.de_block4 = nn.Sequential( # in_ch=160 out_ch=96
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
nn.Conv2d(160, 96, kernel_size=3, padding=1),
|
||||
nn.InstanceNorm2d(96),
|
||||
nn.LeakyReLU(0.1),
|
||||
nn.Dropout(p=0.2)
|
||||
)
|
||||
self.de_block3 = nn.Sequential( # in_ch=96x2 out_ch=32
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
nn.Conv2d(96*2, 32, kernel_size=3, padding=1),
|
||||
nn.InstanceNorm2d(32),
|
||||
nn.LeakyReLU(0.1),
|
||||
nn.Dropout(p=0.2)
|
||||
)
|
||||
self.de_block2 = nn.Sequential( # in_ch=32x2 out_ch=24
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
nn.Conv2d(32*2, 24, kernel_size=3, padding=1),
|
||||
nn.InstanceNorm2d(24),
|
||||
nn.LeakyReLU(0.1),
|
||||
nn.Dropout(p=0.2)
|
||||
)
|
||||
self.de_block1 = nn.Sequential( # in_ch=24x2 out_ch=16
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
nn.Conv2d(24*2, 16, kernel_size=3, padding=1),
|
||||
nn.InstanceNorm2d(16),
|
||||
nn.LeakyReLU(0.1),
|
||||
nn.Dropout(p=0.2)
|
||||
)
|
||||
|
||||
self.de_block0 = nn.Sequential( # in_ch=16x2 out_ch=7
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
nn.Conv2d(16*2, self.NUM_SEG_CLASSES, kernel_size=3, padding=1),
|
||||
nn.Softmax2d()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
e0 = self.en_block0(x)
|
||||
e1 = self.en_block1(e0)
|
||||
e2 = self.en_block2(e1)
|
||||
e3 = self.en_block3(e2)
|
||||
e4 = self.en_block4(e3)
|
||||
|
||||
d4 = self.de_block4(e4)
|
||||
c4 = torch.cat((d4,e3),1)
|
||||
d3 = self.de_block3(c4)
|
||||
c3 = torch.cat((d3,e2),1)
|
||||
d2 = self.de_block2(c3)
|
||||
c2 =torch.cat((d2,e1),1)
|
||||
d1 = self.de_block1(c2)
|
||||
c1 = torch.cat((d1,e0),1)
|
||||
y = self.de_block0(c1)
|
||||
|
||||
return y
|
||||
@@ -0,0 +1,40 @@
|
||||
#https://github.com/siyeong0/Anime-Face-Segmentation/blob/main/util.py
|
||||
#The color palette is changed according to https://github.com/Mikubill/sd-webui-controlnet/blob/91f67ddcc7bc47537a6285864abfc12590f46c3f/annotator/anime_face_segment/__init__.py
|
||||
import cv2 as cv
|
||||
import glob
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
"""
|
||||
COLOR_BACKGROUND = (0,255,255)
|
||||
COLOR_HAIR = (255,0,0)
|
||||
COLOR_EYE = (0,0,255)
|
||||
COLOR_MOUTH = (255,255,255)
|
||||
COLOR_FACE = (0,255,0)
|
||||
COLOR_SKIN = (255,255,0)
|
||||
COLOR_CLOTHES = (255,0,255)
|
||||
"""
|
||||
COLOR_BACKGROUND = (255,255,0)
|
||||
COLOR_HAIR = (0,0,255)
|
||||
COLOR_EYE = (255,0,0)
|
||||
COLOR_MOUTH = (255,255,255)
|
||||
COLOR_FACE = (0,255,0)
|
||||
COLOR_SKIN = (0,255,255)
|
||||
COLOR_CLOTHES = (255,0,255)
|
||||
PALETTE = [COLOR_BACKGROUND,COLOR_HAIR,COLOR_EYE,COLOR_MOUTH,COLOR_FACE,COLOR_SKIN,COLOR_CLOTHES]
|
||||
|
||||
def img2seg(path):
|
||||
src = cv.imread(path)
|
||||
src = src.reshape(-1, 3)
|
||||
seg_list = []
|
||||
for color in PALETTE:
|
||||
seg_list.append(np.where(np.all(src==color, axis=1), 1.0, 0.0))
|
||||
dst = np.stack(seg_list,axis=1).reshape(512,512,7)
|
||||
|
||||
return dst.astype(np.float32)
|
||||
|
||||
def seg2img(src):
|
||||
src = np.moveaxis(src,0,2)
|
||||
dst = [[PALETTE[np.argmax(val)] for val in buf]for buf in src]
|
||||
|
||||
return np.array(dst).astype(np.uint8)
|
||||
@@ -0,0 +1,38 @@
|
||||
import warnings
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from custom_controlnet_aux.util import HWC3, resize_image_with_pad
|
||||
|
||||
class BinaryDetector:
|
||||
def __call__(self, input_image=None, bin_threshold=0, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
|
||||
if "img" in kwargs:
|
||||
warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
|
||||
input_image = kwargs.pop("img")
|
||||
|
||||
if input_image is None:
|
||||
raise ValueError("input_image must be defined.")
|
||||
|
||||
if not isinstance(input_image, np.ndarray):
|
||||
input_image = np.array(input_image, dtype=np.uint8)
|
||||
output_type = output_type or "pil"
|
||||
else:
|
||||
output_type = output_type or "np"
|
||||
|
||||
detected_map, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
|
||||
img_gray = cv2.cvtColor(detected_map, cv2.COLOR_RGB2GRAY)
|
||||
if bin_threshold == 0 or bin_threshold == 255:
|
||||
# Otsu's threshold
|
||||
otsu_threshold, img_bin = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||
print("Otsu threshold:", otsu_threshold)
|
||||
else:
|
||||
_, img_bin = cv2.threshold(img_gray, bin_threshold, 255, cv2.THRESH_BINARY_INV)
|
||||
|
||||
detected_map = cv2.cvtColor(img_bin, cv2.COLOR_GRAY2RGB)
|
||||
detected_map = HWC3(remove_pad(255 - detected_map))
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,17 @@
|
||||
import warnings
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from custom_controlnet_aux.util import resize_image_with_pad, common_input_validate, HWC3
|
||||
|
||||
class CannyDetector:
|
||||
def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
detected_map, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
detected_map = cv2.Canny(detected_map, low_threshold, high_threshold)
|
||||
detected_map = HWC3(remove_pad(detected_map))
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,37 @@
|
||||
import cv2
|
||||
import warnings
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from custom_controlnet_aux.util import HWC3, safer_memory, common_input_validate
|
||||
|
||||
def cv2_resize_shortest_edge(image, size):
|
||||
h, w = image.shape[:2]
|
||||
if h < w:
|
||||
new_h = size
|
||||
new_w = int(round(w / h * size))
|
||||
else:
|
||||
new_w = size
|
||||
new_h = int(round(h / w * size))
|
||||
resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||
return resized_image
|
||||
|
||||
def apply_color(img, res=512):
|
||||
img = cv2_resize_shortest_edge(img, res)
|
||||
h, w = img.shape[:2]
|
||||
|
||||
input_img_color = cv2.resize(img, (w//64, h//64), interpolation=cv2.INTER_CUBIC)
|
||||
input_img_color = cv2.resize(input_img_color, (w, h), interpolation=cv2.INTER_NEAREST)
|
||||
return input_img_color
|
||||
|
||||
#Color T2I like multiples-of-64, upscale methods are fixed.
|
||||
class ColorDetector:
|
||||
def __call__(self, input_image=None, detect_resolution=512, output_type=None, **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
input_image = HWC3(input_image)
|
||||
detected_map = HWC3(apply_color(input_image, detect_resolution))
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,66 @@
|
||||
import torchvision # Fix issue Unknown builtin op: torchvision::nms
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
|
||||
from custom_controlnet_aux.util import HWC3, resize_image_with_pad, common_input_validate, custom_hf_download, DENSEPOSE_MODEL_NAME
|
||||
from .densepose import DensePoseMaskedColormapResultsVisualizer, _extract_i_from_iuvarr, densepose_chart_predictor_output_to_result_with_confidences
|
||||
|
||||
N_PART_LABELS = 24
|
||||
|
||||
class DenseposeDetector:
|
||||
def __init__(self, model):
|
||||
self.dense_pose_estimation = model
|
||||
self.device = "cpu"
|
||||
self.result_visualizer = DensePoseMaskedColormapResultsVisualizer(
|
||||
alpha=1,
|
||||
data_extractor=_extract_i_from_iuvarr,
|
||||
segm_extractor=_extract_i_from_iuvarr,
|
||||
val_scale = 255.0 / N_PART_LABELS
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_or_path=DENSEPOSE_MODEL_NAME, filename="densepose_r50_fpn_dl.torchscript"):
|
||||
torchscript_model_path = custom_hf_download(pretrained_model_or_path, filename)
|
||||
densepose = torch.jit.load(torchscript_model_path, map_location="cpu")
|
||||
return cls(densepose)
|
||||
|
||||
def to(self, device):
|
||||
self.dense_pose_estimation.to(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def __call__(self, input_image, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", cmap="viridis", **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
H, W = input_image.shape[:2]
|
||||
|
||||
hint_image_canvas = np.zeros([H, W], dtype=np.uint8)
|
||||
hint_image_canvas = np.tile(hint_image_canvas[:, :, np.newaxis], [1, 1, 3])
|
||||
|
||||
input_image = rearrange(torch.from_numpy(input_image).to(self.device), 'h w c -> c h w')
|
||||
|
||||
pred_boxes, corase_segm, fine_segm, u, v = self.dense_pose_estimation(input_image)
|
||||
|
||||
extractor = densepose_chart_predictor_output_to_result_with_confidences
|
||||
densepose_results = [extractor(pred_boxes[i:i+1], corase_segm[i:i+1], fine_segm[i:i+1], u[i:i+1], v[i:i+1]) for i in range(len(pred_boxes))]
|
||||
|
||||
if cmap=="viridis":
|
||||
self.result_visualizer.mask_visualizer.cmap = cv2.COLORMAP_VIRIDIS
|
||||
hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
|
||||
hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
|
||||
hint_image[:, :, 0][hint_image[:, :, 0] == 0] = 68
|
||||
hint_image[:, :, 1][hint_image[:, :, 1] == 0] = 1
|
||||
hint_image[:, :, 2][hint_image[:, :, 2] == 0] = 84
|
||||
else:
|
||||
self.result_visualizer.mask_visualizer.cmap = cv2.COLORMAP_PARULA
|
||||
hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
|
||||
hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
|
||||
|
||||
detected_map = remove_pad(HWC3(hint_image))
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
return detected_map
|
||||
@@ -0,0 +1,347 @@
|
||||
from typing import Tuple
|
||||
import math
|
||||
import numpy as np
|
||||
from enum import IntEnum
|
||||
from typing import List, Tuple, Union
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
import logging
|
||||
import cv2
|
||||
|
||||
Image = np.ndarray
|
||||
Boxes = torch.Tensor
|
||||
ImageSizeType = Tuple[int, int]
|
||||
_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]
|
||||
IntTupleBox = Tuple[int, int, int, int]
|
||||
|
||||
class BoxMode(IntEnum):
|
||||
"""
|
||||
Enum of different ways to represent a box.
|
||||
"""
|
||||
|
||||
XYXY_ABS = 0
|
||||
"""
|
||||
(x0, y0, x1, y1) in absolute floating points coordinates.
|
||||
The coordinates in range [0, width or height].
|
||||
"""
|
||||
XYWH_ABS = 1
|
||||
"""
|
||||
(x0, y0, w, h) in absolute floating points coordinates.
|
||||
"""
|
||||
XYXY_REL = 2
|
||||
"""
|
||||
Not yet supported!
|
||||
(x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image.
|
||||
"""
|
||||
XYWH_REL = 3
|
||||
"""
|
||||
Not yet supported!
|
||||
(x0, y0, w, h) in range [0, 1]. They are relative to the size of the image.
|
||||
"""
|
||||
XYWHA_ABS = 4
|
||||
"""
|
||||
(xc, yc, w, h, a) in absolute floating points coordinates.
|
||||
(xc, yc) is the center of the rotated box, and the angle a is in degrees ccw.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def convert(box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode") -> _RawBoxType:
|
||||
"""
|
||||
Args:
|
||||
box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5
|
||||
from_mode, to_mode (BoxMode)
|
||||
|
||||
Returns:
|
||||
The converted box of the same type.
|
||||
"""
|
||||
if from_mode == to_mode:
|
||||
return box
|
||||
|
||||
original_type = type(box)
|
||||
is_numpy = isinstance(box, np.ndarray)
|
||||
single_box = isinstance(box, (list, tuple))
|
||||
if single_box:
|
||||
assert len(box) == 4 or len(box) == 5, (
|
||||
"BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
|
||||
" where k == 4 or 5"
|
||||
)
|
||||
arr = torch.tensor(box)[None, :]
|
||||
else:
|
||||
# avoid modifying the input box
|
||||
if is_numpy:
|
||||
arr = torch.from_numpy(np.asarray(box)).clone()
|
||||
else:
|
||||
arr = box.clone()
|
||||
|
||||
assert to_mode not in [BoxMode.XYXY_REL, BoxMode.XYWH_REL] and from_mode not in [
|
||||
BoxMode.XYXY_REL,
|
||||
BoxMode.XYWH_REL,
|
||||
], "Relative mode not yet supported!"
|
||||
|
||||
if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
|
||||
assert (
|
||||
arr.shape[-1] == 5
|
||||
), "The last dimension of input shape must be 5 for XYWHA format"
|
||||
original_dtype = arr.dtype
|
||||
arr = arr.double()
|
||||
|
||||
w = arr[:, 2]
|
||||
h = arr[:, 3]
|
||||
a = arr[:, 4]
|
||||
c = torch.abs(torch.cos(a * math.pi / 180.0))
|
||||
s = torch.abs(torch.sin(a * math.pi / 180.0))
|
||||
# This basically computes the horizontal bounding rectangle of the rotated box
|
||||
new_w = c * w + s * h
|
||||
new_h = c * h + s * w
|
||||
|
||||
# convert center to top-left corner
|
||||
arr[:, 0] -= new_w / 2.0
|
||||
arr[:, 1] -= new_h / 2.0
|
||||
# bottom-right corner
|
||||
arr[:, 2] = arr[:, 0] + new_w
|
||||
arr[:, 3] = arr[:, 1] + new_h
|
||||
|
||||
arr = arr[:, :4].to(dtype=original_dtype)
|
||||
elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:
|
||||
original_dtype = arr.dtype
|
||||
arr = arr.double()
|
||||
arr[:, 0] += arr[:, 2] / 2.0
|
||||
arr[:, 1] += arr[:, 3] / 2.0
|
||||
angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
|
||||
arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
|
||||
else:
|
||||
if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
|
||||
arr[:, 2] += arr[:, 0]
|
||||
arr[:, 3] += arr[:, 1]
|
||||
elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:
|
||||
arr[:, 2] -= arr[:, 0]
|
||||
arr[:, 3] -= arr[:, 1]
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Conversion from BoxMode {} to {} is not supported yet".format(
|
||||
from_mode, to_mode
|
||||
)
|
||||
)
|
||||
|
||||
if single_box:
|
||||
return original_type(arr.flatten().tolist())
|
||||
if is_numpy:
|
||||
return arr.numpy()
|
||||
else:
|
||||
return arr
|
||||
|
||||
class MatrixVisualizer:
|
||||
"""
|
||||
Base visualizer for matrix data
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inplace=True,
|
||||
cmap=cv2.COLORMAP_PARULA,
|
||||
val_scale=1.0,
|
||||
alpha=0.7,
|
||||
interp_method_matrix=cv2.INTER_LINEAR,
|
||||
interp_method_mask=cv2.INTER_NEAREST,
|
||||
):
|
||||
self.inplace = inplace
|
||||
self.cmap = cmap
|
||||
self.val_scale = val_scale
|
||||
self.alpha = alpha
|
||||
self.interp_method_matrix = interp_method_matrix
|
||||
self.interp_method_mask = interp_method_mask
|
||||
|
||||
def visualize(self, image_bgr, mask, matrix, bbox_xywh):
|
||||
self._check_image(image_bgr)
|
||||
self._check_mask_matrix(mask, matrix)
|
||||
if self.inplace:
|
||||
image_target_bgr = image_bgr
|
||||
else:
|
||||
image_target_bgr = image_bgr * 0
|
||||
x, y, w, h = [int(v) for v in bbox_xywh]
|
||||
if w <= 0 or h <= 0:
|
||||
return image_bgr
|
||||
mask, matrix = self._resize(mask, matrix, w, h)
|
||||
mask_bg = np.tile((mask == 0)[:, :, np.newaxis], [1, 1, 3])
|
||||
matrix_scaled = matrix.astype(np.float32) * self.val_scale
|
||||
_EPSILON = 1e-6
|
||||
if np.any(matrix_scaled > 255 + _EPSILON):
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(
|
||||
f"Matrix has values > {255 + _EPSILON} after " f"scaling, clipping to [0..255]"
|
||||
)
|
||||
matrix_scaled_8u = matrix_scaled.clip(0, 255).astype(np.uint8)
|
||||
matrix_vis = cv2.applyColorMap(matrix_scaled_8u, self.cmap)
|
||||
matrix_vis[mask_bg] = image_target_bgr[y : y + h, x : x + w, :][mask_bg]
|
||||
image_target_bgr[y : y + h, x : x + w, :] = (
|
||||
image_target_bgr[y : y + h, x : x + w, :] * (1.0 - self.alpha) + matrix_vis * self.alpha
|
||||
)
|
||||
return image_target_bgr.astype(np.uint8)
|
||||
|
||||
def _resize(self, mask, matrix, w, h):
|
||||
if (w != mask.shape[1]) or (h != mask.shape[0]):
|
||||
mask = cv2.resize(mask, (w, h), self.interp_method_mask)
|
||||
if (w != matrix.shape[1]) or (h != matrix.shape[0]):
|
||||
matrix = cv2.resize(matrix, (w, h), self.interp_method_matrix)
|
||||
return mask, matrix
|
||||
|
||||
def _check_image(self, image_rgb):
|
||||
assert len(image_rgb.shape) == 3
|
||||
assert image_rgb.shape[2] == 3
|
||||
assert image_rgb.dtype == np.uint8
|
||||
|
||||
def _check_mask_matrix(self, mask, matrix):
|
||||
assert len(matrix.shape) == 2
|
||||
assert len(mask.shape) == 2
|
||||
assert mask.dtype == np.uint8
|
||||
|
||||
class DensePoseResultsVisualizer:
|
||||
def visualize(
|
||||
self,
|
||||
image_bgr: Image,
|
||||
results,
|
||||
) -> Image:
|
||||
context = self.create_visualization_context(image_bgr)
|
||||
for i, result in enumerate(results):
|
||||
boxes_xywh, labels, uv = result
|
||||
iuv_array = torch.cat(
|
||||
(labels[None].type(torch.float32), uv * 255.0)
|
||||
).type(torch.uint8)
|
||||
self.visualize_iuv_arr(context, iuv_array.cpu().numpy(), boxes_xywh)
|
||||
image_bgr = self.context_to_image_bgr(context)
|
||||
return image_bgr
|
||||
|
||||
def create_visualization_context(self, image_bgr: Image):
|
||||
return image_bgr
|
||||
|
||||
def visualize_iuv_arr(self, context, iuv_arr: np.ndarray, bbox_xywh) -> None:
|
||||
pass
|
||||
|
||||
def context_to_image_bgr(self, context):
|
||||
return context
|
||||
|
||||
def get_image_bgr_from_context(self, context):
|
||||
return context
|
||||
|
||||
class DensePoseMaskedColormapResultsVisualizer(DensePoseResultsVisualizer):
|
||||
def __init__(
|
||||
self,
|
||||
data_extractor,
|
||||
segm_extractor,
|
||||
inplace=True,
|
||||
cmap=cv2.COLORMAP_PARULA,
|
||||
alpha=0.7,
|
||||
val_scale=1.0,
|
||||
**kwargs,
|
||||
):
|
||||
self.mask_visualizer = MatrixVisualizer(
|
||||
inplace=inplace, cmap=cmap, val_scale=val_scale, alpha=alpha
|
||||
)
|
||||
self.data_extractor = data_extractor
|
||||
self.segm_extractor = segm_extractor
|
||||
|
||||
def context_to_image_bgr(self, context):
|
||||
return context
|
||||
|
||||
def visualize_iuv_arr(self, context, iuv_arr: np.ndarray, bbox_xywh) -> None:
|
||||
image_bgr = self.get_image_bgr_from_context(context)
|
||||
matrix = self.data_extractor(iuv_arr)
|
||||
segm = self.segm_extractor(iuv_arr)
|
||||
mask = np.zeros(matrix.shape, dtype=np.uint8)
|
||||
mask[segm > 0] = 1
|
||||
image_bgr = self.mask_visualizer.visualize(image_bgr, mask, matrix, bbox_xywh)
|
||||
|
||||
|
||||
def _extract_i_from_iuvarr(iuv_arr):
|
||||
return iuv_arr[0, :, :]
|
||||
|
||||
|
||||
def _extract_u_from_iuvarr(iuv_arr):
|
||||
return iuv_arr[1, :, :]
|
||||
|
||||
|
||||
def _extract_v_from_iuvarr(iuv_arr):
|
||||
return iuv_arr[2, :, :]
|
||||
|
||||
def make_int_box(box: torch.Tensor) -> IntTupleBox:
|
||||
int_box = [0, 0, 0, 0]
|
||||
int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
|
||||
return int_box[0], int_box[1], int_box[2], int_box[3]
|
||||
|
||||
def densepose_chart_predictor_output_to_result_with_confidences(
|
||||
boxes: Boxes,
|
||||
coarse_segm,
|
||||
fine_segm,
|
||||
u, v
|
||||
|
||||
):
|
||||
boxes_xyxy_abs = boxes.clone()
|
||||
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
||||
box_xywh = make_int_box(boxes_xywh_abs[0])
|
||||
|
||||
labels = resample_fine_and_coarse_segm_tensors_to_bbox(fine_segm, coarse_segm, box_xywh).squeeze(0)
|
||||
uv = resample_uv_tensors_to_bbox(u, v, labels, box_xywh)
|
||||
confidences = []
|
||||
return box_xywh, labels, uv
|
||||
|
||||
def resample_fine_and_coarse_segm_tensors_to_bbox(
|
||||
fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
|
||||
):
|
||||
"""
|
||||
Resample fine and coarse segmentation tensors to the given
|
||||
bounding box and derive labels for each pixel of the bounding box
|
||||
|
||||
Args:
|
||||
fine_segm: float tensor of shape [1, C, Hout, Wout]
|
||||
coarse_segm: float tensor of shape [1, K, Hout, Wout]
|
||||
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
|
||||
corner coordinates, width (W) and height (H)
|
||||
Return:
|
||||
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
|
||||
"""
|
||||
x, y, w, h = box_xywh_abs
|
||||
w = max(int(w), 1)
|
||||
h = max(int(h), 1)
|
||||
# coarse segmentation
|
||||
coarse_segm_bbox = F.interpolate(
|
||||
coarse_segm,
|
||||
(h, w),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
).argmax(dim=1)
|
||||
# combined coarse and fine segmentation
|
||||
labels = (
|
||||
F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
|
||||
* (coarse_segm_bbox > 0).long()
|
||||
)
|
||||
return labels
|
||||
|
||||
def resample_uv_tensors_to_bbox(
|
||||
u: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
labels: torch.Tensor,
|
||||
box_xywh_abs: IntTupleBox,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Resamples U and V coordinate estimates for the given bounding box
|
||||
|
||||
Args:
|
||||
u (tensor [1, C, H, W] of float): U coordinates
|
||||
v (tensor [1, C, H, W] of float): V coordinates
|
||||
labels (tensor [H, W] of long): labels obtained by resampling segmentation
|
||||
outputs for the given bounding box
|
||||
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
|
||||
Return:
|
||||
Resampled U and V coordinates - a tensor [2, H, W] of float
|
||||
"""
|
||||
x, y, w, h = box_xywh_abs
|
||||
w = max(int(w), 1)
|
||||
h = max(int(h), 1)
|
||||
u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
|
||||
v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
|
||||
uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
|
||||
for part_id in range(1, u_bbox.size(1)):
|
||||
uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
|
||||
uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
|
||||
return uv
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .transformers import DepthAnythingDetector
|
||||
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
Modern DepthAnything implementation using HuggingFace transformers.
|
||||
Replaces legacy torch.hub.load DINOv2 backbone with transformers pipeline.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import pipeline
|
||||
|
||||
from custom_controlnet_aux.util import HWC3, common_input_validate, resize_image_with_pad
|
||||
|
||||
class DepthAnythingDetector:
|
||||
"""DepthAnything depth estimation using HuggingFace transformers."""
|
||||
|
||||
def __init__(self, model_name="LiheYoung/depth-anything-large-hf"):
|
||||
"""Initialize DepthAnything with specified model."""
|
||||
self.pipe = pipeline(task="depth-estimation", model=model_name)
|
||||
self.device = "cpu"
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_or_path=None, filename="depth_anything_vitl14.pth"):
|
||||
"""Create DepthAnything from pretrained model, mapping legacy names to HuggingFace models."""
|
||||
|
||||
# Map legacy checkpoint names to modern HuggingFace models
|
||||
model_mapping = {
|
||||
"depth_anything_vitl14.pth": "LiheYoung/depth-anything-large-hf",
|
||||
"depth_anything_vitb14.pth": "LiheYoung/depth-anything-base-hf",
|
||||
"depth_anything_vits14.pth": "LiheYoung/depth-anything-small-hf"
|
||||
}
|
||||
|
||||
model_name = model_mapping.get(filename, "LiheYoung/depth-anything-large-hf")
|
||||
return cls(model_name=model_name)
|
||||
|
||||
def to(self, device):
|
||||
"""Move model to specified device."""
|
||||
self.pipe.model = self.pipe.model.to(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def __call__(self, input_image, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
|
||||
"""Perform depth estimation on input image."""
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
|
||||
if isinstance(input_image, np.ndarray):
|
||||
pil_image = Image.fromarray(input_image)
|
||||
else:
|
||||
pil_image = input_image
|
||||
|
||||
with torch.no_grad():
|
||||
result = self.pipe(pil_image)
|
||||
depth = result["depth"]
|
||||
|
||||
if isinstance(depth, Image.Image):
|
||||
depth_array = np.array(depth, dtype=np.float32)
|
||||
else:
|
||||
depth_array = np.array(depth)
|
||||
|
||||
# Normalize depth values to 0-255 range
|
||||
depth_min = depth_array.min()
|
||||
depth_max = depth_array.max()
|
||||
if depth_max > depth_min:
|
||||
depth_array = (depth_array - depth_min) / (depth_max - depth_min) * 255.0
|
||||
else:
|
||||
depth_array = np.zeros_like(depth_array)
|
||||
|
||||
depth_image = depth_array.astype(np.uint8)
|
||||
|
||||
detected_map = remove_pad(HWC3(depth_image))
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,56 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import repeat
|
||||
from PIL import Image
|
||||
from custom_controlnet_aux.util import HWC3, common_input_validate, resize_image_with_pad, custom_hf_download, DEPTH_ANYTHING_V2_MODEL_NAME_DICT
|
||||
from custom_controlnet_aux.depth_anything_v2.dpt import DepthAnythingV2
|
||||
import cv2
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# https://github.com/DepthAnything/Depth-Anything-V2/blob/main/app.py
|
||||
model_configs = {
|
||||
'depth_anything_v2_vits.pth': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
|
||||
'depth_anything_v2_vitb.pth': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
|
||||
'depth_anything_v2_vitl.pth': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
|
||||
'depth_anything_v2_vitg.pth': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]},
|
||||
'depth_anything_v2_metric_vkitti_vitl.pth': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
|
||||
'depth_anything_v2_metric_hypersim_vitl.pth': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
|
||||
}
|
||||
|
||||
class DepthAnythingV2Detector:
|
||||
def __init__(self, model, filename):
|
||||
self.model = model
|
||||
self.device = "cpu"
|
||||
self.filename = filename
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_or_path=None, filename="depth_anything_v2_vits.pth"):
|
||||
if pretrained_model_or_path is None:
|
||||
pretrained_model_or_path = DEPTH_ANYTHING_V2_MODEL_NAME_DICT[filename]
|
||||
model_path = custom_hf_download(pretrained_model_or_path, filename)
|
||||
model = DepthAnythingV2(**model_configs[filename])
|
||||
model.load_state_dict(torch.load(model_path, map_location="cpu"))
|
||||
model = model.eval()
|
||||
return cls(model, filename)
|
||||
|
||||
def to(self, device):
|
||||
self.model.to(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def __call__(self, input_image, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", max_depth=20.0, **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
|
||||
depth = self.model.infer_image(cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR), input_size=518, max_depth=max_depth)
|
||||
depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
|
||||
depth = depth.astype(np.uint8)
|
||||
if 'metric' in self.filename:
|
||||
depth = 255 - depth
|
||||
|
||||
detected_map = repeat(depth, "h w -> h w 3")
|
||||
detected_map, remove_pad = resize_image_with_pad(detected_map, detect_resolution, upscale_method)
|
||||
detected_map = remove_pad(detected_map)
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
#
|
||||
# This source code is licensed under the Apache License, Version 2.0
|
||||
# found in the LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
from functools import partial
|
||||
import math
|
||||
import logging
|
||||
from typing import Sequence, Tuple, Union, Callable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
from torch.nn.init import trunc_normal_
|
||||
|
||||
from custom_controlnet_aux.depth_anything_v2.dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
|
||||
if not depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
for child_name, child_module in module.named_children():
|
||||
child_name = ".".join((name, child_name)) if name else child_name
|
||||
named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
|
||||
if depth_first and include_root:
|
||||
fn(module=module, name=name)
|
||||
return module
|
||||
|
||||
|
||||
class BlockChunk(nn.ModuleList):
|
||||
def forward(self, x):
|
||||
for b in self:
|
||||
x = b(x)
|
||||
return x
|
||||
|
||||
|
||||
class DinoVisionTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
img_size=224,
|
||||
patch_size=16,
|
||||
in_chans=3,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4.0,
|
||||
qkv_bias=True,
|
||||
ffn_bias=True,
|
||||
proj_bias=True,
|
||||
drop_path_rate=0.0,
|
||||
drop_path_uniform=False,
|
||||
init_values=None, # for layerscale: None or 0 => no layerscale
|
||||
embed_layer=PatchEmbed,
|
||||
act_layer=nn.GELU,
|
||||
block_fn=Block,
|
||||
ffn_layer="mlp",
|
||||
block_chunks=1,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
img_size (int, tuple): input image size
|
||||
patch_size (int, tuple): patch size
|
||||
in_chans (int): number of input channels
|
||||
embed_dim (int): embedding dimension
|
||||
depth (int): depth of transformer
|
||||
num_heads (int): number of attention heads
|
||||
mlp_ratio (int): ratio of mlp hidden dim to embedding dim
|
||||
qkv_bias (bool): enable bias for qkv if True
|
||||
proj_bias (bool): enable bias for proj in attn if True
|
||||
ffn_bias (bool): enable bias for ffn if True
|
||||
drop_path_rate (float): stochastic depth rate
|
||||
drop_path_uniform (bool): apply uniform drop rate across blocks
|
||||
weight_init (str): weight init scheme
|
||||
init_values (float): layer-scale init values
|
||||
embed_layer (nn.Module): patch embedding layer
|
||||
act_layer (nn.Module): MLP activation layer
|
||||
block_fn (nn.Module): transformer block class
|
||||
ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
|
||||
block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
|
||||
num_register_tokens: (int) number of extra cls tokens (so-called "registers")
|
||||
interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
|
||||
interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
|
||||
"""
|
||||
super().__init__()
|
||||
norm_layer = partial(nn.LayerNorm, eps=1e-6)
|
||||
|
||||
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
|
||||
self.num_tokens = 1
|
||||
self.n_blocks = depth
|
||||
self.num_heads = num_heads
|
||||
self.patch_size = patch_size
|
||||
self.num_register_tokens = num_register_tokens
|
||||
self.interpolate_antialias = interpolate_antialias
|
||||
self.interpolate_offset = interpolate_offset
|
||||
|
||||
self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
|
||||
num_patches = self.patch_embed.num_patches
|
||||
|
||||
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
|
||||
assert num_register_tokens >= 0
|
||||
self.register_tokens = (
|
||||
nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
|
||||
)
|
||||
|
||||
if drop_path_uniform is True:
|
||||
dpr = [drop_path_rate] * depth
|
||||
else:
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
|
||||
|
||||
if ffn_layer == "mlp":
|
||||
logger.info("using MLP layer as FFN")
|
||||
ffn_layer = Mlp
|
||||
elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
|
||||
logger.info("using SwiGLU layer as FFN")
|
||||
ffn_layer = SwiGLUFFNFused
|
||||
elif ffn_layer == "identity":
|
||||
logger.info("using Identity layer as FFN")
|
||||
|
||||
def f(*args, **kwargs):
|
||||
return nn.Identity()
|
||||
|
||||
ffn_layer = f
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
blocks_list = [
|
||||
block_fn(
|
||||
dim=embed_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
ffn_bias=ffn_bias,
|
||||
drop_path=dpr[i],
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
ffn_layer=ffn_layer,
|
||||
init_values=init_values,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
if block_chunks > 0:
|
||||
self.chunked_blocks = True
|
||||
chunked_blocks = []
|
||||
chunksize = depth // block_chunks
|
||||
for i in range(0, depth, chunksize):
|
||||
# this is to keep the block index consistent if we chunk the block list
|
||||
chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
|
||||
self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
|
||||
else:
|
||||
self.chunked_blocks = False
|
||||
self.blocks = nn.ModuleList(blocks_list)
|
||||
|
||||
self.norm = norm_layer(embed_dim)
|
||||
self.head = nn.Identity()
|
||||
|
||||
self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
|
||||
|
||||
self.init_weights()
|
||||
|
||||
def init_weights(self):
|
||||
trunc_normal_(self.pos_embed, std=0.02)
|
||||
nn.init.normal_(self.cls_token, std=1e-6)
|
||||
if self.register_tokens is not None:
|
||||
nn.init.normal_(self.register_tokens, std=1e-6)
|
||||
named_apply(init_weights_vit_timm, self)
|
||||
|
||||
def interpolate_pos_encoding(self, x, w, h):
|
||||
previous_dtype = x.dtype
|
||||
npatch = x.shape[1] - 1
|
||||
N = self.pos_embed.shape[1] - 1
|
||||
if npatch == N and w == h:
|
||||
return self.pos_embed
|
||||
pos_embed = self.pos_embed.float()
|
||||
class_pos_embed = pos_embed[:, 0]
|
||||
patch_pos_embed = pos_embed[:, 1:]
|
||||
dim = x.shape[-1]
|
||||
w0 = w // self.patch_size
|
||||
h0 = h // self.patch_size
|
||||
# we add a small number to avoid floating point error in the interpolation
|
||||
# see discussion at https://github.com/facebookresearch/dino/issues/8
|
||||
# DINOv2 with register modify the interpolate_offset from 0.1 to 0.0
|
||||
w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
|
||||
# w0, h0 = w0 + 0.1, h0 + 0.1
|
||||
|
||||
sqrt_N = math.sqrt(N)
|
||||
sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
|
||||
patch_pos_embed = nn.functional.interpolate(
|
||||
patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
|
||||
scale_factor=(sx, sy),
|
||||
# (int(w0), int(h0)), # to solve the upsampling shape issue
|
||||
mode="bicubic",
|
||||
antialias=self.interpolate_antialias
|
||||
)
|
||||
|
||||
assert int(w0) == patch_pos_embed.shape[-2]
|
||||
assert int(h0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
|
||||
|
||||
def prepare_tokens_with_masks(self, x, masks=None):
|
||||
B, nc, w, h = x.shape
|
||||
x = self.patch_embed(x)
|
||||
if masks is not None:
|
||||
x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
|
||||
|
||||
x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
|
||||
x = x + self.interpolate_pos_encoding(x, w, h)
|
||||
|
||||
if self.register_tokens is not None:
|
||||
x = torch.cat(
|
||||
(
|
||||
x[:, :1],
|
||||
self.register_tokens.expand(x.shape[0], -1, -1),
|
||||
x[:, 1:],
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_features_list(self, x_list, masks_list):
|
||||
x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
all_x = x
|
||||
output = []
|
||||
for x, masks in zip(all_x, masks_list):
|
||||
x_norm = self.norm(x)
|
||||
output.append(
|
||||
{
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def forward_features(self, x, masks=None):
|
||||
if isinstance(x, list):
|
||||
return self.forward_features_list(x, masks)
|
||||
|
||||
x = self.prepare_tokens_with_masks(x, masks)
|
||||
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
x_norm = self.norm(x)
|
||||
return {
|
||||
"x_norm_clstoken": x_norm[:, 0],
|
||||
"x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
|
||||
"x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
|
||||
"x_prenorm": x,
|
||||
"masks": masks,
|
||||
}
|
||||
|
||||
def _get_intermediate_layers_not_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
output, total_block_len = [], len(self.blocks)
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for i, blk in enumerate(self.blocks):
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def _get_intermediate_layers_chunked(self, x, n=1):
|
||||
x = self.prepare_tokens_with_masks(x)
|
||||
output, i, total_block_len = [], 0, len(self.blocks[-1])
|
||||
# If n is an int, take the n last blocks. If it's a list, take them
|
||||
blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
|
||||
for block_chunk in self.blocks:
|
||||
for blk in block_chunk[i:]: # Passing the nn.Identity()
|
||||
x = blk(x)
|
||||
if i in blocks_to_take:
|
||||
output.append(x)
|
||||
i += 1
|
||||
assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
|
||||
return output
|
||||
|
||||
def get_intermediate_layers(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
n: Union[int, Sequence] = 1, # Layers or n last layers to take
|
||||
reshape: bool = False,
|
||||
return_class_token: bool = False,
|
||||
norm=True
|
||||
) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
|
||||
if self.chunked_blocks:
|
||||
outputs = self._get_intermediate_layers_chunked(x, n)
|
||||
else:
|
||||
outputs = self._get_intermediate_layers_not_chunked(x, n)
|
||||
if norm:
|
||||
outputs = [self.norm(out) for out in outputs]
|
||||
class_tokens = [out[:, 0] for out in outputs]
|
||||
outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
|
||||
if reshape:
|
||||
B, _, w, h = x.shape
|
||||
outputs = [
|
||||
out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
|
||||
for out in outputs
|
||||
]
|
||||
if return_class_token:
|
||||
return tuple(zip(outputs, class_tokens))
|
||||
return tuple(outputs)
|
||||
|
||||
def forward(self, *args, is_training=False, **kwargs):
|
||||
ret = self.forward_features(*args, **kwargs)
|
||||
if is_training:
|
||||
return ret
|
||||
else:
|
||||
return self.head(ret["x_norm_clstoken"])
|
||||
|
||||
|
||||
def init_weights_vit_timm(module: nn.Module, name: str = ""):
|
||||
"""ViT weight initialization, original timm impl (for reproducibility)"""
|
||||
if isinstance(module, nn.Linear):
|
||||
trunc_normal_(module.weight, std=0.02)
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
|
||||
def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=384,
|
||||
depth=12,
|
||||
num_heads=6,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=768,
|
||||
depth=12,
|
||||
num_heads=12,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1024,
|
||||
depth=24,
|
||||
num_heads=16,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
|
||||
"""
|
||||
Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
|
||||
"""
|
||||
model = DinoVisionTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=1536,
|
||||
depth=40,
|
||||
num_heads=24,
|
||||
mlp_ratio=4,
|
||||
block_fn=partial(Block, attn_class=MemEffAttention),
|
||||
num_register_tokens=num_register_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def DINOv2(model_name):
|
||||
model_zoo = {
|
||||
"vits": vit_small,
|
||||
"vitb": vit_base,
|
||||
"vitl": vit_large,
|
||||
"vitg": vit_giant2
|
||||
}
|
||||
|
||||
return model_zoo[model_name](
|
||||
img_size=518,
|
||||
patch_size=14,
|
||||
init_values=1.0,
|
||||
ffn_layer="mlp" if model_name != "vitg" else "swiglufused",
|
||||
block_chunks=0,
|
||||
num_register_tokens=0,
|
||||
interpolate_antialias=False,
|
||||
interpolate_offset=0.1
|
||||
)
|
||||
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .mlp import Mlp
|
||||
from .patch_embed import PatchEmbed
|
||||
from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
|
||||
from .block import NestedTensorBlock
|
||||
from .attention import MemEffAttention
|
||||
@@ -0,0 +1,83 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
|
||||
|
||||
import logging
|
||||
|
||||
from torch import Tensor
|
||||
from torch import nn
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import memory_efficient_attention, unbind, fmha
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
logger.warning("xFormers not available")
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int = 8,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
attn_drop: float = 0.0,
|
||||
proj_drop: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj = nn.Linear(dim, dim, bias=proj_bias)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
|
||||
q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
|
||||
attn = q @ k.transpose(-2, -1)
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = self.attn_drop(attn)
|
||||
|
||||
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
class MemEffAttention(Attention):
|
||||
def forward(self, x: Tensor, attn_bias=None) -> Tensor:
|
||||
if not XFORMERS_AVAILABLE:
|
||||
assert attn_bias is None, "xFormers is required for nested tensors usage"
|
||||
return super().forward(x)
|
||||
|
||||
B, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
|
||||
|
||||
q, k, v = unbind(qkv, 2)
|
||||
|
||||
x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
|
||||
x = x.reshape([B, N, C])
|
||||
|
||||
x = self.proj(x)
|
||||
x = self.proj_drop(x)
|
||||
return x
|
||||
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
import logging
|
||||
from typing import Callable, List, Any, Tuple, Dict
|
||||
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
|
||||
from .attention import Attention, MemEffAttention
|
||||
from .drop_path import DropPath
|
||||
from .layer_scale import LayerScale
|
||||
from .mlp import Mlp
|
||||
|
||||
|
||||
logger = logging.getLogger("dinov2")
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import fmha
|
||||
from xformers.ops import scaled_index_add, index_select_cat
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
logger.warning("xFormers not available")
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qkv_bias: bool = False,
|
||||
proj_bias: bool = True,
|
||||
ffn_bias: bool = True,
|
||||
drop: float = 0.0,
|
||||
attn_drop: float = 0.0,
|
||||
init_values=None,
|
||||
drop_path: float = 0.0,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
|
||||
attn_class: Callable[..., nn.Module] = Attention,
|
||||
ffn_layer: Callable[..., nn.Module] = Mlp,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
# print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = attn_class(
|
||||
dim,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=qkv_bias,
|
||||
proj_bias=proj_bias,
|
||||
attn_drop=attn_drop,
|
||||
proj_drop=drop,
|
||||
)
|
||||
self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.norm2 = norm_layer(dim)
|
||||
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
self.mlp = ffn_layer(
|
||||
in_features=dim,
|
||||
hidden_features=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
drop=drop,
|
||||
bias=ffn_bias,
|
||||
)
|
||||
self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
|
||||
self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
self.sample_drop_ratio = drop_path
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
def attn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x)))
|
||||
|
||||
def ffn_residual_func(x: Tensor) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.1:
|
||||
# the overhead is compensated only for a drop path rate larger than 0.1
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
x = drop_add_residual_stochastic_depth(
|
||||
x,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
)
|
||||
elif self.training and self.sample_drop_ratio > 0.0:
|
||||
x = x + self.drop_path1(attn_residual_func(x))
|
||||
x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
|
||||
else:
|
||||
x = x + attn_residual_func(x)
|
||||
x = x + ffn_residual_func(x)
|
||||
return x
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth(
|
||||
x: Tensor,
|
||||
residual_func: Callable[[Tensor], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
) -> Tensor:
|
||||
# 1) extract subset using permutation
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
x_subset = x[brange]
|
||||
|
||||
# 2) apply residual_func to get residual
|
||||
residual = residual_func(x_subset)
|
||||
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
|
||||
# 3) add the residual
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
return x_plus_residual.view_as(x)
|
||||
|
||||
|
||||
def get_branges_scales(x, sample_drop_ratio=0.0):
|
||||
b, n, d = x.shape
|
||||
sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
|
||||
brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
|
||||
residual_scale_factor = b / sample_subset_size
|
||||
return brange, residual_scale_factor
|
||||
|
||||
|
||||
def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
|
||||
if scaling_vector is None:
|
||||
x_flat = x.flatten(1)
|
||||
residual = residual.flatten(1)
|
||||
x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
|
||||
else:
|
||||
x_plus_residual = scaled_index_add(
|
||||
x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
|
||||
)
|
||||
return x_plus_residual
|
||||
|
||||
|
||||
attn_bias_cache: Dict[Tuple, Any] = {}
|
||||
|
||||
|
||||
def get_attn_bias_and_cat(x_list, branges=None):
|
||||
"""
|
||||
this will perform the index select, cat the tensors, and provide the attn_bias from cache
|
||||
"""
|
||||
batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
|
||||
all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
|
||||
if all_shapes not in attn_bias_cache.keys():
|
||||
seqlens = []
|
||||
for b, x in zip(batch_sizes, x_list):
|
||||
for _ in range(b):
|
||||
seqlens.append(x.shape[1])
|
||||
attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
|
||||
attn_bias._batch_sizes = batch_sizes
|
||||
attn_bias_cache[all_shapes] = attn_bias
|
||||
|
||||
if branges is not None:
|
||||
cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
|
||||
else:
|
||||
tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
|
||||
cat_tensors = torch.cat(tensors_bs1, dim=1)
|
||||
|
||||
return attn_bias_cache[all_shapes], cat_tensors
|
||||
|
||||
|
||||
def drop_add_residual_stochastic_depth_list(
|
||||
x_list: List[Tensor],
|
||||
residual_func: Callable[[Tensor, Any], Tensor],
|
||||
sample_drop_ratio: float = 0.0,
|
||||
scaling_vector=None,
|
||||
) -> Tensor:
|
||||
# 1) generate random set of indices for dropping samples in the batch
|
||||
branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
|
||||
branges = [s[0] for s in branges_scales]
|
||||
residual_scale_factors = [s[1] for s in branges_scales]
|
||||
|
||||
# 2) get attention bias and index+concat the tensors
|
||||
attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
|
||||
|
||||
# 3) apply residual_func to get residual, and split the result
|
||||
residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
|
||||
|
||||
outputs = []
|
||||
for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
|
||||
outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
|
||||
return outputs
|
||||
|
||||
|
||||
class NestedTensorBlock(Block):
|
||||
def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
x_list contains a list of tensors to nest together and run
|
||||
"""
|
||||
assert isinstance(self.attn, MemEffAttention)
|
||||
|
||||
if self.training and self.sample_drop_ratio > 0.0:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.attn(self.norm1(x), attn_bias=attn_bias)
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.mlp(self.norm2(x))
|
||||
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=attn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
|
||||
)
|
||||
x_list = drop_add_residual_stochastic_depth_list(
|
||||
x_list,
|
||||
residual_func=ffn_residual_func,
|
||||
sample_drop_ratio=self.sample_drop_ratio,
|
||||
scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
|
||||
)
|
||||
return x_list
|
||||
else:
|
||||
|
||||
def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
|
||||
|
||||
def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
|
||||
return self.ls2(self.mlp(self.norm2(x)))
|
||||
|
||||
attn_bias, x = get_attn_bias_and_cat(x_list)
|
||||
x = x + attn_residual_func(x, attn_bias=attn_bias)
|
||||
x = x + ffn_residual_func(x)
|
||||
return attn_bias.split(x)
|
||||
|
||||
def forward(self, x_or_x_list):
|
||||
if isinstance(x_or_x_list, Tensor):
|
||||
return super().forward(x_or_x_list)
|
||||
elif isinstance(x_or_x_list, list):
|
||||
assert XFORMERS_AVAILABLE, "Please install xFormers for nested tensors usage"
|
||||
return self.forward_nested(x_or_x_list)
|
||||
else:
|
||||
raise AssertionError
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
|
||||
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0:
|
||||
random_tensor.div_(keep_prob)
|
||||
output = x * random_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
||||
|
||||
def __init__(self, drop_prob=None):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training)
|
||||
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
|
||||
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch import nn
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
init_values: Union[float, Tensor] = 1e-5,
|
||||
inplace: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.inplace = inplace
|
||||
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x.mul_(self.gamma) if self.inplace else x * self.gamma
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
|
||||
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = nn.GELU,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
|
||||
self.act = act_layer()
|
||||
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
self.drop = nn.Dropout(drop)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop(x)
|
||||
return x
|
||||
@@ -0,0 +1,89 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# References:
|
||||
# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
|
||||
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
|
||||
|
||||
from typing import Callable, Optional, Tuple, Union
|
||||
|
||||
from torch import Tensor
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def make_2tuple(x):
|
||||
if isinstance(x, tuple):
|
||||
assert len(x) == 2
|
||||
return x
|
||||
|
||||
assert isinstance(x, int)
|
||||
return (x, x)
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""
|
||||
2D image to patch embedding: (B,C,H,W) -> (B,N,D)
|
||||
|
||||
Args:
|
||||
img_size: Image size.
|
||||
patch_size: Patch token size.
|
||||
in_chans: Number of input image channels.
|
||||
embed_dim: Number of linear projection output channels.
|
||||
norm_layer: Normalization layer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size: Union[int, Tuple[int, int]] = 224,
|
||||
patch_size: Union[int, Tuple[int, int]] = 16,
|
||||
in_chans: int = 3,
|
||||
embed_dim: int = 768,
|
||||
norm_layer: Optional[Callable] = None,
|
||||
flatten_embedding: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
image_HW = make_2tuple(img_size)
|
||||
patch_HW = make_2tuple(patch_size)
|
||||
patch_grid_size = (
|
||||
image_HW[0] // patch_HW[0],
|
||||
image_HW[1] // patch_HW[1],
|
||||
)
|
||||
|
||||
self.img_size = image_HW
|
||||
self.patch_size = patch_HW
|
||||
self.patches_resolution = patch_grid_size
|
||||
self.num_patches = patch_grid_size[0] * patch_grid_size[1]
|
||||
|
||||
self.in_chans = in_chans
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
self.flatten_embedding = flatten_embedding
|
||||
|
||||
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
|
||||
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
_, _, H, W = x.shape
|
||||
patch_H, patch_W = self.patch_size
|
||||
|
||||
assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
|
||||
assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
|
||||
|
||||
x = self.proj(x) # B C H W
|
||||
H, W = x.size(2), x.size(3)
|
||||
x = x.flatten(2).transpose(1, 2) # B HW C
|
||||
x = self.norm(x)
|
||||
if not self.flatten_embedding:
|
||||
x = x.reshape(-1, H, W, self.embed_dim) # B H W C
|
||||
return x
|
||||
|
||||
def flops(self) -> float:
|
||||
Ho, Wo = self.patches_resolution
|
||||
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
|
||||
if self.norm is not None:
|
||||
flops += Ho * Wo * self.embed_dim
|
||||
return flops
|
||||
@@ -0,0 +1,63 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from typing import Callable, Optional
|
||||
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SwiGLUFFN(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = None,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
|
||||
self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
x12 = self.w12(x)
|
||||
x1, x2 = x12.chunk(2, dim=-1)
|
||||
hidden = F.silu(x1) * x2
|
||||
return self.w3(hidden)
|
||||
|
||||
|
||||
try:
|
||||
from xformers.ops import SwiGLU
|
||||
|
||||
XFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
SwiGLU = SwiGLUFFN
|
||||
XFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class SwiGLUFFNFused(SwiGLU):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: Optional[int] = None,
|
||||
out_features: Optional[int] = None,
|
||||
act_layer: Callable[..., nn.Module] = None,
|
||||
drop: float = 0.0,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
out_features = out_features or in_features
|
||||
hidden_features = hidden_features or in_features
|
||||
hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
|
||||
super().__init__(
|
||||
in_features=in_features,
|
||||
hidden_features=hidden_features,
|
||||
out_features=out_features,
|
||||
bias=bias,
|
||||
)
|
||||
@@ -0,0 +1,220 @@
|
||||
import cv2
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import Compose
|
||||
|
||||
from custom_controlnet_aux.depth_anything_v2.dinov2 import DINOv2
|
||||
from custom_controlnet_aux.depth_anything_v2.util.blocks import FeatureFusionBlock, _make_scratch
|
||||
from custom_controlnet_aux.depth_anything_v2.util.transform import Resize, NormalizeImage, PrepareForNet
|
||||
|
||||
|
||||
def _make_fusion_block(features, use_bn, size=None):
|
||||
return FeatureFusionBlock(
|
||||
features,
|
||||
nn.ReLU(False),
|
||||
deconv=False,
|
||||
bn=use_bn,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=size,
|
||||
)
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_feature, out_feature):
|
||||
super().__init__()
|
||||
|
||||
self.conv_block = nn.Sequential(
|
||||
nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),
|
||||
nn.BatchNorm2d(out_feature),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv_block(x)
|
||||
|
||||
|
||||
class DPTHead(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
features=256,
|
||||
use_bn=False,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_clstoken=False
|
||||
):
|
||||
super(DPTHead, self).__init__()
|
||||
|
||||
self.use_clstoken = use_clstoken
|
||||
|
||||
self.projects = nn.ModuleList([
|
||||
nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channel,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
) for out_channel in out_channels
|
||||
])
|
||||
|
||||
self.resize_layers = nn.ModuleList([
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[0],
|
||||
out_channels=out_channels[0],
|
||||
kernel_size=4,
|
||||
stride=4,
|
||||
padding=0),
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=out_channels[1],
|
||||
out_channels=out_channels[1],
|
||||
kernel_size=2,
|
||||
stride=2,
|
||||
padding=0),
|
||||
nn.Identity(),
|
||||
nn.Conv2d(
|
||||
in_channels=out_channels[3],
|
||||
out_channels=out_channels[3],
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1)
|
||||
])
|
||||
|
||||
if use_clstoken:
|
||||
self.readout_projects = nn.ModuleList()
|
||||
for _ in range(len(self.projects)):
|
||||
self.readout_projects.append(
|
||||
nn.Sequential(
|
||||
nn.Linear(2 * in_channels, in_channels),
|
||||
nn.GELU()))
|
||||
|
||||
self.scratch = _make_scratch(
|
||||
out_channels,
|
||||
features,
|
||||
groups=1,
|
||||
expand=False,
|
||||
)
|
||||
|
||||
self.scratch.stem_transpose = None
|
||||
|
||||
self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
|
||||
self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
|
||||
|
||||
head_features_1 = features
|
||||
head_features_2 = 32
|
||||
|
||||
self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
|
||||
self.scratch.output_conv2 = nn.Sequential(
|
||||
nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
|
||||
nn.ReLU(True),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
def forward(self, out_features, patch_h, patch_w):
|
||||
out = []
|
||||
for i, x in enumerate(out_features):
|
||||
if self.use_clstoken:
|
||||
x, cls_token = x[0], x[1]
|
||||
readout = cls_token.unsqueeze(1).expand_as(x)
|
||||
x = self.readout_projects[i](torch.cat((x, readout), -1))
|
||||
else:
|
||||
x = x[0]
|
||||
|
||||
x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
|
||||
|
||||
x = self.projects[i](x)
|
||||
x = self.resize_layers[i](x)
|
||||
|
||||
out.append(x)
|
||||
|
||||
layer_1, layer_2, layer_3, layer_4 = out
|
||||
|
||||
layer_1_rn = self.scratch.layer1_rn(layer_1)
|
||||
layer_2_rn = self.scratch.layer2_rn(layer_2)
|
||||
layer_3_rn = self.scratch.layer3_rn(layer_3)
|
||||
layer_4_rn = self.scratch.layer4_rn(layer_4)
|
||||
|
||||
path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
|
||||
path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
|
||||
path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
|
||||
path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
|
||||
|
||||
out = self.scratch.output_conv1(path_1)
|
||||
out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
|
||||
out = self.scratch.output_conv2(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class DepthAnythingV2(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
encoder='vitl',
|
||||
features=256,
|
||||
out_channels=[256, 512, 1024, 1024],
|
||||
use_bn=False,
|
||||
use_clstoken=False
|
||||
):
|
||||
super(DepthAnythingV2, self).__init__()
|
||||
|
||||
self.intermediate_layer_idx = {
|
||||
'vits': [2, 5, 8, 11],
|
||||
'vitb': [2, 5, 8, 11],
|
||||
'vitl': [4, 11, 17, 23],
|
||||
'vitg': [9, 19, 29, 39]
|
||||
}
|
||||
|
||||
self.encoder = encoder
|
||||
self.pretrained = DINOv2(model_name=encoder)
|
||||
|
||||
self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
|
||||
|
||||
def forward(self, x, max_depth):
|
||||
patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
|
||||
|
||||
features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)
|
||||
|
||||
depth = self.depth_head(features, patch_h, patch_w) * max_depth
|
||||
|
||||
return depth.squeeze(1)
|
||||
|
||||
@torch.no_grad()
|
||||
def infer_image(self, raw_image, input_size=518, max_depth=20.0):
|
||||
image, (h, w) = self.image2tensor(raw_image, input_size)
|
||||
|
||||
depth = self.forward(image, max_depth)
|
||||
|
||||
depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
|
||||
|
||||
return depth.cpu().numpy()
|
||||
|
||||
def image2tensor(self, raw_image, input_size=518):
|
||||
transform = Compose([
|
||||
Resize(
|
||||
width=input_size,
|
||||
height=input_size,
|
||||
resize_target=False,
|
||||
keep_aspect_ratio=True,
|
||||
ensure_multiple_of=14,
|
||||
resize_method='lower_bound',
|
||||
image_interpolation_method=cv2.INTER_CUBIC,
|
||||
),
|
||||
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
PrepareForNet(),
|
||||
])
|
||||
|
||||
h, w = raw_image.shape[:2]
|
||||
|
||||
image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0
|
||||
|
||||
image = transform({'image': image})['image']
|
||||
image = torch.from_numpy(image).unsqueeze(0)
|
||||
|
||||
DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
|
||||
image = image.to(DEVICE)
|
||||
|
||||
return image, (h, w)
|
||||
@@ -0,0 +1,148 @@
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
|
||||
scratch = nn.Module()
|
||||
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape
|
||||
out_shape3 = out_shape
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape
|
||||
|
||||
if expand:
|
||||
out_shape1 = out_shape
|
||||
out_shape2 = out_shape * 2
|
||||
out_shape3 = out_shape * 4
|
||||
if len(in_shape) >= 4:
|
||||
out_shape4 = out_shape * 8
|
||||
|
||||
scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
if len(in_shape) >= 4:
|
||||
scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)
|
||||
|
||||
return scratch
|
||||
|
||||
|
||||
class ResidualConvUnit(nn.Module):
|
||||
"""Residual convolution module.
|
||||
"""
|
||||
|
||||
def __init__(self, features, activation, bn):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.bn = bn
|
||||
|
||||
self.groups=1
|
||||
|
||||
self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
|
||||
|
||||
self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
|
||||
|
||||
if self.bn == True:
|
||||
self.bn1 = nn.BatchNorm2d(features)
|
||||
self.bn2 = nn.BatchNorm2d(features)
|
||||
|
||||
self.activation = activation
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass.
|
||||
|
||||
Args:
|
||||
x (tensor): input
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
|
||||
out = self.activation(x)
|
||||
out = self.conv1(out)
|
||||
if self.bn == True:
|
||||
out = self.bn1(out)
|
||||
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
if self.bn == True:
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.groups > 1:
|
||||
out = self.conv_merge(out)
|
||||
|
||||
return self.skip_add.add(out, x)
|
||||
|
||||
|
||||
class FeatureFusionBlock(nn.Module):
|
||||
"""Feature fusion block.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
features,
|
||||
activation,
|
||||
deconv=False,
|
||||
bn=False,
|
||||
expand=False,
|
||||
align_corners=True,
|
||||
size=None
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
features (int): number of features
|
||||
"""
|
||||
super(FeatureFusionBlock, self).__init__()
|
||||
|
||||
self.deconv = deconv
|
||||
self.align_corners = align_corners
|
||||
|
||||
self.groups=1
|
||||
|
||||
self.expand = expand
|
||||
out_features = features
|
||||
if self.expand == True:
|
||||
out_features = features // 2
|
||||
|
||||
self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
|
||||
|
||||
self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
|
||||
self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
|
||||
|
||||
self.skip_add = nn.quantized.FloatFunctional()
|
||||
|
||||
self.size=size
|
||||
|
||||
def forward(self, *xs, size=None):
|
||||
"""Forward pass.
|
||||
|
||||
Returns:
|
||||
tensor: output
|
||||
"""
|
||||
output = xs[0]
|
||||
|
||||
if len(xs) == 2:
|
||||
res = self.resConfUnit1(xs[1])
|
||||
output = self.skip_add.add(output, res)
|
||||
|
||||
output = self.resConfUnit2(output)
|
||||
|
||||
if (size is None) and (self.size is None):
|
||||
modifier = {"scale_factor": 2}
|
||||
elif size is None:
|
||||
modifier = {"size": self.size}
|
||||
else:
|
||||
modifier = {"size": size}
|
||||
|
||||
output = nn.functional.interpolate(output, **modifier, mode="bilinear", align_corners=self.align_corners)
|
||||
|
||||
output = self.out_conv(output)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,158 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
class Resize(object):
|
||||
"""Resize sample to given size (width, height).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width,
|
||||
height,
|
||||
resize_target=True,
|
||||
keep_aspect_ratio=False,
|
||||
ensure_multiple_of=1,
|
||||
resize_method="lower_bound",
|
||||
image_interpolation_method=cv2.INTER_AREA,
|
||||
):
|
||||
"""Init.
|
||||
|
||||
Args:
|
||||
width (int): desired output width
|
||||
height (int): desired output height
|
||||
resize_target (bool, optional):
|
||||
True: Resize the full sample (image, mask, target).
|
||||
False: Resize image only.
|
||||
Defaults to True.
|
||||
keep_aspect_ratio (bool, optional):
|
||||
True: Keep the aspect ratio of the input sample.
|
||||
Output sample might not have the given width and height, and
|
||||
resize behaviour depends on the parameter 'resize_method'.
|
||||
Defaults to False.
|
||||
ensure_multiple_of (int, optional):
|
||||
Output width and height is constrained to be multiple of this parameter.
|
||||
Defaults to 1.
|
||||
resize_method (str, optional):
|
||||
"lower_bound": Output will be at least as large as the given size.
|
||||
"upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
|
||||
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
|
||||
Defaults to "lower_bound".
|
||||
"""
|
||||
self.__width = width
|
||||
self.__height = height
|
||||
|
||||
self.__resize_target = resize_target
|
||||
self.__keep_aspect_ratio = keep_aspect_ratio
|
||||
self.__multiple_of = ensure_multiple_of
|
||||
self.__resize_method = resize_method
|
||||
self.__image_interpolation_method = image_interpolation_method
|
||||
|
||||
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
|
||||
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if max_val is not None and y > max_val:
|
||||
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
if y < min_val:
|
||||
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
||||
|
||||
return y
|
||||
|
||||
def get_size(self, width, height):
|
||||
# determine new height and width
|
||||
scale_height = self.__height / height
|
||||
scale_width = self.__width / width
|
||||
|
||||
if self.__keep_aspect_ratio:
|
||||
if self.__resize_method == "lower_bound":
|
||||
# scale such that output size is lower bound
|
||||
if scale_width > scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "upper_bound":
|
||||
# scale such that output size is upper bound
|
||||
if scale_width < scale_height:
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
elif self.__resize_method == "minimal":
|
||||
# scale as least as possbile
|
||||
if abs(1 - scale_width) < abs(1 - scale_height):
|
||||
# fit width
|
||||
scale_height = scale_width
|
||||
else:
|
||||
# fit height
|
||||
scale_width = scale_height
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
if self.__resize_method == "lower_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
|
||||
elif self.__resize_method == "upper_bound":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
|
||||
elif self.__resize_method == "minimal":
|
||||
new_height = self.constrain_to_multiple_of(scale_height * height)
|
||||
new_width = self.constrain_to_multiple_of(scale_width * width)
|
||||
else:
|
||||
raise ValueError(f"resize_method {self.__resize_method} not implemented")
|
||||
|
||||
return (new_width, new_height)
|
||||
|
||||
def __call__(self, sample):
|
||||
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
|
||||
|
||||
# resize sample
|
||||
sample["image"] = cv2.resize(sample["image"], (width, height), interpolation=self.__image_interpolation_method)
|
||||
|
||||
if self.__resize_target:
|
||||
if "depth" in sample:
|
||||
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = cv2.resize(sample["mask"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
"""Normlize image by given mean and std.
|
||||
"""
|
||||
|
||||
def __init__(self, mean, std):
|
||||
self.__mean = mean
|
||||
self.__std = std
|
||||
|
||||
def __call__(self, sample):
|
||||
sample["image"] = (sample["image"] - self.__mean) / self.__std
|
||||
|
||||
return sample
|
||||
|
||||
|
||||
class PrepareForNet(object):
|
||||
"""Prepare sample for usage as network input.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, sample):
|
||||
image = np.transpose(sample["image"], (2, 0, 1))
|
||||
sample["image"] = np.ascontiguousarray(image).astype(np.float32)
|
||||
|
||||
if "depth" in sample:
|
||||
depth = sample["depth"].astype(np.float32)
|
||||
sample["depth"] = np.ascontiguousarray(depth)
|
||||
|
||||
if "mask" in sample:
|
||||
sample["mask"] = sample["mask"].astype(np.float32)
|
||||
sample["mask"] = np.ascontiguousarray(sample["mask"])
|
||||
|
||||
return sample
|
||||
@@ -0,0 +1,40 @@
|
||||
from custom_controlnet_aux.diffusion_edge.model import DiffusionEdge, prepare_args
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
from custom_controlnet_aux.util import HWC3, common_input_validate, resize_image_with_pad, custom_hf_download, DIFFUSION_EDGE_MODEL_NAME
|
||||
|
||||
class DiffusionEdgeDetector:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
self.device = "cpu"
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_or_path=DIFFUSION_EDGE_MODEL_NAME, filename="diffusion_edge_indoor.pt"):
|
||||
model_path = custom_hf_download(pretrained_model_or_path, filename)
|
||||
model = DiffusionEdge(prepare_args(model_path))
|
||||
return cls(model)
|
||||
|
||||
def to(self, device):
|
||||
self.model.to(device)
|
||||
self.device = device
|
||||
return self
|
||||
|
||||
def __call__(self, input_image, detect_resolution=512, patch_batch_size=8, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
|
||||
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
|
||||
input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
|
||||
|
||||
with torch.no_grad():
|
||||
input_image = rearrange(torch.from_numpy(input_image), "h w c -> 1 c h w")
|
||||
input_image = input_image.float() / 255.
|
||||
line = self.model(input_image, patch_batch_size)
|
||||
line = rearrange(line, "1 c h w -> h w c")
|
||||
|
||||
detected_map = line.cpu().numpy().__mul__(255.).astype(np.uint8)
|
||||
detected_map = remove_pad(HWC3(detected_map))
|
||||
|
||||
if output_type == "pil":
|
||||
detected_map = Image.fromarray(detected_map)
|
||||
|
||||
return detected_map
|
||||
@@ -0,0 +1,74 @@
|
||||
model:
|
||||
model_type: const_sde
|
||||
model_name: cond_unet
|
||||
image_size: [320, 320]
|
||||
input_keys: ['image', 'cond']
|
||||
ckpt_path:
|
||||
ignore_keys: [ ]
|
||||
only_model: False
|
||||
timesteps: 1000
|
||||
train_sample: -1
|
||||
sampling_timesteps: 1
|
||||
loss_type: l2
|
||||
objective: pred_noise
|
||||
start_dist: normal
|
||||
perceptual_weight: 0
|
||||
scale_factor: 0.3
|
||||
scale_by_std: True
|
||||
default_scale: True
|
||||
scale_by_softsign: False
|
||||
eps: !!float 1e-4
|
||||
weighting_loss: False
|
||||
first_stage:
|
||||
embed_dim: 3
|
||||
lossconfig:
|
||||
disc_start: 50001
|
||||
kl_weight: 0.000001
|
||||
disc_weight: 0.5
|
||||
disc_in_channels: 1
|
||||
ddconfig:
|
||||
double_z: True
|
||||
z_channels: 3
|
||||
resolution: [ 320, 320 ]
|
||||
in_channels: 1
|
||||
out_ch: 1
|
||||
ch: 128
|
||||
ch_mult: [ 1,2,4 ] # num_down = len(ch_mult)-1
|
||||
num_res_blocks: 2
|
||||
attn_resolutions: [ ]
|
||||
dropout: 0.0
|
||||
ckpt_path:
|
||||
unet:
|
||||
dim: 128
|
||||
cond_net: swin
|
||||
without_pretrain: False
|
||||
channels: 3
|
||||
out_mul: 1
|
||||
dim_mults: [ 1, 2, 4, 4, ] # num_down = len(dim_mults)
|
||||
cond_in_dim: 3
|
||||
cond_dim: 128
|
||||
cond_dim_mults: [ 2, 4 ] # num_down = len(cond_dim_mults)
|
||||
# window_sizes1: [ [4, 4], [2, 2], [1, 1], [1, 1] ]
|
||||
# window_sizes2: [ [4, 4], [2, 2], [1, 1], [1, 1] ]
|
||||
window_sizes1: [ [ 8, 8 ], [ 4, 4 ], [ 2, 2 ], [ 1, 1 ] ]
|
||||
window_sizes2: [ [ 8, 8 ], [ 4, 4 ], [ 2, 2 ], [ 1, 1 ] ]
|
||||
fourier_scale: 16
|
||||
cond_pe: False
|
||||
num_pos_feats: 128
|
||||
cond_feature_size: [ 80, 80 ]
|
||||
|
||||
data:
|
||||
name: edge
|
||||
img_folder: '/data/yeyunfan/edge_detection_datasets/datasets/BSDS_test'
|
||||
augment_horizontal_flip: True
|
||||
batch_size: 8
|
||||
num_workers: 4
|
||||
|
||||
sampler:
|
||||
sample_type: "slide"
|
||||
stride: [240, 240]
|
||||
batch_size: 1
|
||||
sample_num: 300
|
||||
use_ema: True
|
||||
save_folder:
|
||||
ckpt_path:
|
||||
@@ -0,0 +1 @@
|
||||
# from custom_controlnet_aux.diffusion_edge.denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet, Trainer
|
||||
@@ -0,0 +1,598 @@
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
import torch.utils.data as data
|
||||
import torch.nn as nn
|
||||
from pathlib import Path
|
||||
from functools import partial
|
||||
from custom_controlnet_aux.diffusion_edge.denoising_diffusion_pytorch.utils import exists, convert_image_to_fn, normalize_to_neg_one_to_one
|
||||
from PIL import Image, ImageDraw
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
import torchvision.transforms.functional as F2
|
||||
import torchvision.transforms as transforms
|
||||
import torchvision.datasets as datasets
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
import os
|
||||
import pickle
|
||||
import numpy as np
|
||||
import copy
|
||||
import albumentations
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
def get_imgs_list(imgs_dir):
|
||||
imgs_list = os.listdir(imgs_dir)
|
||||
imgs_list.sort()
|
||||
return [os.path.join(imgs_dir, f) for f in imgs_list if f.endswith('.jpg') or f.endswith('.JPG')or f.endswith('.png') or f.endswith('.pgm') or f.endswith('.ppm')]
|
||||
|
||||
|
||||
def fit_img_postfix(img_path):
|
||||
if not os.path.exists(img_path) and img_path.endswith(".jpg"):
|
||||
img_path = img_path[:-4] + ".png"
|
||||
if not os.path.exists(img_path) and img_path.endswith(".png"):
|
||||
img_path = img_path[:-4] + ".jpg"
|
||||
return img_path
|
||||
|
||||
|
||||
class AdaptEdgeDataset(data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
# mask_folder,
|
||||
image_size,
|
||||
exts = ['png', 'jpg'],
|
||||
augment_horizontal_flip = False,
|
||||
convert_image_to = None,
|
||||
normalize_to_neg_one_to_one=True,
|
||||
split='train',
|
||||
# inter_type='bicubic',
|
||||
# down=4,
|
||||
threshold=0.3, use_uncertainty=False
|
||||
):
|
||||
super().__init__()
|
||||
# self.img_folder = Path(img_folder)
|
||||
# self.edge_folder = Path(os.path.join(data_root, f'gt_imgs'))
|
||||
# self.img_folder = Path(os.path.join(data_root, f'imgs'))
|
||||
# self.edge_folder = Path(os.path.join(data_root, "edge", "aug"))
|
||||
# self.img_folder = Path(os.path.join(data_root, "image", "aug"))
|
||||
self.data_root = data_root
|
||||
self.image_size = image_size
|
||||
|
||||
# self.edge_paths = [p for ext in exts for p in self.edge_folder.rglob(f'*.{ext}')]
|
||||
# self.img_paths = [(self.img_folder / item.parent.name / f'{item.stem}.jpg') for item in self.edge_paths]
|
||||
# self.img_paths = [(self.img_folder / f'{item.stem}.jpg') for item in self.edge_paths]
|
||||
|
||||
self.threshold = threshold * 256
|
||||
self.use_uncertainty = use_uncertainty
|
||||
self.normalize_to_neg_one_to_one = normalize_to_neg_one_to_one
|
||||
|
||||
maybe_convert_fn = partial(convert_image_to_fn, convert_image_to) if exists(convert_image_to) else Identity()
|
||||
|
||||
# self.normalize_to_neg_one_to_one = normalize_to_neg_one_to_one
|
||||
# self.random_crop = RandomCrop(size=image_size)
|
||||
# self.transform = Compose([
|
||||
# # Lambda(maybe_convert_fn),
|
||||
# # Resize(image_size, interpolation=3, interpolation2=0),
|
||||
# Resize(image_size, interpolation=InterpolationMode.BILINEAR, interpolation2=InterpolationMode.NEAREST),
|
||||
# RandomHorizontalFlip() if augment_horizontal_flip else Identity(),
|
||||
# # RandomCrop(image_size),
|
||||
# ToTensor()
|
||||
# ])
|
||||
self.data_list = self.build_list()
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
# Resize(self.image_size, interpolation=InterpolationMode.BILINEAR, interpolation2=InterpolationMode.NEAREST),
|
||||
transforms.ToTensor()])
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_list)
|
||||
|
||||
|
||||
def read_img(self, image_path):
|
||||
with open(image_path, 'rb') as f:
|
||||
img = Image.open(f)
|
||||
img = img.convert('RGB')
|
||||
|
||||
raw_width, raw_height = img.size
|
||||
# width = int(raw_width / 32) * 32
|
||||
# height = int(raw_height / 32) * 32
|
||||
# img = img.resize((width, height), Image.Resampling.BILINEAR)
|
||||
# # print("img.size:", img.size)
|
||||
# img = self.transform(img)
|
||||
|
||||
return img, (raw_width, raw_height)
|
||||
|
||||
def read_lb(self, lb_path):
|
||||
lb_data = Image.open(lb_path)
|
||||
|
||||
width, height = lb_data.size
|
||||
width = int(width / 32) * 32
|
||||
height = int(height / 32) * 32
|
||||
lb_data = lb_data.resize((width, height), Image.Resampling.BILINEAR)
|
||||
# print("lb_data.size:", lb_data.size)
|
||||
lb = np.array(lb_data, dtype=np.float32)
|
||||
if lb.ndim == 3:
|
||||
lb = np.squeeze(lb[:, :, 0])
|
||||
assert lb.ndim == 2
|
||||
threshold = self.threshold
|
||||
lb = lb[np.newaxis, :, :]
|
||||
|
||||
lb[lb == 0] = 0
|
||||
|
||||
# ---------- important ----------
|
||||
if self.use_uncertainty:
|
||||
lb[np.logical_and(lb > 0, lb < threshold)] = 2
|
||||
else:
|
||||
lb[np.logical_and(lb > 0, lb < threshold)] /= 255.
|
||||
|
||||
lb[lb >= threshold] = 1
|
||||
return lb
|
||||
|
||||
def build_list(self):
|
||||
data_root = os.path.abspath(self.data_root)
|
||||
images_path = os.path.join(data_root, 'image', "raw")
|
||||
labels_path = os.path.join(data_root, 'edge', "raw")
|
||||
|
||||
samples = []
|
||||
for directory_name in os.listdir(images_path):
|
||||
image_directories = os.path.join(images_path, directory_name)
|
||||
for file_name_ext in os.listdir(image_directories):
|
||||
file_name = os.path.basename(file_name_ext)
|
||||
image_path = fit_img_postfix(os.path.join(images_path, directory_name, file_name))
|
||||
lb_path = fit_img_postfix(os.path.join(labels_path, directory_name, file_name))
|
||||
samples.append((image_path, lb_path))
|
||||
return samples
|
||||
|
||||
def __getitem__(self, index):
|
||||
img_path, edge_path = self.data_list[index]
|
||||
# edge_path = self.edge_paths[index]
|
||||
# img_path = self.img_paths[index]
|
||||
img_name = os.path.basename(img_path)
|
||||
|
||||
img, raw_size = self.read_img(img_path)
|
||||
edge = self.read_lb(edge_path)
|
||||
|
||||
# print("-------hhhhhhhhhhhhh--------:", img.shape, edge.shape)
|
||||
# edge = Image.open(edge_path).convert('L')
|
||||
# # default to score-sde preprocessing
|
||||
# mask = Image.open(img_path).convert('RGB')
|
||||
# edge, img = self.transform(edge, mask)
|
||||
if self.normalize_to_neg_one_to_one: # transform to [-1, 1]
|
||||
edge = normalize_to_neg_one_to_one(edge)
|
||||
img = normalize_to_neg_one_to_one(img)
|
||||
return {'image': edge, 'cond': img, 'raw_size': raw_size, 'img_name': img_name}
|
||||
|
||||
class EdgeDataset(data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
# mask_folder,
|
||||
image_size,
|
||||
exts = ['png', 'jpg'],
|
||||
augment_horizontal_flip = True,
|
||||
convert_image_to = None,
|
||||
normalize_to_neg_one_to_one=True,
|
||||
split='train',
|
||||
# inter_type='bicubic',
|
||||
# down=4,
|
||||
threshold=0.3, use_uncertainty=False, cfg={}
|
||||
):
|
||||
super().__init__()
|
||||
# self.img_folder = Path(img_folder)
|
||||
# self.edge_folder = Path(os.path.join(data_root, f'gt_imgs'))
|
||||
# self.img_folder = Path(os.path.join(data_root, f'imgs'))
|
||||
# self.edge_folder = Path(os.path.join(data_root, "edge", "aug"))
|
||||
# self.img_folder = Path(os.path.join(data_root, "image", "aug"))
|
||||
self.data_root = data_root
|
||||
self.image_size = image_size
|
||||
|
||||
# self.edge_paths = [p for ext in exts for p in self.edge_folder.rglob(f'*.{ext}')]
|
||||
# self.img_paths = [(self.img_folder / item.parent.name / f'{item.stem}.jpg') for item in self.edge_paths]
|
||||
# self.img_paths = [(self.img_folder / f'{item.stem}.jpg') for item in self.edge_paths]
|
||||
|
||||
self.threshold = threshold * 255
|
||||
self.use_uncertainty = use_uncertainty
|
||||
self.normalize_to_neg_one_to_one = normalize_to_neg_one_to_one
|
||||
|
||||
maybe_convert_fn = partial(convert_image_to_fn, convert_image_to) if exists(convert_image_to) else Identity()
|
||||
|
||||
self.data_list = self.build_list()
|
||||
|
||||
# self.transform = Compose([
|
||||
# Resize(image_size),
|
||||
# RandomHorizontalFlip() if augment_horizontal_flip else Identity(),
|
||||
# ToTensor()
|
||||
# ])
|
||||
crop_type = cfg.get('crop_type') if 'crop_type' in cfg else 'rand_crop'
|
||||
if crop_type == 'rand_crop':
|
||||
self.transform = Compose([
|
||||
RandomCrop(image_size),
|
||||
RandomHorizontalFlip() if augment_horizontal_flip else Identity(),
|
||||
ToTensor()
|
||||
])
|
||||
elif crop_type == 'rand_resize_crop':
|
||||
self.transform = Compose([
|
||||
RandomResizeCrop(image_size),
|
||||
RandomHorizontalFlip() if augment_horizontal_flip else Identity(),
|
||||
ToTensor()
|
||||
])
|
||||
print("crop_type:", crop_type)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_list)
|
||||
|
||||
|
||||
def read_img(self, image_path):
|
||||
with open(image_path, 'rb') as f:
|
||||
img = Image.open(f)
|
||||
img = img.convert('RGB')
|
||||
|
||||
raw_width, raw_height = img.size
|
||||
# width = int(raw_width / 32) * 32
|
||||
# height = int(raw_height / 32) * 32
|
||||
# img = img.resize((width, height), Image.Resampling.BILINEAR)
|
||||
# # print("img.size:", img.size)
|
||||
# img = self.transform(img)
|
||||
|
||||
return img, (raw_width, raw_height)
|
||||
|
||||
def read_lb(self, lb_path):
|
||||
lb_data = Image.open(lb_path).convert('L')
|
||||
lb = np.array(lb_data).astype(np.float32)
|
||||
# width, height = lb_data.size
|
||||
# width = int(width / 32) * 32
|
||||
# height = int(height / 32) * 32
|
||||
# lb_data = lb_data.resize((width, height), Image.Resampling.BILINEAR)
|
||||
# print("lb_data.size:", lb_data.size)
|
||||
# lb = np.array(lb_data, dtype=np.float32)
|
||||
# if lb.ndim == 3:
|
||||
# lb = np.squeeze(lb[:, :, 0])
|
||||
# assert lb.ndim == 2
|
||||
threshold = self.threshold
|
||||
# lb = lb[np.newaxis, :, :]
|
||||
# lb[lb == 0] = 0
|
||||
|
||||
# ---------- important ----------
|
||||
# if self.use_uncertainty:
|
||||
# lb[np.logical_and(lb > 0, lb < threshold)] = 2
|
||||
# else:
|
||||
# lb[np.logical_and(lb > 0, lb < threshold)] /= 255.
|
||||
|
||||
lb[lb >= threshold] = 255
|
||||
lb = Image.fromarray(lb.astype(np.uint8))
|
||||
return lb
|
||||
|
||||
def build_list(self):
|
||||
data_root = os.path.abspath(self.data_root)
|
||||
images_path = os.path.join(data_root, 'image')
|
||||
labels_path = os.path.join(data_root, 'edge')
|
||||
|
||||
samples = []
|
||||
for directory_name in os.listdir(images_path):
|
||||
image_directories = os.path.join(images_path, directory_name)
|
||||
for file_name_ext in os.listdir(image_directories):
|
||||
file_name = os.path.basename(file_name_ext)
|
||||
image_path = fit_img_postfix(os.path.join(images_path, directory_name, file_name))
|
||||
lb_path = fit_img_postfix(os.path.join(labels_path, directory_name, file_name))
|
||||
samples.append((image_path, lb_path))
|
||||
return samples
|
||||
|
||||
def __getitem__(self, index):
|
||||
img_path, edge_path = self.data_list[index]
|
||||
# edge_path = self.edge_paths[index]
|
||||
# img_path = self.img_paths[index]
|
||||
img_name = os.path.basename(img_path)
|
||||
|
||||
img, raw_size = self.read_img(img_path)
|
||||
edge = self.read_lb(edge_path)
|
||||
img, edge = self.transform(img, edge)
|
||||
|
||||
# print("-------hhhhhhhhhhhhh--------:", img.shape, edge.shape)
|
||||
# edge = Image.open(edge_path).convert('L')
|
||||
# # default to score-sde preprocessing
|
||||
# mask = Image.open(img_path).convert('RGB')
|
||||
# edge, img = self.transform(edge, mask)
|
||||
if self.normalize_to_neg_one_to_one: # transform to [-1, 1]
|
||||
edge = normalize_to_neg_one_to_one(edge)
|
||||
img = normalize_to_neg_one_to_one(img)
|
||||
return {'image': edge, 'cond': img, 'raw_size': raw_size, 'img_name': img_name}
|
||||
|
||||
class EdgeDatasetTest(data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_root,
|
||||
# mask_folder,
|
||||
image_size,
|
||||
exts = ['png', 'jpg'],
|
||||
convert_image_to = None,
|
||||
normalize_to_neg_one_to_one=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.data_root = data_root
|
||||
self.image_size = image_size
|
||||
self.normalize_to_neg_one_to_one = normalize_to_neg_one_to_one
|
||||
|
||||
maybe_convert_fn = partial(convert_image_to_fn, convert_image_to) if exists(convert_image_to) else Identity()
|
||||
|
||||
self.data_list = self.build_list()
|
||||
|
||||
self.transform = Compose([
|
||||
ToTensor()
|
||||
])
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_list)
|
||||
|
||||
|
||||
def read_img(self, image_path):
|
||||
with open(image_path, 'rb') as f:
|
||||
img = Image.open(f)
|
||||
img = img.convert('RGB')
|
||||
|
||||
raw_width, raw_height = img.size
|
||||
|
||||
|
||||
return img, (raw_width, raw_height)
|
||||
|
||||
def read_lb(self, lb_path):
|
||||
lb_data = Image.open(lb_path).convert('L')
|
||||
lb = np.array(lb_data).astype(np.float32)
|
||||
|
||||
threshold = self.threshold
|
||||
|
||||
|
||||
lb[lb >= threshold] = 255
|
||||
lb = Image.fromarray(lb.astype(np.uint8))
|
||||
return lb
|
||||
|
||||
def build_list(self):
|
||||
data_root = os.path.abspath(self.data_root)
|
||||
# images_path = os.path.join(data_root)
|
||||
images_path = data_root
|
||||
samples = get_imgs_list(images_path)
|
||||
return samples
|
||||
|
||||
def __getitem__(self, index):
|
||||
img_path = self.data_list[index]
|
||||
# edge_path = self.edge_paths[index]
|
||||
# img_path = self.img_paths[index]
|
||||
img_name = os.path.basename(img_path)
|
||||
|
||||
img, raw_size = self.read_img(img_path)
|
||||
|
||||
img = self.transform(img)
|
||||
if self.normalize_to_neg_one_to_one: # transform to [-1, 1]
|
||||
img = normalize_to_neg_one_to_one(img)
|
||||
return {'cond': img, 'raw_size': raw_size, 'img_name': img_name}
|
||||
|
||||
|
||||
class Identity(nn.Identity):
|
||||
r"""A placeholder identity operator that is argument-insensitive.
|
||||
|
||||
Args:
|
||||
args: any argument (unused)
|
||||
kwargs: any keyword argument (unused)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*)`, where :math:`*` means any number of dimensions.
|
||||
- Output: :math:`(*)`, same shape as the input.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 20])
|
||||
|
||||
"""
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(Identity, self).__init__(*args, **kwargs)
|
||||
|
||||
def forward(self, input, target):
|
||||
return input, target
|
||||
|
||||
class Resize(T.Resize):
|
||||
def __init__(self, size, interpolation2=None, **kwargs):
|
||||
super().__init__(size, **kwargs)
|
||||
if interpolation2 is None:
|
||||
self.interpolation2 = self.interpolation
|
||||
else:
|
||||
self.interpolation2 = interpolation2
|
||||
|
||||
def forward(self, img, target=None):
|
||||
if target is None:
|
||||
img = F2.resize(img, self.size, self.interpolation, self.max_size, self.antialias)
|
||||
return img
|
||||
else:
|
||||
img = F2.resize(img, self.size, self.interpolation, self.max_size, self.antialias)
|
||||
target = F2.resize(target, self.size, self.interpolation2, self.max_size, self.antialias)
|
||||
return img, target
|
||||
|
||||
class RandomHorizontalFlip(T.RandomHorizontalFlip):
|
||||
def __init__(self, p=0.5):
|
||||
super().__init__(p)
|
||||
|
||||
def forward(self, img, target=None):
|
||||
if target is None:
|
||||
if torch.rand(1) < self.p:
|
||||
img = F2.hflip(img)
|
||||
return img
|
||||
else:
|
||||
if torch.rand(1) < self.p:
|
||||
img = F2.hflip(img)
|
||||
target = F2.hflip(target)
|
||||
return img, target
|
||||
|
||||
class CenterCrop(T.CenterCrop):
|
||||
def __init__(self, size):
|
||||
super().__init__(size)
|
||||
|
||||
def forward(self, img, target=None):
|
||||
if target is None:
|
||||
img = F2.center_crop(img, self.size)
|
||||
return img
|
||||
else:
|
||||
img = F2.center_crop(img, self.size)
|
||||
target = F2.center_crop(target, self.size)
|
||||
return img, target
|
||||
|
||||
class RandomCrop(T.RandomCrop):
|
||||
def __init__(self, size, **kwargs):
|
||||
super().__init__(size, **kwargs)
|
||||
|
||||
def single_forward(self, img, i, j, h, w):
|
||||
if self.padding is not None:
|
||||
img = F2.pad(img, self.padding, self.fill, self.padding_mode)
|
||||
width, height = F2.get_image_size(img)
|
||||
# pad the width if needed
|
||||
if self.pad_if_needed and width < self.size[1]:
|
||||
padding = [self.size[1] - width, 0]
|
||||
img = F2.pad(img, padding, self.fill, self.padding_mode)
|
||||
# pad the height if needed
|
||||
if self.pad_if_needed and height < self.size[0]:
|
||||
padding = [0, self.size[0] - height]
|
||||
img = F2.pad(img, padding, self.fill, self.padding_mode)
|
||||
|
||||
return F2.crop(img, i, j, h, w)
|
||||
|
||||
def forward(self, img, target=None):
|
||||
i, j, h, w = self.get_params(img, self.size)
|
||||
if target is None:
|
||||
img = self.single_forward(img, i, j, h, w)
|
||||
return img
|
||||
else:
|
||||
img = self.single_forward(img, i, j, h, w)
|
||||
target = self.single_forward(target, i, j, h, w)
|
||||
return img, target
|
||||
|
||||
class RandomResizeCrop(T.RandomResizedCrop):
|
||||
def __init__(self, size, scale=(0.25, 1.0), **kwargs):
|
||||
super().__init__(size, scale, **kwargs)
|
||||
|
||||
# def single_forward(self, img, i, j, h, w):
|
||||
# if self.padding is not None:
|
||||
# img = F2.pad(img, self.padding, self.fill, self.padding_mode)
|
||||
# width, height = F2.get_image_size(img)
|
||||
# # pad the width if needed
|
||||
# if self.pad_if_needed and width < self.size[1]:
|
||||
# padding = [self.size[1] - width, 0]
|
||||
# img = F2.pad(img, padding, self.fill, self.padding_mode)
|
||||
# # pad the height if needed
|
||||
# if self.pad_if_needed and height < self.size[0]:
|
||||
# padding = [0, self.size[0] - height]
|
||||
# img = F2.pad(img, padding, self.fill, self.padding_mode)
|
||||
#
|
||||
# return F2.crop(img, i, j, h, w)
|
||||
|
||||
def single_forward(self, img, i, j, h, w, interpolation=InterpolationMode.BILINEAR):
|
||||
"""
|
||||
Args:
|
||||
img (PIL Image or Tensor): Image to be cropped and resized.
|
||||
|
||||
Returns:
|
||||
PIL Image or Tensor: Randomly cropped and resized image.
|
||||
"""
|
||||
# i, j, h, w = self.get_params(img, self.scale, self.ratio)
|
||||
return F2.resized_crop(img, i, j, h, w, self.size, interpolation)
|
||||
|
||||
def forward(self, img, target=None):
|
||||
i, j, h, w = self.get_params(img, self.scale, self.ratio)
|
||||
if target is None:
|
||||
img = self.single_forward(img, i, j, h, w)
|
||||
return img
|
||||
else:
|
||||
img = self.single_forward(img, i, j, h, w)
|
||||
target = self.single_forward(target, i, j, h, w, interpolation=InterpolationMode.NEAREST)
|
||||
return img, target
|
||||
|
||||
class ToTensor(T.ToTensor):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def __call__(self, img, target=None):
|
||||
if target is None:
|
||||
img = F2.to_tensor(img)
|
||||
return img
|
||||
else:
|
||||
img = F2.to_tensor(img)
|
||||
target = F2.to_tensor(target)
|
||||
return img, target
|
||||
|
||||
class Lambda(T.Lambda):
|
||||
"""Apply a user-defined lambda as a transform. This transform does not support torchscript.
|
||||
|
||||
Args:
|
||||
lambd (function): Lambda/function to be used for transform.
|
||||
"""
|
||||
|
||||
def __init__(self, lambd):
|
||||
super().__init__(lambd)
|
||||
|
||||
def __call__(self, img, target=None):
|
||||
if target is None:
|
||||
return self.lambd(img)
|
||||
else:
|
||||
return self.lambd(img), self.lambd(target)
|
||||
|
||||
class Compose(T.Compose):
|
||||
def __init__(self, transforms):
|
||||
super().__init__(transforms)
|
||||
|
||||
def __call__(self, img, target=None):
|
||||
if target is None:
|
||||
for t in self.transforms:
|
||||
img = t(img)
|
||||
return img
|
||||
else:
|
||||
for t in self.transforms:
|
||||
img, target = t(img, target)
|
||||
return img, target
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset = CIFAR10(
|
||||
img_folder='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/cifar-10-python',
|
||||
augment_horizontal_flip=False
|
||||
)
|
||||
# dataset = CityscapesDataset(
|
||||
# # img_folder='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/CelebAHQ/celeba_hq_256',
|
||||
# data_root='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/Cityscapes/',
|
||||
# # data_root='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/ADEChallengeData2016/',
|
||||
# image_size=[512, 1024],
|
||||
# exts = ['png'],
|
||||
# augment_horizontal_flip = False,
|
||||
# convert_image_to = None,
|
||||
# normalize_to_neg_one_to_one=True,
|
||||
# )
|
||||
# dataset = SRDataset(
|
||||
# img_folder='/media/huang/ZX3 512G/data/DIV2K/DIV2K_train_HR',
|
||||
# image_size=[512, 512],
|
||||
# )
|
||||
# dataset = InpaintDataset(
|
||||
# img_folder='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/CelebAHQ/celeba_hq_256',
|
||||
# image_size=[256, 256],
|
||||
# augment_horizontal_flip = True
|
||||
# )
|
||||
dataset = EdgeDataset(
|
||||
data_root='/media/huang/2da18d46-7cba-4259-9abd-0df819bb104c/data/BSDS',
|
||||
image_size=[320, 320],
|
||||
)
|
||||
for i in range(len(dataset)):
|
||||
d = dataset[i]
|
||||
mask = d['cond']
|
||||
print(mask.max())
|
||||
dl = data.DataLoader(dataset, batch_size=2, shuffle=False, pin_memory=True, num_workers=0)
|
||||
|
||||
|
||||
dataset_builder = tfds.builder('cifar10')
|
||||
split = 'train'
|
||||
dataset_options = tf.data.Options()
|
||||
dataset_options.experimental_optimization.map_parallelization = True
|
||||
dataset_options.experimental_threading.private_threadpool_size = 48
|
||||
dataset_options.experimental_threading.max_intra_op_parallelism = 1
|
||||
read_config = tfds.ReadConfig(options=dataset_options)
|
||||
dataset_builder.download_and_prepare()
|
||||
ds = dataset_builder.as_dataset(
|
||||
split=split, shuffle_files=True, read_config=read_config)
|
||||
pause = 0
|
||||
+992
@@ -0,0 +1,992 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.cuda.amp import custom_bwd, custom_fwd
|
||||
import math
|
||||
import torch.nn.functional as F
|
||||
# import torchvision.transforms.functional as F2
|
||||
from .utils import default, identity, normalize_to_neg_one_to_one, unnormalize_to_zero_to_one
|
||||
from tqdm.auto import tqdm
|
||||
from einops import rearrange, reduce
|
||||
from functools import partial
|
||||
from collections import namedtuple
|
||||
from random import random, randint, sample, choice
|
||||
from .encoder_decoder import DiagonalGaussianDistribution
|
||||
import random
|
||||
from custom_controlnet_aux.diffusion_edge.taming.modules.losses.vqperceptual import *
|
||||
|
||||
# gaussian diffusion trainer class
|
||||
ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
|
||||
|
||||
def extract(a, t, x_shape):
|
||||
b, *_ = t.shape
|
||||
out = a.gather(-1, t)
|
||||
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
||||
|
||||
def linear_beta_schedule(timesteps):
|
||||
scale = 1000 / timesteps
|
||||
beta_start = scale * 0.0001
|
||||
beta_end = scale * 0.02
|
||||
return torch.linspace(beta_start, beta_end, timesteps, dtype = torch.float64)
|
||||
|
||||
def cosine_beta_schedule(timesteps, s = 0.008):
|
||||
"""
|
||||
cosine schedule
|
||||
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
|
||||
"""
|
||||
steps = timesteps + 1
|
||||
x = torch.linspace(0, timesteps, steps, dtype = torch.float64)
|
||||
alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
||||
return torch.clip(betas, 0, 0.999)
|
||||
|
||||
class DDPM(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
*,
|
||||
image_size,
|
||||
timesteps = 1000,
|
||||
sampling_timesteps = None,
|
||||
loss_type = 'l2',
|
||||
objective = 'pred_noise',
|
||||
beta_schedule = 'cosine',
|
||||
p2_loss_weight_gamma = 0., # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
|
||||
p2_loss_weight_k = 1,
|
||||
original_elbo_weight=0.,
|
||||
ddim_sampling_eta = 1.,
|
||||
clip_x_start=True,
|
||||
train_sample=-1,
|
||||
input_keys=['image'],
|
||||
start_dist='normal',
|
||||
sample_type='ddim',
|
||||
perceptual_weight=1.,
|
||||
use_l1=False,
|
||||
**kwargs
|
||||
):
|
||||
ckpt_path = kwargs.pop("ckpt_path", None)
|
||||
ignore_keys = kwargs.pop("ignore_keys", [])
|
||||
only_model = kwargs.pop("only_model", False)
|
||||
cfg = kwargs.pop("cfg", None)
|
||||
super().__init__(**kwargs)
|
||||
# assert not (type(self) == DDPM and model.channels != model.out_dim)
|
||||
# assert not model.random_or_learned_sinusoidal_cond
|
||||
|
||||
self.model = model
|
||||
self.channels = self.model.channels
|
||||
self.self_condition = self.model.self_condition
|
||||
self.input_keys = input_keys
|
||||
self.cfg = cfg
|
||||
self.eps = cfg.get('eps', 1e-4) if cfg is not None else 1e-4
|
||||
self.weighting_loss = cfg.get("weighting_loss", False) if cfg is not None else False
|
||||
if self.weighting_loss:
|
||||
print('#### WEIGHTING LOSS ####')
|
||||
|
||||
self.clip_x_start = clip_x_start
|
||||
self.image_size = image_size
|
||||
self.train_sample = train_sample
|
||||
self.objective = objective
|
||||
self.start_dist = start_dist
|
||||
assert start_dist in ['normal', 'uniform']
|
||||
|
||||
assert objective in {'pred_noise', 'pred_x0', 'pred_v', 'pred_delta', 'pred_KC'}, 'objective must be either pred_noise (predict noise) or pred_x0 (predict image start) or pred_v (predict v [v-parameterization as defined in appendix D of progressive distillation paper, used in imagen-video successfully])'
|
||||
|
||||
if beta_schedule == 'linear':
|
||||
betas = linear_beta_schedule(timesteps)
|
||||
elif beta_schedule == 'cosine':
|
||||
betas = cosine_beta_schedule(timesteps, s=1e-4)
|
||||
else:
|
||||
raise ValueError(f'unknown beta schedule {beta_schedule}')
|
||||
# betas[0] = 2e-3 * betas[0]
|
||||
alphas = 1. - betas
|
||||
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||||
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.)
|
||||
|
||||
timesteps, = betas.shape
|
||||
self.num_timesteps = int(timesteps)
|
||||
self.time_range = list(range(self.num_timesteps + 1))
|
||||
self.loss_type = loss_type
|
||||
self.original_elbo_weight = original_elbo_weight
|
||||
|
||||
# sampling related parameters
|
||||
|
||||
self.sampling_timesteps = default(sampling_timesteps, timesteps) # default num sampling timesteps to number of timesteps at training
|
||||
|
||||
# assert self.sampling_timesteps <= timesteps
|
||||
self.is_ddim_sampling = self.sampling_timesteps < timesteps
|
||||
self.ddim_sampling_eta = ddim_sampling_eta
|
||||
|
||||
# helper function to register buffer from float64 to float32
|
||||
|
||||
register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))
|
||||
|
||||
register_buffer('betas', betas)
|
||||
register_buffer('alphas_cumprod', alphas_cumprod)
|
||||
register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
|
||||
register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
|
||||
register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
|
||||
register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
|
||||
register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
|
||||
register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))
|
||||
|
||||
# calculations for posterior q(x_{t-1} | x_t, x_0)
|
||||
|
||||
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
|
||||
|
||||
# above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
|
||||
|
||||
register_buffer('posterior_variance', posterior_variance)
|
||||
|
||||
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
|
||||
|
||||
register_buffer('posterior_log_variance_clipped', torch.log(posterior_variance.clamp(min =1e-20)))
|
||||
register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
|
||||
register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
|
||||
|
||||
# calculate p2 reweighting
|
||||
|
||||
register_buffer('p2_loss_weight', (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod)) ** -p2_loss_weight_gamma)
|
||||
assert not torch.isnan(self.p2_loss_weight).all()
|
||||
if self.objective == "pred_noise":
|
||||
lvlb_weights = self.betas ** 2 / (
|
||||
2 * (self.posterior_variance+1e-5) * alphas * (1 - self.alphas_cumprod))
|
||||
elif self.objective == "pred_x0":
|
||||
lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
|
||||
elif self.objective == "pred_delta":
|
||||
lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
|
||||
elif self.objective == "pred_KC":
|
||||
lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
|
||||
elif self.objective == "pred_v":
|
||||
lvlb_weights = 0.5 * torch.sqrt(alphas_cumprod) / (2. * 1 - alphas_cumprod)
|
||||
self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
|
||||
assert not torch.isnan(self.lvlb_weights).all()
|
||||
self.use_l1 = use_l1
|
||||
|
||||
self.perceptual_weight = perceptual_weight
|
||||
if self.perceptual_weight > 0:
|
||||
self.perceptual_loss = LPIPS().eval()
|
||||
|
||||
if ckpt_path is not None:
|
||||
self.init_from_ckpt(ckpt_path, ignore_keys, only_model)
|
||||
|
||||
def init_from_ckpt(self, path, ignore_keys=list(), only_model=False, use_ema=False):
|
||||
sd = torch.load(path, map_location="cpu")
|
||||
if 'ema' in list(sd.keys()) and use_ema:
|
||||
sd = sd['ema']
|
||||
new_sd = {}
|
||||
for k in sd.keys():
|
||||
if k.startswith("ema_model."):
|
||||
new_k = k[10:] # remove ema_model.
|
||||
new_sd[new_k] = sd[k]
|
||||
sd = new_sd
|
||||
else:
|
||||
if "model" in list(sd.keys()):
|
||||
sd = sd["model"]
|
||||
keys = list(sd.keys())
|
||||
for k in keys:
|
||||
for ik in ignore_keys:
|
||||
if k.startswith(ik):
|
||||
print("Deleting key {} from state_dict.".format(k))
|
||||
del sd[k]
|
||||
missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
|
||||
sd, strict=False)
|
||||
print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
|
||||
if len(missing) > 0:
|
||||
print(f"Missing Keys: {missing}")
|
||||
if len(unexpected) > 0:
|
||||
print(f"Unexpected Keys: {unexpected}")
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample(self, x, mask, t: int, x_self_cond = None, clip_denoised = True):
|
||||
b, *_, device = *x.shape, x.device
|
||||
batched_times = torch.full((x.shape[0],), t, device = x.device, dtype = torch.long)
|
||||
model_mean, _, model_log_variance, x_start = self.p_mean_variance(x = x, mask=mask, t = batched_times, x_self_cond = x_self_cond, clip_denoised = clip_denoised)
|
||||
noise = torch.randn_like(x) if t > 0 else 0. # no noise if t == 0
|
||||
pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
|
||||
return pred_img, x_start
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample_loop(self, shape, mask, up_scale=1, unnormalize=True):
|
||||
batch, device = shape[0], self.betas.device
|
||||
|
||||
img = torch.randn(shape, device=device)
|
||||
img = F.interpolate(img, scale_factor=up_scale, mode='bilinear', align_corners=True)
|
||||
|
||||
x_start = None
|
||||
|
||||
for t in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):
|
||||
self_cond = x_start if self.self_condition else None
|
||||
img, x_start = self.p_sample(img, mask, t, self_cond)
|
||||
if unnormalize:
|
||||
img = unnormalize_to_zero_to_one(img)
|
||||
return img
|
||||
|
||||
@torch.no_grad()
|
||||
def ddim_sample(self, shape, mask, up_scale=1, unnormalize=True):
|
||||
batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
|
||||
|
||||
times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1) # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
|
||||
times = list(reversed(times.int().tolist()))
|
||||
time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
|
||||
|
||||
img = torch.randn(shape, device = device)
|
||||
img = F.interpolate(img, scale_factor=up_scale, mode='bilinear', align_corners=True)
|
||||
|
||||
x_start = None
|
||||
|
||||
for time, time_next in tqdm(time_pairs, desc = 'sampling loop time step', total=len(time_pairs)):
|
||||
time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
|
||||
self_cond = x_start if self.self_condition else None
|
||||
pred_noise, x_start, *_ = self.model_predictions(img, time_cond, mask, self_cond)
|
||||
|
||||
if time_next < 0:
|
||||
img = x_start
|
||||
continue
|
||||
|
||||
alpha = self.alphas_cumprod[time]
|
||||
alpha_next = self.alphas_cumprod[time_next]
|
||||
|
||||
sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
|
||||
c = (1 - alpha_next - sigma ** 2).sqrt()
|
||||
|
||||
noise = torch.randn_like(img)
|
||||
|
||||
img = x_start * alpha_next.sqrt() + \
|
||||
c * pred_noise + \
|
||||
sigma * noise
|
||||
if unnormalize:
|
||||
img = unnormalize_to_zero_to_one(img)
|
||||
return img
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def interpolate(self, x1, x2, mask, t = None, lam = 0.5):
|
||||
b, *_, device = *x1.shape, x1.device
|
||||
t = default(t, self.num_timesteps - 1)
|
||||
|
||||
assert x1.shape == x2.shape
|
||||
|
||||
t_batched = torch.stack([torch.tensor(t, device = device)] * b)
|
||||
xt1, xt2 = map(lambda x: self.q_sample(x, t = t_batched), (x1, x2))
|
||||
|
||||
img = (1 - lam) * xt1 + lam * xt2
|
||||
for i in tqdm(reversed(range(0, t)), desc = 'interpolation sample time step', total = t):
|
||||
img = self.p_sample(img, mask, torch.full((b,), i, device=device, dtype=torch.long))
|
||||
return img
|
||||
|
||||
def get_input(self, batch, return_first_stage_outputs=False, return_original_cond=False):
|
||||
assert 'image' in self.input_keys;
|
||||
if len(self.input_keys) > len(batch.keys()):
|
||||
x, *_ = batch.values()
|
||||
else:
|
||||
x = batch.values()
|
||||
return x
|
||||
|
||||
def training_step(self, batch):
|
||||
z, *_ = self.get_input(batch)
|
||||
cond = batch['cond'] if 'cond' in batch else None
|
||||
loss, loss_dict = self(z, cond)
|
||||
return loss, loss_dict
|
||||
|
||||
def forward(self, x, *args, **kwargs):
|
||||
# continuous time, t in [0, 1]
|
||||
# t = []
|
||||
# for _ in range(x.shape[0]):
|
||||
# if self.train_sample <= 0:
|
||||
# t.append(torch.tensor(sample(self.time_range, 2), device=x.device).long())
|
||||
# else:
|
||||
# sl = choice(self.time_range)
|
||||
# sl_range = list(range(sl - self.train_sample, sl + self.train_sample))
|
||||
# sl_range = list(set(sl_range) & set(self.time_range))
|
||||
# sl_range.pop(sl_range.index(sl))
|
||||
# sl2 = choice(sl_range)
|
||||
# t.append(torch.tensor([sl, sl2], device=x.device).long())
|
||||
# t = torch.stack(t, dim=0)
|
||||
# t = torch.randint(0, self.num_timesteps+1, (x.shape[0],), device=x.device).long()
|
||||
eps = self.eps # smallest time step
|
||||
# t = torch.rand((x.shape[0],), device=x.device) * (self.num_timesteps / eps)
|
||||
# t = t.round() * eps
|
||||
# t[t < eps] = eps
|
||||
t = torch.rand(x.shape[0], device=x.device) * (1. - eps) + eps
|
||||
return self.p_losses(x, t, *args, **kwargs)
|
||||
|
||||
def q_sample2(self, x_start, t, noise=None):
|
||||
b, c, h, w = x_start.shape
|
||||
noise = default(noise, lambda: torch.randn_like(x_start))
|
||||
_, nt = t.shape
|
||||
param_x = self.sqrt_alphas_cumprod.repeat(b, 1).gather(-1, t) # (b, nt)
|
||||
x = x_start.expand(nt, b, c, h, w).transpose(1, 0) * param_x.reshape(b, nt, 1, 1, 1).repeat(1, 1, c, h, w)
|
||||
param_noise = self.sqrt_one_minus_alphas_cumprod.repeat(b, 1).gather(-1, t)
|
||||
n = noise.expand(nt, b, c, h, w).transpose(1, 0) * param_noise.reshape(b, nt, 1, 1, 1).repeat(1, 1, c, h, w)
|
||||
return x + n # (b, nt, c, h, w)
|
||||
|
||||
def q_sample3(self, x_start, t, C):
|
||||
b, c, h, w = x_start.shape
|
||||
_, nt = t.shape
|
||||
# K_ = K.unsqueeze(1).repeat(1, nt, 1, 1, 1)
|
||||
C_ = C.unsqueeze(1).repeat(1, nt, 1, 1, 1)
|
||||
x_noisy = x_start.expand(nt, b, c, h, w).transpose(1, 0) + \
|
||||
+ C_ * t.reshape(b, nt, 1, 1, 1).repeat(1, 1, c, h, w) / self.num_timesteps
|
||||
return x_noisy # (b, nt, c, h, w)
|
||||
|
||||
# def q_sample(self, x_start, t, C):
|
||||
# x_noisy = x_start + C * t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1))) / self.num_timesteps
|
||||
# return x_noisy
|
||||
def q_sample(self, x_start, noise, t, C):
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
x_noisy = x_start + C * time + torch.sqrt(time) * noise
|
||||
return x_noisy
|
||||
|
||||
def q_sample2(self, x_start, noise, t, C):
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
x_noisy = x_start + C / 2 * time ** 2 + torch.sqrt(time) * noise
|
||||
return x_noisy
|
||||
|
||||
def pred_x0_from_xt(self, xt, noise, C, t):
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
x0 = xt - C * time - torch.sqrt(time) * noise
|
||||
return x0
|
||||
|
||||
def pred_x0_from_xt2(self, xt, noise, C, t):
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
x0 = xt - C / 2 * time ** 2 - torch.sqrt(time) * noise
|
||||
return x0
|
||||
|
||||
def pred_xtms_from_xt(self, xt, noise, C, t, s):
|
||||
# noise = noise / noise.std(dim=[1, 2, 3]).reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
s = s.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
mean = xt + C * (time-s) - C * time - s / torch.sqrt(time) * noise
|
||||
epsilon = torch.randn_like(mean, device=xt.device)
|
||||
sigma = torch.sqrt(s * (time-s) / time)
|
||||
xtms = mean + sigma * epsilon
|
||||
return xtms
|
||||
|
||||
def pred_xtms_from_xt2(self, xt, noise, C, t, s):
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
s = s.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
mean = xt + C / 2 * (time-s) ** 2 - C / 2 * time ** 2 - s / torch.sqrt(time) * noise
|
||||
epsilon = torch.randn_like(mean, device=xt.device)
|
||||
sigma = torch.sqrt(s * (time-s) / time)
|
||||
xtms = mean + sigma * epsilon
|
||||
return xtms
|
||||
|
||||
def WCE_loss(self, prediction, labelf, beta=1.1):
|
||||
label = labelf.long()
|
||||
mask = labelf.clone()
|
||||
|
||||
num_positive = torch.sum(label == 1).float()
|
||||
num_negative = torch.sum(label == 0).float()
|
||||
|
||||
mask[label == 1] = 1.0 * num_negative / (num_positive + num_negative)
|
||||
mask[label == 0] = beta * num_positive / (num_positive + num_negative)
|
||||
mask[label == 2] = 0
|
||||
cost = F.binary_cross_entropy(prediction, labelf, weight=mask, reduction='sum')
|
||||
|
||||
return cost
|
||||
|
||||
def Dice_Loss(self, pred, label):
|
||||
# pred = torch.sigmoid(pred)
|
||||
smooth = 1
|
||||
pred_flat = pred.view(-1)
|
||||
label_flat = label.view(-1)
|
||||
|
||||
intersecion = pred_flat * label_flat
|
||||
unionsection = pred_flat.pow(2).sum() + label_flat.pow(2).sum() + smooth
|
||||
loss = unionsection / (2 * intersecion.sum() + smooth)
|
||||
loss = loss.sum()
|
||||
return loss
|
||||
|
||||
def p_losses(self, x_start, t, *args, **kwargs):
|
||||
if self.start_dist == 'normal':
|
||||
noise = torch.randn_like(x_start)
|
||||
elif self.start_dist == 'uniform':
|
||||
noise = 2 * torch.rand_like(x_start) - 1.
|
||||
else:
|
||||
raise NotImplementedError(f'{self.start_dist} is not supported !')
|
||||
# K = -1. * torch.ones_like(x_start)
|
||||
# C = noise - x_start # t = 1000 / 1000
|
||||
C = -1 * x_start # U(t) = Ct, U(1) = -x0
|
||||
# C = -2 * x_start # U(t) = 1/2 * C * t**2, U(1) = 1/2 * C = -x0
|
||||
x_noisy = self.q_sample(x_start=x_start, noise=noise, t=t, C=C) # (b, 2, c, h, w)
|
||||
C_pred, noise_pred = self.model(x_noisy, t, **kwargs)
|
||||
# C_pred = C_pred / torch.sqrt(t)
|
||||
# noise_pred = noise_pred / torch.sqrt(1 - t)
|
||||
x_rec = self.pred_x0_from_xt(x_noisy, noise_pred, C_pred, t) # x_rec:(B, 1, H, W)
|
||||
loss_dict = {}
|
||||
prefix = 'train'
|
||||
|
||||
# elif self.objective == 'pred_KC':
|
||||
# target1 = C
|
||||
# target2 = noise
|
||||
# target3 = x_start
|
||||
|
||||
target1 = C
|
||||
target2 = noise
|
||||
target3 = x_start
|
||||
|
||||
loss_simple = 0.
|
||||
loss_vlb = 0.
|
||||
# use l1 + l2
|
||||
if self.weighting_loss:
|
||||
simple_weight1 = 2*torch.exp(1-t)
|
||||
simple_weight2 = torch.exp(torch.sqrt(t))
|
||||
if self.cfg.model_name == 'ncsnpp9':
|
||||
simple_weight1 = (t + 1) / t.sqrt()
|
||||
simple_weight2 = (2 - t).sqrt() / (1 - t + self.eps).sqrt()
|
||||
else:
|
||||
simple_weight1 = 1
|
||||
simple_weight2 = 1
|
||||
|
||||
loss_simple += simple_weight1 * self.get_loss(C_pred, target1, mean=False).mean([1, 2, 3]) + \
|
||||
simple_weight2 * self.get_loss(noise_pred, target2, mean=False).mean([1, 2, 3])
|
||||
if self.use_l1:
|
||||
loss_simple += simple_weight1 * (C_pred - target1).abs().mean([1, 2, 3]) + \
|
||||
simple_weight2 * (noise_pred - target2).abs().mean([1, 2, 3])
|
||||
loss_simple = loss_simple / 2
|
||||
# rec_weight = (1 - t.reshape(C.shape[0], 1)) ** 2
|
||||
rec_weight = 1 - t.reshape(C.shape[0], 1) # (B, 1)
|
||||
loss_simple = loss_simple.mean()
|
||||
loss_dict.update({f'{prefix}/loss_simple': loss_simple})
|
||||
|
||||
# loss_vlb += torch.abs(x_rec - target3).mean([1, 2, 3]) * rec_weight: (B, 1)
|
||||
loss_vlb += self.Dice_Loss(x_rec, target3)
|
||||
loss_vlb = loss_vlb.mean()
|
||||
loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
|
||||
|
||||
loss = loss_simple + loss_vlb
|
||||
loss_dict.update({f'{prefix}/loss': loss})
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
def get_loss(self, pred, target, mean=True):
|
||||
if self.loss_type == 'l1':
|
||||
loss = (target - pred).abs()
|
||||
if mean:
|
||||
loss = loss.mean()
|
||||
elif self.loss_type == 'l2':
|
||||
if mean:
|
||||
loss = torch.nn.functional.mse_loss(target, pred)
|
||||
else:
|
||||
loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
|
||||
else:
|
||||
raise NotImplementedError("unknown loss type '{loss_type}'")
|
||||
|
||||
return loss
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self, batch_size=16, up_scale=1, cond=None, denoise=True):
|
||||
image_size, channels = self.image_size, self.channels
|
||||
if cond is not None:
|
||||
batch_size = cond.shape[0]
|
||||
return self.sample_fn((batch_size, channels, image_size[0], image_size[1]),
|
||||
up_scale=up_scale, unnormalize=True, cond=cond, denoise=denoise)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_fn(self, shape, up_scale=1, unnormalize=True, cond=None, denoise=False):
|
||||
batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], \
|
||||
self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
|
||||
|
||||
# times = torch.linspace(-1, total_timesteps, steps=self.sampling_timesteps + 1).int()
|
||||
# times = list(reversed(times.int().tolist()))
|
||||
# time_pairs = list(zip(times[:-1], times[1:]))
|
||||
# time_steps = torch.tensor([0.25, 0.15, 0.1, 0.1, 0.1, 0.09, 0.075, 0.06, 0.045, 0.03])
|
||||
step = 1. / self.sampling_timesteps
|
||||
# time_steps = torch.tensor([0.1]).repeat(10)
|
||||
time_steps = torch.tensor([step]).repeat(self.sampling_timesteps)
|
||||
if denoise:
|
||||
eps = self.eps
|
||||
time_steps = torch.cat((time_steps[:-1], torch.tensor([step - eps]), torch.tensor([eps])), dim=0)
|
||||
|
||||
if self.start_dist == 'normal':
|
||||
img = torch.randn(shape, device=device)
|
||||
elif self.start_dist == 'uniform':
|
||||
img = 2 * torch.rand(shape, device=device) - 1.
|
||||
else:
|
||||
raise NotImplementedError(f'{self.start_dist} is not supported !')
|
||||
img = F.interpolate(img, scale_factor=up_scale, mode='bilinear', align_corners=True)
|
||||
# K = -1 * torch.ones_like(img)
|
||||
cur_time = torch.ones((batch,), device=device)
|
||||
for i, time_step in enumerate(time_steps):
|
||||
s = torch.full((batch,), time_step, device=device)
|
||||
if i == time_steps.shape[0] - 1:
|
||||
s = cur_time
|
||||
if cond is not None:
|
||||
pred = self.model(img, cur_time, cond)
|
||||
else:
|
||||
pred = self.model(img, cur_time)
|
||||
# C, noise = pred.chunk(2, dim=1)
|
||||
C, noise = pred[:2]
|
||||
# correct C
|
||||
x0 = self.pred_x0_from_xt(img, noise, C, cur_time)
|
||||
if self.clip_x_start:
|
||||
x0.clamp_(-1., 1.)
|
||||
# C.clamp_(-2., 2.)
|
||||
C = -1 * x0
|
||||
img = self.pred_xtms_from_xt(img, noise, C, cur_time, s)
|
||||
# img = self.pred_xtms_from_xt2(img, noise, C, cur_time, s)
|
||||
cur_time = cur_time - s
|
||||
img.clamp_(-1., 1.)
|
||||
if unnormalize:
|
||||
img = unnormalize_to_zero_to_one(img)
|
||||
return img
|
||||
|
||||
|
||||
|
||||
class LatentDiffusion(DDPM):
|
||||
def __init__(self,
|
||||
auto_encoder,
|
||||
scale_factor=1.0,
|
||||
scale_by_std=True,
|
||||
scale_by_softsign=False,
|
||||
input_keys=['image'],
|
||||
sample_type='ddim',
|
||||
num_timesteps_cond=1,
|
||||
train_sample=-1,
|
||||
default_scale=False,
|
||||
*args,
|
||||
**kwargs
|
||||
):
|
||||
self.scale_by_std = scale_by_std
|
||||
self.scale_by_softsign = scale_by_softsign
|
||||
self.default_scale = default_scale
|
||||
self.num_timesteps_cond = num_timesteps_cond
|
||||
self.train_sample = train_sample
|
||||
self.perceptual_weight = 0
|
||||
ckpt_path = kwargs.pop("ckpt_path", None)
|
||||
ignore_keys = kwargs.pop("ignore_keys", [])
|
||||
only_model = kwargs.pop("only_model", False)
|
||||
super().__init__(*args, **kwargs)
|
||||
assert self.num_timesteps_cond <= self.num_timesteps
|
||||
if not scale_by_std:
|
||||
self.scale_factor = scale_factor
|
||||
else:
|
||||
self.register_buffer('scale_factor', torch.tensor(scale_factor))
|
||||
if self.scale_by_softsign:
|
||||
self.scale_by_std = False
|
||||
print('### USING SOFTSIGN RESCALING')
|
||||
assert (self.scale_by_std and self.scale_by_softsign) is False;
|
||||
|
||||
self.init_first_stage(auto_encoder)
|
||||
# self.instantiate_cond_stage(cond_stage_config)
|
||||
self.input_keys = input_keys
|
||||
self.clip_denoised = False
|
||||
assert sample_type in ['p_loop', 'ddim', 'dpm', 'transformer'] ### 'dpm' is not availible now, suggestion 'ddim'
|
||||
self.sample_type = sample_type
|
||||
|
||||
if ckpt_path is not None:
|
||||
self.init_from_ckpt(ckpt_path, ignore_keys, only_model)
|
||||
|
||||
def init_first_stage(self, first_stage_model):
|
||||
self.first_stage_model = first_stage_model.eval()
|
||||
# self.first_stage_model.train = disabled_train
|
||||
for param in self.first_stage_model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
'''
|
||||
def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
|
||||
sd = torch.load(path, map_location="cpu")
|
||||
if "state_dict" in list(sd.keys()):
|
||||
sd = sd["state_dict"]
|
||||
keys = list(sd.keys())
|
||||
for k in keys:
|
||||
for ik in ignore_keys:
|
||||
if k.startswith(ik):
|
||||
print("Deleting key {} from state_dict.".format(k))
|
||||
del sd[k]
|
||||
missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
|
||||
sd, strict=False)
|
||||
print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
|
||||
if len(missing) > 0:
|
||||
print(f"Missing Keys: {missing}")
|
||||
if len(unexpected) > 0:
|
||||
print(f"Unexpected Keys: {unexpected}")
|
||||
'''
|
||||
|
||||
def get_first_stage_encoding(self, encoder_posterior):
|
||||
if isinstance(encoder_posterior, DiagonalGaussianDistribution):
|
||||
z = encoder_posterior.sample()
|
||||
elif isinstance(encoder_posterior, torch.Tensor):
|
||||
z = encoder_posterior
|
||||
else:
|
||||
raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
|
||||
# return self.scale_factor * z.detach() + self.scale_bias
|
||||
return z.detach()
|
||||
|
||||
@torch.no_grad()
|
||||
def on_train_batch_start(self, batch):
|
||||
# only for the first batch
|
||||
if self.scale_by_std and (not self.scale_by_softsign):
|
||||
if not self.default_scale:
|
||||
assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
|
||||
# set rescale weight to 1./std of encodings
|
||||
print("### USING STD-RESCALING ###")
|
||||
x, *_ = batch.values()
|
||||
encoder_posterior = self.first_stage_model.encode(x)
|
||||
z = self.get_first_stage_encoding(encoder_posterior)
|
||||
del self.scale_factor
|
||||
self.register_buffer('scale_factor', 1. / z.flatten().std())
|
||||
print(f"setting self.scale_factor to {self.scale_factor}")
|
||||
# print("### USING STD-RESCALING ###")
|
||||
else:
|
||||
print(f'### USING DEFAULT SCALE {self.scale_factor}')
|
||||
else:
|
||||
print(f'### USING SOFTSIGN SCALE !')
|
||||
|
||||
@torch.no_grad()
|
||||
def get_input(self, batch, return_first_stage_outputs=False, return_original_cond=False):
|
||||
assert 'image' in self.input_keys;
|
||||
# if len(self.input_keys) > len(batch.keys()):
|
||||
# x, cond, *_ = batch.values()
|
||||
# else:
|
||||
# x, cond = batch.values()
|
||||
x = batch['image']
|
||||
cond = batch['cond'] if 'cond' in batch else None
|
||||
z = self.first_stage_model.encode(x)
|
||||
# print('zzzz', z.shape)
|
||||
z = self.get_first_stage_encoding(z)
|
||||
out = [z, cond, x]
|
||||
if return_first_stage_outputs:
|
||||
xrec = self.first_stage_model.decode(z)
|
||||
out.extend([x, xrec])
|
||||
if return_original_cond:
|
||||
out.append(cond)
|
||||
return out
|
||||
|
||||
def training_step(self, batch):
|
||||
z, c, *_ = self.get_input(batch)
|
||||
# print(_[0].shape)
|
||||
if self.scale_by_softsign:
|
||||
z = F.softsign(z)
|
||||
elif self.scale_by_std:
|
||||
z = self.scale_factor * z
|
||||
# print('grad', self.scale_bias.grad)
|
||||
loss, loss_dict = self(z, c, edge=_[0])
|
||||
return loss, loss_dict
|
||||
|
||||
def q_sample3(self, x_start, t, K, C):
|
||||
b, c, h, w = x_start.shape
|
||||
_, nt = t.shape
|
||||
K_ = K.unsqueeze(1).repeat(1, nt, 1, 1, 1)
|
||||
C_ = C.unsqueeze(1).repeat(1, nt, 1, 1, 1)
|
||||
x_noisy = x_start.expand(nt, b, c, h, w).transpose(1, 0) + K_ / 2 * (t.reshape(b, nt, 1, 1, 1).repeat(1, 1, c, h, w) / self.num_timesteps) ** 2 \
|
||||
+ C_ * t.reshape(b, nt, 1, 1, 1).repeat(1, 1, c, h, w) / self.num_timesteps
|
||||
return x_noisy # (b, nt, c, h, w)
|
||||
|
||||
def pred_xtms_from_xt(self, xt, noise, C, t, s):
|
||||
# noise = noise / noise.std(dim=[1, 2, 3]).reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
time = t.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
s = s.reshape(C.shape[0], *((1,) * (len(C.shape) - 1)))
|
||||
mean = xt - C * s - s / torch.sqrt(time) * noise
|
||||
epsilon = torch.randn_like(mean, device=xt.device)
|
||||
sigma = torch.sqrt(s * (time-s) / time)
|
||||
xtms = mean + sigma * epsilon
|
||||
return xtms
|
||||
|
||||
def WCE_loss(self, prediction, labelf, beta=1.1):
|
||||
label = labelf.long()
|
||||
mask = labelf.clone()
|
||||
|
||||
num_positive = torch.sum(label == 1).float()
|
||||
num_negative = torch.sum(label == 0).float()
|
||||
|
||||
mask[label == 1] = 1.0 * num_negative / (num_positive + num_negative)
|
||||
mask[label == 0] = beta * num_positive / (num_positive + num_negative)
|
||||
mask[label == 2] = 0
|
||||
cost = F.binary_cross_entropy(prediction, labelf, weight=mask, reduction='sum')
|
||||
|
||||
return cost
|
||||
|
||||
def Dice_Loss(self, pred, label):
|
||||
# pred = torch.sigmoid(pred)
|
||||
B = pred.shape[0]
|
||||
smooth = 1
|
||||
pred_flat = pred.view(B, -1)
|
||||
label_flat = label.view(B, -1)
|
||||
|
||||
intersecion = pred_flat * label_flat
|
||||
unionsection = pred_flat.pow(2).sum(dim=-1) + label_flat.pow(2).sum(dim=-1) + smooth
|
||||
loss = unionsection / (2 * intersecion.sum(dim=-1) + smooth)
|
||||
loss = loss.reshape(B, 1)
|
||||
return loss
|
||||
|
||||
def p_losses(self, x_start, t, *args, **kwargs):
|
||||
if self.start_dist == 'normal':
|
||||
noise = torch.randn_like(x_start)
|
||||
elif self.start_dist == 'uniform':
|
||||
noise = 2 * torch.rand_like(x_start) - 1.
|
||||
else:
|
||||
raise NotImplementedError(f'{self.start_dist} is not supported !')
|
||||
# K = -1. * torch.ones_like(x_start)
|
||||
# C = noise - x_start # t = 1000 / 1000
|
||||
C = -1 * x_start # U(t) = Ct, U(1) = -x0
|
||||
# C = -2 * x_start # U(t) = 1/2 * C * t**2, U(1) = 1/2 * C = -x0
|
||||
x_noisy = self.q_sample(x_start=x_start, noise=noise, t=t, C=C) # (b, 2, c, h, w)
|
||||
if self.cfg.model_name == 'cond_unet8':
|
||||
C_pred, noise_pred, (e1, e2) = self.model(x_noisy, t, *args, **kwargs)
|
||||
if self.cfg.model_name == 'cond_unet13':
|
||||
C_pred, noise_pred, aux_C = self.model(x_noisy, t, *args, **kwargs)
|
||||
else:
|
||||
C_pred, noise_pred = self.model(x_noisy, t, *args, **kwargs)
|
||||
# C_pred = C_pred / torch.sqrt(t)
|
||||
# noise_pred = noise_pred / torch.sqrt(1 - t)
|
||||
x_rec = self.pred_x0_from_xt(x_noisy, noise_pred, C_pred, t) # x_rec:(B, C, H, W)
|
||||
loss_dict = {}
|
||||
prefix = 'train'
|
||||
|
||||
# elif self.objective == 'pred_KC':
|
||||
# target1 = C
|
||||
# target2 = noise
|
||||
# target3 = x_start
|
||||
|
||||
target1 = C
|
||||
target2 = noise
|
||||
target3 = x_start
|
||||
|
||||
loss_simple = 0.
|
||||
loss_vlb = 0.
|
||||
|
||||
simple_weight1 = (t + 1) / t.sqrt()
|
||||
simple_weight2 = (2 - t).sqrt() / (1 - t + self.eps).sqrt()
|
||||
|
||||
# if self.weighting_loss:
|
||||
# simple_weight1 = 2 * torch.exp(1 - t)
|
||||
# simple_weight2 = torch.exp(torch.sqrt(t))
|
||||
# if self.cfg.model_name == 'ncsnpp9':
|
||||
# simple_weight1 = (t + 1) / t.sqrt()
|
||||
# simple_weight2 = (2 - t).sqrt() / (1 - t + self.eps).sqrt()
|
||||
# else:
|
||||
# simple_weight1 = 1
|
||||
# simple_weight2 = 1
|
||||
|
||||
loss_simple += simple_weight1 * self.get_loss(C_pred, target1, mean=False).mean([1, 2, 3]) + \
|
||||
simple_weight2 * self.get_loss(noise_pred, target2, mean=False).mean([1, 2, 3])
|
||||
|
||||
# loss_simple += self.Dice_Loss(C_pred, target1) * simple_weight1
|
||||
|
||||
if self.use_l1:
|
||||
loss_simple += simple_weight1 * (C_pred - target1).abs().mean([1, 2, 3]) + \
|
||||
simple_weight2 * (noise_pred - target2).abs().mean([1, 2, 3])
|
||||
loss_simple = loss_simple / 2
|
||||
|
||||
if self.cfg.model_name == 'cond_unet8':
|
||||
loss_simple += 0.05*(self.Dice_Loss(e1, (kwargs['edge'] + 1)/2) + self.Dice_Loss(e2, (kwargs['edge'] + 1)/2))
|
||||
elif self.cfg.model_name == 'cond_unet13':
|
||||
loss_simple += 0.5 * (simple_weight1 * self.get_loss(aux_C, target1, mean=False).mean([1, 2, 3]) + \
|
||||
simple_weight1 * (aux_C - target1).abs().mean([1, 2, 3]))
|
||||
|
||||
rec_weight = (1 - t.reshape(C.shape[0], 1)) ** 2
|
||||
# rec_weight = 1 - t.reshape(C.shape[0], 1) # (B, 1)
|
||||
loss_simple = loss_simple.mean()
|
||||
loss_dict.update({f'{prefix}/loss_simple': loss_simple})
|
||||
|
||||
loss_vlb += torch.abs(x_rec - target3).mean([1, 2, 3]) * rec_weight # : (B, 1)
|
||||
# loss_vlb += self.Dice_Loss(x_rec, target3) * rec_weight
|
||||
|
||||
# loss_vlb = loss_vlb
|
||||
loss_vlb = loss_vlb.mean()
|
||||
|
||||
if self.cfg.get('use_disloss', False):
|
||||
with torch.no_grad():
|
||||
edge_rec = self.first_stage_model.decode(x_rec / self.scale_factor)
|
||||
edge_rec = unnormalize_to_zero_to_one(edge_rec)
|
||||
edge_rec = torch.clamp(edge_rec, min=0., max=1.) # B, 1, 320, 320
|
||||
loss_tmp = self.cross_entropy_loss_RCF(edge_rec, (kwargs['edge'] + 1)/2) * rec_weight # B, 1
|
||||
loss_ce = SpecifyGradient.apply(x_rec, loss_tmp.mean())
|
||||
# print(loss_ce.shape)
|
||||
# print(loss_vlb.shape)
|
||||
loss_vlb += loss_ce.mean()
|
||||
loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
|
||||
|
||||
loss = loss_simple + loss_vlb
|
||||
loss_dict.update({f'{prefix}/loss': loss})
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
def get_loss(self, pred, target, mean=True):
|
||||
if self.loss_type == 'l1':
|
||||
loss = (target - pred).abs()
|
||||
if mean:
|
||||
loss = loss.mean()
|
||||
elif self.loss_type == 'l2':
|
||||
if mean:
|
||||
loss = torch.nn.functional.mse_loss(target, pred)
|
||||
else:
|
||||
loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
|
||||
else:
|
||||
raise NotImplementedError("unknown loss type '{loss_type}'")
|
||||
|
||||
return loss
|
||||
|
||||
def cross_entropy_loss_RCF(self, prediction, labelf, beta=1.1):
|
||||
# label = labelf.long()
|
||||
label = labelf
|
||||
mask = labelf.clone()
|
||||
|
||||
num_positive = torch.sum(label == 1).float()
|
||||
num_negative = torch.sum(label == 0).float()
|
||||
|
||||
mask_temp = (label > 0) & (label <= 0.3)
|
||||
mask[mask_temp] = 0.
|
||||
|
||||
mask[label == 1] = 1.0 * num_negative / (num_positive + num_negative)
|
||||
mask[label == 0] = beta * num_positive / (num_positive + num_negative)
|
||||
|
||||
# mask[label == 2] = 0
|
||||
cost = F.binary_cross_entropy(prediction, labelf, weight=mask, reduction='none')
|
||||
return cost.mean([1, 2, 3])
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self, batch_size=16, up_scale=1, cond=None, mask=None, denoise=True):
|
||||
# image_size, channels = self.image_size, self.channels
|
||||
channels = self.channels
|
||||
image_size = cond.shape[-2:]
|
||||
if cond is not None:
|
||||
batch_size = cond.shape[0]
|
||||
down_ratio = self.first_stage_model.down_ratio
|
||||
if self.cfg.model_name == 'cond_unet8' or self.cfg.model_name == 'cond_unet13':
|
||||
z, aux_out = self.sample_fn((batch_size, channels, image_size[0] // down_ratio, image_size[1] // down_ratio),
|
||||
up_scale=up_scale, unnormalize=False, cond=cond, denoise=denoise)
|
||||
else:
|
||||
z = self.sample_fn((batch_size, channels, image_size[0]//down_ratio, image_size[1]//down_ratio),
|
||||
up_scale=up_scale, unnormalize=False, cond=cond, denoise=denoise)
|
||||
aux_out = None
|
||||
|
||||
if self.scale_by_std:
|
||||
z = 1. / self.scale_factor * z.detach()
|
||||
if self.cfg.model_name == 'cond_unet13':
|
||||
aux_out = 1. / self.scale_factor * aux_out.detach()
|
||||
elif self.scale_by_softsign:
|
||||
z = z / (1 - z.abs())
|
||||
z = z.detach()
|
||||
#print(z.shape)
|
||||
x_rec = self.first_stage_model.decode(z)
|
||||
x_rec = unnormalize_to_zero_to_one(x_rec)
|
||||
x_rec = torch.clamp(x_rec, min=0., max=1.)
|
||||
if self.cfg.model_name == 'cond_unet13':
|
||||
aux_out = self.first_stage_model.decode(aux_out)
|
||||
aux_out = unnormalize_to_zero_to_one(aux_out)
|
||||
aux_out = torch.clamp(aux_out, min=0., max=1.)
|
||||
if mask is not None:
|
||||
x_rec = mask * unnormalize_to_zero_to_one(cond) + (1 - mask) * x_rec
|
||||
if aux_out is not None:
|
||||
return x_rec, aux_out
|
||||
return x_rec
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_fn(self, shape, up_scale=1, unnormalize=True, cond=None, denoise=False):
|
||||
batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], \
|
||||
self.betas.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
|
||||
|
||||
# times = torch.linspace(-1, total_timesteps, steps=self.sampling_timesteps + 1).int()
|
||||
# times = list(reversed(times.int().tolist()))
|
||||
# time_pairs = list(zip(times[:-1], times[1:]))
|
||||
# time_steps = torch.tensor([0.25, 0.15, 0.1, 0.1, 0.1, 0.09, 0.075, 0.06, 0.045, 0.03])
|
||||
step = 1. / self.sampling_timesteps
|
||||
# time_steps = torch.tensor([0.1]).repeat(10)
|
||||
time_steps = torch.tensor([step]).repeat(self.sampling_timesteps)
|
||||
if denoise:
|
||||
eps = self.eps
|
||||
time_steps = torch.cat((time_steps[:-1], torch.tensor([step - eps]), torch.tensor([eps])), dim=0)
|
||||
|
||||
if self.start_dist == 'normal':
|
||||
img = torch.randn(shape, device=device)
|
||||
elif self.start_dist == 'uniform':
|
||||
img = 2 * torch.rand(shape, device=device) - 1.
|
||||
else:
|
||||
raise NotImplementedError(f'{self.start_dist} is not supported !')
|
||||
img = F.interpolate(img, scale_factor=up_scale, mode='bilinear', align_corners=True)
|
||||
img_aux = F.interpolate(img.clone(), scale_factor=up_scale, mode='bilinear', align_corners=True)
|
||||
# img_aux = img.clone()
|
||||
# K = -1 * torch.ones_like(img)
|
||||
cur_time = torch.ones((batch,), device=device)
|
||||
for i, time_step in enumerate(time_steps):
|
||||
s = torch.full((batch,), time_step, device=device)
|
||||
if i == time_steps.shape[0] - 1:
|
||||
s = cur_time
|
||||
if cond is not None:
|
||||
pred = self.model(img, cur_time, cond)
|
||||
else:
|
||||
pred = self.model(img, cur_time)
|
||||
# C, noise = pred.chunk(2, dim=1)
|
||||
C, noise = pred[:2]
|
||||
if self.cfg.model_name == 'cond_unet8' or self.cfg.model_name == 'cond_unet13':
|
||||
aux_out = pred[-1]
|
||||
else:
|
||||
aux_out = None
|
||||
# if self.scale_by_softsign:
|
||||
# # correct the C for softsign
|
||||
# x0 = self.pred_x0_from_xt(img, noise, C, cur_time)
|
||||
# x0 = torch.clamp(x0, min=-0.987654321, max=0.987654321)
|
||||
# C = -x0
|
||||
# correct C
|
||||
x0 = self.pred_x0_from_xt(img, noise, C, cur_time)
|
||||
C = -1 * x0
|
||||
img = self.pred_xtms_from_xt(img, noise, C, cur_time, s)
|
||||
# if self.cfg.model_name == 'cond_unet13' and i == len(time_steps) - 2:
|
||||
# img_aux = img
|
||||
# if self.cfg.model_name == 'cond_unet13' and i in [len(time_steps)-2, len(time_steps)-1]:
|
||||
# x0_aux = self.pred_x0_from_xt(img_aux, noise, aux_out, cur_time)
|
||||
# C_aux = -1 * x0_aux
|
||||
# img_aux = self.pred_xtms_from_xt(img_aux, noise, C_aux, cur_time, s)
|
||||
if self.cfg.model_name == 'cond_unet13':
|
||||
for _ in range(1):
|
||||
x0_aux = self.pred_x0_from_xt(img_aux, noise, aux_out, cur_time)
|
||||
C_aux = -1 * x0_aux
|
||||
img_aux = self.pred_xtms_from_xt(img_aux, noise, C_aux, cur_time, s)
|
||||
cur_time = cur_time - s
|
||||
if self.scale_by_softsign:
|
||||
img.clamp_(-0.987654321, 0.987654321)
|
||||
if unnormalize:
|
||||
img = unnormalize_to_zero_to_one(img)
|
||||
if self.cfg.model_name == 'cond_unet13':
|
||||
aux_out = img_aux
|
||||
if aux_out is not None:
|
||||
return img, aux_out
|
||||
return img
|
||||
|
||||
class SpecifyGradient(torch.autograd.Function):
|
||||
@staticmethod
|
||||
@custom_fwd
|
||||
def forward(ctx, input_tensor, gt_grad):
|
||||
ctx.save_for_backward(gt_grad)
|
||||
# we return a dummy value 1, which will be scaled by amp's scaler so we get the scale in backward.
|
||||
return torch.ones(input_tensor.shape, device=input_tensor.device, dtype=input_tensor.dtype)
|
||||
|
||||
@staticmethod
|
||||
@custom_bwd
|
||||
def backward(ctx, grad_scale):
|
||||
(gt_grad,) = ctx.saved_tensors
|
||||
gt_grad = gt_grad * grad_scale
|
||||
return gt_grad, None
|
||||
|
||||
if __name__ == "__main__":
|
||||
ddconfig = {'double_z': True,
|
||||
'z_channels': 4,
|
||||
'resolution': (240, 960),
|
||||
'in_channels': 3,
|
||||
'out_ch': 3,
|
||||
'ch': 128,
|
||||
'ch_mult': [1, 2, 4, 4], # num_down = len(ch_mult)-1
|
||||
'num_res_blocks': 2,
|
||||
'attn_resolutions': [],
|
||||
'dropout': 0.0}
|
||||
lossconfig = {'disc_start': 50001,
|
||||
'kl_weight': 0.000001,
|
||||
'disc_weight': 0.5}
|
||||
from encoder_decoder import AutoencoderKL
|
||||
auto_encoder = AutoencoderKL(ddconfig, lossconfig, embed_dim=4,
|
||||
)
|
||||
from mask_cond_unet import Unet
|
||||
unet = Unet(dim=64, dim_mults=(1, 2, 4, 8), channels=4, cond_in_dim=1,)
|
||||
ldm = LatentDiffusion(auto_encoder=auto_encoder, model=unet, image_size=ddconfig['resolution'])
|
||||
image = torch.rand(1, 3, 128, 128)
|
||||
mask = torch.rand(1, 1, 128, 128)
|
||||
input = {'image': image, 'cond': mask}
|
||||
time = torch.tensor([1])
|
||||
with torch.no_grad():
|
||||
y = ldm.training_step(input)
|
||||
pass
|
||||
+1130
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,191 @@
|
||||
import copy
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def exists(val):
|
||||
return val is not None
|
||||
|
||||
|
||||
def clamp(value, min_value=None, max_value=None):
|
||||
assert exists(min_value) or exists(max_value)
|
||||
if exists(min_value):
|
||||
value = max(value, min_value)
|
||||
|
||||
if exists(max_value):
|
||||
value = min(value, max_value)
|
||||
|
||||
return value
|
||||
|
||||
|
||||
class EMA(nn.Module):
|
||||
"""
|
||||
Implements exponential moving average shadowing for your model.
|
||||
|
||||
Utilizes an inverse decay schedule to manage longer term training runs.
|
||||
By adjusting the power, you can control how fast EMA will ramp up to your specified beta.
|
||||
|
||||
@crowsonkb's notes on EMA Warmup:
|
||||
|
||||
If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are
|
||||
good values for models you plan to train for a million or more steps (reaches decay
|
||||
factor 0.999 at 31.6K steps, 0.9999 at 1M steps), gamma=1, power=3/4 for models
|
||||
you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
|
||||
215.4k steps).
|
||||
|
||||
Args:
|
||||
inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
|
||||
power (float): Exponential factor of EMA warmup. Default: 1.
|
||||
min_value (float): The minimum EMA decay rate. Default: 0.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
ema_model=None,
|
||||
# if your model has lazylinears or other types of non-deepcopyable modules, you can pass in your own ema model
|
||||
beta=0.9999,
|
||||
update_after_step=100,
|
||||
update_every=10,
|
||||
inv_gamma=1.0,
|
||||
power=2 / 3,
|
||||
min_value=0.0,
|
||||
param_or_buffer_names_no_ema=set(),
|
||||
ignore_names=set(),
|
||||
ignore_startswith_names=set(),
|
||||
include_online_model=True
|
||||
# set this to False if you do not wish for the online model to be saved along with the ema model (managed externally)
|
||||
):
|
||||
super().__init__()
|
||||
self.beta = beta
|
||||
|
||||
# whether to include the online model within the module tree, so that state_dict also saves it
|
||||
|
||||
self.include_online_model = include_online_model
|
||||
|
||||
if include_online_model:
|
||||
self.online_model = model
|
||||
else:
|
||||
self.online_model = [model] # hack
|
||||
|
||||
# ema model
|
||||
|
||||
self.ema_model = ema_model
|
||||
|
||||
if not exists(self.ema_model):
|
||||
try:
|
||||
self.ema_model = copy.deepcopy(model)
|
||||
except:
|
||||
print('Your model was not copyable. Please make sure you are not using any LazyLinear')
|
||||
exit()
|
||||
|
||||
self.ema_model.requires_grad_(False)
|
||||
|
||||
self.parameter_names = {name for name, param in self.ema_model.named_parameters() if param.dtype == torch.float}
|
||||
self.buffer_names = {name for name, buffer in self.ema_model.named_buffers() if buffer.dtype == torch.float}
|
||||
|
||||
self.update_every = update_every
|
||||
self.update_after_step = update_after_step
|
||||
|
||||
self.inv_gamma = inv_gamma
|
||||
self.power = power
|
||||
self.min_value = min_value
|
||||
|
||||
assert isinstance(param_or_buffer_names_no_ema, (set, list))
|
||||
self.param_or_buffer_names_no_ema = param_or_buffer_names_no_ema # parameter or buffer
|
||||
|
||||
self.ignore_names = ignore_names
|
||||
self.ignore_startswith_names = ignore_startswith_names
|
||||
|
||||
self.register_buffer('initted', torch.Tensor([False]))
|
||||
self.register_buffer('step', torch.tensor([0]))
|
||||
|
||||
@property
|
||||
def model(self):
|
||||
return self.online_model if self.include_online_model else self.online_model[0]
|
||||
|
||||
def restore_ema_model_device(self):
|
||||
device = self.initted.device
|
||||
self.ema_model.to(device)
|
||||
|
||||
def get_params_iter(self, model):
|
||||
for name, param in model.named_parameters():
|
||||
if name not in self.parameter_names:
|
||||
continue
|
||||
yield name, param
|
||||
|
||||
def get_buffers_iter(self, model):
|
||||
for name, buffer in model.named_buffers():
|
||||
if name not in self.buffer_names:
|
||||
continue
|
||||
yield name, buffer
|
||||
|
||||
def copy_params_from_model_to_ema(self):
|
||||
for (_, ma_params), (_, current_params) in zip(self.get_params_iter(self.ema_model),
|
||||
self.get_params_iter(self.model)):
|
||||
ma_params.data.copy_(current_params.data)
|
||||
|
||||
for (_, ma_buffers), (_, current_buffers) in zip(self.get_buffers_iter(self.ema_model),
|
||||
self.get_buffers_iter(self.model)):
|
||||
ma_buffers.data.copy_(current_buffers.data)
|
||||
|
||||
def get_current_decay(self):
|
||||
epoch = clamp(self.step.item() - self.update_after_step - 1, min_value=0.)
|
||||
value = 1 - (1 + epoch / self.inv_gamma) ** - self.power
|
||||
|
||||
if epoch <= 0:
|
||||
return 0.
|
||||
|
||||
return clamp(value, min_value=self.min_value, max_value=self.beta)
|
||||
|
||||
def update(self):
|
||||
step = self.step.item()
|
||||
self.step += 1
|
||||
|
||||
if (step % self.update_every) != 0:
|
||||
return
|
||||
|
||||
if step <= self.update_after_step:
|
||||
self.copy_params_from_model_to_ema()
|
||||
return
|
||||
|
||||
if not self.initted.item():
|
||||
self.copy_params_from_model_to_ema()
|
||||
self.initted.data.copy_(torch.Tensor([True]))
|
||||
|
||||
self.update_moving_average(self.ema_model, self.model)
|
||||
|
||||
@torch.no_grad()
|
||||
def update_moving_average(self, ma_model, current_model):
|
||||
current_decay = self.get_current_decay()
|
||||
|
||||
for (name, current_params), (_, ma_params) in zip(self.get_params_iter(current_model),
|
||||
self.get_params_iter(ma_model)):
|
||||
if name in self.ignore_names:
|
||||
continue
|
||||
|
||||
if any([name.startswith(prefix) for prefix in self.ignore_startswith_names]):
|
||||
continue
|
||||
|
||||
if name in self.param_or_buffer_names_no_ema:
|
||||
ma_params.data.copy_(current_params.data)
|
||||
continue
|
||||
|
||||
ma_params.data.lerp_(current_params.data, 1. - current_decay)
|
||||
|
||||
for (name, current_buffer), (_, ma_buffer) in zip(self.get_buffers_iter(current_model),
|
||||
self.get_buffers_iter(ma_model)):
|
||||
if name in self.ignore_names:
|
||||
continue
|
||||
|
||||
if any([name.startswith(prefix) for prefix in self.ignore_startswith_names]):
|
||||
continue
|
||||
|
||||
if name in self.param_or_buffer_names_no_ema:
|
||||
ma_buffer.data.copy_(current_buffer.data)
|
||||
continue
|
||||
|
||||
ma_buffer.data.lerp_(current_buffer.data, 1. - current_decay)
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self.ema_model(*args, **kwargs)
|
||||
+1086
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,395 @@
|
||||
import os, yaml, pickle, shutil, tarfile, glob
|
||||
import cv2
|
||||
import albumentations
|
||||
import PIL
|
||||
import numpy as np
|
||||
import torchvision.transforms.functional as TF
|
||||
# from omegaconf import OmegaConf
|
||||
from functools import partial
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
from torch.utils.data import Dataset, Subset
|
||||
|
||||
import taming.data.utils as tdu
|
||||
from custom_controlnet_aux.diffusion_edge.taming.data.imagenet import str_to_indices, give_synsets_from_indices, download, retrieve
|
||||
from custom_controlnet_aux.diffusion_edge.taming.data.imagenet import ImagePaths
|
||||
|
||||
# from ldm.modules.image_degradation import degradation_fn_bsr, degradation_fn_bsr_light
|
||||
|
||||
|
||||
def synset2idx(path_to_yaml="data/index_synset.yaml"):
|
||||
with open(path_to_yaml) as f:
|
||||
di2s = yaml.load(f)
|
||||
return dict((v,k) for k,v in di2s.items())
|
||||
|
||||
|
||||
class ImageNetBase(Dataset):
|
||||
def __init__(self, config=None):
|
||||
self.config = config
|
||||
# if not type(self.config)==dict:
|
||||
# self.config = OmegaConf.to_container(self.config)
|
||||
self.keep_orig_class_label = self.config.get("keep_orig_class_label", False)
|
||||
self.process_images = True # if False we skip loading & processing images and self.data contains filepaths
|
||||
self._prepare()
|
||||
self._prepare_synset_to_human()
|
||||
self._prepare_idx_to_synset()
|
||||
self._prepare_human_to_integer_label()
|
||||
self._load()
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, i):
|
||||
return self.data[i]
|
||||
|
||||
def _prepare(self):
|
||||
raise NotImplementedError()
|
||||
|
||||
def _filter_relpaths(self, relpaths):
|
||||
ignore = set([
|
||||
"n06596364_9591.JPEG",
|
||||
])
|
||||
relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore]
|
||||
if "sub_indices" in self.config:
|
||||
indices = str_to_indices(self.config["sub_indices"])
|
||||
synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn) # returns a list of strings
|
||||
self.synset2idx = synset2idx(path_to_yaml=self.idx2syn)
|
||||
files = []
|
||||
for rpath in relpaths:
|
||||
syn = rpath.split("/")[0]
|
||||
if syn in synsets:
|
||||
files.append(rpath)
|
||||
return files
|
||||
else:
|
||||
return relpaths
|
||||
|
||||
def _prepare_synset_to_human(self):
|
||||
SIZE = 2655750
|
||||
URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1"
|
||||
self.human_dict = os.path.join(self.root, "synset_human.txt")
|
||||
if (not os.path.exists(self.human_dict) or
|
||||
not os.path.getsize(self.human_dict)==SIZE):
|
||||
download(URL, self.human_dict)
|
||||
|
||||
def _prepare_idx_to_synset(self):
|
||||
URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1"
|
||||
self.idx2syn = os.path.join(self.root, "index_synset.yaml")
|
||||
if (not os.path.exists(self.idx2syn)):
|
||||
download(URL, self.idx2syn)
|
||||
|
||||
def _prepare_human_to_integer_label(self):
|
||||
URL = "https://heibox.uni-heidelberg.de/f/2362b797d5be43b883f6/?dl=1"
|
||||
self.human2integer = os.path.join(self.root, "imagenet1000_clsidx_to_labels.txt")
|
||||
if (not os.path.exists(self.human2integer)):
|
||||
download(URL, self.human2integer)
|
||||
with open(self.human2integer, "r") as f:
|
||||
lines = f.read().splitlines()
|
||||
assert len(lines) == 1000
|
||||
self.human2integer_dict = dict()
|
||||
for line in lines:
|
||||
value, key = line.split(":")
|
||||
self.human2integer_dict[key] = int(value)
|
||||
|
||||
def _load(self):
|
||||
with open(self.txt_filelist, "r") as f:
|
||||
self.relpaths = f.read().splitlines()
|
||||
l1 = len(self.relpaths)
|
||||
self.relpaths = self._filter_relpaths(self.relpaths)
|
||||
print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths)))
|
||||
|
||||
self.synsets = [p.split("/")[0] for p in self.relpaths]
|
||||
self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths]
|
||||
|
||||
unique_synsets = np.unique(self.synsets)
|
||||
class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets))
|
||||
if not self.keep_orig_class_label:
|
||||
self.class_labels = [class_dict[s] for s in self.synsets]
|
||||
else:
|
||||
self.class_labels = [self.synset2idx[s] for s in self.synsets]
|
||||
|
||||
with open(self.human_dict, "r") as f:
|
||||
human_dict = f.read().splitlines()
|
||||
human_dict = dict(line.split(maxsplit=1) for line in human_dict)
|
||||
|
||||
self.human_labels = [human_dict[s] for s in self.synsets]
|
||||
|
||||
labels = {
|
||||
"relpath": np.array(self.relpaths),
|
||||
"synsets": np.array(self.synsets),
|
||||
"class_label": np.array(self.class_labels),
|
||||
"human_label": np.array(self.human_labels),
|
||||
}
|
||||
|
||||
if self.process_images:
|
||||
# self.size = retrieve(self.config, "size", default=256)
|
||||
self.size = self.config.get("size", default=256)
|
||||
self.data = ImagePaths(self.abspaths,
|
||||
labels=labels,
|
||||
size=self.size,
|
||||
random_crop=self.random_crop,
|
||||
)
|
||||
else:
|
||||
self.data = self.abspaths
|
||||
|
||||
|
||||
class ImageNetTrain(ImageNetBase):
|
||||
NAME = "ILSVRC2012_train"
|
||||
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
||||
AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2"
|
||||
FILES = [
|
||||
"ILSVRC2012_img_train.tar",
|
||||
]
|
||||
SIZES = [
|
||||
147897477120,
|
||||
]
|
||||
|
||||
def __init__(self, process_images=True, data_root=None, **kwargs):
|
||||
self.process_images = process_images
|
||||
self.data_root = data_root
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def _prepare(self):
|
||||
if self.data_root:
|
||||
self.root = os.path.join(self.data_root, self.NAME)
|
||||
else:
|
||||
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
||||
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
||||
|
||||
self.datadir = os.path.join(self.root, "data")
|
||||
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
||||
self.expected_length = 1281167
|
||||
# self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop", default=True)
|
||||
self.random_crop = self.config.get("random_crop", default=True)
|
||||
if not tdu.is_prepared(self.root):
|
||||
# prep
|
||||
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
||||
|
||||
datadir = self.datadir
|
||||
if not os.path.exists(datadir):
|
||||
path = os.path.join(self.root, self.FILES[0])
|
||||
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
||||
import academictorrents as at
|
||||
atpath = at.get(self.AT_HASH, datastore=self.root)
|
||||
assert atpath == path
|
||||
|
||||
print("Extracting {} to {}".format(path, datadir))
|
||||
os.makedirs(datadir, exist_ok=True)
|
||||
with tarfile.open(path, "r:") as tar:
|
||||
tar.extractall(path=datadir)
|
||||
|
||||
print("Extracting sub-tars.")
|
||||
subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar")))
|
||||
for subpath in tqdm(subpaths):
|
||||
subdir = subpath[:-len(".tar")]
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
with tarfile.open(subpath, "r:") as tar:
|
||||
tar.extractall(path=subdir)
|
||||
|
||||
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
||||
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
||||
filelist = sorted(filelist)
|
||||
filelist = "\n".join(filelist)+"\n"
|
||||
with open(self.txt_filelist, "w") as f:
|
||||
f.write(filelist)
|
||||
|
||||
tdu.mark_prepared(self.root)
|
||||
|
||||
|
||||
class ImageNetValidation(ImageNetBase):
|
||||
NAME = "ILSVRC2012_validation"
|
||||
URL = "http://www.image-net.org/challenges/LSVRC/2012/"
|
||||
AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5"
|
||||
VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1"
|
||||
FILES = [
|
||||
"ILSVRC2012_img_val.tar",
|
||||
"validation_synset.txt",
|
||||
]
|
||||
SIZES = [
|
||||
6744924160,
|
||||
1950000,
|
||||
]
|
||||
|
||||
def __init__(self, process_images=True, data_root=None, **kwargs):
|
||||
self.data_root = data_root
|
||||
self.process_images = process_images
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def _prepare(self):
|
||||
if self.data_root:
|
||||
self.root = os.path.join(self.data_root, self.NAME)
|
||||
else:
|
||||
cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
|
||||
self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
|
||||
self.datadir = os.path.join(self.root, "data")
|
||||
self.txt_filelist = os.path.join(self.root, "filelist.txt")
|
||||
self.expected_length = 50000
|
||||
# self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop", default=False)
|
||||
self.random_crop = self.config.get("random_crop", default=False)
|
||||
if not tdu.is_prepared(self.root):
|
||||
# prep
|
||||
print("Preparing dataset {} in {}".format(self.NAME, self.root))
|
||||
|
||||
datadir = self.datadir
|
||||
if not os.path.exists(datadir):
|
||||
path = os.path.join(self.root, self.FILES[0])
|
||||
if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
|
||||
import academictorrents as at
|
||||
atpath = at.get(self.AT_HASH, datastore=self.root)
|
||||
assert atpath == path
|
||||
|
||||
print("Extracting {} to {}".format(path, datadir))
|
||||
os.makedirs(datadir, exist_ok=True)
|
||||
with tarfile.open(path, "r:") as tar:
|
||||
tar.extractall(path=datadir)
|
||||
|
||||
vspath = os.path.join(self.root, self.FILES[1])
|
||||
if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]:
|
||||
download(self.VS_URL, vspath)
|
||||
|
||||
with open(vspath, "r") as f:
|
||||
synset_dict = f.read().splitlines()
|
||||
synset_dict = dict(line.split() for line in synset_dict)
|
||||
|
||||
print("Reorganizing into synset folders")
|
||||
synsets = np.unique(list(synset_dict.values()))
|
||||
for s in synsets:
|
||||
os.makedirs(os.path.join(datadir, s), exist_ok=True)
|
||||
for k, v in synset_dict.items():
|
||||
src = os.path.join(datadir, k)
|
||||
dst = os.path.join(datadir, v)
|
||||
shutil.move(src, dst)
|
||||
|
||||
filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
|
||||
filelist = [os.path.relpath(p, start=datadir) for p in filelist]
|
||||
filelist = sorted(filelist)
|
||||
filelist = "\n".join(filelist)+"\n"
|
||||
with open(self.txt_filelist, "w") as f:
|
||||
f.write(filelist)
|
||||
|
||||
tdu.mark_prepared(self.root)
|
||||
|
||||
|
||||
|
||||
class ImageNetSR(Dataset):
|
||||
def __init__(self, size=None,
|
||||
degradation=None, downscale_f=4, min_crop_f=0.5, max_crop_f=1.,
|
||||
random_crop=True):
|
||||
"""
|
||||
Imagenet Superresolution Dataloader
|
||||
Performs following ops in order:
|
||||
1. crops a crop of size s from image either as random or center crop
|
||||
2. resizes crop to size with cv2.area_interpolation
|
||||
3. degrades resized crop with degradation_fn
|
||||
|
||||
:param size: resizing to size after cropping
|
||||
:param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light
|
||||
:param downscale_f: Low Resolution Downsample factor
|
||||
:param min_crop_f: determines crop size s,
|
||||
where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f)
|
||||
:param max_crop_f: ""
|
||||
:param data_root:
|
||||
:param random_crop:
|
||||
"""
|
||||
self.base = self.get_base()
|
||||
assert size
|
||||
assert (size / downscale_f).is_integer()
|
||||
self.size = size
|
||||
self.LR_size = int(size / downscale_f)
|
||||
self.min_crop_f = min_crop_f
|
||||
self.max_crop_f = max_crop_f
|
||||
assert(max_crop_f <= 1.)
|
||||
self.center_crop = not random_crop
|
||||
|
||||
self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)
|
||||
|
||||
self.pil_interpolation = False # gets reset later if incase interp_op is from pillow
|
||||
|
||||
if degradation == "bsrgan":
|
||||
self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f)
|
||||
|
||||
elif degradation == "bsrgan_light":
|
||||
self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f)
|
||||
|
||||
else:
|
||||
interpolation_fn = {
|
||||
"cv_nearest": cv2.INTER_NEAREST,
|
||||
"cv_bilinear": cv2.INTER_LINEAR,
|
||||
"cv_bicubic": cv2.INTER_CUBIC,
|
||||
"cv_area": cv2.INTER_AREA,
|
||||
"cv_lanczos": cv2.INTER_LANCZOS4,
|
||||
"pil_nearest": PIL.Image.NEAREST,
|
||||
"pil_bilinear": PIL.Image.BILINEAR,
|
||||
"pil_bicubic": PIL.Image.BICUBIC,
|
||||
"pil_box": PIL.Image.BOX,
|
||||
"pil_hamming": PIL.Image.HAMMING,
|
||||
"pil_lanczos": PIL.Image.LANCZOS,
|
||||
}[degradation]
|
||||
|
||||
self.pil_interpolation = degradation.startswith("pil_")
|
||||
|
||||
if self.pil_interpolation:
|
||||
self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn)
|
||||
|
||||
else:
|
||||
self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size,
|
||||
interpolation=interpolation_fn)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.base)
|
||||
|
||||
def __getitem__(self, i):
|
||||
example = self.base[i]
|
||||
image = Image.open(example["file_path_"])
|
||||
|
||||
if not image.mode == "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
image = np.array(image).astype(np.uint8)
|
||||
|
||||
min_side_len = min(image.shape[:2])
|
||||
crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None)
|
||||
crop_side_len = int(crop_side_len)
|
||||
|
||||
if self.center_crop:
|
||||
self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len)
|
||||
|
||||
else:
|
||||
self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len)
|
||||
|
||||
image = self.cropper(image=image)["image"]
|
||||
image = self.image_rescaler(image=image)["image"]
|
||||
|
||||
if self.pil_interpolation:
|
||||
image_pil = PIL.Image.fromarray(image)
|
||||
LR_image = self.degradation_process(image_pil)
|
||||
LR_image = np.array(LR_image).astype(np.uint8)
|
||||
|
||||
else:
|
||||
LR_image = self.degradation_process(image=image)["image"]
|
||||
|
||||
example["image"] = (image/127.5 - 1.0).astype(np.float32)
|
||||
example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32)
|
||||
|
||||
return example
|
||||
|
||||
|
||||
class ImageNetSRTrain(ImageNetSR):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def get_base(self):
|
||||
with open("data/imagenet_train_hr_indices.p", "rb") as f:
|
||||
indices = pickle.load(f)
|
||||
dset = ImageNetTrain(process_images=False,)
|
||||
return Subset(dset, indices)
|
||||
|
||||
|
||||
class ImageNetSRValidation(ImageNetSR):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def get_base(self):
|
||||
with open("data/imagenet_val_hr_indices.p", "rb") as f:
|
||||
indices = pickle.load(f)
|
||||
dset = ImageNetValidation(process_images=False,)
|
||||
return Subset(dset, indices)
|
||||
@@ -0,0 +1,113 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import sys
|
||||
# .path.append()
|
||||
from custom_controlnet_aux.diffusion_edge.taming.modules.losses.vqperceptual import *
|
||||
|
||||
|
||||
class LPIPSWithDiscriminator(nn.Module):
|
||||
def __init__(self, *, disc_start, logvar_init=0.0, kl_weight=1.0, pixelloss_weight=1.0,
|
||||
disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
|
||||
perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
|
||||
disc_loss="hinge"):
|
||||
|
||||
super().__init__()
|
||||
assert disc_loss in ["hinge", "vanilla"]
|
||||
self.kl_weight = kl_weight
|
||||
self.pixel_weight = pixelloss_weight
|
||||
self.perceptual_loss = LPIPS().eval()
|
||||
self.perceptual_weight = perceptual_weight
|
||||
# output log variance
|
||||
self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init)
|
||||
|
||||
self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels,
|
||||
n_layers=disc_num_layers,
|
||||
use_actnorm=use_actnorm
|
||||
).apply(weights_init)
|
||||
self.discriminator_iter_start = disc_start
|
||||
self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
|
||||
self.disc_factor = disc_factor
|
||||
self.discriminator_weight = disc_weight
|
||||
self.disc_conditional = disc_conditional
|
||||
|
||||
def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
|
||||
if last_layer is not None:
|
||||
nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
|
||||
g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
|
||||
else:
|
||||
nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
|
||||
g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
|
||||
|
||||
d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
|
||||
d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
|
||||
d_weight = d_weight * self.discriminator_weight
|
||||
return d_weight
|
||||
|
||||
def forward(self, inputs, reconstructions, posteriors, optimizer_idx,
|
||||
global_step, last_layer=None, cond=None, split="train",
|
||||
weights=None):
|
||||
rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous()) + \
|
||||
F.mse_loss(inputs, reconstructions, reduction="none")
|
||||
if self.perceptual_weight > 0:
|
||||
p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
|
||||
rec_loss = rec_loss + self.perceptual_weight * p_loss
|
||||
|
||||
nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
|
||||
weighted_nll_loss = nll_loss
|
||||
if weights is not None:
|
||||
weighted_nll_loss = weights*nll_loss
|
||||
weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
|
||||
nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
|
||||
kl_loss = posteriors.kl()
|
||||
kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
|
||||
|
||||
# now the GAN part
|
||||
if optimizer_idx == 0:
|
||||
# generator update
|
||||
if cond is None:
|
||||
assert not self.disc_conditional
|
||||
logits_fake = self.discriminator(reconstructions.contiguous())
|
||||
else:
|
||||
assert self.disc_conditional
|
||||
logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
|
||||
g_loss = -torch.mean(logits_fake)
|
||||
|
||||
if self.disc_factor > 0.0:
|
||||
try:
|
||||
d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
|
||||
except RuntimeError:
|
||||
assert not self.training
|
||||
d_weight = torch.tensor(0.0)
|
||||
else:
|
||||
d_weight = torch.tensor(0.0)
|
||||
|
||||
disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
|
||||
loss = weighted_nll_loss + self.kl_weight * kl_loss + d_weight * disc_factor * g_loss
|
||||
|
||||
log = {"{}/total_loss".format(split): loss.clone().detach().mean(), "{}/logvar".format(split): self.logvar.detach(),
|
||||
"{}/kl_loss".format(split): kl_loss.detach().mean(), "{}/nll_loss".format(split): nll_loss.detach().mean(),
|
||||
"{}/rec_loss".format(split): rec_loss.detach().mean(),
|
||||
"{}/d_weight".format(split): d_weight.detach(),
|
||||
"{}/disc_factor".format(split): torch.tensor(disc_factor),
|
||||
"{}/g_loss".format(split): g_loss.detach().mean(),
|
||||
}
|
||||
return loss, log
|
||||
|
||||
if optimizer_idx == 1:
|
||||
# second pass for discriminator update
|
||||
if cond is None:
|
||||
logits_real = self.discriminator(inputs.contiguous().detach())
|
||||
logits_fake = self.discriminator(reconstructions.contiguous().detach())
|
||||
else:
|
||||
logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
|
||||
logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))
|
||||
|
||||
disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
|
||||
d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
|
||||
|
||||
log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(),
|
||||
"{}/logits_real".format(split): logits_real.detach().mean(),
|
||||
"{}/logits_fake".format(split): logits_fake.detach().mean()
|
||||
}
|
||||
return d_loss, log
|
||||
|
||||
+1009
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,103 @@
|
||||
import torch
|
||||
from torch import nn as nn
|
||||
from torch.nn import Parameter
|
||||
|
||||
|
||||
def weight_quantization(b):
|
||||
def uniform_quant(x, b):
|
||||
xdiv = x.mul((2 ** b - 1))
|
||||
xhard = xdiv.round().div(2 ** b - 1)
|
||||
return xhard
|
||||
|
||||
class _pq(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, input, alpha):
|
||||
input.div_(alpha) # weights are first divided by alpha
|
||||
input_c = input.clamp(min=-1, max=1) # then clipped to [-1,1]
|
||||
sign = input_c.sign()
|
||||
input_abs = input_c.abs()
|
||||
input_q = uniform_quant(input_abs, b).mul(sign)
|
||||
ctx.save_for_backward(input, input_q)
|
||||
input_q = input_q.mul(alpha) # rescale to the original range
|
||||
return input_q
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
grad_input = grad_output.clone() # grad for weights will not be clipped
|
||||
input, input_q = ctx.saved_tensors
|
||||
i = (input.abs() > 1.).float()
|
||||
sign = input.sign()
|
||||
grad_alpha = (grad_output * (sign * i + (input_q - input) * (1 - i))).sum()
|
||||
return grad_input, grad_alpha
|
||||
|
||||
return _pq().apply
|
||||
|
||||
|
||||
class weight_quantize_fn(nn.Module):
|
||||
def __init__(self, bit_w):
|
||||
super(weight_quantize_fn, self).__init__()
|
||||
assert bit_w > 0
|
||||
|
||||
self.bit_w = bit_w - 1
|
||||
self.weight_q = weight_quantization(b=self.bit_w)
|
||||
self.register_parameter('w_alpha', Parameter(torch.tensor(3.0), requires_grad=True))
|
||||
|
||||
def forward(self, weight):
|
||||
mean = weight.data.mean()
|
||||
std = weight.data.std()
|
||||
weight = weight.add(-mean).div(std) # weights normalization
|
||||
weight_q = self.weight_q(weight, self.w_alpha)
|
||||
return weight_q
|
||||
|
||||
def change_bit(self, bit_w):
|
||||
self.bit_w = bit_w - 1
|
||||
self.weight_q = weight_quantization(b=self.bit_w)
|
||||
|
||||
def act_quantization(b, signed=False):
|
||||
def uniform_quant(x, b=3):
|
||||
xdiv = x.mul(2 ** b - 1)
|
||||
xhard = xdiv.round().div(2 ** b - 1)
|
||||
return xhard
|
||||
|
||||
class _uq(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, input, alpha):
|
||||
input = input.div(alpha)
|
||||
input_c = input.clamp(min=-1, max=1) if signed else input.clamp(max=1)
|
||||
input_q = uniform_quant(input_c, b)
|
||||
ctx.save_for_backward(input, input_q)
|
||||
input_q = input_q.mul(alpha)
|
||||
return input_q
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
grad_input = grad_output.clone()
|
||||
input, input_q = ctx.saved_tensors
|
||||
i = (input.abs() > 1.).float()
|
||||
sign = input.sign()
|
||||
grad_alpha = (grad_output * (sign * i + (input_q - input) * (1 - i))).sum()
|
||||
grad_input = grad_input * (1 - i)
|
||||
return grad_input, grad_alpha
|
||||
|
||||
return _uq().apply
|
||||
|
||||
class act_quantize_fn(nn.Module):
|
||||
def __init__(self, bit_a, signed=False):
|
||||
super(act_quantize_fn, self).__init__()
|
||||
self.bit_a = bit_a
|
||||
self.signed = signed
|
||||
if signed:
|
||||
self.bit_a -= 1
|
||||
assert bit_a > 0
|
||||
|
||||
self.act_q = act_quantization(b=self.bit_a, signed=signed)
|
||||
self.register_parameter('a_alpha', Parameter(torch.tensor(8.0), requires_grad=True))
|
||||
|
||||
def forward(self, x):
|
||||
return self.act_q(x, self.a_alpha)
|
||||
|
||||
def change_bit(self, bit_a):
|
||||
self.bit_a = bit_a
|
||||
if self.signed:
|
||||
self.bit_a -= 1
|
||||
self.act_q = act_quantization(b=self.bit_a, signed=self.signed)
|
||||
@@ -0,0 +1,963 @@
|
||||
from functools import partial
|
||||
from typing import Type, Any, Callable, Union, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
from torchvision.transforms._presets import ImageClassification
|
||||
from torchvision.utils import _log_api_usage_once
|
||||
from torchvision.models._api import WeightsEnum, Weights
|
||||
from torchvision.models._meta import _IMAGENET_CATEGORIES
|
||||
from torchvision.models._utils import handle_legacy_interface, _ovewrite_named_param
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ResNet",
|
||||
"ResNet18_Weights",
|
||||
"ResNet34_Weights",
|
||||
"ResNet50_Weights",
|
||||
"ResNet101_Weights",
|
||||
"ResNet152_Weights",
|
||||
"ResNeXt50_32X4D_Weights",
|
||||
"ResNeXt101_32X8D_Weights",
|
||||
"ResNeXt101_64X4D_Weights",
|
||||
"Wide_ResNet50_2_Weights",
|
||||
"Wide_ResNet101_2_Weights",
|
||||
"resnet18",
|
||||
"resnet34",
|
||||
"resnet50",
|
||||
"resnet101",
|
||||
"resnet152",
|
||||
"resnext50_32x4d",
|
||||
"resnext101_32x8d",
|
||||
"resnext101_64x4d",
|
||||
"wide_resnet50_2",
|
||||
"wide_resnet101_2",
|
||||
]
|
||||
|
||||
|
||||
def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
padding=dilation,
|
||||
groups=groups,
|
||||
bias=False,
|
||||
dilation=dilation,
|
||||
)
|
||||
|
||||
|
||||
def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion: int = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inplanes: int,
|
||||
planes: int,
|
||||
stride: int = 1,
|
||||
downsample: Optional[nn.Module] = None,
|
||||
groups: int = 1,
|
||||
base_width: int = 64,
|
||||
dilation: int = 1,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
if groups != 1 or base_width != 64:
|
||||
raise ValueError("BasicBlock only supports groups=1 and base_width=64")
|
||||
if dilation > 1:
|
||||
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
# Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
|
||||
# while original implementation places the stride at the first 1x1 convolution(self.conv1)
|
||||
# according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
|
||||
# This variant is also known as ResNet V1.5 and improves accuracy according to
|
||||
# https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
|
||||
|
||||
expansion: int = 4
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inplanes: int,
|
||||
planes: int,
|
||||
stride: int = 1,
|
||||
downsample: Optional[nn.Module] = None,
|
||||
groups: int = 1,
|
||||
base_width: int = 64,
|
||||
dilation: int = 1,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
width = int(planes * (base_width / 64.0)) * groups
|
||||
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv1x1(inplanes, width)
|
||||
self.bn1 = norm_layer(width)
|
||||
self.conv2 = conv3x3(width, width, stride, groups, dilation)
|
||||
self.bn2 = norm_layer(width)
|
||||
self.conv3 = conv1x1(width, planes * self.expansion)
|
||||
self.bn3 = norm_layer(planes * self.expansion)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block: Type[Union[BasicBlock, Bottleneck]],
|
||||
layers: List[int],
|
||||
num_classes: int = 1000,
|
||||
zero_init_residual: bool = False,
|
||||
groups: int = 1,
|
||||
width_per_group: int = 64,
|
||||
replace_stride_with_dilation: Optional[List[bool]] = None,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
_log_api_usage_once(self)
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.dilation = 1
|
||||
if replace_stride_with_dilation is None:
|
||||
# each element in the tuple indicates if we should replace
|
||||
# the 2x2 stride with a dilated convolution instead
|
||||
replace_stride_with_dilation = [False, False, False]
|
||||
if len(replace_stride_with_dilation) != 3:
|
||||
raise ValueError(
|
||||
"replace_stride_with_dilation should be None "
|
||||
f"or a 3-element tuple, got {replace_stride_with_dilation}"
|
||||
)
|
||||
self.groups = groups
|
||||
self.base_width = width_per_group
|
||||
self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.bn1 = norm_layer(self.inplanes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
|
||||
# self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
# self.fc = nn.Linear(512 * block.expansion, num_classes)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck) and m.bn3.weight is not None:
|
||||
nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type]
|
||||
elif isinstance(m, BasicBlock) and m.bn2.weight is not None:
|
||||
nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type]
|
||||
|
||||
def _make_layer(
|
||||
self,
|
||||
block: Type[Union[BasicBlock, Bottleneck]],
|
||||
planes: int,
|
||||
blocks: int,
|
||||
stride: int = 1,
|
||||
dilate: bool = False,
|
||||
) -> nn.Sequential:
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
previous_dilation = self.dilation
|
||||
if dilate:
|
||||
self.dilation *= stride
|
||||
stride = 1
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
conv1x1(self.inplanes, planes * block.expansion, stride),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(
|
||||
block(
|
||||
self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer
|
||||
)
|
||||
)
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(
|
||||
block(
|
||||
self.inplanes,
|
||||
planes,
|
||||
groups=self.groups,
|
||||
base_width=self.base_width,
|
||||
dilation=self.dilation,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
)
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _forward_impl(self, x: Tensor) -> Tensor:
|
||||
# See note [TorchScript super()]
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
feats = []
|
||||
x = self.layer1(x)
|
||||
feats.append(x)
|
||||
x = self.layer2(x)
|
||||
feats.append(x)
|
||||
x = self.layer3(x)
|
||||
feats.append(x)
|
||||
x = self.layer4(x)
|
||||
feats.append(x)
|
||||
|
||||
# x = self.avgpool(x)
|
||||
# x = torch.flatten(x, 1)
|
||||
# x = self.fc(x)
|
||||
|
||||
return feats
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def _resnet(
|
||||
block: Type[Union[BasicBlock, Bottleneck]],
|
||||
layers: List[int],
|
||||
weights: Optional[WeightsEnum],
|
||||
progress: bool,
|
||||
**kwargs: Any,
|
||||
) -> ResNet:
|
||||
if weights is not None:
|
||||
_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
|
||||
|
||||
model = ResNet(block, layers, **kwargs)
|
||||
|
||||
if weights is not None:
|
||||
model.load_state_dict(weights.get_state_dict(progress=progress), strict=False)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
_COMMON_META = {
|
||||
"min_size": (1, 1),
|
||||
"categories": _IMAGENET_CATEGORIES,
|
||||
}
|
||||
|
||||
|
||||
class ResNet18_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet18-f37072fd.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 11689512,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 69.758,
|
||||
"acc@5": 89.078,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class ResNet34_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet34-b627a593.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 21797672,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 73.314,
|
||||
"acc@5": 91.420,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class ResNet50_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet50-0676ba61.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 25557032,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 76.130,
|
||||
"acc@5": 92.862,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet50-11ad3fa6.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 25557032,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#issuecomment-1013906621",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 80.858,
|
||||
"acc@5": 95.434,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class ResNet101_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet101-63fe2227.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 44549160,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 77.374,
|
||||
"acc@5": 93.546,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet101-cd907fc2.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 44549160,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 81.886,
|
||||
"acc@5": 95.780,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class ResNet152_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet152-394f9c45.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 60192808,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 78.312,
|
||||
"acc@5": 94.046,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/resnet152-f82ba261.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 60192808,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 82.284,
|
||||
"acc@5": 96.002,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class ResNeXt50_32X4D_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 25028904,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 77.618,
|
||||
"acc@5": 93.698,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/resnext50_32x4d-1a0047aa.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 25028904,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 81.198,
|
||||
"acc@5": 95.340,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class ResNeXt101_32X8D_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 88791336,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 79.312,
|
||||
"acc@5": 94.526,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/resnext101_32x8d-110c445d.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 88791336,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 82.834,
|
||||
"acc@5": 96.228,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class ResNeXt101_64X4D_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/resnext101_64x4d-173b62eb.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 83455272,
|
||||
"recipe": "https://github.com/pytorch/vision/pull/5935",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 83.246,
|
||||
"acc@5": 96.454,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights were trained from scratch by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class Wide_ResNet50_2_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 68883240,
|
||||
"recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 78.468,
|
||||
"acc@5": 94.086,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/wide_resnet50_2-9ba9bcbe.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 68883240,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 81.602,
|
||||
"acc@5": 95.758,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
class Wide_ResNet101_2_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 126886696,
|
||||
"recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 78.848,
|
||||
"acc@5": 94.284,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a simple training recipe.""",
|
||||
},
|
||||
)
|
||||
IMAGENET1K_V2 = Weights(
|
||||
url="https://download.pytorch.org/models/wide_resnet101_2-d733dc28.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224, resize_size=232),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 126886696,
|
||||
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 82.510,
|
||||
"acc@5": 96.020,
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights improve upon the results of the original paper by using TorchVision's `new training recipe
|
||||
<https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V2
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNet18_Weights.IMAGENET1K_V1))
|
||||
def resnet18(*, weights: Optional[ResNet18_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet:
|
||||
"""ResNet-18 from `Deep Residual Learning for Image Recognition <https://arxiv.org/pdf/1512.03385.pdf>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNet18_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNet18_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.ResNet18_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNet18_Weights.verify(weights)
|
||||
|
||||
return _resnet(BasicBlock, [2, 2, 2, 2], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNet34_Weights.IMAGENET1K_V1))
|
||||
def resnet34(*, weights: Optional[ResNet34_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet:
|
||||
"""ResNet-34 from `Deep Residual Learning for Image Recognition <https://arxiv.org/pdf/1512.03385.pdf>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNet34_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNet34_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.ResNet34_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNet34_Weights.verify(weights)
|
||||
|
||||
return _resnet(BasicBlock, [3, 4, 6, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNet50_Weights.IMAGENET1K_V1))
|
||||
def resnet50(*, weights: Optional[ResNet50_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet:
|
||||
"""ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/pdf/1512.03385.pdf>`__.
|
||||
|
||||
.. note::
|
||||
The bottleneck of TorchVision places the stride for downsampling to the second 3x3
|
||||
convolution while the original paper places it to the first 1x1 convolution.
|
||||
This variant improves the accuracy and is known as `ResNet V1.5
|
||||
<https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNet50_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNet50_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.ResNet50_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNet50_Weights.verify(weights)
|
||||
|
||||
return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNet101_Weights.IMAGENET1K_V1))
|
||||
def resnet101(*, weights: Optional[ResNet101_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet:
|
||||
"""ResNet-101 from `Deep Residual Learning for Image Recognition <https://arxiv.org/pdf/1512.03385.pdf>`__.
|
||||
|
||||
.. note::
|
||||
The bottleneck of TorchVision places the stride for downsampling to the second 3x3
|
||||
convolution while the original paper places it to the first 1x1 convolution.
|
||||
This variant improves the accuracy and is known as `ResNet V1.5
|
||||
<https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNet101_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNet101_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.ResNet101_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNet101_Weights.verify(weights)
|
||||
|
||||
return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNet152_Weights.IMAGENET1K_V1))
|
||||
def resnet152(*, weights: Optional[ResNet152_Weights] = None, progress: bool = True, **kwargs: Any) -> ResNet:
|
||||
"""ResNet-152 from `Deep Residual Learning for Image Recognition <https://arxiv.org/pdf/1512.03385.pdf>`__.
|
||||
|
||||
.. note::
|
||||
The bottleneck of TorchVision places the stride for downsampling to the second 3x3
|
||||
convolution while the original paper places it to the first 1x1 convolution.
|
||||
This variant improves the accuracy and is known as `ResNet V1.5
|
||||
<https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNet152_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNet152_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.ResNet152_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNet152_Weights.verify(weights)
|
||||
|
||||
return _resnet(Bottleneck, [3, 8, 36, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNeXt50_32X4D_Weights.IMAGENET1K_V1))
|
||||
def resnext50_32x4d(
|
||||
*, weights: Optional[ResNeXt50_32X4D_Weights] = None, progress: bool = True, **kwargs: Any
|
||||
) -> ResNet:
|
||||
"""ResNeXt-50 32x4d model from
|
||||
`Aggregated Residual Transformation for Deep Neural Networks <https://arxiv.org/abs/1611.05431>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNeXt50_32X4D_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNext50_32X4D_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
.. autoclass:: torchvision.models.ResNeXt50_32X4D_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNeXt50_32X4D_Weights.verify(weights)
|
||||
|
||||
_ovewrite_named_param(kwargs, "groups", 32)
|
||||
_ovewrite_named_param(kwargs, "width_per_group", 4)
|
||||
return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", ResNeXt101_32X8D_Weights.IMAGENET1K_V1))
|
||||
def resnext101_32x8d(
|
||||
*, weights: Optional[ResNeXt101_32X8D_Weights] = None, progress: bool = True, **kwargs: Any
|
||||
) -> ResNet:
|
||||
"""ResNeXt-101 32x8d model from
|
||||
`Aggregated Residual Transformation for Deep Neural Networks <https://arxiv.org/abs/1611.05431>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNeXt101_32X8D_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNeXt101_32X8D_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
.. autoclass:: torchvision.models.ResNeXt101_32X8D_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNeXt101_32X8D_Weights.verify(weights)
|
||||
|
||||
_ovewrite_named_param(kwargs, "groups", 32)
|
||||
_ovewrite_named_param(kwargs, "width_per_group", 8)
|
||||
return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
def resnext101_64x4d(
|
||||
*, weights: Optional[ResNeXt101_64X4D_Weights] = None, progress: bool = True, **kwargs: Any
|
||||
) -> ResNet:
|
||||
"""ResNeXt-101 64x4d model from
|
||||
`Aggregated Residual Transformation for Deep Neural Networks <https://arxiv.org/abs/1611.05431>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.ResNeXt101_64X4D_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.ResNeXt101_64X4D_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
.. autoclass:: torchvision.models.ResNeXt101_64X4D_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = ResNeXt101_64X4D_Weights.verify(weights)
|
||||
|
||||
_ovewrite_named_param(kwargs, "groups", 64)
|
||||
_ovewrite_named_param(kwargs, "width_per_group", 4)
|
||||
return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", Wide_ResNet50_2_Weights.IMAGENET1K_V1))
|
||||
def wide_resnet50_2(
|
||||
*, weights: Optional[Wide_ResNet50_2_Weights] = None, progress: bool = True, **kwargs: Any
|
||||
) -> ResNet:
|
||||
"""Wide ResNet-50-2 model from
|
||||
`Wide Residual Networks <https://arxiv.org/abs/1605.07146>`_.
|
||||
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
|
||||
channels, and in Wide ResNet-50-2 has 2048-1024-2048.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.Wide_ResNet50_2_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.Wide_ResNet50_2_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
.. autoclass:: torchvision.models.Wide_ResNet50_2_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = Wide_ResNet50_2_Weights.verify(weights)
|
||||
|
||||
_ovewrite_named_param(kwargs, "width_per_group", 64 * 2)
|
||||
return _resnet(Bottleneck, [3, 4, 6, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", Wide_ResNet101_2_Weights.IMAGENET1K_V1))
|
||||
def wide_resnet101_2(
|
||||
*, weights: Optional[Wide_ResNet101_2_Weights] = None, progress: bool = True, **kwargs: Any
|
||||
) -> ResNet:
|
||||
"""Wide ResNet-101-2 model from
|
||||
`Wide Residual Networks <https://arxiv.org/abs/1605.07146>`_.
|
||||
|
||||
The model is the same as ResNet except for the bottleneck number of channels
|
||||
which is twice larger in every block. The number of channels in outer 1x1
|
||||
convolutions is the same, e.g. last block in ResNet-101 has 2048-512-2048
|
||||
channels, and in Wide ResNet-101-2 has 2048-1024-2048.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.Wide_ResNet101_2_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.Wide_ResNet101_2_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.resnet.ResNet``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py>`_
|
||||
for more details about this class.
|
||||
.. autoclass:: torchvision.models.Wide_ResNet101_2_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = Wide_ResNet101_2_Weights.verify(weights)
|
||||
|
||||
_ovewrite_named_param(kwargs, "width_per_group", 64 * 2)
|
||||
return _resnet(Bottleneck, [3, 4, 23, 3], weights, progress, **kwargs)
|
||||
|
||||
|
||||
# The dictionary below is internal implementation detail and will be removed in v0.15
|
||||
from torchvision.models._utils import _ModelURLs
|
||||
|
||||
|
||||
model_urls = _ModelURLs(
|
||||
{
|
||||
"resnet18": ResNet18_Weights.IMAGENET1K_V1.url,
|
||||
"resnet34": ResNet34_Weights.IMAGENET1K_V1.url,
|
||||
"resnet50": ResNet50_Weights.IMAGENET1K_V1.url,
|
||||
"resnet101": ResNet101_Weights.IMAGENET1K_V1.url,
|
||||
"resnet152": ResNet152_Weights.IMAGENET1K_V1.url,
|
||||
"resnext50_32x4d": ResNeXt50_32X4D_Weights.IMAGENET1K_V1.url,
|
||||
"resnext101_32x8d": ResNeXt101_32X8D_Weights.IMAGENET1K_V1.url,
|
||||
"wide_resnet50_2": Wide_ResNet50_2_Weights.IMAGENET1K_V1.url,
|
||||
"wide_resnet101_2": Wide_ResNet101_2_Weights.IMAGENET1K_V1.url,
|
||||
}
|
||||
)
|
||||
+651
@@ -0,0 +1,651 @@
|
||||
from functools import partial
|
||||
from typing import Optional, Callable, List, Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn, Tensor
|
||||
|
||||
from torchvision.ops.misc import MLP, Permute
|
||||
from torchvision.ops.stochastic_depth import StochasticDepth
|
||||
from torchvision.transforms._presets import ImageClassification, InterpolationMode
|
||||
from torchvision.utils import _log_api_usage_once
|
||||
from torchvision.models._api import WeightsEnum, Weights
|
||||
from torchvision.models._meta import _IMAGENET_CATEGORIES
|
||||
from torchvision.models._utils import _ovewrite_named_param
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SwinTransformer",
|
||||
"Swin_T_Weights",
|
||||
"Swin_S_Weights",
|
||||
"Swin_B_Weights",
|
||||
"swin_t",
|
||||
"swin_s",
|
||||
"swin_b",
|
||||
]
|
||||
|
||||
|
||||
def _patch_merging_pad(x):
|
||||
H, W, _ = x.shape[-3:]
|
||||
x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
|
||||
return x
|
||||
|
||||
|
||||
torch.fx.wrap("_patch_merging_pad")
|
||||
|
||||
|
||||
class PatchMerging(nn.Module):
|
||||
"""Patch Merging Layer.
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, norm_layer: Callable[..., nn.Module] = nn.LayerNorm):
|
||||
super().__init__()
|
||||
_log_api_usage_once(self)
|
||||
self.dim = dim
|
||||
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
|
||||
self.norm = norm_layer(4 * dim)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
"""
|
||||
Args:
|
||||
x (Tensor): input tensor with expected layout of [..., H, W, C]
|
||||
Returns:
|
||||
Tensor with layout of [..., H/2, W/2, 2*C]
|
||||
"""
|
||||
x = _patch_merging_pad(x)
|
||||
|
||||
x0 = x[..., 0::2, 0::2, :] # ... H/2 W/2 C
|
||||
x1 = x[..., 1::2, 0::2, :] # ... H/2 W/2 C
|
||||
x2 = x[..., 0::2, 1::2, :] # ... H/2 W/2 C
|
||||
x3 = x[..., 1::2, 1::2, :] # ... H/2 W/2 C
|
||||
x = torch.cat([x0, x1, x2, x3], -1) # ... H/2 W/2 4*C
|
||||
|
||||
x = self.norm(x)
|
||||
x = self.reduction(x) # ... H/2 W/2 2*C
|
||||
return x
|
||||
|
||||
|
||||
def shifted_window_attention(
|
||||
input: Tensor,
|
||||
qkv_weight: Tensor,
|
||||
proj_weight: Tensor,
|
||||
relative_position_bias: Tensor,
|
||||
window_size: List[int],
|
||||
num_heads: int,
|
||||
shift_size: List[int],
|
||||
attention_dropout: float = 0.0,
|
||||
dropout: float = 0.0,
|
||||
qkv_bias: Optional[Tensor] = None,
|
||||
proj_bias: Optional[Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Window based multi-head self attention (W-MSA) module with relative position bias.
|
||||
It supports both of shifted and non-shifted window.
|
||||
Args:
|
||||
input (Tensor[N, H, W, C]): The input tensor or 4-dimensions.
|
||||
qkv_weight (Tensor[in_dim, out_dim]): The weight tensor of query, key, value.
|
||||
proj_weight (Tensor[out_dim, out_dim]): The weight tensor of projection.
|
||||
relative_position_bias (Tensor): The learned relative position bias added to attention.
|
||||
window_size (List[int]): Window size.
|
||||
num_heads (int): Number of attention heads.
|
||||
shift_size (List[int]): Shift size for shifted window attention.
|
||||
attention_dropout (float): Dropout ratio of attention weight. Default: 0.0.
|
||||
dropout (float): Dropout ratio of output. Default: 0.0.
|
||||
qkv_bias (Tensor[out_dim], optional): The bias tensor of query, key, value. Default: None.
|
||||
proj_bias (Tensor[out_dim], optional): The bias tensor of projection. Default: None.
|
||||
Returns:
|
||||
Tensor[N, H, W, C]: The output tensor after shifted window attention.
|
||||
"""
|
||||
B, H, W, C = input.shape
|
||||
# pad feature maps to multiples of window size
|
||||
pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
|
||||
pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
|
||||
x = F.pad(input, (0, 0, 0, pad_r, 0, pad_b))
|
||||
_, pad_H, pad_W, _ = x.shape
|
||||
|
||||
# If window size is larger than feature size, there is no need to shift window
|
||||
if window_size[0] >= pad_H:
|
||||
shift_size[0] = 0
|
||||
if window_size[1] >= pad_W:
|
||||
shift_size[1] = 0
|
||||
|
||||
# cyclic shift
|
||||
if sum(shift_size) > 0:
|
||||
x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2))
|
||||
|
||||
# partition windows
|
||||
num_windows = (pad_H // window_size[0]) * (pad_W // window_size[1])
|
||||
x = x.view(B, pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1], C)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).reshape(B * num_windows, window_size[0] * window_size[1], C) # B*nW, Ws*Ws, C
|
||||
|
||||
# multi-head attention
|
||||
qkv = F.linear(x, qkv_weight, qkv_bias)
|
||||
qkv = qkv.reshape(x.size(0), x.size(1), 3, num_heads, C // num_heads).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv[0], qkv[1], qkv[2]
|
||||
q = q * (C // num_heads) ** -0.5
|
||||
attn = q.matmul(k.transpose(-2, -1))
|
||||
# add relative position bias
|
||||
attn = attn + relative_position_bias
|
||||
|
||||
if sum(shift_size) > 0:
|
||||
# generate attention mask
|
||||
attn_mask = x.new_zeros((pad_H, pad_W))
|
||||
h_slices = ((0, -window_size[0]), (-window_size[0], -shift_size[0]), (-shift_size[0], None))
|
||||
w_slices = ((0, -window_size[1]), (-window_size[1], -shift_size[1]), (-shift_size[1], None))
|
||||
count = 0
|
||||
for h in h_slices:
|
||||
for w in w_slices:
|
||||
attn_mask[h[0] : h[1], w[0] : w[1]] = count
|
||||
count += 1
|
||||
attn_mask = attn_mask.view(pad_H // window_size[0], window_size[0], pad_W // window_size[1], window_size[1])
|
||||
attn_mask = attn_mask.permute(0, 2, 1, 3).reshape(num_windows, window_size[0] * window_size[1])
|
||||
attn_mask = attn_mask.unsqueeze(1) - attn_mask.unsqueeze(2)
|
||||
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
|
||||
attn = attn.view(x.size(0) // num_windows, num_windows, num_heads, x.size(1), x.size(1))
|
||||
attn = attn + attn_mask.unsqueeze(1).unsqueeze(0)
|
||||
attn = attn.view(-1, num_heads, x.size(1), x.size(1))
|
||||
|
||||
attn = F.softmax(attn, dim=-1)
|
||||
attn = F.dropout(attn, p=attention_dropout)
|
||||
|
||||
x = attn.matmul(v).transpose(1, 2).reshape(x.size(0), x.size(1), C)
|
||||
x = F.linear(x, proj_weight, proj_bias)
|
||||
x = F.dropout(x, p=dropout)
|
||||
|
||||
# reverse windows
|
||||
x = x.view(B, pad_H // window_size[0], pad_W // window_size[1], window_size[0], window_size[1], C)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, pad_H, pad_W, C)
|
||||
|
||||
# reverse cyclic shift
|
||||
if sum(shift_size) > 0:
|
||||
x = torch.roll(x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2))
|
||||
|
||||
# unpad features
|
||||
x = x[:, :H, :W, :].contiguous()
|
||||
return x
|
||||
|
||||
|
||||
torch.fx.wrap("shifted_window_attention")
|
||||
|
||||
|
||||
class ShiftedWindowAttention(nn.Module):
|
||||
"""
|
||||
See :func:`shifted_window_attention`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
window_size: List[int],
|
||||
shift_size: List[int],
|
||||
num_heads: int,
|
||||
qkv_bias: bool = True,
|
||||
proj_bias: bool = True,
|
||||
attention_dropout: float = 0.0,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
if len(window_size) != 2 or len(shift_size) != 2:
|
||||
raise ValueError("window_size and shift_size must be of length 2")
|
||||
self.window_size = window_size
|
||||
self.shift_size = shift_size
|
||||
self.num_heads = num_heads
|
||||
self.attention_dropout = attention_dropout
|
||||
self.dropout = dropout
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.proj = nn.Linear(dim, dim, bias=proj_bias)
|
||||
|
||||
# define a parameter table of relative position bias
|
||||
self.relative_position_bias_table = nn.Parameter(
|
||||
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
|
||||
) # 2*Wh-1 * 2*Ww-1, nH
|
||||
|
||||
# get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(self.window_size[0])
|
||||
coords_w = torch.arange(self.window_size[1])
|
||||
# coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij")) # 2, Wh, Ww
|
||||
coords = torch.stack(torch.meshgrid(coords_h, coords_w)) # 2, Wh, Ww
|
||||
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
|
||||
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
|
||||
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
|
||||
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
|
||||
relative_coords[:, :, 1] += self.window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
|
||||
relative_position_index = relative_coords.sum(-1).view(-1) # Wh*Ww*Wh*Ww
|
||||
self.register_buffer("relative_position_index", relative_position_index)
|
||||
|
||||
nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
"""
|
||||
Args:
|
||||
x (Tensor): Tensor with layout of [B, H, W, C]
|
||||
Returns:
|
||||
Tensor with same layout as input, i.e. [B, H, W, C]
|
||||
"""
|
||||
|
||||
N = self.window_size[0] * self.window_size[1]
|
||||
relative_position_bias = self.relative_position_bias_table[self.relative_position_index] # type: ignore[index]
|
||||
relative_position_bias = relative_position_bias.view(N, N, -1)
|
||||
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous().unsqueeze(0)
|
||||
|
||||
return shifted_window_attention(
|
||||
x,
|
||||
self.qkv.weight,
|
||||
self.proj.weight,
|
||||
relative_position_bias,
|
||||
self.window_size,
|
||||
self.num_heads,
|
||||
shift_size=self.shift_size,
|
||||
attention_dropout=self.attention_dropout,
|
||||
dropout=self.dropout,
|
||||
qkv_bias=self.qkv.bias,
|
||||
proj_bias=self.proj.bias,
|
||||
)
|
||||
|
||||
|
||||
class SwinTransformerBlock(nn.Module):
|
||||
"""
|
||||
Swin Transformer Block.
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads.
|
||||
window_size (List[int]): Window size.
|
||||
shift_size (List[int]): Shift size for shifted window attention.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
|
||||
dropout (float): Dropout rate. Default: 0.0.
|
||||
attention_dropout (float): Attention dropout rate. Default: 0.0.
|
||||
stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0.
|
||||
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
|
||||
attn_layer (nn.Module): Attention layer. Default: ShiftedWindowAttention
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
window_size: List[int],
|
||||
shift_size: List[int],
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.0,
|
||||
attention_dropout: float = 0.0,
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
|
||||
attn_layer: Callable[..., nn.Module] = ShiftedWindowAttention,
|
||||
):
|
||||
super().__init__()
|
||||
_log_api_usage_once(self)
|
||||
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = attn_layer(
|
||||
dim,
|
||||
window_size,
|
||||
shift_size,
|
||||
num_heads,
|
||||
attention_dropout=attention_dropout,
|
||||
dropout=dropout,
|
||||
)
|
||||
self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
|
||||
self.norm2 = norm_layer(dim)
|
||||
self.mlp = MLP(dim, [int(dim * mlp_ratio), dim], activation_layer=nn.GELU, inplace=None, dropout=dropout)
|
||||
|
||||
for m in self.mlp.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.normal_(m.bias, std=1e-6)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = x + self.stochastic_depth(self.attn(self.norm1(x)))
|
||||
x = x + self.stochastic_depth(self.mlp(self.norm2(x)))
|
||||
return x
|
||||
|
||||
|
||||
class SwinTransformer(nn.Module):
|
||||
"""
|
||||
Implements Swin Transformer from the `"Swin Transformer: Hierarchical Vision Transformer using
|
||||
Shifted Windows" <https://arxiv.org/pdf/2103.14030>`_ paper.
|
||||
Args:
|
||||
patch_size (List[int]): Patch size.
|
||||
embed_dim (int): Patch embedding dimension.
|
||||
depths (List(int)): Depth of each Swin Transformer layer.
|
||||
num_heads (List(int)): Number of attention heads in different layers.
|
||||
window_size (List[int]): Window size.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
|
||||
dropout (float): Dropout rate. Default: 0.0.
|
||||
attention_dropout (float): Attention dropout rate. Default: 0.0.
|
||||
stochastic_depth_prob (float): Stochastic depth rate. Default: 0.0.
|
||||
num_classes (int): Number of classes for classification head. Default: 1000.
|
||||
block (nn.Module, optional): SwinTransformer Block. Default: None.
|
||||
norm_layer (nn.Module, optional): Normalization layer. Default: None.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: List[int],
|
||||
embed_dim: int,
|
||||
depths: List[int],
|
||||
num_heads: List[int],
|
||||
window_size: List[int],
|
||||
mlp_ratio: float = 4.0,
|
||||
dropout: float = 0.0,
|
||||
attention_dropout: float = 0.0,
|
||||
stochastic_depth_prob: float = 0.0,
|
||||
num_classes: int = 1000,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
_log_api_usage_once(self)
|
||||
self.num_classes = num_classes
|
||||
|
||||
if block is None:
|
||||
block = SwinTransformerBlock
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(nn.LayerNorm, eps=1e-5)
|
||||
|
||||
layers: List[nn.Module] = []
|
||||
# split image into non-overlapping patches
|
||||
# layers.append(
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(
|
||||
# 3, embed_dim, kernel_size=(patch_size[0], patch_size[1]), stride=(patch_size[0], patch_size[1])
|
||||
# ),
|
||||
# Permute([0, 2, 3, 1]),
|
||||
# norm_layer(embed_dim),
|
||||
# )
|
||||
# )
|
||||
self.first_coonv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
3, embed_dim, kernel_size=(patch_size[0], patch_size[1]), stride=(patch_size[0], patch_size[1])
|
||||
),
|
||||
Permute([0, 2, 3, 1]),
|
||||
norm_layer(embed_dim),
|
||||
)
|
||||
|
||||
total_stage_blocks = sum(depths)
|
||||
stage_block_id = 0
|
||||
# build SwinTransformer blocks
|
||||
for i_stage in range(len(depths)):
|
||||
stage: List[nn.Module] = []
|
||||
dim = embed_dim * 2 ** i_stage
|
||||
for i_layer in range(depths[i_stage]):
|
||||
# adjust stochastic depth probability based on the depth of the stage block
|
||||
sd_prob = stochastic_depth_prob * float(stage_block_id) / (total_stage_blocks - 1)
|
||||
stage.append(
|
||||
block(
|
||||
dim,
|
||||
num_heads[i_stage],
|
||||
window_size=window_size,
|
||||
shift_size=[0 if i_layer % 2 == 0 else w // 2 for w in window_size],
|
||||
mlp_ratio=mlp_ratio,
|
||||
dropout=dropout,
|
||||
attention_dropout=attention_dropout,
|
||||
stochastic_depth_prob=sd_prob,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
)
|
||||
stage_block_id += 1
|
||||
layers.append(nn.Sequential(*stage))
|
||||
# add patch merging layer
|
||||
if i_stage < (len(depths) - 1):
|
||||
layers.append(PatchMerging(dim, norm_layer))
|
||||
# self.features = nn.Sequential(*layers)
|
||||
self.features = nn.ModuleList(layers)
|
||||
|
||||
num_features = embed_dim * 2 ** (len(depths) - 1)
|
||||
self.norm = norm_layer(num_features)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.head = nn.Linear(num_features, num_classes)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def forward(self, x):
|
||||
feats = []
|
||||
x = self.first_coonv(x)
|
||||
for i, layer in enumerate(self.features):
|
||||
x = layer(x)
|
||||
if i in [0, 2, 4, 6]:
|
||||
feats.append(x.permute(0, 3, 1, 2).contiguous())
|
||||
# x = self.features(x)
|
||||
# x = self.norm(x)
|
||||
# x = x.permute(0, 3, 1, 2)
|
||||
# x = self.avgpool(x)
|
||||
# x = torch.flatten(x, 1)
|
||||
# x = self.head(x)
|
||||
return feats
|
||||
|
||||
|
||||
def _swin_transformer(
|
||||
patch_size: List[int],
|
||||
embed_dim: int,
|
||||
depths: List[int],
|
||||
num_heads: List[int],
|
||||
window_size: List[int],
|
||||
stochastic_depth_prob: float,
|
||||
weights: Optional[WeightsEnum],
|
||||
progress: bool,
|
||||
**kwargs: Any,
|
||||
) -> SwinTransformer:
|
||||
if weights is not None:
|
||||
_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
|
||||
|
||||
model = SwinTransformer(
|
||||
patch_size=patch_size,
|
||||
embed_dim=embed_dim,
|
||||
depths=depths,
|
||||
num_heads=num_heads,
|
||||
window_size=window_size,
|
||||
stochastic_depth_prob=stochastic_depth_prob,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if weights is not None:
|
||||
ckpt1 = weights.get_state_dict(progress=progress)
|
||||
ckpt2 = model.state_dict()
|
||||
kl1 = list(ckpt1.keys())
|
||||
for i, k in enumerate(list(ckpt2.keys())):
|
||||
ckpt2[k] = ckpt1[kl1[i]]
|
||||
msg = model.load_state_dict(ckpt2, strict=False)
|
||||
print(f'Load swin_transformer: {msg}')
|
||||
|
||||
return model
|
||||
|
||||
|
||||
_COMMON_META = {
|
||||
"categories": _IMAGENET_CATEGORIES,
|
||||
}
|
||||
|
||||
|
||||
class Swin_T_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/swin_t-704ceda3.pth",
|
||||
transforms=partial(
|
||||
ImageClassification, crop_size=224, resize_size=232, interpolation=InterpolationMode.BICUBIC
|
||||
),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 28288354,
|
||||
"min_size": (224, 224),
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 81.474,
|
||||
"acc@5": 95.776,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class Swin_S_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/swin_s-5e29d889.pth",
|
||||
transforms=partial(
|
||||
ImageClassification, crop_size=224, resize_size=246, interpolation=InterpolationMode.BICUBIC
|
||||
),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 49606258,
|
||||
"min_size": (224, 224),
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 83.196,
|
||||
"acc@5": 96.360,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class Swin_B_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/swin_b-68c6b09e.pth",
|
||||
transforms=partial(
|
||||
ImageClassification, crop_size=224, resize_size=238, interpolation=InterpolationMode.BICUBIC
|
||||
),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 87768224,
|
||||
"min_size": (224, 224),
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#swintransformer",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 83.582,
|
||||
"acc@5": 96.640,
|
||||
}
|
||||
},
|
||||
"_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
def swin_t(*, weights: Optional[Swin_T_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer:
|
||||
"""
|
||||
Constructs a swin_tiny architecture from
|
||||
`Swin Transformer: Hierarchical Vision Transformer using Shifted Windows <https://arxiv.org/pdf/2103.14030>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.Swin_T_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.Swin_T_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.Swin_T_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = Swin_T_Weights.verify(weights)
|
||||
|
||||
return _swin_transformer(
|
||||
patch_size=[4, 4],
|
||||
embed_dim=96,
|
||||
depths=[2, 2, 6, 2],
|
||||
num_heads=[3, 6, 12, 24],
|
||||
window_size=[7, 7],
|
||||
stochastic_depth_prob=0.2,
|
||||
weights=weights,
|
||||
progress=progress,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def swin_s(*, weights: Optional[Swin_S_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer:
|
||||
"""
|
||||
Constructs a swin_small architecture from
|
||||
`Swin Transformer: Hierarchical Vision Transformer using Shifted Windows <https://arxiv.org/pdf/2103.14030>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.Swin_S_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.Swin_S_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.Swin_S_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = Swin_S_Weights.verify(weights)
|
||||
|
||||
return _swin_transformer(
|
||||
patch_size=[4, 4],
|
||||
embed_dim=96,
|
||||
depths=[2, 2, 18, 2],
|
||||
num_heads=[3, 6, 12, 24],
|
||||
window_size=[7, 7],
|
||||
stochastic_depth_prob=0.3,
|
||||
weights=weights,
|
||||
progress=progress,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
from torchvision.models._utils import handle_legacy_interface
|
||||
@handle_legacy_interface(weights=("pretrained", Swin_B_Weights.IMAGENET1K_V1))
|
||||
def swin_b(*, weights: Optional[Swin_B_Weights] = None, progress: bool = True, **kwargs: Any) -> SwinTransformer:
|
||||
"""
|
||||
Constructs a swin_base architecture from
|
||||
`Swin Transformer: Hierarchical Vision Transformer using Shifted Windows <https://arxiv.org/pdf/2103.14030>`_.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.Swin_B_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.Swin_B_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.swin_transformer.SwinTransformer``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.Swin_B_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = Swin_B_Weights.verify(weights)
|
||||
|
||||
return _swin_transformer(
|
||||
patch_size=[4, 4],
|
||||
embed_dim=128,
|
||||
depths=[2, 2, 18, 2],
|
||||
num_heads=[4, 8, 16, 32],
|
||||
window_size=[7, 7],
|
||||
stochastic_depth_prob=0.5,
|
||||
weights=weights,
|
||||
progress=progress,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = swin_b(weights=Swin_B_Weights)
|
||||
x = torch.rand(1, 3, 320, 320)
|
||||
y = model(x)
|
||||
pause = 0
|
||||
@@ -0,0 +1,376 @@
|
||||
import math
|
||||
import torch
|
||||
from torch import nn, einsum
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange, reduce
|
||||
from functools import partial
|
||||
|
||||
|
||||
def exists(x):
|
||||
return x is not None
|
||||
|
||||
def default(val, d):
|
||||
if exists(val):
|
||||
return val
|
||||
return d() if callable(d) else d
|
||||
|
||||
def identity(t, *args, **kwargs):
|
||||
return t
|
||||
|
||||
def cycle(dl):
|
||||
while True:
|
||||
for data in dl:
|
||||
yield data
|
||||
|
||||
def has_int_squareroot(num):
|
||||
return (math.sqrt(num) ** 2) == num
|
||||
|
||||
def num_to_groups(num, divisor):
|
||||
groups = num // divisor
|
||||
remainder = num % divisor
|
||||
arr = [divisor] * groups
|
||||
if remainder > 0:
|
||||
arr.append(remainder)
|
||||
return arr
|
||||
|
||||
def convert_image_to_fn(img_type, image):
|
||||
if image.mode != img_type:
|
||||
return image.convert(img_type)
|
||||
return image
|
||||
|
||||
# normalization functions
|
||||
|
||||
def normalize_to_neg_one_to_one(img):
|
||||
return img * 2 - 1
|
||||
|
||||
def unnormalize_to_zero_to_one(t):
|
||||
return (t + 1) * 0.5
|
||||
|
||||
# small helper modules
|
||||
|
||||
class Residual(nn.Module):
|
||||
def __init__(self, fn):
|
||||
super().__init__()
|
||||
self.fn = fn
|
||||
|
||||
def forward(self, x, *args, **kwargs):
|
||||
return self.fn(x, *args, **kwargs) + x
|
||||
|
||||
def Upsample(dim, dim_out = None):
|
||||
return nn.Sequential(
|
||||
nn.Upsample(scale_factor = 2, mode = 'nearest'),
|
||||
nn.Conv2d(dim, default(dim_out, dim), 3, padding = 1)
|
||||
)
|
||||
|
||||
def Downsample(dim, dim_out = None):
|
||||
return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
|
||||
|
||||
class WeightStandardizedConv2d(nn.Conv2d):
|
||||
"""
|
||||
https://arxiv.org/abs/1903.10520
|
||||
weight standardization purportedly works synergistically with group normalization
|
||||
"""
|
||||
def forward(self, x):
|
||||
eps = 1e-5 if x.dtype == torch.float32 else 1e-3
|
||||
|
||||
weight = self.weight
|
||||
mean = reduce(weight, 'o ... -> o 1 1 1', 'mean')
|
||||
var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased = False))
|
||||
normalized_weight = (weight - mean) * (var + eps).rsqrt()
|
||||
|
||||
return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
|
||||
|
||||
def forward(self, x):
|
||||
eps = 1e-5 if x.dtype == torch.float32 else 1e-3
|
||||
var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
|
||||
mean = torch.mean(x, dim = 1, keepdim = True)
|
||||
return (x - mean) * (var + eps).rsqrt() * self.g
|
||||
|
||||
class PreNorm(nn.Module):
|
||||
def __init__(self, dim, fn):
|
||||
super().__init__()
|
||||
self.fn = fn
|
||||
self.norm = LayerNorm(dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm(x)
|
||||
return self.fn(x)
|
||||
|
||||
# sinusoidal positional embeds
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
device = x.device
|
||||
half_dim = self.dim // 2
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
|
||||
emb = x[:, None] * emb[None, :]
|
||||
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
||||
return emb
|
||||
|
||||
class RandomOrLearnedSinusoidalPosEmb(nn.Module):
|
||||
""" following @crowsonkb 's lead with random (learned optional) sinusoidal pos emb """
|
||||
""" https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """
|
||||
|
||||
def __init__(self, dim, is_random = False):
|
||||
super().__init__()
|
||||
assert (dim % 2) == 0
|
||||
half_dim = dim // 2
|
||||
self.weights = nn.Parameter(torch.randn(half_dim), requires_grad = not is_random)
|
||||
|
||||
def forward(self, x):
|
||||
x = rearrange(x, 'b -> b 1')
|
||||
freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi
|
||||
fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1)
|
||||
fouriered = torch.cat((x, fouriered), dim = -1)
|
||||
return fouriered
|
||||
|
||||
# building block modules
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, dim, dim_out, groups = 8):
|
||||
super().__init__()
|
||||
self.proj = WeightStandardizedConv2d(dim, dim_out, 3, padding = 1)
|
||||
self.norm = nn.GroupNorm(groups, dim_out)
|
||||
self.act = nn.SiLU()
|
||||
|
||||
def forward(self, x, scale_shift = None):
|
||||
x = self.proj(x)
|
||||
x = self.norm(x)
|
||||
|
||||
if exists(scale_shift):
|
||||
scale, shift = scale_shift
|
||||
x = x * (scale + 1) + shift
|
||||
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(time_emb_dim, dim_out * 2)
|
||||
) if exists(time_emb_dim) else None
|
||||
|
||||
self.block1 = Block(dim, dim_out, groups = groups)
|
||||
self.block2 = Block(dim_out, dim_out, groups = groups)
|
||||
self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
|
||||
|
||||
def forward(self, x, time_emb = None):
|
||||
|
||||
scale_shift = None
|
||||
if exists(self.mlp) and exists(time_emb):
|
||||
time_emb = self.mlp(time_emb)
|
||||
time_emb = rearrange(time_emb, 'b c -> b c 1 1')
|
||||
scale_shift = time_emb.chunk(2, dim = 1)
|
||||
|
||||
h = self.block1(x, scale_shift = scale_shift)
|
||||
|
||||
h = self.block2(h)
|
||||
|
||||
return h + self.res_conv(x)
|
||||
|
||||
class LinearAttention(nn.Module):
|
||||
def __init__(self, dim, heads = 4, dim_head = 32):
|
||||
super().__init__()
|
||||
self.scale = dim_head ** -0.5
|
||||
self.heads = heads
|
||||
hidden_dim = dim_head * heads
|
||||
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
|
||||
|
||||
self.to_out = nn.Sequential(
|
||||
nn.Conv2d(hidden_dim, dim, 1),
|
||||
LayerNorm(dim)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, h, w = x.shape
|
||||
qkv = self.to_qkv(x).chunk(3, dim = 1)
|
||||
q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
|
||||
|
||||
q = q.softmax(dim = -2)
|
||||
k = k.softmax(dim = -1)
|
||||
|
||||
q = q * self.scale
|
||||
v = v / (h * w)
|
||||
|
||||
context = torch.einsum('b h d n, b h e n -> b h d e', k, v)
|
||||
|
||||
out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
|
||||
out = rearrange(out, 'b h c (x y) -> b (h c) x y', h = self.heads, x = h, y = w)
|
||||
return self.to_out(out)
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, dim, heads = 4, dim_head = 32):
|
||||
super().__init__()
|
||||
self.scale = dim_head ** -0.5
|
||||
self.heads = heads
|
||||
hidden_dim = dim_head * heads
|
||||
|
||||
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
|
||||
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
|
||||
|
||||
def forward(self, x):
|
||||
b, c, h, w = x.shape
|
||||
qkv = self.to_qkv(x).chunk(3, dim = 1)
|
||||
q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv)
|
||||
|
||||
q = q * self.scale
|
||||
|
||||
sim = einsum('b h d i, b h d j -> b h i j', q, k)
|
||||
attn = sim.softmax(dim = -1)
|
||||
out = einsum('b h i j, b h d j -> b h i d', attn, v)
|
||||
|
||||
out = rearrange(out, 'b h (x y) d -> b (h d) x y', x = h, y = w)
|
||||
return self.to_out(out)
|
||||
|
||||
# model
|
||||
|
||||
class Unet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
init_dim = None,
|
||||
out_dim = None,
|
||||
dim_mults=(1, 2, 4, 8),
|
||||
channels = 3,
|
||||
self_condition = False,
|
||||
resnet_block_groups = 8,
|
||||
heads=8,
|
||||
learned_variance = False,
|
||||
learned_sinusoidal_cond = False,
|
||||
random_fourier_features = False,
|
||||
learned_sinusoidal_dim = 16,
|
||||
out_mul=1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# determine dimensions
|
||||
|
||||
self.channels = channels
|
||||
self.self_condition = self_condition
|
||||
input_channels = channels * (2 if self_condition else 1)
|
||||
|
||||
init_dim = default(init_dim, dim)
|
||||
self.init_conv = nn.Conv2d(input_channels, init_dim, 7, padding = 3)
|
||||
|
||||
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
|
||||
in_out = list(zip(dims[:-1], dims[1:]))
|
||||
|
||||
block_klass = partial(ResnetBlock, groups = resnet_block_groups)
|
||||
|
||||
# time embeddings
|
||||
|
||||
time_dim = dim * 4
|
||||
|
||||
self.random_or_learned_sinusoidal_cond = learned_sinusoidal_cond or random_fourier_features
|
||||
|
||||
if self.random_or_learned_sinusoidal_cond:
|
||||
sinu_pos_emb = RandomOrLearnedSinusoidalPosEmb(learned_sinusoidal_dim, random_fourier_features)
|
||||
fourier_dim = learned_sinusoidal_dim + 1
|
||||
else:
|
||||
sinu_pos_emb = SinusoidalPosEmb(dim)
|
||||
fourier_dim = dim
|
||||
|
||||
self.time_mlp = nn.Sequential(
|
||||
sinu_pos_emb,
|
||||
nn.Linear(fourier_dim, time_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(time_dim, time_dim)
|
||||
)
|
||||
|
||||
# layers
|
||||
|
||||
self.downs = nn.ModuleList([])
|
||||
self.ups = nn.ModuleList([])
|
||||
num_resolutions = len(in_out)
|
||||
|
||||
for ind, (dim_in, dim_out) in enumerate(in_out):
|
||||
is_last = ind >= (num_resolutions - 1)
|
||||
|
||||
self.downs.append(nn.ModuleList([
|
||||
block_klass(dim_in, dim_in, time_emb_dim = time_dim),
|
||||
block_klass(dim_in, dim_in, time_emb_dim = time_dim),
|
||||
Residual(PreNorm(dim_in, LinearAttention(dim_in, heads=heads))),
|
||||
Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding = 1)
|
||||
]))
|
||||
|
||||
mid_dim = dims[-1]
|
||||
self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
|
||||
self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim, heads=heads)))
|
||||
self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
|
||||
|
||||
for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
|
||||
is_last = ind == (len(in_out) - 1)
|
||||
|
||||
self.ups.append(nn.ModuleList([
|
||||
block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
|
||||
block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
|
||||
Residual(PreNorm(dim_out, LinearAttention(dim_out))),
|
||||
Upsample(dim_out, dim_in) if not is_last else nn.Conv2d(dim_out, dim_in, 3, padding = 1)
|
||||
]))
|
||||
|
||||
default_out_dim = channels * out_mul
|
||||
self.out_dim = default(out_dim, default_out_dim)
|
||||
|
||||
self.final_res_block = block_klass(dim * 2, dim, time_emb_dim = time_dim)
|
||||
self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
|
||||
|
||||
def forward(self, x, time, cond=None, x_self_cond=None): ## cond is always None for unconditional model
|
||||
if self.self_condition:
|
||||
x_self_cond = default(x_self_cond, lambda: torch.zeros_like(x))
|
||||
x = torch.cat((x_self_cond, x), dim = 1)
|
||||
|
||||
x = self.init_conv(x)
|
||||
r = x.clone()
|
||||
|
||||
t = self.time_mlp(time)
|
||||
|
||||
h = []
|
||||
|
||||
for block1, block2, attn, downsample in self.downs:
|
||||
x = block1(x, t)
|
||||
h.append(x)
|
||||
|
||||
x = block2(x, t)
|
||||
x = attn(x)
|
||||
h.append(x)
|
||||
|
||||
x = downsample(x)
|
||||
|
||||
x = self.mid_block1(x, t)
|
||||
x = self.mid_attn(x)
|
||||
x = self.mid_block2(x, t)
|
||||
|
||||
for block1, block2, attn, upsample in self.ups:
|
||||
x = torch.cat((x, h.pop()), dim = 1)
|
||||
x = block1(x, t)
|
||||
|
||||
x = torch.cat((x, h.pop()), dim = 1)
|
||||
x = block2(x, t)
|
||||
x = attn(x)
|
||||
|
||||
x = upsample(x)
|
||||
|
||||
x = torch.cat((x, r), dim = 1)
|
||||
|
||||
x = self.final_res_block(x, t)
|
||||
return self.final_conv(x)
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = Unet(96, out_mul=2, dim_mults=[1,2,4,8], heads=8)
|
||||
x = torch.rand(2, 3, 8, 8)
|
||||
time = torch.tensor([2, 5])
|
||||
with torch.no_grad():
|
||||
y = model(x, time)
|
||||
pass
|
||||
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
import time
|
||||
import logging
|
||||
import math
|
||||
|
||||
def create_logger(root_dir, des=''):
|
||||
root_output_dir = Path(root_dir)
|
||||
# set up logger
|
||||
if not root_output_dir.exists():
|
||||
print('=> creating {}'.format(root_output_dir))
|
||||
root_output_dir.mkdir(exist_ok=True, parents=True)
|
||||
time_str = time.strftime('%Y-%m-%d-%H-%M')
|
||||
log_file = '{}_{}.log'.format(time_str, des)
|
||||
final_log_file = root_output_dir / log_file
|
||||
head = '%(asctime)-15s %(message)s'
|
||||
logging.basicConfig(filename=str(final_log_file), format=head)
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.INFO)
|
||||
console = logging.StreamHandler()
|
||||
logging.getLogger('').addHandler(console)
|
||||
return logger
|
||||
|
||||
def exists(x):
|
||||
return x is not None
|
||||
|
||||
def default(val, d):
|
||||
if exists(val):
|
||||
return val
|
||||
return d() if callable(d) else d
|
||||
|
||||
def identity(t, *args, **kwargs):
|
||||
return t
|
||||
|
||||
def cycle(dl):
|
||||
while True:
|
||||
for data in dl:
|
||||
yield data
|
||||
|
||||
def has_int_squareroot(num):
|
||||
return (math.sqrt(num) ** 2) == num
|
||||
|
||||
def num_to_groups(num, divisor):
|
||||
groups = num // divisor
|
||||
remainder = num % divisor
|
||||
arr = [divisor] * groups
|
||||
if remainder > 0:
|
||||
arr.append(remainder)
|
||||
return arr
|
||||
|
||||
def convert_image_to_fn(img_type, image):
|
||||
if image.mode != img_type:
|
||||
return image.convert(img_type)
|
||||
return image
|
||||
|
||||
# normalization functions
|
||||
|
||||
def normalize_to_neg_one_to_one(img):
|
||||
return img * 2 - 1
|
||||
|
||||
def unnormalize_to_zero_to_one(t):
|
||||
return (t + 1) * 0.5
|
||||
|
||||
def dict2str(dict):
|
||||
s = ''
|
||||
for k, v in dict.items():
|
||||
s += "{}: {:.5f}, ".format(k, v)
|
||||
return s
|
||||
@@ -0,0 +1,517 @@
|
||||
from functools import partial
|
||||
from typing import Union, List, Dict, Any, Optional, cast
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from torchvision.transforms._presets import ImageClassification
|
||||
from torchvision.utils import _log_api_usage_once
|
||||
from torchvision.models._api import WeightsEnum, Weights
|
||||
from torchvision.models._meta import _IMAGENET_CATEGORIES
|
||||
from torchvision.models._utils import handle_legacy_interface, _ovewrite_named_param
|
||||
|
||||
|
||||
__all__ = [
|
||||
"VGG",
|
||||
"VGG11_Weights",
|
||||
"VGG11_BN_Weights",
|
||||
"VGG13_Weights",
|
||||
"VGG13_BN_Weights",
|
||||
"VGG16_Weights",
|
||||
"VGG16_BN_Weights",
|
||||
"VGG19_Weights",
|
||||
"VGG19_BN_Weights",
|
||||
"vgg11",
|
||||
"vgg11_bn",
|
||||
"vgg13",
|
||||
"vgg13_bn",
|
||||
"vgg16",
|
||||
"vgg16_bn",
|
||||
"vgg19",
|
||||
"vgg19_bn",
|
||||
]
|
||||
|
||||
|
||||
class VGG(nn.Module):
|
||||
def __init__(
|
||||
self, features: nn.Module, num_classes: int = 1000, init_weights: bool = True, dropout: float = 0.5
|
||||
) -> None:
|
||||
super().__init__()
|
||||
_log_api_usage_once(self)
|
||||
self.features = features
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(512 * 7 * 7, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(p=dropout),
|
||||
nn.Linear(4096, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(p=dropout),
|
||||
nn.Linear(4096, num_classes),
|
||||
)
|
||||
if init_weights:
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
feats = []
|
||||
# x = self.features(x)
|
||||
# x = self.avgpool(x)
|
||||
# x = torch.flatten(x, 1)
|
||||
# x = self.classifier(x)
|
||||
for i, layer in enumerate(self.features):
|
||||
x = layer(x)
|
||||
if i in [9, 16, 23, 30]:
|
||||
feats.append(x)
|
||||
return feats
|
||||
|
||||
|
||||
def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequential:
|
||||
layers: List[nn.Module] = []
|
||||
in_channels = 3
|
||||
for v in cfg:
|
||||
if v == "M":
|
||||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||||
else:
|
||||
v = cast(int, v)
|
||||
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = v
|
||||
return nn.ModuleList(layers)
|
||||
|
||||
|
||||
cfgs: Dict[str, List[Union[str, int]]] = {
|
||||
"A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
|
||||
"B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
|
||||
"D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"],
|
||||
"E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"],
|
||||
}
|
||||
|
||||
|
||||
def _vgg(cfg: str, batch_norm: bool, weights: Optional[WeightsEnum], progress: bool, **kwargs: Any) -> VGG:
|
||||
if weights is not None:
|
||||
kwargs["init_weights"] = False
|
||||
if weights.meta["categories"] is not None:
|
||||
_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
|
||||
model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
|
||||
|
||||
if weights is not None:
|
||||
ckpt1 = weights.get_state_dict(progress=progress)
|
||||
ckpt2 = model.state_dict()
|
||||
kl1 = list(ckpt1.keys())
|
||||
for i, k in enumerate(list(ckpt2.keys())):
|
||||
ckpt2[k] = ckpt1[kl1[i]]
|
||||
msg = model.load_state_dict(ckpt2, strict=False)
|
||||
print(f'Load VGG: {msg}')
|
||||
else:
|
||||
print('No pretrained weight loaded!')
|
||||
return model
|
||||
|
||||
|
||||
_COMMON_META = {
|
||||
"min_size": (32, 32),
|
||||
"categories": _IMAGENET_CATEGORIES,
|
||||
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg",
|
||||
"_docs": """These weights were trained from scratch by using a simplified training recipe.""",
|
||||
}
|
||||
|
||||
|
||||
class VGG11_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg11-8a719046.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 132863336,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 69.020,
|
||||
"acc@5": 88.628,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG11_BN_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg11_bn-6002323d.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 132868840,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 70.370,
|
||||
"acc@5": 89.810,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG13_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg13-19584684.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 133047848,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 69.928,
|
||||
"acc@5": 89.246,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG13_BN_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg13_bn-abd245e5.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 133053736,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 71.586,
|
||||
"acc@5": 90.374,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG16_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg16-397923af.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 138357544,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 71.592,
|
||||
"acc@5": 90.382,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
IMAGENET1K_FEATURES = Weights(
|
||||
# Weights ported from https://github.com/amdegroot/ssd.pytorch/
|
||||
url="https://download.pytorch.org/models/vgg16_features-amdegroot-88682ab5.pth",
|
||||
transforms=partial(
|
||||
ImageClassification,
|
||||
crop_size=224,
|
||||
mean=(0.48235, 0.45882, 0.40784),
|
||||
std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0),
|
||||
),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 138357544,
|
||||
"categories": None,
|
||||
"recipe": "https://github.com/amdegroot/ssd.pytorch#training-ssd",
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": float("nan"),
|
||||
"acc@5": float("nan"),
|
||||
}
|
||||
},
|
||||
"_docs": """
|
||||
These weights can't be used for classification because they are missing values in the `classifier`
|
||||
module. Only the `features` module has valid values and can be used for feature extraction. The weights
|
||||
were trained using the original input standardization method as described in the paper.
|
||||
""",
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG16_BN_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg16_bn-6c64b313.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 138365992,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 73.360,
|
||||
"acc@5": 91.516,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG19_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg19-dcbb9e9d.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 143667240,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 72.376,
|
||||
"acc@5": 90.876,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
class VGG19_BN_Weights(WeightsEnum):
|
||||
IMAGENET1K_V1 = Weights(
|
||||
url="https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",
|
||||
transforms=partial(ImageClassification, crop_size=224),
|
||||
meta={
|
||||
**_COMMON_META,
|
||||
"num_params": 143678248,
|
||||
"_metrics": {
|
||||
"ImageNet-1K": {
|
||||
"acc@1": 74.218,
|
||||
"acc@5": 91.842,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
DEFAULT = IMAGENET1K_V1
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG11_Weights.IMAGENET1K_V1))
|
||||
def vgg11(*, weights: Optional[VGG11_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-11 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG11_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG11_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG11_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG11_Weights.verify(weights)
|
||||
|
||||
return _vgg("A", False, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG11_BN_Weights.IMAGENET1K_V1))
|
||||
def vgg11_bn(*, weights: Optional[VGG11_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-11-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG11_BN_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG11_BN_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG11_BN_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG11_BN_Weights.verify(weights)
|
||||
|
||||
return _vgg("A", True, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG13_Weights.IMAGENET1K_V1))
|
||||
def vgg13(*, weights: Optional[VGG13_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-13 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG13_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG13_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG13_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG13_Weights.verify(weights)
|
||||
|
||||
return _vgg("B", False, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG13_BN_Weights.IMAGENET1K_V1))
|
||||
def vgg13_bn(*, weights: Optional[VGG13_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-13-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG13_BN_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG13_BN_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG13_BN_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG13_BN_Weights.verify(weights)
|
||||
|
||||
return _vgg("B", True, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG16_Weights.IMAGENET1K_V1))
|
||||
def vgg16(*, weights: Optional[VGG16_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-16 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG16_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG16_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG16_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG16_Weights.verify(weights)
|
||||
|
||||
return _vgg("D", False, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG16_BN_Weights.IMAGENET1K_V1))
|
||||
def vgg16_bn(*, weights: Optional[VGG16_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-16-BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG16_BN_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG16_BN_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG16_BN_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG16_BN_Weights.verify(weights)
|
||||
|
||||
return _vgg("D", True, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG19_Weights.IMAGENET1K_V1))
|
||||
def vgg19(*, weights: Optional[VGG19_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-19 from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG19_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG19_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG19_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG19_Weights.verify(weights)
|
||||
|
||||
return _vgg("E", False, weights, progress, **kwargs)
|
||||
|
||||
|
||||
@handle_legacy_interface(weights=("pretrained", VGG19_BN_Weights.IMAGENET1K_V1))
|
||||
def vgg19_bn(*, weights: Optional[VGG19_BN_Weights] = None, progress: bool = True, **kwargs: Any) -> VGG:
|
||||
"""VGG-19_BN from `Very Deep Convolutional Networks for Large-Scale Image Recognition <https://arxiv.org/abs/1409.1556>`__.
|
||||
|
||||
Args:
|
||||
weights (:class:`~torchvision.models.VGG19_BN_Weights`, optional): The
|
||||
pretrained weights to use. See
|
||||
:class:`~torchvision.models.VGG19_BN_Weights` below for
|
||||
more details, and possible values. By default, no pre-trained
|
||||
weights are used.
|
||||
progress (bool, optional): If True, displays a progress bar of the
|
||||
download to stderr. Default is True.
|
||||
**kwargs: parameters passed to the ``torchvision.models.vgg.VGG``
|
||||
base class. Please refer to the `source code
|
||||
<https://github.com/pytorch/vision/blob/main/torchvision/models/vgg.py>`_
|
||||
for more details about this class.
|
||||
|
||||
.. autoclass:: torchvision.models.VGG19_BN_Weights
|
||||
:members:
|
||||
"""
|
||||
weights = VGG19_BN_Weights.verify(weights)
|
||||
|
||||
return _vgg("E", True, weights, progress, **kwargs)
|
||||
|
||||
|
||||
# The dictionary below is internal implementation detail and will be removed in v0.15
|
||||
from torchvision.models._utils import _ModelURLs
|
||||
|
||||
|
||||
model_urls = _ModelURLs(
|
||||
{
|
||||
"vgg11": VGG11_Weights.IMAGENET1K_V1.url,
|
||||
"vgg13": VGG13_Weights.IMAGENET1K_V1.url,
|
||||
"vgg16": VGG16_Weights.IMAGENET1K_V1.url,
|
||||
"vgg19": VGG19_Weights.IMAGENET1K_V1.url,
|
||||
"vgg11_bn": VGG11_BN_Weights.IMAGENET1K_V1.url,
|
||||
"vgg13_bn": VGG13_BN_Weights.IMAGENET1K_V1.url,
|
||||
"vgg16_bn": VGG16_BN_Weights.IMAGENET1K_V1.url,
|
||||
"vgg19_bn": VGG19_BN_Weights.IMAGENET1K_V1.url,
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,83 @@
|
||||
import pywt
|
||||
import pywt.data
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.autograd import Function
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def create_wavelet_filter(wave, in_size, out_size, type=torch.float):
|
||||
w = pywt.Wavelet(wave)
|
||||
dec_hi = torch.tensor(w.dec_hi[::-1], dtype=type)
|
||||
dec_lo = torch.tensor(w.dec_lo[::-1], dtype=type)
|
||||
dec_filters = torch.stack([dec_lo.unsqueeze(0) * dec_lo.unsqueeze(1),
|
||||
dec_lo.unsqueeze(0) * dec_hi.unsqueeze(1),
|
||||
dec_hi.unsqueeze(0) * dec_lo.unsqueeze(1),
|
||||
dec_hi.unsqueeze(0) * dec_hi.unsqueeze(1)], dim=0)
|
||||
|
||||
dec_filters = dec_filters[:, None].repeat(in_size, 1, 1, 1)
|
||||
|
||||
rec_hi = torch.tensor(w.rec_hi[::-1], dtype=type).flip(dims=[0])
|
||||
rec_lo = torch.tensor(w.rec_lo[::-1], dtype=type).flip(dims=[0])
|
||||
rec_filters = torch.stack([rec_lo.unsqueeze(0) * rec_lo.unsqueeze(1),
|
||||
rec_lo.unsqueeze(0) * rec_hi.unsqueeze(1),
|
||||
rec_hi.unsqueeze(0) * rec_lo.unsqueeze(1),
|
||||
rec_hi.unsqueeze(0) * rec_hi.unsqueeze(1)], dim=0)
|
||||
|
||||
rec_filters = rec_filters[:, None].repeat(out_size, 1, 1, 1)
|
||||
|
||||
return dec_filters, rec_filters
|
||||
|
||||
|
||||
def wt(x, filters, in_size, level):
|
||||
_, _, h, w = x.shape
|
||||
pad = (filters.shape[2] // 2 - 1, filters.shape[3] // 2 - 1)
|
||||
res = F.conv2d(x, filters, stride=2, groups=in_size, padding=pad)
|
||||
if level > 1:
|
||||
res[:, ::4] = wt(res[:, ::4], filters, in_size, level - 1)
|
||||
res = res.reshape(-1, 2, h // 2, w // 2).transpose(1, 2).reshape(-1, in_size, h, w)
|
||||
return res
|
||||
|
||||
|
||||
def iwt(x, inv_filters, in_size, level):
|
||||
_, _, h, w = x.shape
|
||||
pad = (inv_filters.shape[2] // 2 - 1, inv_filters.shape[3] // 2 - 1)
|
||||
res = x.reshape(-1, h // 2, 2, w // 2).transpose(1, 2).reshape(-1, 4 * in_size, h // 2, w // 2)
|
||||
if level > 1:
|
||||
res[:, ::4] = iwt(res[:, ::4], inv_filters, in_size, level - 1)
|
||||
res = F.conv_transpose2d(res, inv_filters, stride=2, groups=in_size, padding=pad)
|
||||
return res
|
||||
|
||||
|
||||
def get_inverse_transform(weights, in_size, level):
|
||||
class InverseWaveletTransform(Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
with torch.no_grad():
|
||||
x = iwt(input, weights, in_size, level)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
grad = wt(grad_output, weights, in_size, level)
|
||||
return grad, None
|
||||
|
||||
return InverseWaveletTransform().apply
|
||||
|
||||
|
||||
def get_transform(weights, in_size, level):
|
||||
class WaveletTransform(Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
with torch.no_grad():
|
||||
x = wt(input, weights, in_size, level)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
grad = iwt(grad_output, weights, in_size, level)
|
||||
return grad, None
|
||||
|
||||
return WaveletTransform().apply
|
||||
@@ -0,0 +1,101 @@
|
||||
from typing import Union, Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from custom_controlnet_aux.diffusion_edge.denoising_diffusion_pytorch.quantization import weight_quantize_fn, act_quantize_fn
|
||||
from custom_controlnet_aux.diffusion_edge.denoising_diffusion_pytorch import wavelet
|
||||
|
||||
|
||||
class WCC(nn.Conv1d):
|
||||
def __init__(self, in_channels: int,
|
||||
out_channels: int,
|
||||
stride: Union[int, Tuple] = 1,
|
||||
padding: Union[int, Tuple] = 0,
|
||||
dilation: Union[int, Tuple] = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = False,
|
||||
levels: int = 3,
|
||||
compress_rate: float = 0.25,
|
||||
bit_w: int = 8,
|
||||
bit_a: int = 8,
|
||||
wt_type: str = "db1"):
|
||||
super(WCC, self).__init__(in_channels, out_channels, 1, stride, padding, dilation, groups, bias)
|
||||
self.layer_type = 'WCC'
|
||||
self.bit_w = bit_w
|
||||
self.bit_a = bit_a
|
||||
|
||||
self.weight_quant = weight_quantize_fn(self.bit_w)
|
||||
self.act_quant = act_quantize_fn(self.bit_a, signed=True)
|
||||
|
||||
self.levels = levels
|
||||
self.wt_type = wt_type
|
||||
self.compress_rate = compress_rate
|
||||
|
||||
dec_filters, rec_filters = wavelet.create_wavelet_filter(wave=self.wt_type,
|
||||
in_size=in_channels,
|
||||
out_size=out_channels)
|
||||
self.wt_filters = nn.Parameter(dec_filters, requires_grad=False)
|
||||
self.iwt_filters = nn.Parameter(rec_filters, requires_grad=False)
|
||||
self.wt = wavelet.get_transform(self.wt_filters, in_channels, levels)
|
||||
self.iwt = wavelet.get_inverse_transform(self.iwt_filters, out_channels, levels)
|
||||
|
||||
self.get_pad = lambda n: ((2 ** levels) - n) % (2 ** levels)
|
||||
|
||||
def forward(self, x):
|
||||
in_shape = x.shape
|
||||
pads = (0, self.get_pad(in_shape[2]), 0, self.get_pad(in_shape[3]))
|
||||
x = F.pad(x, pads) # pad to match 2^(levels)
|
||||
|
||||
weight_q = self.weight_quant(self.weight) # quantize weights
|
||||
x = self.wt(x) # H
|
||||
topk, ids = self.compress(x) # T
|
||||
topk_q = self.act_quant(topk) # quantize activations
|
||||
topk_q = F.conv1d(topk_q, weight_q, self.bias, self.stride, self.padding, self.dilation, self.groups) # K_1x1
|
||||
x = self.decompress(topk_q, ids, x.shape) # T^T
|
||||
x = self.iwt(x) # H^T
|
||||
|
||||
x = x[:, :, :in_shape[2], :in_shape[3]] # remove pads
|
||||
return x
|
||||
|
||||
def compress(self, x):
|
||||
b, c, h, w = x.shape
|
||||
acc = x.norm(dim=1).pow(2)
|
||||
acc = acc.view(b, h * w)
|
||||
k = int(h * w * self.compress_rate)
|
||||
ids = acc.topk(k, dim=1, sorted=False)[1]
|
||||
ids.unsqueeze_(dim=1)
|
||||
topk = x.reshape((b, c, h * w)).gather(dim=2, index=ids.repeat(1, c, 1))
|
||||
return topk, ids
|
||||
|
||||
def decompress(self, topk, ids, shape):
|
||||
b, _, h, w = shape
|
||||
ids = ids.repeat(1, self.out_channels, 1)
|
||||
x = torch.zeros(size=(b, self.out_channels, h * w), requires_grad=True, device=topk.device)
|
||||
x = x.scatter(dim=2, index=ids, src=topk)
|
||||
x = x.reshape((b, self.out_channels, h, w))
|
||||
return x
|
||||
|
||||
def change_wt_params(self, compress_rate, levels, wt_type="db1"):
|
||||
self.compress_rate = compress_rate
|
||||
self.levels = levels
|
||||
dec_filters, rec_filters = wavelet.create_wavelet_filter(wave=self.wt_type,
|
||||
in_size=self.in_channels,
|
||||
out_size=self.out_channels)
|
||||
self.wt_filters = nn.Parameter(dec_filters, requires_grad=False)
|
||||
self.iwt_filters = nn.Parameter(rec_filters, requires_grad=False)
|
||||
self.wt = wavelet.get_transform(self.wt_filters, self.in_channels, levels)
|
||||
self.iwt = wavelet.get_inverse_transform(self.iwt_filters, self.out_channels, levels)
|
||||
|
||||
def change_bit(self, bit_w, bit_a):
|
||||
self.bit_w = bit_w
|
||||
self.bit_a = bit_a
|
||||
self.weight_quant.change_bit(bit_w)
|
||||
self.act_quant.change_bit(bit_a)
|
||||
|
||||
if __name__ == '__main__':
|
||||
wcc = WCC(80, 80)
|
||||
x = torch.rand(1, 80, 80, 80)
|
||||
y = wcc(x)
|
||||
pause = 0
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user