Showing preview only (2,398K chars total). Download the full file or copy to clipboard to get everything.
Repository: Sanster/IOPaint
Branch: main
Commit: 61a759fb3f33
Files: 302
Total size: 2.2 MB
Directory structure:
gitextract_t_2nmh0z/
├── .github/
│ ├── FUNDING.yml
│ └── ISSUE_TEMPLATE/
│ ├── config.yml
│ ├── 🐛-bug-report.md
│ └── 🚀-feature-request.md
├── .gitignore
├── LICENSE
├── README.md
├── build_docker.sh
├── docker/
│ ├── CPUDockerfile
│ └── GPUDockerfile
├── iopaint/
│ ├── __init__.py
│ ├── __main__.py
│ ├── api.py
│ ├── batch_processing.py
│ ├── benchmark.py
│ ├── cli.py
│ ├── const.py
│ ├── download.py
│ ├── file_manager/
│ │ ├── __init__.py
│ │ ├── file_manager.py
│ │ ├── storage_backends.py
│ │ └── utils.py
│ ├── helper.py
│ ├── installer.py
│ ├── model/
│ │ ├── __init__.py
│ │ ├── anytext/
│ │ │ ├── __init__.py
│ │ │ ├── anytext_model.py
│ │ │ ├── anytext_pipeline.py
│ │ │ ├── anytext_sd15.yaml
│ │ │ ├── cldm/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── cldm.py
│ │ │ │ ├── ddim_hacked.py
│ │ │ │ ├── embedding_manager.py
│ │ │ │ ├── hack.py
│ │ │ │ ├── model.py
│ │ │ │ └── recognizer.py
│ │ │ ├── ldm/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── models/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── autoencoder.py
│ │ │ │ │ └── diffusion/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── ddim.py
│ │ │ │ │ ├── ddpm.py
│ │ │ │ │ ├── dpm_solver/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ ├── dpm_solver.py
│ │ │ │ │ │ └── sampler.py
│ │ │ │ │ ├── plms.py
│ │ │ │ │ └── sampling_util.py
│ │ │ │ ├── modules/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── attention.py
│ │ │ │ │ ├── diffusionmodules/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ ├── model.py
│ │ │ │ │ │ ├── openaimodel.py
│ │ │ │ │ │ ├── upscaling.py
│ │ │ │ │ │ └── util.py
│ │ │ │ │ ├── distributions/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ └── distributions.py
│ │ │ │ │ ├── ema.py
│ │ │ │ │ └── encoders/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ └── modules.py
│ │ │ │ └── util.py
│ │ │ ├── main.py
│ │ │ ├── ocr_recog/
│ │ │ │ ├── RNN.py
│ │ │ │ ├── RecCTCHead.py
│ │ │ │ ├── RecModel.py
│ │ │ │ ├── RecMv1_enhance.py
│ │ │ │ ├── RecSVTR.py
│ │ │ │ ├── __init__.py
│ │ │ │ ├── common.py
│ │ │ │ ├── en_dict.txt
│ │ │ │ └── ppocr_keys_v1.txt
│ │ │ └── utils.py
│ │ ├── base.py
│ │ ├── brushnet/
│ │ │ ├── __init__.py
│ │ │ ├── brushnet.py
│ │ │ ├── brushnet_unet_forward.py
│ │ │ ├── brushnet_wrapper.py
│ │ │ ├── brushnet_xl_wrapper.py
│ │ │ ├── pipeline_brushnet.py
│ │ │ ├── pipeline_brushnet_sd_xl.py
│ │ │ └── unet_2d_blocks.py
│ │ ├── controlnet.py
│ │ ├── ddim_sampler.py
│ │ ├── fcf.py
│ │ ├── helper/
│ │ │ ├── __init__.py
│ │ │ ├── controlnet_preprocess.py
│ │ │ ├── cpu_text_encoder.py
│ │ │ └── g_diffuser_bot.py
│ │ ├── instruct_pix2pix.py
│ │ ├── kandinsky.py
│ │ ├── lama.py
│ │ ├── ldm.py
│ │ ├── manga.py
│ │ ├── mat.py
│ │ ├── mi_gan.py
│ │ ├── opencv2.py
│ │ ├── original_sd_configs/
│ │ │ ├── __init__.py
│ │ │ ├── sd_xl_base.yaml
│ │ │ ├── sd_xl_refiner.yaml
│ │ │ ├── v1-inference.yaml
│ │ │ └── v2-inference-v.yaml
│ │ ├── paint_by_example.py
│ │ ├── plms_sampler.py
│ │ ├── power_paint/
│ │ │ ├── __init__.py
│ │ │ ├── pipeline_powerpaint.py
│ │ │ ├── power_paint.py
│ │ │ ├── power_paint_v2.py
│ │ │ ├── powerpaint_tokenizer.py
│ │ │ └── v2/
│ │ │ ├── BrushNet_CA.py
│ │ │ ├── __init__.py
│ │ │ ├── pipeline_PowerPaint_Brushnet_CA.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ └── unet_2d_condition.py
│ │ ├── sd.py
│ │ ├── sdxl.py
│ │ ├── utils.py
│ │ └── zits.py
│ ├── model_manager.py
│ ├── plugins/
│ │ ├── __init__.py
│ │ ├── anime_seg.py
│ │ ├── base_plugin.py
│ │ ├── basicsr/
│ │ │ ├── LICENSE
│ │ │ ├── __init__.py
│ │ │ ├── arch_util.py
│ │ │ ├── img_util.py
│ │ │ └── rrdbnet_arch.py
│ │ ├── briarmbg.py
│ │ ├── briarmbg2.py
│ │ ├── facexlib/
│ │ │ ├── .gitignore
│ │ │ ├── __init__.py
│ │ │ ├── detection/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── align_trans.py
│ │ │ │ ├── matlab_cp2tform.py
│ │ │ │ ├── retinaface.py
│ │ │ │ ├── retinaface_net.py
│ │ │ │ └── retinaface_utils.py
│ │ │ ├── parsing/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── bisenet.py
│ │ │ │ ├── parsenet.py
│ │ │ │ └── resnet.py
│ │ │ └── utils/
│ │ │ ├── __init__.py
│ │ │ ├── face_restoration_helper.py
│ │ │ ├── face_utils.py
│ │ │ └── misc.py
│ │ ├── gfpgan/
│ │ │ ├── __init__.py
│ │ │ └── archs/
│ │ │ ├── __init__.py
│ │ │ ├── gfpganv1_clean_arch.py
│ │ │ ├── restoreformer_arch.py
│ │ │ └── stylegan2_clean_arch.py
│ │ ├── gfpgan_plugin.py
│ │ ├── gfpganer.py
│ │ ├── interactive_seg.py
│ │ ├── realesrgan.py
│ │ ├── remove_bg.py
│ │ ├── restoreformer.py
│ │ ├── segment_anything/
│ │ │ ├── __init__.py
│ │ │ ├── build_sam.py
│ │ │ ├── modeling/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── common.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ ├── image_encoder_hq.py
│ │ │ │ ├── mask_decoder.py
│ │ │ │ ├── prompt_encoder.py
│ │ │ │ ├── sam.py
│ │ │ │ ├── sam_hq.py
│ │ │ │ ├── tiny_vit_sam.py
│ │ │ │ └── transformer.py
│ │ │ ├── predictor.py
│ │ │ ├── predictor_hq.py
│ │ │ └── utils/
│ │ │ ├── __init__.py
│ │ │ └── transforms.py
│ │ └── segment_anything2/
│ │ ├── __init__.py
│ │ ├── build_sam.py
│ │ ├── modeling/
│ │ │ ├── __init__.py
│ │ │ ├── backbones/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── hieradet.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── utils.py
│ │ │ ├── memory_attention.py
│ │ │ ├── memory_encoder.py
│ │ │ ├── position_encoding.py
│ │ │ ├── sam/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mask_decoder.py
│ │ │ │ ├── prompt_encoder.py
│ │ │ │ └── transformer.py
│ │ │ ├── sam2_base.py
│ │ │ └── sam2_utils.py
│ │ ├── sam2_image_predictor.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── misc.py
│ │ └── transforms.py
│ ├── runtime.py
│ ├── schema.py
│ ├── tests/
│ │ ├── .gitignore
│ │ ├── __init__.py
│ │ ├── test_adjust_mask.py
│ │ ├── test_anytext.py
│ │ ├── test_brushnet.py
│ │ ├── test_controlnet.py
│ │ ├── test_instruct_pix2pix.py
│ │ ├── test_load_img.py
│ │ ├── test_low_mem.py
│ │ ├── test_match_histograms.py
│ │ ├── test_model.py
│ │ ├── test_model_md5.py
│ │ ├── test_model_switch.py
│ │ ├── test_outpainting.py
│ │ ├── test_paint_by_example.py
│ │ ├── test_plugins.py
│ │ ├── test_save_exif.py
│ │ ├── test_save_quality.py
│ │ ├── test_sd_model.py
│ │ ├── test_sdxl.py
│ │ └── utils.py
│ └── web_config.py
├── main.py
├── publish.sh
├── requirements-dev.txt
├── requirements.txt
├── scripts/
│ ├── .gitignore
│ ├── README.md
│ ├── convert_vae_pt_to_diffusers.py
│ ├── environment.yaml
│ ├── pack.bat
│ ├── pack.sh
│ ├── tool.py
│ └── user_scripts/
│ ├── win_config.bat
│ ├── win_setup.bat
│ ├── win_setup_cn.bat
│ ├── win_start.bat
│ ├── win_start_cn.bat
│ └── win_update.bat
├── setup.py
└── web_app/
├── .eslintrc.cjs
├── .gitignore
├── README.md
├── components.json
├── index.html
├── package.json
├── postcss.config.js
├── src/
│ ├── App.tsx
│ ├── components/
│ │ ├── Coffee.tsx
│ │ ├── Cropper.tsx
│ │ ├── DiffusionProgress.tsx
│ │ ├── Editor.tsx
│ │ ├── Extender.tsx
│ │ ├── FileManager.tsx
│ │ ├── FileSelect.tsx
│ │ ├── Header.tsx
│ │ ├── ImageSize.tsx
│ │ ├── InteractiveSeg.tsx
│ │ ├── Plugins.tsx
│ │ ├── PromptInput.tsx
│ │ ├── Settings.tsx
│ │ ├── Shortcuts.tsx
│ │ ├── SidePanel/
│ │ │ ├── CV2Options.tsx
│ │ │ ├── DiffusionOptions.tsx
│ │ │ ├── LDMOptions.tsx
│ │ │ ├── LabelTitle.tsx
│ │ │ └── index.tsx
│ │ ├── Workspace.tsx
│ │ └── ui/
│ │ ├── accordion.tsx
│ │ ├── alert-dialog.tsx
│ │ ├── button.tsx
│ │ ├── context-menu.tsx
│ │ ├── dialog.tsx
│ │ ├── dropdown-menu.tsx
│ │ ├── form.tsx
│ │ ├── input.tsx
│ │ ├── label.tsx
│ │ ├── popover.tsx
│ │ ├── progress.tsx
│ │ ├── radio-group.tsx
│ │ ├── scroll-area.tsx
│ │ ├── select.tsx
│ │ ├── separator.tsx
│ │ ├── sheet.tsx
│ │ ├── slider.tsx
│ │ ├── switch.tsx
│ │ ├── tabs.tsx
│ │ ├── textarea.tsx
│ │ ├── toast.tsx
│ │ ├── toaster.tsx
│ │ ├── toggle.tsx
│ │ ├── tooltip.tsx
│ │ └── use-toast.ts
│ ├── globals.css
│ ├── hooks/
│ │ ├── useAsyncMemo.tsx
│ │ ├── useHotkey.tsx
│ │ ├── useImage.tsx
│ │ ├── useInputImage.tsx
│ │ └── useResolution.tsx
│ ├── lib/
│ │ ├── api.ts
│ │ ├── const.ts
│ │ ├── states.ts
│ │ ├── types.ts
│ │ └── utils.ts
│ ├── main.tsx
│ └── vite-env.d.ts
├── tailwind.config.js
├── tsconfig.json
├── tsconfig.node.json
└── vite.config.ts
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/FUNDING.yml
================================================
# These are supported funding model platforms
ko_fi: Z8Z1CZJGY
================================================
FILE: .github/ISSUE_TEMPLATE/config.yml
================================================
contact_links:
- name: Blank issue
url: https://github.com/Sanster/lama-cleaner/issues/new
about: Other
================================================
FILE: .github/ISSUE_TEMPLATE/🐛-bug-report.md
================================================
---
name: "\U0001F41B Bug Report"
about: Create a report to help us improve
title: "[BUG]"
labels: bug
assignees: ''
---
**Model**
Which model are you using?
**Describe the bug**
A clear and concise description of what the bug is.
**Screenshots**
If applicable, add screenshots to help explain your problem.
**System Info**
Software version used
- iopaint:
- pytorch:
- CUDA:
================================================
FILE: .github/ISSUE_TEMPLATE/🚀-feature-request.md
================================================
---
name: "\U0001F680 Feature request"
about: Suggest an idea for this project
title: "[Feature Request]"
labels: ''
assignees: ''
---
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
**Additional context**
Add any other context or screenshots about the feature request here.
================================================
FILE: .gitignore
================================================
.DS_Store
**/__pycache__
examples/
.idea/
.vscode/
build
!iopaint/app/build
dist/
IOPaint.egg-info/
venv/
tmp/
iopaint/web_app/
================================================
FILE: LICENSE
================================================
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.
================================================
FILE: README.md
================================================
<h1 align="center">IOPaint</h1>
<p align="center">A free and open-source inpainting & outpainting tool powered by SOTA AI model.</p>
<p align="center">
<a href="https://github.com/Sanster/IOPaint">
<img alt="total download" src="https://pepy.tech/badge/iopaint" />
</a>
<a href="https://pypi.org/project/iopaint">
<img alt="version" src="https://img.shields.io/pypi/v/iopaint" />
</a>
<a href="">
<img alt="python version" src="https://img.shields.io/pypi/pyversions/iopaint" />
</a>
<a href="https://huggingface.co/spaces/Sanster/iopaint-lama">
<img alt="HuggingFace Spaces" src="https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Spaces-blue" />
</a>
<a href="https://colab.research.google.com/drive/1TKVlDZiE3MIZnAUMpv2t_S4hLr6TUY1d?usp=sharing">
<img alt="Open in Colab" src="https://colab.research.google.com/assets/colab-badge.svg" />
</a>
</p>
|Erase([LaMa](https://www.iopaint.com/models/erase/lama))|Replace Object([PowerPaint](https://www.iopaint.com/models/diffusion/powerpaint))|
|-----|----|
|<video src="https://github.com/Sanster/IOPaint/assets/3998421/264bc27c-0abd-4d8b-bb1e-0078ab264c4a"> | <video src="https://github.com/Sanster/IOPaint/assets/3998421/1de5c288-e0e1-4f32-926d-796df0655846">|
|Draw Text([AnyText](https://www.iopaint.com/models/diffusion/anytext))|Out-painting([PowerPaint](https://www.iopaint.com/models/diffusion/powerpaint))|
|---------|-----------|
|<video src="https://github.com/Sanster/IOPaint/assets/3998421/ffd4eda4-f7d4-4693-93d8-d2cd5aa7c6d6">|<video src="https://github.com/Sanster/IOPaint/assets/3998421/c4af8aef-8c29-49e0-96eb-0aae2f768da2">|
## Features
- Completely free and open-source, fully self-hosted, support CPU & GPU & Apple Silicon
- [Windows 1-Click Installer](https://www.iopaint.com/install/windows_1click_installer)
- [OptiClean](https://apps.apple.com/ca/app/opticlean/id6452387177): macOS & iOS App for object erase
- Supports various AI [models](https://www.iopaint.com/models) to perform erase, inpainting or outpainting task.
- [Erase models](https://www.iopaint.com/models#erase-models): These models can be used to remove unwanted object, defect, watermarks, people from image.
- Diffusion models: These models can be used to replace objects or perform outpainting. Some popular used models include:
- [runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting)
- [diffusers/stable-diffusion-xl-1.0-inpainting-0.1](https://huggingface.co/diffusers/stable-diffusion-xl-1.0-inpainting-0.1)
- [andregn/Realistic_Vision_V3.0-inpainting](https://huggingface.co/andregn/Realistic_Vision_V3.0-inpainting)
- [Lykon/dreamshaper-8-inpainting](https://huggingface.co/Lykon/dreamshaper-8-inpainting)
- [Sanster/anything-4.0-inpainting](https://huggingface.co/Sanster/anything-4.0-inpainting)
- [BrushNet](https://www.iopaint.com/models/diffusion/brushnet)
- [PowerPaintV2](https://www.iopaint.com/models/diffusion/powerpaint_v2)
- [Sanster/AnyText](https://huggingface.co/Sanster/AnyText)
- [Fantasy-Studio/Paint-by-Example](https://huggingface.co/Fantasy-Studio/Paint-by-Example)
- [Plugins](https://www.iopaint.com/plugins):
- [Segment Anything](https://iopaint.com/plugins/interactive_seg): Accurate and fast Interactive Object Segmentation
- [RemoveBG](https://iopaint.com/plugins/rembg): Remove image background or generate masks for foreground objects
- [Anime Segmentation](https://iopaint.com/plugins/anime_seg): Similar to RemoveBG, the model is specifically trained for anime images.
- [RealESRGAN](https://iopaint.com/plugins/RealESRGAN): Super Resolution
- [GFPGAN](https://iopaint.com/plugins/GFPGAN): Face Restoration
- [RestoreFormer](https://iopaint.com/plugins/RestoreFormer): Face Restoration
- [FileManager](https://iopaint.com/file_manager): Browse your pictures conveniently and save them directly to the output directory.
## Quick Start
### Start webui
IOPaint provides a convenient webui for using the latest AI models to edit your images.
You can install and start IOPaint easily by running following command:
```bash
# In order to use GPU, install cuda version of pytorch first.
# pip3 install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu118
# AMD GPU users, please utilize the following command, only works on linux, as pytorch is not yet supported on Windows with ROCm.
# pip3 install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/rocm5.6
pip3 install iopaint
iopaint start --model=lama --device=cpu --port=8080
```
That's it, you can start using IOPaint by visiting http://localhost:8080 in your web browser.
All models will be downloaded automatically at startup. If you want to change the download directory, you can add `--model-dir`. More documentation can be found [here](https://www.iopaint.com/install/download_model)
You can see other supported models at [here](https://www.iopaint.com/models) and how to use local sd ckpt/safetensors file at [here](https://www.iopaint.com/models#load-ckptsafetensors).
### Plugins
You can specify which plugins to use when starting the service, and you can view the commands to enable plugins by using `iopaint start --help`.
More demonstrations of the Plugin can be seen [here](https://www.iopaint.com/plugins)
```bash
iopaint start --enable-interactive-seg --interactive-seg-device=cuda
```
### Batch processing
You can also use IOPaint in the command line to batch process images:
```bash
iopaint run --model=lama --device=cpu \
--image=/path/to/image_folder \
--mask=/path/to/mask_folder \
--output=output_dir
```
`--image` is the folder containing input images, `--mask` is the folder containing corresponding mask images.
When `--mask` is a path to a mask file, all images will be processed using this mask.
You can see more information about the available models and plugins supported by IOPaint below.
## Development
Install [nodejs](https://nodejs.org/en), then install the frontend dependencies.
```bash
git clone https://github.com/Sanster/IOPaint.git
cd IOPaint/web_app
npm install
npm run build
cp -r dist/ ../iopaint/web_app
```
Create a `.env.local` file in `web_app` and fill in the backend IP and port.
```
VITE_BACKEND=http://127.0.0.1:8080
```
Start front-end development environment
```bash
npm run dev
```
Install back-end requirements and start backend service
```bash
pip install -r requirements.txt
python3 main.py start --model lama --port 8080
```
Then you can visit `http://localhost:5173/` for development.
The frontend code will automatically update after being modified,
but the backend needs to restart the service after modifying the python code.
================================================
FILE: build_docker.sh
================================================
#!/usr/bin/env bash
set -e
GIT_TAG=$1
IMAGE_DESC="Image inpainting tool powered by SOTA AI Model"
GIT_REPO="https://github.com/Sanster/lama-cleaner"
echo "Building cpu docker image..."
docker buildx build \
--platform linux/amd64 \
--file ./docker/CPUDockerfile \
--label org.opencontainers.image.title=lama-cleaner \
--label org.opencontainers.image.description="$IMAGE_DESC" \
--label org.opencontainers.image.url=$GIT_REPO \
--label org.opencontainers.image.source=$GIT_REPO \
--label org.opencontainers.image.version=$GIT_TAG \
--build-arg version=$GIT_TAG \
--tag cwq1913/lama-cleaner:cpu-$GIT_TAG .
# echo "Building NVIDIA GPU docker image..."
docker buildx build \
--platform linux/amd64 \
--file ./docker/GPUDockerfile \
--label org.opencontainers.image.title=lama-cleaner \
--label org.opencontainers.image.description="$IMAGE_DESC" \
--label org.opencontainers.image.url=$GIT_REPO \
--label org.opencontainers.image.source=$GIT_REPO \
--label org.opencontainers.image.version=$GIT_TAG \
--build-arg version=$GIT_TAG \
--tag cwq1913/lama-cleaner:gpu-$GIT_TAG .
================================================
FILE: docker/CPUDockerfile
================================================
FROM python:3.10.11-slim-buster
RUN apt-get update && apt-get install -y --no-install-recommends \
software-properties-common \
libsm6 libxext6 ffmpeg libfontconfig1 libxrender1 libgl1-mesa-glx \
curl gcc build-essential
RUN pip install --upgrade pip && \
pip install torch==1.13.1 torchvision==0.14.1 --extra-index-url https://download.pytorch.org/whl/cpu
ARG version
RUN pip install lama-cleaner==$version
RUN lama-cleaner --install-plugins-package
ENV LD_PRELOAD=/usr/local/lib/python3.10/site-packages/skimage/_shared/../../scikit_image.libs/libgomp-d22c30c5.so.1.0.0
EXPOSE 8080
CMD ["bash"]
================================================
FILE: docker/GPUDockerfile
================================================
FROM nvidia/cuda:11.7.1-runtime-ubuntu20.04
RUN apt-get update && apt-get install -y --no-install-recommends \
software-properties-common \
libsm6 libxext6 ffmpeg libfontconfig1 libxrender1 libgl1-mesa-glx \
curl python3-pip
RUN pip3 install --upgrade pip
RUN pip3 install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118
RUN pip3 install xformers==0.0.22.post4 --index-url https://download.pytorch.org/whl/cu118
ARG version
RUN pip3 install lama-cleaner==$version
RUN lama-cleaner --install-plugins-package
EXPOSE 8080
CMD ["bash"]
================================================
FILE: iopaint/__init__.py
================================================
import os
import importlib.util
import shutil
import ctypes
import logging
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
# https://github.com/pytorch/pytorch/issues/27971#issuecomment-1768868068
os.environ["ONEDNN_PRIMITIVE_CACHE_CAPACITY"] = "1"
os.environ["LRU_CACHE_CAPACITY"] = "1"
# prevent CPU memory leak when run model on GPU
# https://github.com/pytorch/pytorch/issues/98688#issuecomment-1869288431
# https://github.com/pytorch/pytorch/issues/108334#issuecomment-1752763633
os.environ["TORCH_CUDNN_V8_API_LRU_CACHE_LIMIT"] = "1"
import warnings
warnings.simplefilter("ignore", UserWarning)
def fix_window_pytorch():
# copy from: https://github.com/comfyanonymous/ComfyUI/blob/5cbaa9e07c97296b536f240688f5a19300ecf30d/fix_torch.py#L4
import platform
try:
if platform.system() != "Windows":
return
torch_spec = importlib.util.find_spec("torch")
for folder in torch_spec.submodule_search_locations:
lib_folder = os.path.join(folder, "lib")
test_file = os.path.join(lib_folder, "fbgemm.dll")
dest = os.path.join(lib_folder, "libomp140.x86_64.dll")
if os.path.exists(dest):
break
with open(test_file, "rb") as f:
contents = f.read()
if b"libomp140.x86_64.dll" not in contents:
break
try:
mydll = ctypes.cdll.LoadLibrary(test_file)
except FileNotFoundError:
logging.warning("Detected pytorch version with libomp issue, patching.")
shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)
except:
pass
def entry_point():
# To make os.environ["XDG_CACHE_HOME"] = args.model_cache_dir works for diffusers
# https://github.com/huggingface/diffusers/blob/be99201a567c1ccd841dc16fb24e88f7f239c187/src/diffusers/utils/constants.py#L18
from iopaint.cli import typer_app
fix_window_pytorch()
typer_app()
================================================
FILE: iopaint/__main__.py
================================================
from iopaint import entry_point
if __name__ == "__main__":
entry_point()
================================================
FILE: iopaint/api.py
================================================
import asyncio
import os
import threading
import time
import traceback
from pathlib import Path
from typing import Optional, Dict, List
import cv2
import numpy as np
import socketio
import torch
try:
torch._C._jit_override_can_fuse_on_cpu(False)
torch._C._jit_override_can_fuse_on_gpu(False)
torch._C._jit_set_texpr_fuser_enabled(False)
torch._C._jit_set_nvfuser_enabled(False)
torch._C._jit_set_profiling_mode(False)
except:
pass
import uvicorn
from PIL import Image
from fastapi import APIRouter, FastAPI, Request, UploadFile
from fastapi.encoders import jsonable_encoder
from fastapi.exceptions import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse, Response
from fastapi.staticfiles import StaticFiles
from loguru import logger
from socketio import AsyncServer
from iopaint.file_manager import FileManager
from iopaint.helper import (
load_img,
decode_base64_to_image,
pil_to_bytes,
numpy_to_bytes,
concat_alpha_channel,
gen_frontend_mask,
adjust_mask,
)
from iopaint.model.utils import torch_gc
from iopaint.model_manager import ModelManager
from iopaint.plugins import build_plugins, RealESRGANUpscaler, InteractiveSeg
from iopaint.plugins.base_plugin import BasePlugin
from iopaint.plugins.remove_bg import RemoveBG
from iopaint.schema import (
GenInfoResponse,
ApiConfig,
ServerConfigResponse,
SwitchModelRequest,
InpaintRequest,
RunPluginRequest,
SDSampler,
PluginInfo,
AdjustMaskRequest,
RemoveBGModel,
SwitchPluginModelRequest,
ModelInfo,
InteractiveSegModel,
RealESRGANModel,
)
CURRENT_DIR = Path(__file__).parent.absolute().resolve()
WEB_APP_DIR = CURRENT_DIR / "web_app"
def api_middleware(app: FastAPI):
rich_available = False
try:
if os.environ.get("WEBUI_RICH_EXCEPTIONS", None) is not None:
import anyio # importing just so it can be placed on silent list
import starlette # importing just so it can be placed on silent list
from rich.console import Console
console = Console()
rich_available = True
except Exception:
pass
def handle_exception(request: Request, e: Exception):
err = {
"error": type(e).__name__,
"detail": vars(e).get("detail", ""),
"body": vars(e).get("body", ""),
"errors": str(e),
}
if not isinstance(
e, HTTPException
): # do not print backtrace on known httpexceptions
message = f"API error: {request.method}: {request.url} {err}"
if rich_available:
print(message)
console.print_exception(
show_locals=True,
max_frames=2,
extra_lines=1,
suppress=[anyio, starlette],
word_wrap=False,
width=min([console.width, 200]),
)
else:
traceback.print_exc()
return JSONResponse(
status_code=vars(e).get("status_code", 500), content=jsonable_encoder(err)
)
@app.middleware("http")
async def exception_handling(request: Request, call_next):
try:
return await call_next(request)
except Exception as e:
return handle_exception(request, e)
@app.exception_handler(Exception)
async def fastapi_exception_handler(request: Request, e: Exception):
return handle_exception(request, e)
@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, e: HTTPException):
return handle_exception(request, e)
cors_options = {
"allow_methods": ["*"],
"allow_headers": ["*"],
"allow_origins": ["*"],
"allow_credentials": True,
"expose_headers": ["X-Seed"],
}
app.add_middleware(CORSMiddleware, **cors_options)
global_sio: AsyncServer = None
def diffuser_callback(pipe, step: int, timestep: int, callback_kwargs: Dict = {}):
# self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict
# logger.info(f"diffusion callback: step={step}, timestep={timestep}")
# We use asyncio loos for task processing. Perhaps in the future, we can add a processing queue similar to InvokeAI,
# but for now let's just start a separate event loop. It shouldn't make a difference for single person use
asyncio.run(global_sio.emit("diffusion_progress", {"step": step}))
return {}
class Api:
def __init__(self, app: FastAPI, config: ApiConfig):
self.app = app
self.config = config
self.router = APIRouter()
self.queue_lock = threading.Lock()
api_middleware(self.app)
self.file_manager = self._build_file_manager()
self.plugins = self._build_plugins()
self.model_manager = self._build_model_manager()
# fmt: off
self.add_api_route("/api/v1/gen-info", self.api_geninfo, methods=["POST"], response_model=GenInfoResponse)
self.add_api_route("/api/v1/server-config", self.api_server_config, methods=["GET"],
response_model=ServerConfigResponse)
self.add_api_route("/api/v1/model", self.api_current_model, methods=["GET"], response_model=ModelInfo)
self.add_api_route("/api/v1/model", self.api_switch_model, methods=["POST"], response_model=ModelInfo)
self.add_api_route("/api/v1/inputimage", self.api_input_image, methods=["GET"])
self.add_api_route("/api/v1/inpaint", self.api_inpaint, methods=["POST"])
self.add_api_route("/api/v1/switch_plugin_model", self.api_switch_plugin_model, methods=["POST"])
self.add_api_route("/api/v1/run_plugin_gen_mask", self.api_run_plugin_gen_mask, methods=["POST"])
self.add_api_route("/api/v1/run_plugin_gen_image", self.api_run_plugin_gen_image, methods=["POST"])
self.add_api_route("/api/v1/samplers", self.api_samplers, methods=["GET"])
self.add_api_route("/api/v1/adjust_mask", self.api_adjust_mask, methods=["POST"])
self.add_api_route("/api/v1/save_image", self.api_save_image, methods=["POST"])
self.app.mount("/", StaticFiles(directory=WEB_APP_DIR, html=True), name="assets")
# fmt: on
global global_sio
self.sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*")
self.combined_asgi_app = socketio.ASGIApp(self.sio, self.app)
self.app.mount("/ws", self.combined_asgi_app)
global_sio = self.sio
def add_api_route(self, path: str, endpoint, **kwargs):
return self.app.add_api_route(path, endpoint, **kwargs)
def api_save_image(self, file: UploadFile):
# Sanitize filename to prevent path traversal
safe_filename = Path(file.filename).name # Get just the filename component
# Construct the full path within output_dir
output_path = self.config.output_dir / safe_filename
# Ensure output directory exists
if not self.config.output_dir or not self.config.output_dir.exists():
raise HTTPException(
status_code=400,
detail="Output directory not configured or doesn't exist",
)
# Read and write the file
origin_image_bytes = file.file.read()
with open(output_path, "wb") as fw:
fw.write(origin_image_bytes)
def api_current_model(self) -> ModelInfo:
return self.model_manager.current_model
def api_switch_model(self, req: SwitchModelRequest) -> ModelInfo:
if req.name == self.model_manager.name:
return self.model_manager.current_model
self.model_manager.switch(req.name)
return self.model_manager.current_model
def api_switch_plugin_model(self, req: SwitchPluginModelRequest):
if req.plugin_name in self.plugins:
self.plugins[req.plugin_name].switch_model(req.model_name)
if req.plugin_name == RemoveBG.name:
self.config.remove_bg_model = req.model_name
if req.plugin_name == RealESRGANUpscaler.name:
self.config.realesrgan_model = req.model_name
if req.plugin_name == InteractiveSeg.name:
self.config.interactive_seg_model = req.model_name
torch_gc()
def api_server_config(self) -> ServerConfigResponse:
plugins = []
for it in self.plugins.values():
plugins.append(
PluginInfo(
name=it.name,
support_gen_image=it.support_gen_image,
support_gen_mask=it.support_gen_mask,
)
)
return ServerConfigResponse(
plugins=plugins,
modelInfos=self.model_manager.scan_models(),
removeBGModel=self.config.remove_bg_model,
removeBGModels=RemoveBGModel.values(),
realesrganModel=self.config.realesrgan_model,
realesrganModels=RealESRGANModel.values(),
interactiveSegModel=self.config.interactive_seg_model,
interactiveSegModels=InteractiveSegModel.values(),
enableFileManager=self.file_manager is not None,
enableAutoSaving=self.config.output_dir is not None,
enableControlnet=self.model_manager.enable_controlnet,
controlnetMethod=self.model_manager.controlnet_method,
disableModelSwitch=False,
isDesktop=False,
samplers=self.api_samplers(),
)
def api_input_image(self) -> FileResponse:
if self.config.input is None:
raise HTTPException(status_code=200, detail="No input image configured")
if self.config.input.is_file():
return FileResponse(self.config.input)
raise HTTPException(status_code=404, detail="Input image not found")
def api_geninfo(self, file: UploadFile) -> GenInfoResponse:
_, _, info = load_img(file.file.read(), return_info=True)
parts = info.get("parameters", "").split("Negative prompt: ")
prompt = parts[0].strip()
negative_prompt = ""
if len(parts) > 1:
negative_prompt = parts[1].split("\n")[0].strip()
return GenInfoResponse(prompt=prompt, negative_prompt=negative_prompt)
def api_inpaint(self, req: InpaintRequest):
image, alpha_channel, infos, ext = decode_base64_to_image(req.image)
mask, _, _, _ = decode_base64_to_image(req.mask, gray=True)
logger.info(f"image ext: {ext}")
mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
if image.shape[:2] != mask.shape[:2]:
raise HTTPException(
400,
detail=f"Image size({image.shape[:2]}) and mask size({mask.shape[:2]}) not match.",
)
start = time.time()
rgb_np_img = self.model_manager(image, mask, req)
logger.info(f"process time: {(time.time() - start) * 1000:.2f}ms")
torch_gc()
rgb_np_img = cv2.cvtColor(rgb_np_img.astype(np.uint8), cv2.COLOR_BGR2RGB)
rgb_res = concat_alpha_channel(rgb_np_img, alpha_channel)
res_img_bytes = pil_to_bytes(
Image.fromarray(rgb_res),
ext=ext,
quality=self.config.quality,
infos=infos,
)
asyncio.run(self.sio.emit("diffusion_finish"))
return Response(
content=res_img_bytes,
media_type=f"image/{ext}",
headers={"X-Seed": str(req.sd_seed)},
)
def api_run_plugin_gen_image(self, req: RunPluginRequest):
ext = "png"
if req.name not in self.plugins:
raise HTTPException(status_code=422, detail="Plugin not found")
if not self.plugins[req.name].support_gen_image:
raise HTTPException(
status_code=422, detail="Plugin does not support output image"
)
rgb_np_img, alpha_channel, infos, _ = decode_base64_to_image(req.image)
bgr_or_rgba_np_img = self.plugins[req.name].gen_image(rgb_np_img, req)
torch_gc()
if bgr_or_rgba_np_img.shape[2] == 4:
rgba_np_img = bgr_or_rgba_np_img
else:
rgba_np_img = cv2.cvtColor(bgr_or_rgba_np_img, cv2.COLOR_BGR2RGB)
rgba_np_img = concat_alpha_channel(rgba_np_img, alpha_channel)
return Response(
content=pil_to_bytes(
Image.fromarray(rgba_np_img),
ext=ext,
quality=self.config.quality,
infos=infos,
),
media_type=f"image/{ext}",
)
def api_run_plugin_gen_mask(self, req: RunPluginRequest):
if req.name not in self.plugins:
raise HTTPException(status_code=422, detail="Plugin not found")
if not self.plugins[req.name].support_gen_mask:
raise HTTPException(
status_code=422, detail="Plugin does not support output image"
)
rgb_np_img, _, _, _ = decode_base64_to_image(req.image)
bgr_or_gray_mask = self.plugins[req.name].gen_mask(rgb_np_img, req)
torch_gc()
res_mask = gen_frontend_mask(bgr_or_gray_mask)
return Response(
content=numpy_to_bytes(res_mask, "png"),
media_type="image/png",
)
def api_samplers(self) -> List[str]:
return [member.value for member in SDSampler.__members__.values()]
def api_adjust_mask(self, req: AdjustMaskRequest):
mask, _, _, _ = decode_base64_to_image(req.mask, gray=True)
mask = adjust_mask(mask, req.kernel_size, req.operate)
return Response(content=numpy_to_bytes(mask, "png"), media_type="image/png")
def launch(self):
self.app.include_router(self.router)
uvicorn.run(
self.combined_asgi_app,
host=self.config.host,
port=self.config.port,
timeout_keep_alive=999999999,
)
def _build_file_manager(self) -> Optional[FileManager]:
if self.config.input and self.config.input.is_dir():
logger.info(
f"Input is directory, initialize file manager {self.config.input}"
)
return FileManager(
app=self.app,
input_dir=self.config.input,
mask_dir=self.config.mask_dir,
output_dir=self.config.output_dir,
)
return None
def _build_plugins(self) -> Dict[str, BasePlugin]:
return build_plugins(
self.config.enable_interactive_seg,
self.config.interactive_seg_model,
self.config.interactive_seg_device,
self.config.enable_remove_bg,
self.config.remove_bg_device,
self.config.remove_bg_model,
self.config.enable_anime_seg,
self.config.enable_realesrgan,
self.config.realesrgan_device,
self.config.realesrgan_model,
self.config.enable_gfpgan,
self.config.gfpgan_device,
self.config.enable_restoreformer,
self.config.restoreformer_device,
self.config.no_half,
)
def _build_model_manager(self):
return ModelManager(
name=self.config.model,
device=torch.device(self.config.device),
no_half=self.config.no_half,
low_mem=self.config.low_mem,
disable_nsfw=self.config.disable_nsfw_checker,
sd_cpu_textencoder=self.config.cpu_textencoder,
local_files_only=self.config.local_files_only,
cpu_offload=self.config.cpu_offload,
callback=diffuser_callback,
)
================================================
FILE: iopaint/batch_processing.py
================================================
import json
from pathlib import Path
from typing import Dict, Optional
import cv2
import numpy as np
from PIL import Image
from loguru import logger
from rich.console import Console
from rich.progress import (
Progress,
SpinnerColumn,
TimeElapsedColumn,
MofNCompleteColumn,
TextColumn,
BarColumn,
TaskProgressColumn,
)
from iopaint.helper import pil_to_bytes
from iopaint.model.utils import torch_gc
from iopaint.model_manager import ModelManager
from iopaint.schema import InpaintRequest
def glob_images(path: Path) -> Dict[str, Path]:
# png/jpg/jpeg
if path.is_file():
return {path.stem: path}
elif path.is_dir():
res = {}
for it in path.glob("*.*"):
if it.suffix.lower() in [".png", ".jpg", ".jpeg"]:
res[it.stem] = it
return res
def batch_inpaint(
model: str,
device,
image: Path,
mask: Path,
output: Path,
config: Optional[Path] = None,
concat: bool = False,
):
if image.is_dir() and output.is_file():
logger.error(
"invalid --output: when image is a directory, output should be a directory"
)
exit(-1)
output.mkdir(parents=True, exist_ok=True)
image_paths = glob_images(image)
mask_paths = glob_images(mask)
if len(image_paths) == 0:
logger.error("invalid --image: empty image folder")
exit(-1)
if len(mask_paths) == 0:
logger.error("invalid --mask: empty mask folder")
exit(-1)
if config is None:
inpaint_request = InpaintRequest()
logger.info(f"Using default config: {inpaint_request}")
else:
with open(config, "r", encoding="utf-8") as f:
inpaint_request = InpaintRequest(**json.load(f))
logger.info(f"Using config: {inpaint_request}")
model_manager = ModelManager(name=model, device=device)
first_mask = list(mask_paths.values())[0]
console = Console()
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
console=console,
transient=False,
) as progress:
task = progress.add_task("Batch processing...", total=len(image_paths))
for stem, image_p in image_paths.items():
if stem not in mask_paths and mask.is_dir():
progress.log(f"mask for {image_p} not found")
progress.update(task, advance=1)
continue
mask_p = mask_paths.get(stem, first_mask)
infos = Image.open(image_p).info
img = np.array(Image.open(image_p).convert("RGB"))
mask_img = np.array(Image.open(mask_p).convert("L"))
if mask_img.shape[:2] != img.shape[:2]:
progress.log(
f"resize mask {mask_p.name} to image {image_p.name} size: {img.shape[:2]}"
)
mask_img = cv2.resize(
mask_img,
(img.shape[1], img.shape[0]),
interpolation=cv2.INTER_NEAREST,
)
mask_img[mask_img >= 127] = 255
mask_img[mask_img < 127] = 0
# bgr
inpaint_result = model_manager(img, mask_img, inpaint_request)
inpaint_result = cv2.cvtColor(inpaint_result, cv2.COLOR_BGR2RGB)
if concat:
mask_img = cv2.cvtColor(mask_img, cv2.COLOR_GRAY2RGB)
inpaint_result = cv2.hconcat([img, mask_img, inpaint_result])
img_bytes = pil_to_bytes(Image.fromarray(inpaint_result), "png", 100, infos)
save_p = output / f"{stem}.png"
with open(save_p, "wb") as fw:
fw.write(img_bytes)
progress.update(task, advance=1)
torch_gc()
# pid = psutil.Process().pid
# memory_info = psutil.Process(pid).memory_info()
# memory_in_mb = memory_info.rss / (1024 * 1024)
# print(f"原图大小：{img.shape},当前进程的内存占用：{memory_in_mb}MB")
================================================
FILE: iopaint/benchmark.py
================================================
#!/usr/bin/env python3
import argparse
import os
import time
import numpy as np
import nvidia_smi
import psutil
import torch
from iopaint.model_manager import ModelManager
from iopaint.schema import InpaintRequest, HDStrategy, SDSampler
try:
torch._C._jit_override_can_fuse_on_cpu(False)
torch._C._jit_override_can_fuse_on_gpu(False)
torch._C._jit_set_texpr_fuser_enabled(False)
torch._C._jit_set_nvfuser_enabled(False)
except:
pass
NUM_THREADS = str(4)
os.environ["OMP_NUM_THREADS"] = NUM_THREADS
os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS
os.environ["MKL_NUM_THREADS"] = NUM_THREADS
os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS
os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS
if os.environ.get("CACHE_DIR"):
os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"]
def run_model(model, size):
# RGB
image = np.random.randint(0, 256, (size[0], size[1], 3)).astype(np.uint8)
mask = np.random.randint(0, 255, size).astype(np.uint8)
config = InpaintRequest(
ldm_steps=2,
hd_strategy=HDStrategy.ORIGINAL,
hd_strategy_crop_margin=128,
hd_strategy_crop_trigger_size=128,
hd_strategy_resize_limit=128,
prompt="a fox is sitting on a bench",
sd_steps=5,
sd_sampler=SDSampler.ddim,
)
model(image, mask, config)
def benchmark(model, times: int, empty_cache: bool):
sizes = [(512, 512)]
nvidia_smi.nvmlInit()
device_id = 0
handle = nvidia_smi.nvmlDeviceGetHandleByIndex(device_id)
def format(metrics):
return f"{np.mean(metrics):.2f} ± {np.std(metrics):.2f}"
process = psutil.Process(os.getpid())
# 每个 size 给出显存和内存占用的指标
for size in sizes:
torch.cuda.empty_cache()
time_metrics = []
cpu_metrics = []
memory_metrics = []
gpu_memory_metrics = []
for _ in range(times):
start = time.time()
run_model(model, size)
torch.cuda.synchronize()
# cpu_metrics.append(process.cpu_percent())
time_metrics.append((time.time() - start) * 1000)
memory_metrics.append(process.memory_info().rss / 1024 / 1024)
gpu_memory_metrics.append(
nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used / 1024 / 1024
)
print(f"size: {size}".center(80, "-"))
# print(f"cpu: {format(cpu_metrics)}")
print(f"latency: {format(time_metrics)}ms")
print(f"memory: {format(memory_metrics)} MB")
print(f"gpu memory: {format(gpu_memory_metrics)} MB")
nvidia_smi.nvmlShutdown()
def get_args_parser():
parser = argparse.ArgumentParser()
parser.add_argument("--name")
parser.add_argument("--device", default="cuda", type=str)
parser.add_argument("--times", default=10, type=int)
parser.add_argument("--empty-cache", action="store_true")
return parser.parse_args()
if __name__ == "__main__":
args = get_args_parser()
device = torch.device(args.device)
model = ModelManager(
name=args.name,
device=device,
disable_nsfw=True,
sd_cpu_textencoder=True,
)
benchmark(model, args.times, args.empty_cache)
================================================
FILE: iopaint/cli.py
================================================
import webbrowser
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional
import typer
from fastapi import FastAPI
from loguru import logger
from typer import Option
from typer_config import use_json_config
from iopaint.const import *
from iopaint.runtime import setup_model_dir, dump_environment_info, check_device
from iopaint.schema import InteractiveSegModel, Device, RealESRGANModel, RemoveBGModel
typer_app = typer.Typer(pretty_exceptions_show_locals=False, add_completion=False)
@typer_app.command(help="Install all plugins dependencies")
def install_plugins_packages():
from iopaint.installer import install_plugins_package
install_plugins_package()
@typer_app.command(help="Download SD/SDXL normal/inpainting model from HuggingFace")
def download(
model: str = Option(
..., help="Model id on HuggingFace e.g: runwayml/stable-diffusion-inpainting"
),
model_dir: Path = Option(
DEFAULT_MODEL_DIR,
help=MODEL_DIR_HELP,
file_okay=False,
callback=setup_model_dir,
),
):
from iopaint.download import cli_download_model
cli_download_model(model)
@typer_app.command(name="list", help="List downloaded models")
def list_model(
model_dir: Path = Option(
DEFAULT_MODEL_DIR,
help=MODEL_DIR_HELP,
file_okay=False,
callback=setup_model_dir,
),
):
from iopaint.download import scan_models
scanned_models = scan_models()
for it in scanned_models:
print(it.name)
@typer_app.command(help="Batch processing images")
def run(
model: str = Option("lama"),
device: Device = Option(Device.cpu),
image: Path = Option(..., help="Image folders or file path"),
mask: Path = Option(
...,
help="Mask folders or file path. "
"If it is a directory, the mask images in the directory should have the same name as the original image."
"If it is a file, all images will use this mask."
"Mask will automatically resize to the same size as the original image.",
),
output: Path = Option(..., help="Output directory or file path"),
config: Path = Option(
None, help="Config file path. You can use dump command to create a base config."
),
concat: bool = Option(
False, help="Concat original image, mask and output images into one image"
),
model_dir: Path = Option(
DEFAULT_MODEL_DIR,
help=MODEL_DIR_HELP,
file_okay=False,
callback=setup_model_dir,
),
):
from iopaint.download import cli_download_model, scan_models
scanned_models = scan_models()
if model not in [it.name for it in scanned_models]:
logger.info(f"{model} not found in {model_dir}, try to downloading")
cli_download_model(model)
from iopaint.batch_processing import batch_inpaint
batch_inpaint(model, device, image, mask, output, config, concat)
@typer_app.command(help="Start IOPaint server")
@use_json_config()
def start(
host: str = Option("127.0.0.1"),
port: int = Option(8080),
inbrowser: bool = Option(False, help=INBROWSER_HELP),
model: str = Option(
DEFAULT_MODEL,
help=f"Erase models: [{', '.join(AVAILABLE_MODELS)}].\n"
f"Diffusion models: [{', '.join(DIFFUSION_MODELS)}] or any SD/SDXL normal/inpainting models on HuggingFace.",
),
model_dir: Path = Option(
DEFAULT_MODEL_DIR,
help=MODEL_DIR_HELP,
dir_okay=True,
file_okay=False,
callback=setup_model_dir,
),
low_mem: bool = Option(False, help=LOW_MEM_HELP),
no_half: bool = Option(False, help=NO_HALF_HELP),
cpu_offload: bool = Option(False, help=CPU_OFFLOAD_HELP),
disable_nsfw_checker: bool = Option(False, help=DISABLE_NSFW_HELP),
cpu_textencoder: bool = Option(False, help=CPU_TEXTENCODER_HELP),
local_files_only: bool = Option(False, help=LOCAL_FILES_ONLY_HELP),
device: Device = Option(Device.cpu),
input: Optional[Path] = Option(None, help=INPUT_HELP),
mask_dir: Optional[Path] = Option(
None, help=MODEL_DIR_HELP, dir_okay=True, file_okay=False
),
output_dir: Optional[Path] = Option(
None, help=OUTPUT_DIR_HELP, dir_okay=True, file_okay=False
),
quality: int = Option(100, help=QUALITY_HELP),
enable_interactive_seg: bool = Option(False, help=INTERACTIVE_SEG_HELP),
interactive_seg_model: InteractiveSegModel = Option(
InteractiveSegModel.sam2_1_tiny, help=INTERACTIVE_SEG_MODEL_HELP
),
interactive_seg_device: Device = Option(Device.cpu),
enable_remove_bg: bool = Option(False, help=REMOVE_BG_HELP),
remove_bg_device: Device = Option(Device.cpu, help=REMOVE_BG_DEVICE_HELP),
remove_bg_model: RemoveBGModel = Option(RemoveBGModel.briaai_rmbg_1_4),
enable_anime_seg: bool = Option(False, help=ANIMESEG_HELP),
enable_realesrgan: bool = Option(False),
realesrgan_device: Device = Option(Device.cpu),
realesrgan_model: RealESRGANModel = Option(RealESRGANModel.realesr_general_x4v3),
enable_gfpgan: bool = Option(False),
gfpgan_device: Device = Option(Device.cpu),
enable_restoreformer: bool = Option(False),
restoreformer_device: Device = Option(Device.cpu),
):
dump_environment_info()
device = check_device(device)
remove_bg_device = check_device(remove_bg_device)
realesrgan_device = check_device(realesrgan_device)
gfpgan_device = check_device(gfpgan_device)
if input and not input.exists():
logger.error(f"invalid --input: {input} not exists")
exit(-1)
if mask_dir and not mask_dir.exists():
logger.error(f"invalid --mask-dir: {mask_dir} not exists")
exit(-1)
if input and input.is_dir() and not output_dir:
logger.error(
"invalid --output-dir: --output-dir must be set when --input is a directory"
)
exit(-1)
if output_dir:
output_dir = output_dir.expanduser().absolute()
logger.info(f"Image will be saved to {output_dir}")
if not output_dir.exists():
logger.info(f"Create output directory {output_dir}")
output_dir.mkdir(parents=True)
if mask_dir:
mask_dir = mask_dir.expanduser().absolute()
model_dir = model_dir.expanduser().absolute()
if local_files_only:
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"
from iopaint.download import cli_download_model, scan_models
scanned_models = scan_models()
if model not in [it.name for it in scanned_models]:
logger.info(f"{model} not found in {model_dir}, try to downloading")
cli_download_model(model)
from iopaint.api import Api
from iopaint.schema import ApiConfig
@asynccontextmanager
async def lifespan(app: FastAPI):
if inbrowser:
webbrowser.open(f"http://localhost:{port}", new=0, autoraise=True)
yield
app = FastAPI(lifespan=lifespan)
api_config = ApiConfig(
host=host,
port=port,
inbrowser=inbrowser,
model=model,
no_half=no_half,
low_mem=low_mem,
cpu_offload=cpu_offload,
disable_nsfw_checker=disable_nsfw_checker,
local_files_only=local_files_only,
cpu_textencoder=cpu_textencoder if device == Device.cuda else False,
device=device,
input=input,
mask_dir=mask_dir,
output_dir=output_dir,
quality=quality,
enable_interactive_seg=enable_interactive_seg,
interactive_seg_model=interactive_seg_model,
interactive_seg_device=interactive_seg_device,
enable_remove_bg=enable_remove_bg,
remove_bg_device=remove_bg_device,
remove_bg_model=remove_bg_model,
enable_anime_seg=enable_anime_seg,
enable_realesrgan=enable_realesrgan,
realesrgan_device=realesrgan_device,
realesrgan_model=realesrgan_model,
enable_gfpgan=enable_gfpgan,
gfpgan_device=gfpgan_device,
enable_restoreformer=enable_restoreformer,
restoreformer_device=restoreformer_device,
)
print(api_config.model_dump_json(indent=4))
api = Api(app, api_config)
api.launch()
@typer_app.command(help="Start IOPaint web config page")
def start_web_config(
config_file: Path = Option("config.json"),
):
dump_environment_info()
from iopaint.web_config import main
main(config_file)
================================================
FILE: iopaint/const.py
================================================
import os
from typing import List
INSTRUCT_PIX2PIX_NAME = "timbrooks/instruct-pix2pix"
KANDINSKY22_NAME = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
POWERPAINT_NAME = "Sanster/PowerPaint-V1-stable-diffusion-inpainting"
ANYTEXT_NAME = "Sanster/AnyText"
DIFFUSERS_SD_CLASS_NAME = "StableDiffusionPipeline"
DIFFUSERS_SD_INPAINT_CLASS_NAME = "StableDiffusionInpaintPipeline"
DIFFUSERS_SDXL_CLASS_NAME = "StableDiffusionXLPipeline"
DIFFUSERS_SDXL_INPAINT_CLASS_NAME = "StableDiffusionXLInpaintPipeline"
MPS_UNSUPPORT_MODELS = [
"lama",
"ldm",
"zits",
"mat",
"fcf",
"cv2",
"manga",
]
DEFAULT_MODEL = "lama"
AVAILABLE_MODELS = ["lama", "ldm", "zits", "mat", "fcf", "manga", "cv2", "migan"]
DIFFUSION_MODELS = [
"runwayml/stable-diffusion-inpainting",
"Uminosachi/realisticVisionV51_v51VAE-inpainting",
"redstonehero/dreamshaper-inpainting",
"Sanster/anything-4.0-inpainting",
"diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
"Fantasy-Studio/Paint-by-Example",
"RunDiffusion/Juggernaut-XI-v11",
"SG161222/RealVisXL_V5.0",
"eienmojiki/Anything-XL",
POWERPAINT_NAME,
ANYTEXT_NAME,
]
NO_HALF_HELP = """
Using full precision(fp32) model.
If your diffusion model generate result is always black or green, use this argument.
"""
CPU_OFFLOAD_HELP = """
Offloads diffusion model's weight to CPU RAM, significantly reducing vRAM usage.
"""
LOW_MEM_HELP = "Enable attention slicing and vae tiling to save memory."
DISABLE_NSFW_HELP = """
Disable NSFW checker for diffusion model.
"""
CPU_TEXTENCODER_HELP = """
Run diffusion models text encoder on CPU to reduce vRAM usage.
"""
SD_CONTROLNET_CHOICES: List[str] = [
"lllyasviel/control_v11p_sd15_canny",
# "lllyasviel/control_v11p_sd15_seg",
"lllyasviel/control_v11p_sd15_openpose",
"lllyasviel/control_v11p_sd15_inpaint",
"lllyasviel/control_v11f1p_sd15_depth",
]
SD_BRUSHNET_CHOICES: List[str] = [
"Sanster/brushnet_random_mask",
"Sanster/brushnet_segmentation_mask",
]
SD2_CONTROLNET_CHOICES = [
"thibaud/controlnet-sd21-canny-diffusers",
"thibaud/controlnet-sd21-depth-diffusers",
"thibaud/controlnet-sd21-openpose-diffusers",
]
SDXL_CONTROLNET_CHOICES = [
"thibaud/controlnet-openpose-sdxl-1.0",
"destitech/controlnet-inpaint-dreamer-sdxl",
"diffusers/controlnet-canny-sdxl-1.0",
"diffusers/controlnet-canny-sdxl-1.0-mid",
"diffusers/controlnet-canny-sdxl-1.0-small",
"diffusers/controlnet-depth-sdxl-1.0",
"diffusers/controlnet-depth-sdxl-1.0-mid",
"diffusers/controlnet-depth-sdxl-1.0-small",
]
SDXL_BRUSHNET_CHOICES = [
"Regulus0725/random_mask_brushnet_ckpt_sdxl_regulus_v1"
]
LOCAL_FILES_ONLY_HELP = """
When loading diffusion models, using local files only, not connect to HuggingFace server.
"""
DEFAULT_MODEL_DIR = os.path.abspath(
os.getenv("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache"))
)
MODEL_DIR_HELP = f"""
Model download directory (by setting XDG_CACHE_HOME environment variable), by default model download to {DEFAULT_MODEL_DIR}
"""
OUTPUT_DIR_HELP = """
Result images will be saved to output directory automatically.
"""
MASK_DIR_HELP = """
You can view masks in FileManager
"""
INPUT_HELP = """
If input is image, it will be loaded by default.
If input is directory, you can browse and select image in file manager.
"""
GUI_HELP = """
Launch Lama Cleaner as desktop app
"""
QUALITY_HELP = """
Quality of image encoding, 0-100. Default is 95, higher quality will generate larger file size.
"""
INTERACTIVE_SEG_HELP = "Enable interactive segmentation using Segment Anything."
INTERACTIVE_SEG_MODEL_HELP = "Model size: mobile_sam < vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed."
REMOVE_BG_HELP = "Enable remove background plugin."
REMOVE_BG_DEVICE_HELP = "Device for remove background plugin. 'cuda' only supports briaai models(briaai/RMBG-1.4 and briaai/RMBG-2.0)"
ANIMESEG_HELP = "Enable anime segmentation plugin. Always run on CPU"
REALESRGAN_HELP = "Enable realesrgan super resolution"
GFPGAN_HELP = "Enable GFPGAN face restore. To also enhance background, use with --enable-realesrgan"
RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To also enhance background, use with --enable-realesrgan"
GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image"
INBROWSER_HELP = "Automatically launch IOPaint in a new tab on the default browser"
================================================
FILE: iopaint/download.py
================================================
import glob
import json
import os
from functools import lru_cache
from typing import List, Optional
from iopaint.schema import ModelType, ModelInfo
from loguru import logger
from pathlib import Path
from iopaint.const import (
DEFAULT_MODEL_DIR,
DIFFUSERS_SD_CLASS_NAME,
DIFFUSERS_SD_INPAINT_CLASS_NAME,
DIFFUSERS_SDXL_CLASS_NAME,
DIFFUSERS_SDXL_INPAINT_CLASS_NAME,
ANYTEXT_NAME,
)
from iopaint.model.original_sd_configs import get_config_files
def cli_download_model(model: str):
from iopaint.model import models
from iopaint.model.utils import handle_from_pretrained_exceptions
if model in models and models[model].is_erase_model:
logger.info(f"Downloading {model}...")
models[model].download()
logger.info("Done.")
elif model == ANYTEXT_NAME:
logger.info(f"Downloading {model}...")
models[model].download()
logger.info("Done.")
else:
logger.info(f"Downloading model from Huggingface: {model}")
from diffusers import DiffusionPipeline
downloaded_path = handle_from_pretrained_exceptions(
DiffusionPipeline.download, pretrained_model_name=model, variant="fp16"
)
logger.info(f"Done. Downloaded to {downloaded_path}")
def folder_name_to_show_name(name: str) -> str:
return name.replace("models--", "").replace("--", "/")
@lru_cache(maxsize=512)
def get_sd_model_type(model_abs_path: str) -> Optional[ModelType]:
if "inpaint" in Path(model_abs_path).name.lower():
model_type = ModelType.DIFFUSERS_SD_INPAINT
else:
# load once to check num_in_channels
from diffusers import StableDiffusionInpaintPipeline
try:
StableDiffusionInpaintPipeline.from_single_file(
model_abs_path,
load_safety_checker=False,
num_in_channels=9,
original_config_file=get_config_files()["v1"],
)
model_type = ModelType.DIFFUSERS_SD_INPAINT
except ValueError as e:
if "[320, 4, 3, 3]" in str(e):
model_type = ModelType.DIFFUSERS_SD
else:
logger.info(f"Ignore non sdxl file: {model_abs_path}")
return
except Exception as e:
logger.error(f"Failed to load {model_abs_path}: {e}")
return
return model_type
@lru_cache()
def get_sdxl_model_type(model_abs_path: str) -> Optional[ModelType]:
if "inpaint" in model_abs_path:
model_type = ModelType.DIFFUSERS_SDXL_INPAINT
else:
# load once to check num_in_channels
from diffusers import StableDiffusionXLInpaintPipeline
try:
model = StableDiffusionXLInpaintPipeline.from_single_file(
model_abs_path,
load_safety_checker=False,
num_in_channels=9,
original_config_file=get_config_files()["xl"],
)
if model.unet.config.in_channels == 9:
# https://github.com/huggingface/diffusers/issues/6610
model_type = ModelType.DIFFUSERS_SDXL_INPAINT
else:
model_type = ModelType.DIFFUSERS_SDXL
except ValueError as e:
if "[320, 4, 3, 3]" in str(e):
model_type = ModelType.DIFFUSERS_SDXL
else:
logger.info(f"Ignore non sdxl file: {model_abs_path}")
return
except Exception as e:
logger.error(f"Failed to load {model_abs_path}: {e}")
return
return model_type
def scan_single_file_diffusion_models(cache_dir) -> List[ModelInfo]:
cache_dir = Path(cache_dir)
stable_diffusion_dir = cache_dir / "stable_diffusion"
cache_file = stable_diffusion_dir / "iopaint_cache.json"
model_type_cache = {}
if cache_file.exists():
try:
with open(cache_file, "r", encoding="utf-8") as f:
model_type_cache = json.load(f)
assert isinstance(model_type_cache, dict)
except:
pass
res = []
for it in stable_diffusion_dir.glob("*.*"):
if it.suffix not in [".safetensors", ".ckpt"]:
continue
model_abs_path = str(it.absolute())
model_type = model_type_cache.get(it.name)
if model_type is None:
model_type = get_sd_model_type(model_abs_path)
if model_type is None:
continue
model_type_cache[it.name] = model_type
res.append(
ModelInfo(
name=it.name,
path=model_abs_path,
model_type=model_type,
is_single_file_diffusers=True,
)
)
if stable_diffusion_dir.exists():
with open(cache_file, "w", encoding="utf-8") as fw:
json.dump(model_type_cache, fw, indent=2, ensure_ascii=False)
stable_diffusion_xl_dir = cache_dir / "stable_diffusion_xl"
sdxl_cache_file = stable_diffusion_xl_dir / "iopaint_cache.json"
sdxl_model_type_cache = {}
if sdxl_cache_file.exists():
try:
with open(sdxl_cache_file, "r", encoding="utf-8") as f:
sdxl_model_type_cache = json.load(f)
assert isinstance(sdxl_model_type_cache, dict)
except:
pass
for it in stable_diffusion_xl_dir.glob("*.*"):
if it.suffix not in [".safetensors", ".ckpt"]:
continue
model_abs_path = str(it.absolute())
model_type = sdxl_model_type_cache.get(it.name)
if model_type is None:
model_type = get_sdxl_model_type(model_abs_path)
if model_type is None:
continue
sdxl_model_type_cache[it.name] = model_type
if stable_diffusion_xl_dir.exists():
with open(sdxl_cache_file, "w", encoding="utf-8") as fw:
json.dump(sdxl_model_type_cache, fw, indent=2, ensure_ascii=False)
res.append(
ModelInfo(
name=it.name,
path=model_abs_path,
model_type=model_type,
is_single_file_diffusers=True,
)
)
return res
def scan_inpaint_models(model_dir: Path) -> List[ModelInfo]:
res = []
from iopaint.model import models
# logger.info(f"Scanning inpaint models in {model_dir}")
for name, m in models.items():
if m.is_erase_model and m.is_downloaded():
res.append(
ModelInfo(
name=name,
path=name,
model_type=ModelType.INPAINT,
)
)
return res
def scan_diffusers_models() -> List[ModelInfo]:
from huggingface_hub.constants import HF_HUB_CACHE
available_models = []
cache_dir = Path(HF_HUB_CACHE)
# logger.info(f"Scanning diffusers models in {cache_dir}")
diffusers_model_names = []
model_index_files = glob.glob(
os.path.join(cache_dir, "**/*", "model_index.json"), recursive=True
)
for it in model_index_files:
it = Path(it)
try:
with open(it, "r", encoding="utf-8") as f:
data = json.load(f)
except:
continue
_class_name = data["_class_name"]
name = folder_name_to_show_name(it.parent.parent.parent.name)
if name in diffusers_model_names:
continue
if "PowerPaint" in name:
model_type = ModelType.DIFFUSERS_OTHER
elif _class_name == DIFFUSERS_SD_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SD
elif _class_name == DIFFUSERS_SD_INPAINT_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SD_INPAINT
elif _class_name == DIFFUSERS_SDXL_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SDXL
elif _class_name == DIFFUSERS_SDXL_INPAINT_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SDXL_INPAINT
elif _class_name in [
"StableDiffusionInstructPix2PixPipeline",
"PaintByExamplePipeline",
"KandinskyV22InpaintPipeline",
"AnyText",
]:
model_type = ModelType.DIFFUSERS_OTHER
else:
continue
diffusers_model_names.append(name)
available_models.append(
ModelInfo(
name=name,
path=name,
model_type=model_type,
)
)
return available_models
def _scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
cache_dir = Path(cache_dir)
available_models = []
diffusers_model_names = []
model_index_files = glob.glob(
os.path.join(cache_dir, "**/*", "model_index.json"), recursive=True
)
for it in model_index_files:
it = Path(it)
with open(it, "r", encoding="utf-8") as f:
try:
data = json.load(f)
except:
logger.error(
f"Failed to load {it}, please try revert from original model or fix model_index.json by hand."
)
continue
_class_name = data["_class_name"]
name = folder_name_to_show_name(it.parent.name)
if name in diffusers_model_names:
continue
elif _class_name == DIFFUSERS_SD_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SD
elif _class_name == DIFFUSERS_SD_INPAINT_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SD_INPAINT
elif _class_name == DIFFUSERS_SDXL_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SDXL
elif _class_name == DIFFUSERS_SDXL_INPAINT_CLASS_NAME:
model_type = ModelType.DIFFUSERS_SDXL_INPAINT
else:
continue
diffusers_model_names.append(name)
available_models.append(
ModelInfo(
name=name,
path=str(it.parent.absolute()),
model_type=model_type,
)
)
return available_models
def scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
cache_dir = Path(cache_dir)
available_models = []
stable_diffusion_dir = cache_dir / "stable_diffusion"
stable_diffusion_xl_dir = cache_dir / "stable_diffusion_xl"
available_models.extend(_scan_converted_diffusers_models(stable_diffusion_dir))
available_models.extend(_scan_converted_diffusers_models(stable_diffusion_xl_dir))
return available_models
def scan_models() -> List[ModelInfo]:
model_dir = os.getenv("XDG_CACHE_HOME", DEFAULT_MODEL_DIR)
available_models = []
available_models.extend(scan_inpaint_models(model_dir))
available_models.extend(scan_single_file_diffusion_models(model_dir))
available_models.extend(scan_diffusers_models())
available_models.extend(scan_converted_diffusers_models(model_dir))
return available_models
================================================
FILE: iopaint/file_manager/__init__.py
================================================
from .file_manager import FileManager
================================================
FILE: iopaint/file_manager/file_manager.py
================================================
import os
from io import BytesIO
from pathlib import Path
from typing import List
from PIL import Image, ImageOps, PngImagePlugin
from fastapi import FastAPI, HTTPException
from starlette.responses import FileResponse
from ..schema import MediasResponse, MediaTab
LARGE_ENOUGH_NUMBER = 100
PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024**2)
from .storage_backends import FilesystemStorageBackend
from .utils import aspect_to_string, generate_filename, glob_img
class FileManager:
def __init__(self, app: FastAPI, input_dir: Path, mask_dir: Path, output_dir: Path):
self.app = app
self.input_dir: Path = input_dir
self.mask_dir: Path = mask_dir
self.output_dir: Path = output_dir
self.image_dir_filenames = []
self.output_dir_filenames = []
if not self.thumbnail_directory.exists():
self.thumbnail_directory.mkdir(parents=True)
# fmt: off
self.app.add_api_route("/api/v1/medias", self.api_medias, methods=["GET"], response_model=List[MediasResponse])
self.app.add_api_route("/api/v1/media_file", self.api_media_file, methods=["GET"])
self.app.add_api_route("/api/v1/media_thumbnail_file", self.api_media_thumbnail_file, methods=["GET"])
# fmt: on
def api_medias(self, tab: MediaTab) -> List[MediasResponse]:
img_dir = self._get_dir(tab)
return self._media_names(img_dir)
def api_media_file(self, tab: MediaTab, filename: str) -> FileResponse:
file_path = self._get_file(tab, filename)
return FileResponse(file_path, media_type="image/png")
# tab=${tab}?filename=${filename.name}?width=${width}&height=${height}
def api_media_thumbnail_file(
self, tab: MediaTab, filename: str, width: int, height: int
) -> FileResponse:
img_dir = self._get_dir(tab)
thumb_filename, (width, height) = self.get_thumbnail(
img_dir, filename, width=width, height=height
)
thumbnail_filepath = self.thumbnail_directory / thumb_filename
return FileResponse(
thumbnail_filepath,
headers={
"X-Width": str(width),
"X-Height": str(height),
},
media_type="image/jpeg",
)
def _get_dir(self, tab: MediaTab) -> Path:
if tab == "input":
return self.input_dir
elif tab == "output":
return self.output_dir
elif tab == "mask":
return self.mask_dir
else:
raise HTTPException(status_code=422, detail=f"tab not found: {tab}")
def _get_file(self, tab: MediaTab, filename: str) -> Path:
file_path = self._get_dir(tab) / filename
if not file_path.exists():
raise HTTPException(status_code=422, detail=f"file not found: {file_path}")
return file_path
@property
def thumbnail_directory(self) -> Path:
return self.output_dir / "thumbnails"
@staticmethod
def _media_names(directory: Path) -> List[MediasResponse]:
if directory is None:
return []
names = sorted([it.name for it in glob_img(directory)])
res = []
for name in names:
path = os.path.join(directory, name)
img = Image.open(path)
res.append(
MediasResponse(
name=name,
height=img.height,
width=img.width,
ctime=os.path.getctime(path),
mtime=os.path.getmtime(path),
)
)
return res
def get_thumbnail(
self, directory: Path, original_filename: str, width, height, **options
):
directory = Path(directory)
storage = FilesystemStorageBackend(self.app)
crop = options.get("crop", "fit")
background = options.get("background")
quality = options.get("quality", 90)
original_path, original_filename = os.path.split(original_filename)
original_filepath = os.path.join(directory, original_path, original_filename)
image = Image.open(BytesIO(storage.read(original_filepath)))
# keep ratio resize
if not width and not height:
width = 256
if width != 0:
height = int(image.height * width / image.width)
else:
width = int(image.width * height / image.height)
thumbnail_size = (width, height)
thumbnail_filename = generate_filename(
directory,
original_filename,
aspect_to_string(thumbnail_size),
crop,
background,
quality,
)
thumbnail_filepath = os.path.join(
self.thumbnail_directory, original_path, thumbnail_filename
)
if storage.exists(thumbnail_filepath):
return thumbnail_filepath, (width, height)
try:
image.load()
except (IOError, OSError):
self.app.logger.warning("Thumbnail not load image: %s", original_filepath)
return thumbnail_filepath, (width, height)
# get original image format
options["format"] = options.get("format", image.format)
image = self._create_thumbnail(
image, thumbnail_size, crop, background=background
)
raw_data = self.get_raw_data(image, **options)
storage.save(thumbnail_filepath, raw_data)
return thumbnail_filepath, (width, height)
def get_raw_data(self, image, **options):
data = {
"format": self._get_format(image, **options),
"quality": options.get("quality", 90),
}
_file = BytesIO()
image.save(_file, **data)
return _file.getvalue()
@staticmethod
def colormode(image, colormode="RGB"):
if colormode == "RGB" or colormode == "RGBA":
if image.mode == "RGBA":
return image
if image.mode == "LA":
return image.convert("RGBA")
return image.convert(colormode)
if colormode == "GRAY":
return image.convert("L")
return image.convert(colormode)
@staticmethod
def background(original_image, color=0xFF):
size = (max(original_image.size),) * 2
image = Image.new("L", size, color)
image.paste(
original_image,
tuple(map(lambda x: (x[0] - x[1]) / 2, zip(size, original_image.size))),
)
return image
def _get_format(self, image, **options):
if options.get("format"):
return options.get("format")
if image.format:
return image.format
return "JPEG"
def _create_thumbnail(self, image, size, crop="fit", background=None):
try:
resample = Image.Resampling.LANCZOS
except AttributeError: # pylint: disable=raise-missing-from
resample = Image.ANTIALIAS
if crop == "fit":
image = ImageOps.fit(image, size, resample)
else:
image = image.copy()
image.thumbnail(size, resample=resample)
if background is not None:
image = self.background(image)
image = self.colormode(image)
return image
================================================
FILE: iopaint/file_manager/storage_backends.py
================================================
# Copy from https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/storage_backends.py
import errno
import os
from abc import ABC, abstractmethod
class BaseStorageBackend(ABC):
def __init__(self, app=None):
self.app = app
@abstractmethod
def read(self, filepath, mode="rb", **kwargs):
raise NotImplementedError
@abstractmethod
def exists(self, filepath):
raise NotImplementedError
@abstractmethod
def save(self, filepath, data):
raise NotImplementedError
class FilesystemStorageBackend(BaseStorageBackend):
def read(self, filepath, mode="rb", **kwargs):
with open(filepath, mode) as f: # pylint: disable=unspecified-encoding
return f.read()
def exists(self, filepath):
return os.path.exists(filepath)
def save(self, filepath, data):
directory = os.path.dirname(filepath)
if not os.path.exists(directory):
try:
os.makedirs(directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
if not os.path.isdir(directory):
raise IOError("{} is not a directory".format(directory))
with open(filepath, "wb") as f:
f.write(data)
================================================
FILE: iopaint/file_manager/utils.py
================================================
# Copy from: https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/utils.py
import hashlib
from pathlib import Path
from typing import Union
def generate_filename(directory: Path, original_filename, *options) -> str:
text = str(directory.absolute()) + original_filename
for v in options:
text += "%s" % v
md5_hash = hashlib.md5()
md5_hash.update(text.encode("utf-8"))
return md5_hash.hexdigest() + ".jpg"
def parse_size(size):
if isinstance(size, int):
# If the size parameter is a single number, assume square aspect.
return [size, size]
if isinstance(size, (tuple, list)):
if len(size) == 1:
# If single value tuple/list is provided, exand it to two elements
return size + type(size)(size)
return size
try:
thumbnail_size = [int(x) for x in size.lower().split("x", 1)]
except ValueError:
raise ValueError( # pylint: disable=raise-missing-from
"Bad thumbnail size format. Valid format is INTxINT."
)
if len(thumbnail_size) == 1:
# If the size parameter only contains a single integer, assume square aspect.
thumbnail_size.append(thumbnail_size[0])
return thumbnail_size
def aspect_to_string(size):
if isinstance(size, str):
return size
return "x".join(map(str, size))
IMG_SUFFIX = {".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"}
def glob_img(p: Union[Path, str], recursive: bool = False):
p = Path(p)
if p.is_file() and p.suffix in IMG_SUFFIX:
yield p
else:
if recursive:
files = Path(p).glob("**/*.*")
else:
files = Path(p).glob("*.*")
for it in files:
if it.suffix not in IMG_SUFFIX:
continue
yield it
================================================
FILE: iopaint/helper.py
================================================
import base64
import imghdr
import io
import os
import sys
from typing import List, Optional, Dict, Tuple
from urllib.parse import urlparse
import cv2
from PIL import Image, ImageOps, PngImagePlugin
import numpy as np
import torch
from iopaint.const import MPS_UNSUPPORT_MODELS
from loguru import logger
from torch.hub import download_url_to_file, get_dir
import hashlib
def md5sum(filename):
md5 = hashlib.md5()
with open(filename, "rb") as f:
for chunk in iter(lambda: f.read(128 * md5.block_size), b""):
md5.update(chunk)
return md5.hexdigest()
def switch_mps_device(model_name, device):
if model_name in MPS_UNSUPPORT_MODELS and str(device) == "mps":
logger.info(f"{model_name} not support mps, switch to cpu")
return torch.device("cpu")
return device
def get_cache_path_by_url(url):
parts = urlparse(url)
hub_dir = get_dir()
model_dir = os.path.join(hub_dir, "checkpoints")
if not os.path.isdir(model_dir):
os.makedirs(model_dir)
filename = os.path.basename(parts.path)
cached_file = os.path.join(model_dir, filename)
return cached_file
def download_model(url, model_md5: str = None):
if os.path.exists(url):
cached_file = url
else:
cached_file = get_cache_path_by_url(url)
if not os.path.exists(cached_file):
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
hash_prefix = None
download_url_to_file(url, cached_file, hash_prefix, progress=True)
if model_md5:
_md5 = md5sum(cached_file)
if model_md5 == _md5:
logger.info(f"Download model success, md5: {_md5}")
else:
try:
os.remove(cached_file)
logger.error(
f"Model md5: {_md5}, expected md5: {model_md5}, wrong model deleted. Please restart iopaint."
f"If you still have errors, please try download model manually first https://lama-cleaner-docs.vercel.app/install/download_model_manually.\n"
)
except:
logger.error(
f"Model md5: {_md5}, expected md5: {model_md5}, please delete {cached_file} and restart iopaint."
)
exit(-1)
return cached_file
def ceil_modulo(x, mod):
if x % mod == 0:
return x
return (x // mod + 1) * mod
def handle_error(model_path, model_md5, e):
_md5 = md5sum(model_path)
if _md5 != model_md5:
try:
os.remove(model_path)
logger.error(
f"Model md5: {_md5}, expected md5: {model_md5}, wrong model deleted. Please restart iopaint."
f"If you still have errors, please try download model manually first https://lama-cleaner-docs.vercel.app/install/download_model_manually.\n"
)
except:
logger.error(
f"Model md5: {_md5}, expected md5: {model_md5}, please delete {model_path} and restart iopaint."
)
else:
logger.error(
f"Failed to load model {model_path},"
f"please submit an issue at https://github.com/Sanster/lama-cleaner/issues and include a screenshot of the error:\n{e}"
)
exit(-1)
def load_jit_model(url_or_path, device, model_md5: str):
if os.path.exists(url_or_path):
model_path = url_or_path
else:
model_path = download_model(url_or_path, model_md5)
logger.info(f"Loading model from: {model_path}")
try:
model = torch.jit.load(model_path, map_location="cpu").to(device)
except Exception as e:
handle_error(model_path, model_md5, e)
model.eval()
return model
def load_model(model: torch.nn.Module, url_or_path, device, model_md5):
if os.path.exists(url_or_path):
model_path = url_or_path
else:
model_path = download_model(url_or_path, model_md5)
try:
logger.info(f"Loading model from: {model_path}")
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict, strict=True)
model.to(device)
except Exception as e:
handle_error(model_path, model_md5, e)
model.eval()
return model
def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
data = cv2.imencode(
f".{ext}",
image_numpy,
[int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
)[1]
image_bytes = data.tobytes()
return image_bytes
def pil_to_bytes(pil_img, ext: str, quality: int = 95, infos={}) -> bytes:
with io.BytesIO() as output:
kwargs = {k: v for k, v in infos.items() if v is not None}
if ext == "jpg":
ext = "jpeg"
if "png" == ext.lower() and "parameters" in kwargs:
pnginfo_data = PngImagePlugin.PngInfo()
pnginfo_data.add_text("parameters", kwargs["parameters"])
kwargs["pnginfo"] = pnginfo_data
pil_img.save(output, format=ext, quality=quality, **kwargs)
image_bytes = output.getvalue()
return image_bytes
def load_img(img_bytes, gray: bool = False, return_info: bool = False):
alpha_channel = None
image = Image.open(io.BytesIO(img_bytes))
if return_info:
infos = image.info
try:
image = ImageOps.exif_transpose(image)
except:
pass
if gray:
image = image.convert("L")
np_img = np.array(image)
else:
if image.mode == "RGBA":
np_img = np.array(image)
alpha_channel = np_img[:, :, -1]
np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
else:
image = image.convert("RGB")
np_img = np.array(image)
if return_info:
return np_img, alpha_channel, infos
return np_img, alpha_channel
def norm_img(np_img):
if len(np_img.shape) == 2:
np_img = np_img[:, :, np.newaxis]
np_img = np.transpose(np_img, (2, 0, 1))
np_img = np_img.astype("float32") / 255
return np_img
def resize_max_size(
np_img, size_limit: int, interpolation=cv2.INTER_CUBIC
) -> np.ndarray:
# Resize image's longer size to size_limit if longer size larger than size_limit
h, w = np_img.shape[:2]
if max(h, w) > size_limit:
ratio = size_limit / max(h, w)
new_w = int(w * ratio + 0.5)
new_h = int(h * ratio + 0.5)
return cv2.resize(np_img, dsize=(new_w, new_h), interpolation=interpolation)
else:
return np_img
def pad_img_to_modulo(
img: np.ndarray, mod: int, square: bool = False, min_size: Optional[int] = None
):
"""
Args:
img: [H, W, C]
mod:
square: 是否为正方形
min_size:
Returns:
"""
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
height, width = img.shape[:2]
out_height = ceil_modulo(height, mod)
out_width = ceil_modulo(width, mod)
if min_size is not None:
assert min_size % mod == 0
out_width = max(min_size, out_width)
out_height = max(min_size, out_height)
if square:
max_size = max(out_height, out_width)
out_height = max_size
out_width = max_size
return np.pad(
img,
((0, out_height - height), (0, out_width - width), (0, 0)),
mode="symmetric",
)
def boxes_from_mask(mask: np.ndarray) -> List[np.ndarray]:
"""
Args:
mask: (h, w, 1) 0~255
Returns:
"""
height, width = mask.shape[:2]
_, thresh = cv2.threshold(mask, 127, 255, 0)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
boxes = []
for cnt in contours:
x, y, w, h = cv2.boundingRect(cnt)
box = np.array([x, y, x + w, y + h]).astype(int)
box[::2] = np.clip(box[::2], 0, width)
box[1::2] = np.clip(box[1::2], 0, height)
boxes.append(box)
return boxes
def only_keep_largest_contour(mask: np.ndarray) -> List[np.ndarray]:
"""
Args:
mask: (h, w) 0~255
Returns:
"""
_, thresh = cv2.threshold(mask, 127, 255, 0)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
max_area = 0
max_index = -1
for i, cnt in enumerate(contours):
area = cv2.contourArea(cnt)
if area > max_area:
max_area = area
max_index = i
if max_index != -1:
new_mask = np.zeros_like(mask)
return cv2.drawContours(new_mask, contours, max_index, 255, -1)
else:
return mask
def is_mac():
return sys.platform == "darwin"
def get_image_ext(img_bytes):
w = imghdr.what("", img_bytes)
if w is None:
w = "jpeg"
return w
def decode_base64_to_image(
encoding: str, gray=False
) -> Tuple[np.array, Optional[np.array], Dict, str]:
if encoding.startswith("data:image/") or encoding.startswith(
"data:application/octet-stream;base64,"
):
encoding = encoding.split(";")[1].split(",")[1]
image_bytes = base64.b64decode(encoding)
ext = get_image_ext(image_bytes)
image = Image.open(io.BytesIO(image_bytes))
alpha_channel = None
try:
image = ImageOps.exif_transpose(image)
except:
pass
# exif_transpose will remove exif rotate info，we must call image.info after exif_transpose
infos = image.info
if gray:
image = image.convert("L")
np_img = np.array(image)
else:
if image.mode == "RGBA":
np_img = np.array(image)
alpha_channel = np_img[:, :, -1]
np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
else:
image = image.convert("RGB")
np_img = np.array(image)
return np_img, alpha_channel, infos, ext
def encode_pil_to_base64(image: Image, quality: int, infos: Dict) -> bytes:
img_bytes = pil_to_bytes(
image,
"png",
quality=quality,
infos=infos,
)
return base64.b64encode(img_bytes)
def concat_alpha_channel(rgb_np_img, alpha_channel) -> np.ndarray:
if alpha_channel is not None:
if alpha_channel.shape[:2] != rgb_np_img.shape[:2]:
alpha_channel = cv2.resize(
alpha_channel, dsize=(rgb_np_img.shape[1], rgb_np_img.shape[0])
)
rgb_np_img = np.concatenate(
(rgb_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
)
return rgb_np_img
def adjust_mask(mask: np.ndarray, kernel_size: int, operate):
# fronted brush color "ffcc00bb"
# kernel_size = kernel_size*2+1
mask[mask >= 127] = 255
mask[mask < 127] = 0
if operate == "reverse":
mask = 255 - mask
else:
kernel = cv2.getStructuringElement(
cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
)
if operate == "expand":
mask = cv2.dilate(
mask,
kernel,
iterations=1,
)
else:
mask = cv2.erode(
mask,
kernel,
iterations=1,
)
res_mask = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
res_mask[mask > 128] = [255, 203, 0, int(255 * 0.73)]
res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
return res_mask
def gen_frontend_mask(bgr_or_gray_mask):
if len(bgr_or_gray_mask.shape) == 3 and bgr_or_gray_mask.shape[2] != 1:
bgr_or_gray_mask = cv2.cvtColor(bgr_or_gray_mask, cv2.COLOR_BGR2GRAY)
# fronted brush color "ffcc00bb"
# TODO: how to set kernel size?
kernel_size = 9
bgr_or_gray_mask = cv2.dilate(
bgr_or_gray_mask,
np.ones((kernel_size, kernel_size), np.uint8),
iterations=1,
)
res_mask = np.zeros(
(bgr_or_gray_mask.shape[0], bgr_or_gray_mask.shape[1], 4), dtype=np.uint8
)
res_mask[bgr_or_gray_mask > 128] = [255, 203, 0, int(255 * 0.73)]
res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
return res_mask
================================================
FILE: iopaint/installer.py
================================================
import subprocess
import sys
def install(package):
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
def install_plugins_package():
install("onnxruntime<=1.19.2")
install("rembg[cpu]")
================================================
FILE: iopaint/model/__init__.py
================================================
from .anytext.anytext_model import AnyText
from .controlnet import ControlNet
from .fcf import FcF
from .instruct_pix2pix import InstructPix2Pix
from .kandinsky import Kandinsky22
from .lama import LaMa, AnimeLaMa
from .ldm import LDM
from .manga import Manga
from .mat import MAT
from .mi_gan import MIGAN
from .opencv2 import OpenCV2
from .paint_by_example import PaintByExample
from .power_paint.power_paint import PowerPaint
from .sd import SD15, SD2, Anything4, RealisticVision14, SD
from .sdxl import SDXL
from .zits import ZITS
models = {
LaMa.name: LaMa,
AnimeLaMa.name: AnimeLaMa,
LDM.name: LDM,
ZITS.name: ZITS,
MAT.name: MAT,
FcF.name: FcF,
OpenCV2.name: OpenCV2,
Manga.name: Manga,
MIGAN.name: MIGAN,
SD15.name: SD15,
Anything4.name: Anything4,
RealisticVision14.name: RealisticVision14,
SD2.name: SD2,
PaintByExample.name: PaintByExample,
InstructPix2Pix.name: InstructPix2Pix,
Kandinsky22.name: Kandinsky22,
SDXL.name: SDXL,
PowerPaint.name: PowerPaint,
AnyText.name: AnyText,
}
================================================
FILE: iopaint/model/anytext/__init__.py
================================================
================================================
FILE: iopaint/model/anytext/anytext_model.py
================================================
import torch
from huggingface_hub import hf_hub_download
from iopaint.const import ANYTEXT_NAME
from iopaint.model.anytext.anytext_pipeline import AnyTextPipeline
from iopaint.model.base import DiffusionInpaintModel
from iopaint.model.utils import get_torch_dtype, is_local_files_only
from iopaint.schema import InpaintRequest
class AnyText(DiffusionInpaintModel):
name = ANYTEXT_NAME
pad_mod = 64
is_erase_model = False
@staticmethod
def download(local_files_only=False):
hf_hub_download(
repo_id=ANYTEXT_NAME,
filename="model_index.json",
local_files_only=local_files_only,
)
ckpt_path = hf_hub_download(
repo_id=ANYTEXT_NAME,
filename="pytorch_model.fp16.safetensors",
local_files_only=local_files_only,
)
font_path = hf_hub_download(
repo_id=ANYTEXT_NAME,
filename="SourceHanSansSC-Medium.otf",
local_files_only=local_files_only,
)
return ckpt_path, font_path
def init_model(self, device, **kwargs):
local_files_only = is_local_files_only(**kwargs)
ckpt_path, font_path = self.download(local_files_only)
use_gpu, torch_dtype = get_torch_dtype(device, kwargs.get("no_half", False))
self.model = AnyTextPipeline(
ckpt_path=ckpt_path,
font_path=font_path,
device=device,
use_fp16=torch_dtype == torch.float16,
)
self.callback = kwargs.pop("callback", None)
def forward(self, image, mask, config: InpaintRequest):
"""Input image and output image have same size
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to inpainting
return: BGR IMAGE
"""
height, width = image.shape[:2]
mask = mask.astype("float32") / 255.0
masked_image = image * (1 - mask)
# list of rgb ndarray
results, rtn_code, rtn_warning = self.model(
image=image,
masked_image=masked_image,
prompt=config.prompt,
negative_prompt=config.negative_prompt,
num_inference_steps=config.sd_steps,
strength=config.sd_strength,
guidance_scale=config.sd_guidance_scale,
height=height,
width=width,
seed=config.sd_seed,
sort_priority="y",
callback=self.callback
)
inpainted_rgb_image = results[0][..., ::-1]
return inpainted_rgb_image
================================================
FILE: iopaint/model/anytext/anytext_pipeline.py
================================================
"""
AnyText: Multilingual Visual Text Generation And Editing
Paper: https://arxiv.org/abs/2311.03054
Code: https://github.com/tyxsspa/AnyText
Copyright (c) Alibaba, Inc. and its affiliates.
"""
import os
from pathlib import Path
from iopaint.model.utils import set_seed
from safetensors.torch import load_file
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import torch
import re
import numpy as np
import cv2
import einops
from PIL import ImageFont
from iopaint.model.anytext.cldm.model import create_model, load_state_dict
from iopaint.model.anytext.cldm.ddim_hacked import DDIMSampler
from iopaint.model.anytext.utils import (
check_channels,
draw_glyph,
draw_glyph2,
)
BBOX_MAX_NUM = 8
PLACE_HOLDER = "*"
max_chars = 20
ANYTEXT_CFG = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "anytext_sd15.yaml"
)
def check_limits(tensor):
float16_min = torch.finfo(torch.float16).min
float16_max = torch.finfo(torch.float16).max
# 检查张量中是否有值小于float16的最小值或大于float16的最大值
is_below_min = (tensor < float16_min).any()
is_above_max = (tensor > float16_max).any()
return is_below_min or is_above_max
class AnyTextPipeline:
def __init__(self, ckpt_path, font_path, device, use_fp16=True):
self.cfg_path = ANYTEXT_CFG
self.font_path = font_path
self.use_fp16 = use_fp16
self.device = device
self.font = ImageFont.truetype(font_path, size=60)
self.model = create_model(
self.cfg_path,
device=self.device,
use_fp16=self.use_fp16,
)
if self.use_fp16:
self.model = self.model.half()
if Path(ckpt_path).suffix == ".safetensors":
state_dict = load_file(ckpt_path, device="cpu")
else:
state_dict = load_state_dict(ckpt_path, location="cpu")
self.model.load_state_dict(state_dict, strict=False)
self.model = self.model.eval().to(self.device)
self.ddim_sampler = DDIMSampler(self.model, device=self.device)
def __call__(
self,
prompt: str,
negative_prompt: str,
image: np.ndarray,
masked_image: np.ndarray,
num_inference_steps: int,
strength: float,
guidance_scale: float,
height: int,
width: int,
seed: int,
sort_priority: str = "y",
callback=None,
):
"""
Args:
prompt:
negative_prompt:
image:
masked_image:
num_inference_steps:
strength:
guidance_scale:
height:
width:
seed:
sort_priority: x: left-right, y: top-down
Returns:
result: list of images in numpy.ndarray format
rst_code: 0: normal -1: error 1:warning
rst_info: string of error or warning
"""
set_seed(seed)
str_warning = ""
mode = "text-editing"
revise_pos = False
img_count = 1
ddim_steps = num_inference_steps
w = width
h = height
strength = strength
cfg_scale = guidance_scale
eta = 0.0
prompt, texts = self.modify_prompt(prompt)
if prompt is None and texts is None:
return (
None,
-1,
"You have input Chinese prompt but the translator is not loaded!",
"",
)
n_lines = len(texts)
if mode in ["text-generation", "gen"]:
edit_image = np.ones((h, w, 3)) * 127.5 # empty mask image
elif mode in ["text-editing", "edit"]:
if masked_image is None or image is None:
return (
None,
-1,
"Reference image and position image are needed for text editing!",
"",
)
if isinstance(image, str):
image = cv2.imread(image)[..., ::-1]
assert image is not None, f"Can't read ori_image image from{image}!"
elif isinstance(image, torch.Tensor):
image = image.cpu().numpy()
else:
assert isinstance(
image, np.ndarray
), f"Unknown format of ori_image: {type(image)}"
edit_image = image.clip(1, 255) # for mask reason
edit_image = check_channels(edit_image)
# edit_image = resize_image(
# edit_image, max_length=768
# ) # make w h multiple of 64, resize if w or h > max_length
h, w = edit_image.shape[:2] # change h, w by input ref_img
# preprocess pos_imgs(if numpy, make sure it's white pos in black bg)
if masked_image is None:
pos_imgs = np.zeros((w, h, 1))
if isinstance(masked_image, str):
masked_image = cv2.imread(masked_image)[..., ::-1]
assert (
masked_image is not None
), f"Can't read draw_pos image from{masked_image}!"
pos_imgs = 255 - masked_image
elif isinstance(masked_image, torch.Tensor):
pos_imgs = masked_image.cpu().numpy()
else:
assert isinstance(
masked_image, np.ndarray
), f"Unknown format of draw_pos: {type(masked_image)}"
pos_imgs = 255 - masked_image
pos_imgs = pos_imgs[..., 0:1]
pos_imgs = cv2.convertScaleAbs(pos_imgs)
_, pos_imgs = cv2.threshold(pos_imgs, 254, 255, cv2.THRESH_BINARY)
# seprate pos_imgs
pos_imgs = self.separate_pos_imgs(pos_imgs, sort_priority)
if len(pos_imgs) == 0:
pos_imgs = [np.zeros((h, w, 1))]
if len(pos_imgs) < n_lines:
if n_lines == 1 and texts[0] == " ":
pass # text-to-image without text
else:
raise RuntimeError(
f"{n_lines} text line to draw from prompt, not enough mask area({len(pos_imgs)}) on images"
)
elif len(pos_imgs) > n_lines:
str_warning = f"Warning: found {len(pos_imgs)} positions that > needed {n_lines} from prompt."
# get pre_pos, poly_list, hint that needed for anytext
pre_pos = []
poly_list = []
for input_pos in pos_imgs:
if input_pos.mean() != 0:
input_pos = (
input_pos[..., np.newaxis]
if len(input_pos.shape) == 2
else input_pos
)
poly, pos_img = self.find_polygon(input_pos)
pre_pos += [pos_img / 255.0]
poly_list += [poly]
else:
pre_pos += [np.zeros((h, w, 1))]
poly_list += [None]
np_hint = np.sum(pre_pos, axis=0).clip(0, 1)
# prepare info dict
info = {}
info["glyphs"] = []
info["gly_line"] = []
info["positions"] = []
info["n_lines"] = [len(texts)] * img_count
gly_pos_imgs = []
for i in range(len(texts)):
text = texts[i]
if len(text) > max_chars:
str_warning = (
f'"{text}" length > max_chars: {max_chars}, will be cut off...'
)
text = text[:max_chars]
gly_scale = 2
if pre_pos[i].mean() != 0:
gly_line = draw_glyph(self.font, text)
glyphs = draw_glyph2(
self.font,
text,
poly_list[i],
scale=gly_scale,
width=w,
height=h,
add_space=False,
)
gly_pos_img = cv2.drawContours(
glyphs * 255, [poly_list[i] * gly_scale], 0, (255, 255, 255), 1
)
if revise_pos:
resize_gly = cv2.resize(
glyphs, (pre_pos[i].shape[1], pre_pos[i].shape[0])
)
new_pos = cv2.morphologyEx(
(resize_gly * 255).astype(np.uint8),
cv2.MORPH_CLOSE,
kernel=np.ones(
(resize_gly.shape[0] // 10, resize_gly.shape[1] // 10),
dtype=np.uint8,
),
iterations=1,
)
new_pos = (
new_pos[..., np.newaxis] if len(new_pos.shape) == 2 else new_pos
)
contours, _ = cv2.findContours(
new_pos, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
)
if len(contours) != 1:
str_warning = f"Fail to revise position {i} to bounding rect, remain position unchanged..."
else:
rect = cv2.minAreaRect(contours[0])
poly = np.int0(cv2.boxPoints(rect))
pre_pos[i] = (
cv2.drawContours(new_pos, [poly], -1, 255, -1) / 255.0
)
gly_pos_img = cv2.drawContours(
glyphs * 255, [poly * gly_scale], 0, (255, 255, 255), 1
)
gly_pos_imgs += [gly_pos_img] # for show
else:
glyphs = np.zeros((h * gly_scale, w * gly_scale, 1))
gly_line = np.zeros((80, 512, 1))
gly_pos_imgs += [
np.zeros((h * gly_scale, w * gly_scale, 1))
] # for show
pos = pre_pos[i]
info["glyphs"] += [self.arr2tensor(glyphs, img_count)]
info["gly_line"] += [self.arr2tensor(gly_line, img_count)]
info["positions"] += [self.arr2tensor(pos, img_count)]
# get masked_x
masked_img = ((edit_image.astype(np.float32) / 127.5) - 1.0) * (1 - np_hint)
masked_img = np.transpose(masked_img, (2, 0, 1))
masked_img = torch.from_numpy(masked_img.copy()).float().to(self.device)
if self.use_fp16:
masked_img = masked_img.half()
encoder_posterior = self.model.encode_first_stage(masked_img[None, ...])
masked_x = self.model.get_first_stage_encoding(encoder_posterior).detach()
if self.use_fp16:
masked_x = masked_x.half()
info["masked_x"] = torch.cat([masked_x for _ in range(img_count)], dim=0)
hint = self.arr2tensor(np_hint, img_count)
cond = self.model.get_learned_conditioning(
dict(
c_concat=[hint],
c_crossattn=[[prompt] * img_count],
text_info=info,
)
)
un_cond = self.model.get_learned_conditioning(
dict(
c_concat=[hint],
c_crossattn=[[negative_prompt] * img_count],
text_info=info,
)
)
shape = (4, h // 8, w // 8)
self.model.control_scales = [strength] * 13
samples, intermediates = self.ddim_sampler.sample(
ddim_steps,
img_count,
shape,
cond,
verbose=False,
eta=eta,
unconditional_guidance_scale=cfg_scale,
unconditional_conditioning=un_cond,
callback=callback
)
if self.use_fp16:
samples = samples.half()
x_samples = self.model.decode_first_stage(samples)
x_samples = (
(einops.rearrange(x_samples, "b c h w -> b h w c") * 127.5 + 127.5)
.cpu()
.numpy()
.clip(0, 255)
.astype(np.uint8)
)
results = [x_samples[i] for i in range(img_count)]
# if (
# mode == "edit" and False
# ): # replace backgound in text editing but not ideal yet
# results = [r * np_hint + edit_image * (1 - np_hint) for r in results]
# results = [r.clip(0, 255).astype(np.uint8) for r in results]
# if len(gly_pos_imgs) > 0 and show_debug:
# glyph_bs = np.stack(gly_pos_imgs, axis=2)
# glyph_img = np.sum(glyph_bs, axis=2) * 255
# glyph_img = glyph_img.clip(0, 255).astype(np.uint8)
# results += [np.repeat(glyph_img, 3, axis=2)]
rst_code = 1 if str_warning else 0
return results, rst_code, str_warning
def modify_prompt(self, prompt):
prompt = prompt.replace("“", '"')
prompt = prompt.replace("”", '"')
p = '"(.*?)"'
strs = re.findall(p, prompt)
if len(strs) == 0:
strs = [" "]
else:
for s in strs:
prompt = prompt.replace(f'"{s}"', f" {PLACE_HOLDER} ", 1)
# if self.is_chinese(prompt):
# if self.trans_pipe is None:
# return None, None
# old_prompt = prompt
# prompt = self.trans_pipe(input=prompt + " .")["translation"][:-1]
# print(f"Translate: {old_prompt} --> {prompt}")
return prompt, strs
# def is_chinese(self, text):
# text = checker._clean_text(text)
# for char in text:
# cp = ord(char)
# if checker._is_chinese_char(cp):
# return True
# return False
def separate_pos_imgs(self, img, sort_priority, gap=102):
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
components = []
for label in range(1, num_labels):
component = np.zeros_like(img)
component[labels == label] = 255
components.append((component, centroids[label]))
if sort_priority == "y":
fir, sec = 1, 0 # top-down first
elif sort_priority == "x":
fir, sec = 0, 1 # left-right first
components.sort(key=lambda c: (c[1][fir] // gap, c[1][sec] // gap))
sorted_components = [c[0] for c in components]
return sorted_components
def find_polygon(self, image, min_rect=False):
contours, hierarchy = cv2.findContours(
image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
)
max_contour = max(contours, key=cv2.contourArea) # get contour with max area
if min_rect:
# get minimum enclosing rectangle
rect = cv2.minAreaRect(max_contour)
poly = np.int0(cv2.boxPoints(rect))
else:
# get approximate polygon
epsilon = 0.01 * cv2.arcLength(max_contour, True)
poly = cv2.approxPolyDP(max_contour, epsilon, True)
n, _, xy = poly.shape
poly = poly.reshape(n, xy)
cv2.drawContours(image, [poly], -1, 255, -1)
return poly, image
def arr2tensor(self, arr, bs):
arr = np.transpose(arr, (2, 0, 1))
_arr = torch.from_numpy(arr.copy()).float().to(self.device)
if self.use_fp16:
_arr = _arr.half()
_arr = torch.stack([_arr for _ in range(bs)], dim=0)
return _arr
================================================
FILE: iopaint/model/anytext/anytext_sd15.yaml
================================================
model:
target: iopaint.model.anytext.cldm.cldm.ControlLDM
params:
linear_start: 0.00085
linear_end: 0.0120
num_timesteps_cond: 1
log_every_t: 200
timesteps: 1000
first_stage_key: "img"
cond_stage_key: "caption"
control_key: "hint"
glyph_key: "glyphs"
position_key: "positions"
image_size: 64
channels: 4
cond_stage_trainable: true # need be true when embedding_manager is valid
conditioning_key: crossattn
monitor: val/loss_simple_ema
scale_factor: 0.18215
use_ema: False
only_mid_control: False
loss_alpha: 0 # perceptual loss, 0.003
loss_beta: 0 # ctc loss
latin_weight: 1.0 # latin text line may need smaller weigth
with_step_weight: true
use_vae_upsample: true
embedding_manager_config:
target: iopaint.model.anytext.cldm.embedding_manager.EmbeddingManager
params:
valid: true # v6
emb_type: ocr # ocr, vit, conv
glyph_channels: 1
position_channels: 1
add_pos: false
placeholder_string: '*'
control_stage_config:
target: iopaint.model.anytext.cldm.cldm.ControlNet
params:
image_size: 32 # unused
in_channels: 4
model_channels: 320
glyph_channels: 1
position_channels: 1
attention_resolutions: [ 4, 2, 1 ]
num_res_blocks: 2
channel_mult: [ 1, 2, 4, 4 ]
num_heads: 8
use_spatial_transformer: True
transformer_depth: 1
context_dim: 768
use_checkpoint: True
legacy: False
unet_config:
target: iopaint.model.anytext.cldm.cldm.ControlledUnetModel
params:
image_size: 32 # unused
in_channels: 4
out_channels: 4
model_channels: 320
attention_resolutions: [ 4, 2, 1 ]
num_res_blocks: 2
channel_mult: [ 1, 2, 4, 4 ]
num_heads: 8
use_spatial_transformer: True
transformer_depth: 1
context_dim: 768
use_checkpoint: True
legacy: False
first_stage_config:
target: iopaint.model.anytext.ldm.models.autoencoder.AutoencoderKL
params:
embed_dim: 4
monitor: val/rec_loss
ddconfig:
double_z: true
z_channels: 4
resolution: 256
in_channels: 3
out_ch: 3
ch: 128
ch_mult:
- 1
- 2
- 4
- 4
num_res_blocks: 2
attn_resolutions: []
dropout: 0.0
lossconfig:
target: torch.nn.Identity
cond_stage_config:
target: iopaint.model.anytext.ldm.modules.encoders.modules.FrozenCLIPEmbedderT3
params:
version: openai/clip-vit-large-patch14
use_vision: false # v6
================================================
FILE: iopaint/model/anytext/cldm/__init__.py
================================================
================================================
FILE: iopaint/model/anytext/cldm/cldm.py
================================================
import os
from pathlib import Path
import einops
import torch
import torch as th
import torch.nn as nn
import copy
from easydict import EasyDict as edict
from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
conv_nd,
linear,
zero_module,
timestep_embedding,
)
from einops import rearrange, repeat
from iopaint.model.anytext.ldm.modules.attention import SpatialTransformer
from iopaint.model.anytext.ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
from iopaint.model.anytext.ldm.models.diffusion.ddpm import LatentDiffusion
from iopaint.model.anytext.ldm.util import log_txt_as_img, exists, instantiate_from_config
from iopaint.model.anytext.ldm.models.diffusion.ddim import DDIMSampler
from iopaint.model.anytext.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
from .recognizer import TextRecognizer, create_predictor
CURRENT_DIR = Path(os.path.dirname(os.path.abspath(__file__)))
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
class ControlledUnetModel(UNetModel):
def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs):
hs = []
with torch.no_grad():
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
if self.use_fp16:
t_emb = t_emb.half()
emb = self.time_embed(t_emb)
h = x.type(self.dtype)
for module in self.input_blocks:
h = module(h, emb, context)
hs.append(h)
h = self.middle_block(h, emb, context)
if control is not None:
h += control.pop()
for i, module in enumerate(self.output_blocks):
if only_mid_control or control is None:
h = torch.cat([h, hs.pop()], dim=1)
else:
h = torch.cat([h, hs.pop() + control.pop()], dim=1)
h = module(h, emb, context)
h = h.type(x.dtype)
return self.out(h)
class ControlNet(nn.Module):
def __init__(
self,
image_size,
in_channels,
model_channels,
glyph_channels,
position_channels,
num_res_blocks,
attention_resolutions,
dropout=0,
channel_mult=(1, 2, 4, 8),
conv_resample=True,
dims=2,
use_checkpoint=False,
use_fp16=False,
num_heads=-1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
use_spatial_transformer=False, # custom transformer support
transformer_depth=1, # custom transformer support
context_dim=None, # custom transformer support
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
legacy=True,
disable_self_attentions=None,
num_attention_blocks=None,
disable_middle_self_attn=False,
use_linear_in_transformer=False,
):
super().__init__()
if use_spatial_transformer:
assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
if context_dim is not None:
assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
from omegaconf.listconfig import ListConfig
if type(context_dim) == ListConfig:
context_dim = list(context_dim)
if num_heads_upsample == -1:
num_heads_upsample = num_heads
if num_heads == -1:
assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
if num_head_channels == -1:
assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
self.dims = dims
self.image_size = image_size
self.in_channels = in_channels
self.model_channels = model_channels
if isinstance(num_res_blocks, int):
self.num_res_blocks = len(channel_mult) * [num_res_blocks]
else:
if len(num_res_blocks) != len(channel_mult):
raise ValueError("provide num_res_blocks either as an int (globally constant) or "
"as a list/tuple (per-level) with the same length as channel_mult")
self.num_res_blocks = num_res_blocks
if disable_self_attentions is not None:
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
assert len(disable_self_attentions) == len(channel_mult)
if num_attention_blocks is not None:
assert len(num_attention_blocks) == len(self.num_res_blocks)
assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
f"attention will still not be set.")
self.attention_resolutions = attention_resolutions
self.dropout = dropout
self.channel_mult = channel_mult
self.conv_resample = conv_resample
self.use_checkpoint = use_checkpoint
self.use_fp16 = use_fp16
self.dtype = th.float16 if use_fp16 else th.float32
self.num_heads = num_heads
self.num_head_channels = num_head_channels
self.num_heads_upsample = num_heads_upsample
self.predict_codebook_ids = n_embed is not None
time_embed_dim = model_channels * 4
self.time_embed = nn.Sequential(
linear(model_channels, time_embed_dim),
nn.SiLU(),
linear(time_embed_dim, time_embed_dim),
)
self.input_blocks = nn.ModuleList(
[
TimestepEmbedSequential(
conv_nd(dims, in_channels, model_channels, 3, padding=1)
)
]
)
self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])
self.glyph_block = TimestepEmbedSequential(
conv_nd(dims, glyph_channels, 8, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 8, 8, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 8, 16, 3, padding=1, stride=2),
nn.SiLU(),
conv_nd(dims, 16, 16, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 16, 32, 3, padding=1, stride=2),
nn.SiLU(),
conv_nd(dims, 32, 32, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 32, 96, 3, padding=1, stride=2),
nn.SiLU(),
conv_nd(dims, 96, 96, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 96, 256, 3, padding=1, stride=2),
nn.SiLU(),
)
self.position_block = TimestepEmbedSequential(
conv_nd(dims, position_channels, 8, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 8, 8, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 8, 16, 3, padding=1, stride=2),
nn.SiLU(),
conv_nd(dims, 16, 16, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 16, 32, 3, padding=1, stride=2),
nn.SiLU(),
conv_nd(dims, 32, 32, 3, padding=1),
nn.SiLU(),
conv_nd(dims, 32, 64, 3, padding=1, stride=2),
nn.SiLU(),
)
self.fuse_block = zero_module(conv_nd(dims, 256+64+4, model_channels, 3, padding=1))
self._feature_size = model_channels
input_block_chans = [model_channels]
ch = model_channels
ds = 1
for level, mult in enumerate(channel_mult):
for nr in range(self.num_res_blocks[level]):
layers = [
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=mult * model_channels,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
)
]
ch = mult * model_channels
if ds in attention_resolutions:
if num_head_channels == -1:
dim_head = ch // num_heads
else:
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
# num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
if exists(disable_self_attentions):
disabled_sa = disable_self_attentions[level]
else:
disabled_sa = False
if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
layers.append(
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=dim_head,
use_new_attention_order=use_new_attention_order,
) if not use_spatial_transformer else SpatialTransformer(
ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
use_checkpoint=use_checkpoint
)
)
self.input_blocks.append(TimestepEmbedSequential(*layers))
self.zero_convs.append(self.make_zero_conv(ch))
self._feature_size += ch
input_block_chans.append(ch)
if level != len(channel_mult) - 1:
out_ch = ch
self.input_blocks.append(
TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
out_channels=out_ch,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
down=True,
)
if resblock_updown
else Downsample(
ch, conv_resample, dims=dims, out_channels=out_ch
)
)
)
ch = out_ch
input_block_chans.append(ch)
self.zero_convs.append(self.make_zero_conv(ch))
ds *= 2
self._feature_size += ch
if num_head_channels == -1:
dim_head = ch // num_heads
else:
num_heads = ch // num_head_channels
dim_head = num_head_channels
if legacy:
# num_heads = 1
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
self.middle_block = TimestepEmbedSequential(
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
AttentionBlock(
ch,
use_checkpoint=use_checkpoint,
num_heads=num_heads,
num_head_channels=dim_head,
use_new_attention_order=use_new_attention_order,
) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
use_checkpoint=use_checkpoint
),
ResBlock(
ch,
time_embed_dim,
dropout,
dims=dims,
use_checkpoint=use_checkpoint,
use_scale_shift_norm=use_scale_shift_norm,
),
)
self.middle_block_out = self.make_zero_conv(ch)
self._feature_size += ch
def make_zero_conv(self, channels):
return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))
def forward(self, x, hint, text_info, timesteps, context, **kwargs):
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
if self.use_fp16:
t_emb = t_emb.half()
emb = self.time_embed(t_emb)
# guided_hint from text_info
B, C, H, W = x.shape
glyphs = torch.cat(text_info['glyphs'], dim=1).sum(dim=1, keepdim=True)
positions = torch.cat(text_info['positions'], dim=1).sum(dim=1, keepdim=True)
enc_glyph = self.glyph_block(glyphs, emb, context)
enc_pos = self.position_block(positions, emb, context)
guided_hint = self.fuse_block(torch.cat([enc_glyph, enc_pos, text_info['masked_x']], dim=1))
outs = []
h = x.type(self.dtype)
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
if guided_hint is not None:
h = module(h, emb, context)
h += guided_hint
guided_hint = None
else:
h = module(h, emb, context)
outs.append(zero_conv(h, emb, context))
h = self.middle_block(h, emb, context)
outs.append(self.middle_block_out(h, emb, context))
return outs
class ControlLDM(LatentDiffusion):
def __init__(self, control_stage_config, control_key, glyph_key, position_key, only_mid_control, loss_alpha=0, loss_beta=0, with_step_weight=False, use_vae_upsample=False, latin_weight=1.0, embedding_manager_config=None, *args, **kwargs):
self.use_fp16 = kwargs.pop('use_fp16', False)
super().__init__(*args, **kwargs)
self.control_model = instantiate_from_config(control_stage_config)
self.control_key = control_key
self.glyph_key = glyph_key
self.position_key = position_key
self.only_mid_control = only_mid_control
self.control_scales = [1.0] * 13
self.loss_alpha = loss_alpha
self.loss_beta = loss_beta
self.with_step_weight = with_step_weight
self.use_vae_upsample = use_vae_upsample
self.latin_weight = latin_weight
if embedding_manager_config is not None and embedding_manager_config.params.valid:
self.embedding_manager = self.instantiate_embedding_manager(embedding_manager_config, self.cond_stage_model)
for param in self.embedding_manager.embedding_parameters():
param.requires_grad = True
else:
self.embedding_manager = None
if self.loss_alpha > 0 or self.loss_beta > 0 or self.embedding_manager:
if embedding_manager_config.params.emb_type == 'ocr':
self.text_predictor = create_predictor().eval()
args = edict()
args.rec_image_shape = "3, 48, 320"
args.rec_batch_num = 6
args.rec_char_dict_path = str(CURRENT_DIR.parent / "ocr_recog" / "ppocr_keys_v1.txt")
args.use_fp16 = self.use_fp16
self.cn_recognizer = TextRecognizer(args, self.text_predictor)
for param in self.text_predictor.parameters():
param.requires_grad = False
if self.embedding_manager:
self.embedding_manager.recog = self.cn_recognizer
@torch.no_grad()
def get_input(self, batch, k, bs=None, *args, **kwargs):
if self.embedding_manager is None: # fill in full caption
self.fill_caption(batch)
x, c, mx = super().get_input(batch, self.first_stage_key, mask_k='masked_img', *args, **kwargs)
control = batch[self.control_key] # for log_images and loss_alpha, not real control
if bs is not None:
control = control[:bs]
control = control.to(self.device)
control = einops.rearrange(control, 'b h w c -> b c h w')
control = control.to(memory_format=torch.contiguous_format).float()
inv_mask = batch['inv_mask']
if bs is not None:
inv_mask = inv_mask[:bs]
inv_mask = inv_mask.to(self.device)
inv_mask = einops.rearrange(inv_mask, 'b h w c -> b c h w')
inv_mask = inv_mask.to(memory_format=torch.contiguous_format).float()
glyphs = batch[self.glyph_key]
gly_line = batch['gly_line']
positions = batch[self.position_key]
n_lines = batch['n_lines']
language = batch['language']
texts = batch['texts']
assert len(glyphs) == len(positions)
for i in range(len(glyphs)):
if bs is not None:
glyphs[i] = glyphs[i][:bs]
gly_line[i] = gly_line[i][:bs]
positions[i] = positions[i][:bs]
n_lines = n_lines[:bs]
glyphs[i] = glyphs[i].to(self.device)
gly_line[i] = gly_line[i].to(self.device)
positions[i] = positions[i].to(self.device)
glyphs[i] = einops.rearrange(glyphs[i], 'b h w c -> b c h w')
gly_line[i] = einops.rearrange(gly_line[i], 'b h w c -> b c h w')
positions[i] = einops.rearrange(positions[i], 'b h w c -> b c h w')
glyphs[i] = glyphs[i].to(memory_format=torch.contiguous_format).float()
gly_line[i] = gly_line[i].to(memory_format=torch.contiguous_format).float()
positions[i] = positions[i].to(memory_format=torch.contiguous_format).float()
info = {}
info['glyphs'] = glyphs
info['positions'] = positions
info['n_lines'] = n_lines
info['language'] = language
info['texts'] = texts
info['img'] = batch['img'] # nhwc, (-1,1)
info['masked_x'] = mx
info['gly_line'] = gly_line
info['inv_mask'] = inv_mask
return x, dict(c_crossattn=[c], c_concat=[control], text_info=info)
def apply_model(self, x_noisy, t, cond, *args, **kwargs):
assert isinstance(cond, dict)
diffusion_model = self.model.diffusion_model
_cond = torch.cat(cond['c_crossattn'], 1)
_hint = torch.cat(cond['c_concat'], 1)
if self.use_fp16:
x_noisy = x_noisy.half()
control = self.control_model(x=x_noisy, timesteps=t, context=_cond, hint=_hint, text_info=cond['text_info'])
control = [c * scale for c, scale in zip(control, self.control_scales)]
eps = diffusion_model(x=x_noisy, timesteps=t, context=_cond, control=control, only_mid_control=self.only_mid_control)
return eps
def instantiate_embedding_manager(self, config, embedder):
model = instantiate_from_config(config, embedder=embedder)
return model
@torch.no_grad()
def get_unconditional_conditioning(self, N):
return self.get_learned_conditioning(dict(c_crossattn=[[""] * N], text_info=None))
def get_learned_conditioning(self, c):
if self.cond_stage_forward is None:
if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
if self.embedding_manager is not None and c['text_info'] is not None:
self.embedding_manager.encode_text(c['text_info'])
if isinstance(c, dict):
cond_txt = c['c_crossattn'][0]
else:
cond_txt = c
if self.embedding_manager is not None:
cond_txt = self.cond_stage_model.encode(cond_txt, embedding_manager=self.embedding_manager)
else:
cond_txt = self.cond_stage_model.encode(cond_txt)
if isinstance(c, dict):
c['c_crossattn'][0] = cond_txt
else:
c = cond_txt
if isinstance(c, DiagonalGaussianDistribution):
c = c.mode()
else:
c = self.cond_stage_model(c)
else:
assert hasattr(self.cond_stage_model, self.cond_stage_forward)
c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
return c
def fill_caption(self, batch, place_holder='*'):
bs = len(batch['n_lines'])
cond_list = copy.deepcopy(batch[self.cond_stage_key])
for i in range(bs):
n_lines = batch['n_lines'][i]
if n_lines == 0:
continue
cur_cap = cond_list[i]
for j in range(n_lines):
r_txt = batch['texts'][j][i]
cur_cap = cur_cap.replace(place_holder, f'"{r_txt}"', 1)
cond_list[i] = cur_cap
batch[self.cond_stage_key] = cond_list
@torch.no_grad()
def log_images(self, batch, N=4, n_row=2, sample=False, ddim_steps=50, ddim_eta=0.0, return_keys=None,
quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
plot_diffusion_rows=False, unconditional_guidance_scale=9.0, unconditional_guidance_label=None,
use_ema_scope=True,
**kwargs):
use_ddim = ddim_steps is not None
log = dict()
z, c = self.get_input(batch, self.first_stage_key, bs=N)
if self.cond_stage_trainable:
with torch.no_grad():
c = self.get_learned_conditioning(c)
c_crossattn = c["c_crossattn"][0][:N]
c_cat = c["c_concat"][0][:N]
text_info = c["text_info"]
text_info['glyphs'] = [i[:N] for i in text_info['glyphs']]
text_info['gly_line'] = [i[:N] for i in text_info['gly_line']]
text_info['positions'] = [i[:N] for i in text_info['positions']]
text_info['n_lines'] = text_info['n_lines'][:N]
text_info['masked_x'] = text_info['masked_x'][:N]
text_info['img'] = text_info['img'][:N]
N = min(z.shape[0], N)
n_row = min(z.shape[0], n_row)
log["reconstruction"] = self.decode_first_stage(z)
log["masked_image"] = self.decode_first_stage(text_info['masked_x'])
log["control"] = c_cat * 2.0 - 1.0
log["img"] = text_info['img'].permute(0, 3, 1, 2) # log source image if needed
# get glyph
glyph_bs = torch.stack(text_info['glyphs'])
glyph_bs = torch.sum(glyph_bs, dim=0) * 2.0 - 1.0
log["glyph"] = torch.nn.functional.interpolate(glyph_bs, size=(512, 512), mode='bilinear', align_corners=True,)
# fill caption
if not self.embedding_manager:
self.fill_caption(batch)
captions = batch[self.cond_stage_key]
log["conditioning"] = log_txt_as_img((512, 512), captions, size=16)
if plot_diffusion_rows:
# get diffusion row
diffusion_row = list()
z_start = z[:n_row]
for t in range(self.num_timesteps):
if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
t = t.to(self.device).long()
noise = torch.randn_like(z_start)
z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
diffusion_row.append(self.decode_first_stage(z_noisy))
diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
log["diffusion_row"] = diffusion_grid
if sample:
# get denoise row
samples, z_denoise_row = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c], "text_info": text_info},
batch_size=N, ddim=use_ddim,
ddim_steps=ddim_steps, eta=ddim_eta)
x_samples = self.decode_first_stage(samples)
log["samples"] = x_samples
if plot_denoise_rows:
denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
log["denoise_row"] = denoise_grid
if unconditional_guidance_scale > 1.0:
uc_cross = self.get_unconditional_conditioning(N)
uc_cat = c_cat # torch.zeros_like(c_cat)
uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross['c_crossattn'][0]], "text_info": text_info}
samples_cfg, tmps = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c_crossattn], "text_info": text_info},
batch_size=N, ddim=use_ddim,
ddim_steps=ddim_steps, eta=ddim_eta,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=uc_full,
)
x_samples_cfg = self.decode_first_stage(samples_cfg)
log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
pred_x0 = False # wether log pred_x0
if pred_x0:
for idx in range(len(tmps['pred_x0'])):
pred_x0 = self.decode_first_stage(tmps['pred_x0'][idx])
log[f"pred_x0_{tmps['index'][idx]}"] = pred_x0
return log
@torch.no_grad()
def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
ddim_sampler = DDIMSampler(self)
b, c, h, w = cond["c_concat"][0].shape
shape = (self.channels, h // 8, w // 8)
samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size, shape, cond, verbose=False, log_every_t=5, **kwargs)
return samples, intermediates
def configure_optimizers(self):
lr = self.learning_rate
params = list(self.control_model.parameters())
if self.embedding_manager:
params += list(self.embedding_manager.embedding_parameters())
if not self.sd_locked:
# params += list(self.model.diffusion_model.input_blocks.parameters())
# params += list(self.model.diffusion_model.middle_block.parameters())
params += list(self.model.diffusion_model.output_blocks.parameters())
params += list(self.model.diffusion_model.out.parameters())
if self.unlockKV:
nCount = 0
for name, param in self.model.diffusion_model.named_parameters():
if 'attn2.to_k' in name or 'attn2.to_v' in name:
params += [param]
nCount += 1
print(f'Cross attention is unlocked, and {nCount} Wk or Wv are added to potimizers!!!')
opt = torch.optim.AdamW(params, lr=lr)
return opt
def low_vram_shift(self, is_diffusing):
if is_diffusing:
self.model = self.model.cuda()
self.control_model = self.control_model.cuda()
self.first_stage_model = self.first_stage_model.cpu()
self.cond_stage_model = self.cond_stage_model.cpu()
else:
self.model = self.model.cpu()
self.control_model = self.control_model.cpu()
self.first_stage_model = self.first_stage_model.cuda()
self.cond_stage_model = self.cond_stage_model.cuda()
================================================
FILE: iopaint/model/anytext/cldm/ddim_hacked.py
================================================
"""SAMPLING ONLY."""
import torch
import numpy as np
from tqdm import tqdm
from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
make_ddim_sampling_parameters,
make_ddim_timesteps,
noise_like,
extract_into_tensor,
)
class DDIMSampler(object):
def __init__(self, model, device, schedule="linear", **kwargs):
super().__init__()
self.device = device
self.model = model
self.ddpm_num_timesteps = model.num_timesteps
self.schedule = schedule
def register_buffer(self, name, attr):
if type(attr) == torch.Tensor:
if attr.device != torch.device(self.device):
attr = attr.to(torch.device(self.device))
setattr(self, name, attr)
def make_schedule(
self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0.0, verbose=True
):
self.ddim_timesteps = make_ddim_timesteps(
ddim_discr_method=ddim_discretize,
num_ddim_timesteps=ddim_num_steps,
num_ddpm_timesteps=self.ddpm_num_timesteps,
verbose=verbose,
)
alphas_cumprod = self.model.alphas_cumprod
assert (
alphas_cumprod.shape[0] == self.ddpm_num_timesteps
), "alphas have to be defined for each timestep"
to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.device)
self.register_buffer("betas", to_torch(self.model.betas))
self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
self.register_buffer(
"alphas_cumprod_prev", to_torch(self.model.alphas_cumprod_prev)
)
# calculations for diffusion q(x_t | x_{t-1}) and others
self.register_buffer(
"sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod.cpu()))
)
self.register_buffer(
"sqrt_one_minus_alphas_cumprod",
to_torch(np.sqrt(1.0 - alphas_cumprod.cpu())),
)
self.register_buffer(
"log_one_minus_alphas_cumprod", to_torch(np.log(1.0 - alphas_cumprod.cpu()))
)
self.register_buffer(
"sqrt_recip_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod.cpu()))
)
self.register_buffer(
"sqrt_recipm1_alphas_cumprod",
to_torch(np.sqrt(1.0 / alphas_cumprod.cpu() - 1)),
)
# ddim sampling parameters
ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(
alphacums=alphas_cumprod.cpu(),
ddim_timesteps=self.ddim_timesteps,
eta=ddim_eta,
verbose=verbose,
)
self.register_buffer("ddim_sigmas", ddim_sigmas)
self.register_buffer("ddim_alphas", ddim_alphas)
self.register_buffer("ddim_alphas_prev", ddim_alphas_prev)
self.register_buffer("ddim_sqrt_one_minus_alphas", np.sqrt(1.0 - ddim_alphas))
sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
(1 - self.alphas_cumprod_prev)
/ (1 - self.alphas_cumprod)
* (1 - self.alphas_cumprod / self.alphas_cumprod_prev)
)
self.register_buffer(
"ddim_sigmas_for_original_num_steps", sigmas_for_original_sampling_steps
)
@torch.no_grad()
def sample(
self,
S,
batch_size,
shape,
conditioning=None,
callback=None,
normals_sequence=None,
img_callback=None,
quantize_x0=False,
eta=0.0,
mask=None,
x0=None,
temperature=1.0,
noise_dropout=0.0,
score_corrector=None,
corrector_kwargs=None,
verbose=True,
x_T=None,
log_every_t=100,
unconditional_guidance_scale=1.0,
unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
dynamic_threshold=None,
ucg_schedule=None,
**kwargs,
):
if conditioning is not None:
if isinstance(conditioning, dict):
ctmp = conditioning[list(conditioning.keys())[0]]
while isinstance(ctmp, list):
ctmp = ctmp[0]
cbs = ctmp.shape[0]
if cbs != batch_size:
print(
f"Warning: Got {cbs} conditionings but batch-size is {batch_size}"
)
elif isinstance(conditioning, list):
for ctmp in conditioning:
if ctmp.shape[0] != batch_size:
print(
f"Warning: Got {cbs} conditionings but batch-size is {batch_size}"
)
else:
if conditioning.shape[0] != batch_size:
print(
f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}"
)
self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
# sampling
C, H, W = shape
size = (batch_size, C, H, W)
print(f"Data shape for DDIM sampling is {size}, eta {eta}")
samples, intermediates = self.ddim_sampling(
conditioning,
size,
callback=callback,
img_callback=img_callback,
quantize_denoised=quantize_x0,
mask=mask,
x0=x0,
ddim_use_original_steps=False,
noise_dropout=noise_dropout,
temperature=temperature,
score_corrector=score_corrector,
corrector_kwargs=corrector_kwargs,
x_T=x_T,
log_every_t=log_every_t,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=unconditional_conditioning,
dynamic_threshold=dynamic_threshold,
ucg_schedule=ucg_schedule,
)
return samples, intermediates
@torch.no_grad()
def ddim_sampling(
self,
cond,
shape,
x_T=None,
ddim_use_original_steps=False,
callback=None,
timesteps=None,
quantize_denoised=False,
mask=None,
x0=None,
img_callback=None,
log_every_t=100,
temperature=1.0,
noise_dropout=0.0,
score_corrector=None,
corrector_kwargs=None,
unconditional_guidance_scale=1.0,
unconditional_conditioning=None,
dynamic_threshold=None,
ucg_schedule=None,
):
device = self.model.betas.device
b = shape[0]
if x_T is None:
img = torch.randn(shape, device=device)
else:
img = x_T
if timesteps is None:
timesteps = (
self.ddpm_num_timesteps
if ddim_use_original_steps
else self.ddim_timesteps
)
elif timesteps is not None and not ddim_use_original_steps:
subset_end = (
int(
min(timesteps / self.ddim_timesteps.shape[0], 1)
* self.ddim_timesteps.shape[0]
)
- 1
)
timesteps = self.ddim_timesteps[:subset_end]
intermediates = {"x_inter": [img], "pred_x0": [img]}
time_range = (
reversed(range(0, timesteps))
if ddim_use_original_steps
else np.flip(timesteps)
)
total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
print(f"Running DDIM Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc="DDIM Sampler", total=total_steps)
for i, step in enumerate(iterator):
index = total_steps - i - 1
ts = torch.full((b,), step, device=device, dtype=torch.long)
if mask is not None:
assert x0 is not None
img_orig = self.model.q_sample(
x0, ts
) # TODO: deterministic forward pass?
img = img_orig * mask + (1.0 - mask) * img
if ucg_schedule is not None:
assert len(ucg_schedule) == len(time_range)
unconditional_guidance_scale = ucg_schedule[i]
outs = self.p_sample_ddim(
img,
cond,
ts,
index=index,
use_original_steps=ddim_use_original_steps,
quantize_denoised=quantize_denoised,
temperature=temperature,
noise_dropout=noise_dropout,
score_corrector=score_corrector,
corrector_kwargs=corrector_kwargs,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=unconditional_conditioning,
dynamic_threshold=dynamic_threshold,
)
img, pred_x0 = outs
if callback:
callback(None, i, None, None)
if img_callback:
img_callback(pred_x0, i)
if index % log_every_t == 0 or index == total_steps - 1:
intermediates["x_inter"].append(img)
intermediates["pred_x0"].append(pred_x0)
return img, intermediates
@torch.no_grad()
def p_sample_ddim(
self,
x,
c,
t,
index,
repeat_noise=False,
use_original_steps=False,
quantize_denoised=False,
temperature=1.0,
noise_dropout=0.0,
score_corrector=None,
corrector_kwargs=None,
unconditional_guidance_scale=1.0,
unconditional_conditioning=None,
dynamic_threshold=None,
):
b, *_, device = *x.shape, x.device
if unconditional_conditioning is None or unconditional_guidance_scale == 1.0:
model_output = self.model.apply_model(x, t, c)
else:
model_t = self.model.apply_model(x, t, c)
model_uncond = self.model.apply_model(x, t, unconditional_conditioning)
model_output = model_uncond + unconditional_guidance_scale * (
model_t - model_uncond
)
if self.model.parameterization == "v":
e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
else:
e_t = model_output
if score_corrector is not None:
assert self.model.parameterization == "eps", "not implemented"
e_t = score_corrector.modify_score(
self.model, e_t, x, t, c, **corrector_kwargs
)
alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
alphas_prev = (
self.model.alphas_cumprod_prev
if use_original_steps
else self.ddim_alphas_prev
)
sqrt_one_minus_alphas = (
self.model.sqrt_one_minus_alphas_cumprod
if use_original_steps
else self.ddim_sqrt_one_minus_alphas
)
sigmas = (
self.model.ddim_sigmas_for_original_num_steps
if use_original_steps
else self.ddim_sigmas
)
# select parameters corresponding to the currently considered timestep
a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
sqrt_one_minus_at = torch.full(
(b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device
)
# current prediction for x_0
if self.model.parameterization != "v":
pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
else:
pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
if quantize_denoised:
pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
if dynamic_threshold is not None:
raise NotImplementedError()
# direction pointing to x_t
dir_xt = (1.0 - a_prev - sigma_t**2).sqrt() * e_t
noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
if noise_dropout > 0.0:
noise = torch.nn.functional.dropout(noise, p=noise_dropout)
x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
return x_prev, pred_x0
@torch.no_grad()
def encode(
self,
x0,
c,
t_enc,
use_original_steps=False,
return_intermediates=None,
unconditional_guidance_scale=1.0,
unconditional_conditioning=None,
callback=None,
):
timesteps = (
np.arange(self.ddpm_num_timesteps)
if use_original_steps
else self.ddim_timesteps
)
num_reference_steps = timesteps.shape[0]
assert t_enc <= num_reference_steps
num_steps = t_enc
if use_original_steps:
alphas_next = self.alphas_cumprod[:num_steps]
alphas = self.alphas_cumprod_prev[:num_steps]
else:
alphas_next = self.ddim_alphas[:num_steps]
alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
x_next = x0
intermediates = []
inter_steps = []
for i in tqdm(range(num_steps), desc="Encoding Image"):
t = torch.full(
(x0.shape[0],), timesteps[i], device=self.model.device, dtype=torch.long
)
if unconditional_guidance_scale == 1.0:
noise_pred = self.model.apply_model(x_next, t, c)
else:
assert unconditional_conditioning is not None
e_t_uncond, noise_pred = torch.chunk(
self.model.apply_model(
torch.cat((x_next, x_next)),
torch.cat((t, t)),
torch.cat((unconditional_conditioning, c)),
),
2,
)
noise_pred = e_t_uncond + unconditional_guidance_scale * (
noise_pred - e_t_uncond
)
xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
weighted_noise_pred = (
alphas_next[i].sqrt()
* ((1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt())
* noise_pred
)
x_next = xt_weighted + weighted_noise_pred
if (
return_intermediates
and i % (num_steps // return_intermediates) == 0
and i < num_steps - 1
):
intermediates.append(x_next)
inter_steps.append(i)
elif return_intermediates and i >= num_steps - 2:
intermediates.append(x_next)
inter_steps.append(i)
if callback:
callback(i)
out = {"x_encoded": x_next, "intermediate_steps": inter_steps}
if return_intermediates:
out.update({"intermediates": intermediates})
return x_next, out
@torch.no_grad()
def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
# fast, but does not allow for exact reconstruction
# t serves as an index to gather the correct alphas
if use_original_steps:
sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
else:
sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
if noise is None:
noise = torch.randn_like(x0)
return (
extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise
)
@torch.no_grad()
def decode(
self,
x_latent,
cond,
t_start,
unconditional_guidance_scale=1.0,
unconditional_conditioning=None,
use_original_steps=False,
callback=None,
):
timesteps = (
np.arange(self.ddpm_num_timesteps)
if use_original_steps
else self.ddim_timesteps
)
timesteps = timesteps[:t_start]
time_range = np.flip(timesteps)
total_steps = timesteps.shape[0]
print(f"Running DDIM Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc="Decoding image", total=total_steps)
x_dec = x_latent
for i, step in enumerate(iterator):
index = total_steps - i - 1
ts = torch.full(
(x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long
)
x_dec, _ = self.p_sample_ddim(
x_dec,
cond,
ts,
index=index,
use_original_steps=use_original_steps,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=unconditional_conditioning,
)
if callback:
callback(i)
return x_dec
================================================
FILE: iopaint/model/anytext/cldm/embedding_manager.py
================================================
'''
Copyright (c) Alibaba, Inc. and its affiliates.
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from iopaint.model.anytext.ldm.modules.diffusionmodules.util import conv_nd, linear
def get_clip_token_for_string(tokenizer, string):
batch_encoding = tokenizer(string, truncation=True, max_length=77, return_length=True,
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
tokens = batch_encoding["input_ids"]
assert torch.count_nonzero(tokens - 49407) == 2, f"String '{string}' maps to more than a single token. Please use another string"
return tokens[0, 1]
def get_bert_token_for_string(tokenizer, string):
token = tokenizer(string)
assert torch.count_nonzero(token) == 3, f"String '{string}' maps to more than a single token. Please use another string"
token = token[0, 1]
return token
def get_clip_vision_emb(encoder, processor, img):
_img = img.repeat(1, 3, 1, 1)*255
inputs = processor(images=_img, return_tensors="pt")
inputs['pixel_values'] = inputs['pixel_values'].to(img.device)
outputs = encoder(**inputs)
emb = outputs.image_embeds
return emb
def get_recog_emb(encoder, img_list):
_img_list = [(img.repeat(1, 3, 1, 1)*255)[0] for img in img_list]
encoder.predictor.eval()
_, preds_neck = encoder.pred_imglist(_img_list, show_debug=False)
return preds_neck
def pad_H(x):
_, _, H, W = x.shape
p_top = (W - H) // 2
p_bot = W - H - p_top
return F.pad(x, (0, 0, p_top, p_bot))
class EncodeNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(EncodeNet, self).__init__()
chan = 16
n_layer = 4 # downsample
self.conv1 = conv_nd(2, in_channels, chan, 3, padding=1)
self.conv_list = nn.ModuleList([])
_c = chan
for i in range(n_layer):
self.conv_list.append(conv_nd(2, _c, _c*2, 3, padding=1, stride=2))
_c *= 2
self.conv2 = conv_nd(2, _c, out_channels, 3, padding=1)
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.act = nn.SiLU()
def forward(self, x):
x = self.act(self.conv1(x))
for layer in self.conv_list:
x = self.act(layer(x))
x = self.act(self.conv2(x))
x = self.avgpool(x)
x = x.view(x.size(0), -1)
return x
class EmbeddingManager(nn.Module):
def __init__(
self,
embedder,
valid=True,
glyph_channels=20,
position_channels=1,
placeholder_string='*',
add_pos=False,
emb_type='ocr',
**kwargs
):
super().__init__()
if hasattr(embedder, 'tokenizer'): # using Stable Diffusion's CLIP encoder
get_token_for_string = partial(get_clip_token_for_string, embedder.tokenizer)
token_dim = 768
if hasattr(embedder, 'vit'):
assert emb_type == 'vit'
self.get_vision_emb = partial(get_clip_vision_emb, embedder.vit, embedder.processor)
self.get_recog_emb = None
else: # using LDM's BERT encoder
get_token_for_string = partial(get_bert_token_for_string, embedder.tknz_fn)
token_dim = 1280
self.token_dim = token_dim
self.emb_type = emb_type
self.add_pos = add_pos
if add_pos:
self.position_encoder = EncodeNet(position_channels, token_dim)
if emb_type == 'ocr':
self.proj = linear(40*64, token_dim)
if emb_type == 'conv':
self.glyph_encoder = EncodeNet(glyph_channels, token_dim)
self.placeholder_token = get_token_for_string(placeholder_string)
def encode_text(self, text_info):
if self.get_recog_emb is None and self.emb_type == 'ocr':
self.get_recog_emb = partial(get_recog_emb, self.recog)
gline_list = []
pos_list = []
for i in range(len(text_info['n_lines'])): # sample index in a batch
n_lines = text_info['n_lines'][i]
for j in range(n_lines): # line
gline_list += [text_info['gly_line'][j][i:i+1]]
if self.add_pos:
pos_list += [text_info['positions'][j][i:i+1]]
if len(gline_list) > 0:
if self.emb_type == 'ocr':
recog_emb = self.get_recog_emb(gline_list)
enc_glyph = self.proj(recog_emb.reshape(recog_emb.shape[0], -1))
elif self.emb_type == 'vit':
enc_glyph = self.get_vision_emb(pad_H(torch.cat(gline_list, dim=0)))
elif self.emb_type == 'conv':
enc_glyph = self.glyph_encoder(pad_H(torch.cat(gline_list, dim=0)))
if self.add_pos:
enc_pos = self.position_encoder(torch.cat(gline_list, dim=0))
enc_glyph = enc_glyph+enc_pos
self.text_embs_all = []
n_idx = 0
for i in range(len(text_info['n_lines'])): # sample index in a batch
n_lines = text_info['n_lines'][i]
text_embs = []
for j in range(n_lines): # line
text_embs += [enc_glyph[n_idx:n_idx+1]]
n_idx += 1
self.text_embs_all += [text_embs]
def forward(
self,
tokenized_text,
embedded_text,
):
b, device = tokenized_text.shape[0], tokenized_text.device
for i in range(b):
idx = tokenized_text[i] == self.placeholder_token.to(device)
if sum(idx) > 0:
if i >= len(self.text_embs_all):
print('truncation for log images...')
break
text_emb = torch.cat(self.text_embs_all[i], dim=0)
if sum(idx) != len(text_emb):
print('truncation for long caption...')
embedded_text[i][idx] = text_emb[:sum(idx)]
return embedded_text
def embedding_parameters(self):
return self.parameters()
================================================
FILE: iopaint/model/anytext/cldm/hack.py
================================================
import torch
import einops
import iopaint.model.anytext.ldm.modules.encoders.modules
import iopaint.model.anytext.ldm.modules.attention
from transformers import logging
from iopaint.model.anytext.ldm.modules.attention import default
def disable_verbosity():
logging.set_verbosity_error()
print('logging improved.')
return
def enable_sliced_attention():
iopaint.model.anytext.ldm.modules.attention.CrossAttention.forward = _hacked_sliced_attentin_forward
print('Enabled sliced_attention.')
return
def hack_everything(clip_skip=0):
disable_verbosity()
iopaint.model.anytext.ldm.modules.encoders.modules.FrozenCLIPEmbedder.forward = _hacked_clip_forward
iopaint.model.anytext.ldm.modules.encoders.modules.FrozenCLIPEmbedder.clip_skip = clip_skip
print('Enabled clip hacks.')
return
# Written by Lvmin
def _hacked_clip_forward(self, text):
PAD = self.tokenizer.pad_token_id
EOS = self.tokenizer.eos_token_id
BOS = self.tokenizer.bos_token_id
def tokenize(t):
return self.tokenizer(t, truncation=False, add_special_tokens=False)["input_ids"]
def transformer_encode(t):
if self.clip_skip > 1:
rt = self.transformer(input_ids=t, output_hidden_states=True)
return self.transformer.text_model.final_layer_norm(rt.hidden_states[-self.clip_skip])
else:
return self.transformer(input_ids=t, output_hidden_states=False).last_hidden_state
def split(x):
return x[75 * 0: 75 * 1], x[75 * 1: 75 * 2], x[75 * 2: 75 * 3]
def pad(x, p, i):
return x[:i] if len(x) >= i else x + [p] * (i - len(x))
raw_tokens_list = tokenize(text)
tokens_list = []
for raw_tokens in raw_tokens_list:
raw_tokens_123 = split(raw_tokens)
raw_tokens_123 = [[BOS] + raw_tokens_i + [EOS] for raw_tokens_i in raw_tokens_123]
raw_tokens_123 = [pad(raw_tokens_i, PAD, 77) for raw_tokens_i in raw_tokens_123]
tokens_list.append(raw_tokens_123)
tokens_list = torch.IntTensor(tokens_list).to(self.device)
feed = einops.rearrange(tokens_list, 'b f i -> (b f) i')
y = transformer_encode(feed)
z = einops.rearrange(y, '(b f) i c -> b (f i) c', f=3)
return z
# Stolen from https://github.com/basujindal/stable-diffusion/blob/main/optimizedSD/splitAttention.py
def _hacked_sliced_attentin_forward(self, x, context=None, mask=None):
h = self.heads
q = self.to_q(x)
context = default(context, x)
k = self.to_k(context)
v = self.to_v(context)
del context, x
q, k, v = map(lambda t: einops.rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
limit = k.shape[0]
att_step = 1
q_chunks = list(torch.tensor_split(q, limit // att_step, dim=0))
k_chunks = list(torch.tensor_split(k, limit // att_step, dim=0))
v_chunks = list(torch.tensor_split(v, limit // att_step, dim=0))
q_chunks.reverse()
k_chunks.reverse()
v_chunks.reverse()
sim = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
del k, q, v
for i in range(0, limit, att_step):
q_buffer = q_chunks.pop()
k_buffer = k_chunks.pop()
v_buffer = v_chunks.pop()
sim_buffer = torch.einsum('b i d, b j d -> b i j', q_buffer, k_buffer) * self.scale
del k_buffer, q_buffer
# attention, what we cannot get enough of, by chunks
sim_buffer = sim_buffer.softmax(dim=-1)
sim_buffer = torch.einsum('b i j, b j d -> b i d', sim_buffer, v_buffer)
del v_buffer
sim[i:i + att_step, :, :] = sim_buffer
del sim_buffer
sim = einops.rearrange(sim, '(b h) n d -> b n (h d)', h=h)
return self.to_out(sim)
================================================
FILE: iopaint/model/anytext/cldm/model.py
================================================
import os
import torch
from omegaconf import OmegaConf
from iopaint.model.anytext.ldm.util import instantiate_from_config
def get_state_dict(d):
return d.get("state_dict", d)
def load_state_dict(ckpt_path, location="cpu"):
_, extension = os.path.splitext(ckpt_path)
if extension.lower() == ".safetensors":
import safetensors.torch
state_dict = safetensors.torch.load_file(ckpt_path, device=location)
else:
state_dict = get_state_dict(
torch.load(ckpt_path, map_location=torch.device(location))
)
state_dict = get_state_dict(state_dict)
print(f"Loaded state_dict from [{ckpt_path}]")
return state_dict
def create_model(config_path, device, cond_stage_path=None, use_fp16=False):
config = OmegaConf.load(config_path)
# if cond_stage_path:
# config.model.params.cond_stage_config.params.version = (
# cond_stage_path # use pre-downloaded ckpts, in case blocked
# )
config.model.params.cond_stage_config.params.device = str(device)
if use_fp16:
config.model.params.use_fp16 = True
config.model.params.control_stage_config.params.use_fp16 = True
config.model.params.unet_config.params.use_fp16 = True
model = instantiate_from_config(config.model).cpu()
print(f"Loaded model config from [{config_path}]")
return model
================================================
FILE: iopaint/model/anytext/cldm/recognizer.py
================================================
"""
Copyright (c) Alibaba, Inc. and its affiliates.
"""
import os
import cv2
import numpy as np
import math
import traceback
from easydict import EasyDict as edict
import time
from iopaint.model.anytext.ocr_recog.RecModel import RecModel
import torch
import torch.nn.functional as F
def min_bounding_rect(img):
ret, thresh = cv2.threshold(img, 127, 255, 0)
contours, hierarchy = cv2.findContours(
thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
)
if len(contours) == 0:
print("Bad contours, using fake bbox...")
return np.array([[0, 0], [100, 0], [100, 100], [0, 100]])
max_contour = max(contours, key=cv2.contourArea)
rect = cv2.minAreaRect(max_contour)
box = cv2.boxPoints(rect)
box = np.int0(box)
# sort
x_sorted = sorted(box, key=lambda x: x[0])
left = x_sorted[:2]
right = x_sorted[2:]
left = sorted(left, key=lambda x: x[1])
(tl, bl) = left
right = sorted(right, key=lambda x: x[1])
(tr, br) = right
if tl[1] > bl[1]:
(tl, bl) = (bl, tl)
if tr[1] > br[1]:
(tr, br) = (br, tr)
return np.array([tl, tr, br, bl])
def create_predictor(model_dir=None, model_lang="ch", is_onnx=False):
model_file_path = model_dir
if model_file_path is not None and not os.path.exists(model_file_path):
raise ValueError("not find model file path {}".format(model_file_path))
if is_onnx:
import onnxruntime as ort
sess = ort.InferenceSession(
model_file_path, providers=["CPUExecutionProvider"]
) # 'TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'
return sess
else:
if model_lang == "ch":
n_class = 6625
elif model_lang == "en":
n_class = 97
else:
raise ValueError(f"Unsupported OCR recog model_lang: {model_lang}")
rec_config = edict(
in_channels=3,
backbone=edict(
type="MobileNetV1Enhance",
scale=0.5,
last_conv_stride=[1, 2],
last_pool_type="avg",
),
neck=edict(
type="SequenceEncoder",
encoder_type="svtr",
dims=64,
depth=2,
hidden_dims=120,
use_guide=True,
),
head=edict(
type="CTCHead",
fc_decay=0.00001,
out_channels=n_class,
return_feats=True,
),
)
rec_model = RecModel(rec_config)
if model_file_path is not None:
rec_model.load_state_dict(torch.load(model_file_path, map_location="cpu"))
rec_model.eval()
return rec_model.eval()
def _check_image_file(path):
img_end = {"jpg", "bmp", "png", "jpeg", "rgb", "tif", "tiff"}
return any([path.lower().endswith(e) for e in img_end])
def get_image_file_list(img_file):
imgs_lists = []
if img_file is None or not os.path.exists(img_file):
raise Exception("not found any img file in {}".format(img_file))
if os.path.isfile(img_file) and _check_image_file(img_file):
imgs_lists.append(img_file)
elif os.path.isdir(img_file):
for single_file in os.listdir(img_file):
file_path = os.path.join(img_file, single_file)
if os.path.isfile(file_path) and _check_image_file(file_path):
imgs_lists.append(file_path)
if len(imgs_lists) == 0:
raise Exception("not found any img file in {}".format(img_file))
imgs_lists = sorted(imgs_lists)
return imgs_lists
class TextRecognizer(object):
def __init__(self, args, predictor):
self.rec_image_shape = [int(v) for v in args.rec_image_shape.split(",")]
self.rec_batch_num = args.rec_batch_num
self.predictor = predictor
self.chars = self.get_char_dict(args.rec_char_dict_path)
self.char2id = {x: i for i, x in enumerate(self.chars)}
self.is_onnx = not isinstance(self.predictor, torch.nn.Module)
self.use_fp16 = args.use_fp16
# img: CHW
def resize_norm_img(self, img, max_wh_ratio):
imgC, imgH, imgW = self.rec_image_shape
assert imgC == img.shape[0]
imgW = int((imgH * max_wh_ratio))
h, w = img.shape[1:]
ratio = w / float(h)
if math.ceil(imgH * ratio) > imgW:
resized_w = imgW
else:
resized_w = int(math.ceil(imgH * ratio))
resized_image = torch.nn.functional.interpolate(
img.unsqueeze(0),
size=(imgH, resized_w),
mode="bilinear",
align_corners=True,
)
resized_image /= 255.0
resized_image -= 0.5
resized_image /= 0.5
padding_im = torch.zeros((imgC, imgH, imgW), dtype=torch.float32).to(img.device)
padding_im[:, :, 0:resized_w] = resized_image[0]
return padding_im
# img_list: list of tensors with shape chw 0-255
def pred_imglist(self, img_list, show_debug=False, is_ori=False):
img_num = len(img_list)
assert img_num > 0
# Calculate the aspect ratio of all text bars
width_list = []
for img in img_list:
width_list.append(img.shape[2] / float(img.shape[1]))
# Sorting can speed up the recognition process
indices = torch.from_numpy(np.argsort(np.array(width_list)))
batch_num = self.rec_batch_num
preds_all = [None] * img_num
preds_neck_all = [None] * img_num
for beg_img_no in range(0, img_num, batch_num):
end_img_no = min(img_num, beg_img_no + batch_num)
norm_img_batch = []
imgC, imgH, imgW = self.rec_image_shape[:3]
max_wh_ratio = imgW / imgH
for ino in range(beg_img_no, end_img_no):
h, w = img_list[indices[ino]].shape[1:]
if h > w * 1.2:
img = img_list[indices[ino]]
img = torch.transpose(img, 1, 2).flip(dims=[1])
img_list[indices[ino]] = img
h, w = img.shape[1:]
# wh_ratio = w * 1.0 / h
# max_wh_ratio = max(max_wh_ratio, wh_ratio) # comment to not use different ratio
for ino in range(beg_img_no, end_img_no):
norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio)
if self.use_fp16:
norm_img = norm_img.half()
norm_img = norm_img.unsqueeze(0)
norm_img_batch.append(norm_img)
norm_img_batch = torch.cat(norm_img_batch, dim=0)
if show_debug:
for i in range(len(norm_img_batch)):
_img = norm_img_batch[i].permute(1, 2, 0).detach().cpu().numpy()
_img = (_img + 0.5) * 255
_img = _img[:, :, ::-1]
file_name = f"{indices[beg_img_no + i]}"
file_name = file_name + "_ori" if is_ori else file_name
cv2.imwrite(file_name + ".jpg", _img)
if self.is_onnx:
input_dict = {}
input_dict[self.predictor.get_inputs()[0].name] = (
norm_img_batch.detach().cpu().numpy()
)
outputs = self.predictor.run(None, input_dict)
preds = {}
preds["ctc"] = torch.from_numpy(outputs[0])
preds["ctc_neck"] = [torch.zeros(1)] * img_num
else:
preds = self.predictor(norm_img_batch)
for rno in range(preds["ctc"].shape[0]):
preds_all[indices[beg_img_no + rno]] = preds["ctc"][rno]
preds_neck_all[indices[beg_img_no + rno]] = preds["ctc_neck"][rno]
return torch.stack(preds_all, dim=0), torch.stack(preds_neck_all, dim=0)
def get_char_dict(self, character_dict_path):
character_str = []
with open(character_dict_path, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode("utf-8").strip("\n").strip("\r\n")
character_str.append(line)
dict_character = list(character_str)
dict_character = ["sos"] + dict_character + [" "] # eos is space
return dict_character
def get_text(self, order):
char_list = [self.chars[text_id] for text_id in order]
return "".join(char_list)
def decode(self, mat):
text_index = mat.detach().cpu().numpy().argmax(axis=1)
ignored_tokens = [0]
selection = np.ones(len(text_index), dtype=bool)
selection[1:] = text_index[1:] != text_index[:-1]
for ignored_token in ignored_tokens:
selection &= text_index != ignored_token
return text_index[selection], np.where(selection)[0]
def get_ctcloss(self, preds, gt_text, weight):
if not isinstance(weight, torch.Tensor):
weight = torch.tensor(weight).to(preds.device)
ctc_loss = torch.nn.CTCLoss(reduction="none")
log_probs = preds.log_softmax(dim=2).permute(1, 0, 2) # NTC-->TNC
targets = []
target_lengths = []
for t in gt_text:
targets += [self.char2id.get(i, len(self.chars) - 1) for i in t]
target_lengths += [len(t)]
targets = torch.tensor(targets).to(preds.device)
target_lengths = torch.tensor(target_lengths).to(preds.device)
input_lengths = torch.tensor([log_probs.shape[0]] * (log_probs.shape[1])).to(
preds.device
)
loss = ctc_loss(log_probs, targets, input_lengths, target_lengths)
loss = loss / input_lengths * weight
return loss
def main():
rec_model_dir = "./ocr_weights/ppv3_rec.pth"
predictor = create_predictor(rec_model_dir)
args = edict()
args.rec_image_shape = "3, 48, 320"
args.rec_char_dict_path = "./ocr_weights/ppocr_keys_v1.txt"
args.rec_batch_num = 6
text_recognizer = TextRecognizer(args, predictor)
image_dir = "./test_imgs_cn"
gt_text = ["韩国小馆"] * 14
image_file_list = get_image_file_list(image_dir)
valid_image_file_list = []
img_list = []
for image_file in image_file_list:
img = cv2.imread(image_file)
if img is None:
print("error in loading image:{}".format(image_file))
continue
valid_image_file_list.append(image_file)
img_list.append(torch.from_numpy(img).permute(2, 0, 1).float())
try:
tic = time.time()
times = []
for i in range(10):
preds, _ = text_recognizer.pred_imglist(img_list) # get text
preds_all = preds.softmax(dim=2)
times += [(time.time() - tic) * 1000.0]
tic = time.time()
print(times)
print(np.mean(times[1:]) / len(preds_all))
weight = np.ones(len(gt_text))
loss = text_recognizer.get_ctcloss(preds, gt_text, weight)
for i in range(len(valid_image_file_list)):
pred = preds_all[i]
order, idx = text_recognizer.decode(pred)
text = text_recognizer.get_text(order)
print(
f'{valid_image_file_list[i]}: pred/gt="{text}"/"{gt_text[i]}", loss={loss[i]:.2f}'
)
except Exception as E:
print(traceback.format_exc(), E)
if __name__ == "__main__":
main()
================================================
FILE: iopaint/model/anytext/ldm/__init__.py
================================================
================================================
FILE: iopaint/model/anytext/ldm/models/__init__.py
================================================
================================================
FILE: iopaint/model/anytext/ldm/models/autoencoder.py
================================================
import torch
import torch.nn.functional as F
from contextlib import contextmanager
from iopaint.model.anytext.ldm.modules.diffusionmodules.model import Encoder, Decoder
from iopaint.model.anytext.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
from iopaint.model.anytext.ldm.util import instantiate_from_config
from iopaint.model.anytext.ldm.modules.ema import LitEma
class AutoencoderKL(torch.nn.Module):
def __init__(self,
ddconfig,
lossconfig,
embed_dim,
ckpt_path=None,
ignore_keys=[],
image_key="image",
colorize_nlabels=None,
monitor=None,
ema_decay=None,
learn_logvar=False
):
super().__init__()
self.learn_logvar = learn_logvar
self.image_key = image_key
self.encoder = Encoder(**ddconfig)
self.decoder = Decoder(**ddconfig)
self.loss = instantiate_from_config(lossconfig)
assert ddconfig["double_z"]
self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
self.embed_dim = embed_dim
if colorize_nlabels is not None:
assert type(colorize_nlabels)==int
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
if monitor is not None:
self.monitor = monitor
self.use_ema = ema_decay is not None
if self.use_ema:
self.ema_decay = ema_decay
assert 0. < ema_decay < 1.
self.model_ema = LitEma(self, decay=ema_decay)
print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
if ckpt_path is not None:
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
def init_from_ckpt(self, path, ignore_keys=list()):
sd = torch.load(path, map_location="cpu")["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]
self.load_state_dict(sd, strict=False)
print(f"Restored from {path}")
@contextmanager
def ema_scope(self, context=None):
if self.use_ema:
self.model_ema.store(self.parameters())
self.model_ema.copy_to(self)
if context is not None:
print(f"{context}: Switched to EMA weights")
try:
yield None
finally:
if self.use_ema:
self.model_ema.restore(self.parameters())
if context is not None:
print(f"{context}: Restored training weights")
def on_train_batch_end(self, *args, **kwargs):
if self.use_ema:
self.model_ema(self)
def encode(self, x):
h = self.encoder(x)
moments = self.quant_conv(h)
posterior = DiagonalGaussianDistribution(moments)
return posterior
def decode(self, z):
z = self.post_quant_conv(z)
dec = self.decoder(z)
return dec
def forward(self, input, sample_posterior=True):
posterior = self.encode(input)
if sample_posterior:
z = posterior.sample()
else:
z = posterior.mode()
dec = self.decode(z)
return dec, posterior
def get_input(self, batch, k):
x = batch[k]
if len(x.shape) == 3:
x = x[..., None]
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
return x
def training_step(self, batch, batch_idx, optimizer_idx):
inputs = self.get_input(batch, self.image_key)
reconstructions, posterior = self(inputs)
if optimizer_idx == 0:
# train encoder+decoder+logvar
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(), split="train")
self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
return aeloss
if optimizer_idx == 1:
# train the discriminator
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
last_layer=self.get_last_layer(), split="train")
self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
return discloss
def validation_step(self, batch, batch_idx):
log_dict = self._validation_step(batch, batch_idx)
with self.ema_scope():
log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
return log_dict
def _validation_step(self, batch, batch_idx, postfix=""):
inputs = self.get_input(batch, self.image_key)
reconstructions, posterior = self(inputs)
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
last_layer=self.get_last_layer(), split="val"+postfix)
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
last_layer=self.get_last_layer(), split="val"+postfix)
self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
self.log_dict(log_dict_ae)
self.log_dict(log_dict_disc)
return self.log_dict
def configure_optimizers(self):
lr = self.learning_rate
ae_params_list = list(self.encoder.parameters()) + list(self.decoder.parameters()) + list(
self.quant_conv.parameters()) + list(self.post_quant_conv.parameters())
if self.learn_logvar:
print(f"{self.__class__.__name__}: Learning logvar")
ae_params_list.append(self.loss.logvar)
opt_ae = torch.optim.Adam(ae_params_list,
lr=lr, betas=(0.5, 0.9))
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
lr=lr, betas=(0.5, 0.9))
return [opt_ae, opt_disc], []
def get_last_layer(self):
return self.decoder.conv_out.weight
@torch.no_grad()
def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
log = dict()
x = self.get_input(batch, self.image_key)
x = x.to(self.device)
if not only_inputs:
xrec, posterior = self(x)
if x.shape[1] > 3:
# colorize with random projection
assert xrec.shape[1] > 3
x = self.to_rgb(x)
xrec = self.to_rgb(xrec)
log["samples"] = self.decode(torch.randn_like(posterior.sample()))
log["reconstructions"] = xrec
if log_ema or self.use_ema:
with self.ema_scope():
xrec_ema, posterior_ema = self(x)
if x.shape[1] > 3:
# colorize with random projection
assert xrec_ema.shape[1] > 3
xrec_ema = self.to_rgb(xrec_ema)
log["samples_ema"] = self.decode(torch.randn_like(posterior_ema.sample()))
log["reconstructions_ema"] = xrec_ema
log["inputs"] = x
return log
def to_rgb(self, x):
assert self.image_key == "segmentation"
if not hasattr(self, "colorize"):
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
x = F.conv2d(x, weight=self.colorize)
x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
return x
class IdentityFirstStage(torch.nn.Module):
def __init__(self, *args, vq_interface=False, **kwargs):
self.vq_interface = vq_interface
super().__init__()
def encode(self, x, *args, **kwargs):
return x
def decode(self, x, *args, **kwargs):
return x
def quantize(self, x, *args, **kwargs):
if self.vq_interface:
return x, None, [None, None, None]
return x
def forward(self, x, *args, **kwargs):
return x
================================================
FILE: iopaint/model/anytext/ldm/models/diffusion/__init__.py
================================================
================================================
FILE: iopaint/model/anytext/ldm/models/diffusion/ddim.py
================================================
"""SAMPLING ONLY."""
import torch
import numpy as np
from tqdm import tqdm
from iopaint.model.anytext.ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
class DDIMSampler(object):
def __init__(self, model, schedule="linear", **kwargs):
super().__init__()
self.model = model
self.ddpm_num_timesteps = model.num_timesteps
self.schedule = schedule
def register_buffer(self, name, attr):
if type(attr) == torch.Tensor:
if attr.device != torch.device("cuda"):
attr = attr.to(torch.device("cuda"))
setattr(self, name, attr)
def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
alphas_cumprod = self.model.alphas_cumprod
assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
self.register_buffer('betas', to_torch(self.model.betas))
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
# calculations for diffusion q(x_t | x_{t-1}) and others
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
# ddim sampling parameters
ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
ddim_timesteps=self.ddim_timesteps,
eta=ddim_eta,verbose=
gitextract_t_2nmh0z/
├── .github/
│ ├── FUNDING.yml
│ └── ISSUE_TEMPLATE/
│ ├── config.yml
│ ├── 🐛-bug-report.md
│ └── 🚀-feature-request.md
├── .gitignore
├── LICENSE
├── README.md
├── build_docker.sh
├── docker/
│ ├── CPUDockerfile
│ └── GPUDockerfile
├── iopaint/
│ ├── __init__.py
│ ├── __main__.py
│ ├── api.py
│ ├── batch_processing.py
│ ├── benchmark.py
│ ├── cli.py
│ ├── const.py
│ ├── download.py
│ ├── file_manager/
│ │ ├── __init__.py
│ │ ├── file_manager.py
│ │ ├── storage_backends.py
│ │ └── utils.py
│ ├── helper.py
│ ├── installer.py
│ ├── model/
│ │ ├── __init__.py
│ │ ├── anytext/
│ │ │ ├── __init__.py
│ │ │ ├── anytext_model.py
│ │ │ ├── anytext_pipeline.py
│ │ │ ├── anytext_sd15.yaml
│ │ │ ├── cldm/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── cldm.py
│ │ │ │ ├── ddim_hacked.py
│ │ │ │ ├── embedding_manager.py
│ │ │ │ ├── hack.py
│ │ │ │ ├── model.py
│ │ │ │ └── recognizer.py
│ │ │ ├── ldm/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── models/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── autoencoder.py
│ │ │ │ │ └── diffusion/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── ddim.py
│ │ │ │ │ ├── ddpm.py
│ │ │ │ │ ├── dpm_solver/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ ├── dpm_solver.py
│ │ │ │ │ │ └── sampler.py
│ │ │ │ │ ├── plms.py
│ │ │ │ │ └── sampling_util.py
│ │ │ │ ├── modules/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── attention.py
│ │ │ │ │ ├── diffusionmodules/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ ├── model.py
│ │ │ │ │ │ ├── openaimodel.py
│ │ │ │ │ │ ├── upscaling.py
│ │ │ │ │ │ └── util.py
│ │ │ │ │ ├── distributions/
│ │ │ │ │ │ ├── __init__.py
│ │ │ │ │ │ └── distributions.py
│ │ │ │ │ ├── ema.py
│ │ │ │ │ └── encoders/
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ └── modules.py
│ │ │ │ └── util.py
│ │ │ ├── main.py
│ │ │ ├── ocr_recog/
│ │ │ │ ├── RNN.py
│ │ │ │ ├── RecCTCHead.py
│ │ │ │ ├── RecModel.py
│ │ │ │ ├── RecMv1_enhance.py
│ │ │ │ ├── RecSVTR.py
│ │ │ │ ├── __init__.py
│ │ │ │ ├── common.py
│ │ │ │ ├── en_dict.txt
│ │ │ │ └── ppocr_keys_v1.txt
│ │ │ └── utils.py
│ │ ├── base.py
│ │ ├── brushnet/
│ │ │ ├── __init__.py
│ │ │ ├── brushnet.py
│ │ │ ├── brushnet_unet_forward.py
│ │ │ ├── brushnet_wrapper.py
│ │ │ ├── brushnet_xl_wrapper.py
│ │ │ ├── pipeline_brushnet.py
│ │ │ ├── pipeline_brushnet_sd_xl.py
│ │ │ └── unet_2d_blocks.py
│ │ ├── controlnet.py
│ │ ├── ddim_sampler.py
│ │ ├── fcf.py
│ │ ├── helper/
│ │ │ ├── __init__.py
│ │ │ ├── controlnet_preprocess.py
│ │ │ ├── cpu_text_encoder.py
│ │ │ └── g_diffuser_bot.py
│ │ ├── instruct_pix2pix.py
│ │ ├── kandinsky.py
│ │ ├── lama.py
│ │ ├── ldm.py
│ │ ├── manga.py
│ │ ├── mat.py
│ │ ├── mi_gan.py
│ │ ├── opencv2.py
│ │ ├── original_sd_configs/
│ │ │ ├── __init__.py
│ │ │ ├── sd_xl_base.yaml
│ │ │ ├── sd_xl_refiner.yaml
│ │ │ ├── v1-inference.yaml
│ │ │ └── v2-inference-v.yaml
│ │ ├── paint_by_example.py
│ │ ├── plms_sampler.py
│ │ ├── power_paint/
│ │ │ ├── __init__.py
│ │ │ ├── pipeline_powerpaint.py
│ │ │ ├── power_paint.py
│ │ │ ├── power_paint_v2.py
│ │ │ ├── powerpaint_tokenizer.py
│ │ │ └── v2/
│ │ │ ├── BrushNet_CA.py
│ │ │ ├── __init__.py
│ │ │ ├── pipeline_PowerPaint_Brushnet_CA.py
│ │ │ ├── unet_2d_blocks.py
│ │ │ └── unet_2d_condition.py
│ │ ├── sd.py
│ │ ├── sdxl.py
│ │ ├── utils.py
│ │ └── zits.py
│ ├── model_manager.py
│ ├── plugins/
│ │ ├── __init__.py
│ │ ├── anime_seg.py
│ │ ├── base_plugin.py
│ │ ├── basicsr/
│ │ │ ├── LICENSE
│ │ │ ├── __init__.py
│ │ │ ├── arch_util.py
│ │ │ ├── img_util.py
│ │ │ └── rrdbnet_arch.py
│ │ ├── briarmbg.py
│ │ ├── briarmbg2.py
│ │ ├── facexlib/
│ │ │ ├── .gitignore
│ │ │ ├── __init__.py
│ │ │ ├── detection/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── align_trans.py
│ │ │ │ ├── matlab_cp2tform.py
│ │ │ │ ├── retinaface.py
│ │ │ │ ├── retinaface_net.py
│ │ │ │ └── retinaface_utils.py
│ │ │ ├── parsing/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── bisenet.py
│ │ │ │ ├── parsenet.py
│ │ │ │ └── resnet.py
│ │ │ └── utils/
│ │ │ ├── __init__.py
│ │ │ ├── face_restoration_helper.py
│ │ │ ├── face_utils.py
│ │ │ └── misc.py
│ │ ├── gfpgan/
│ │ │ ├── __init__.py
│ │ │ └── archs/
│ │ │ ├── __init__.py
│ │ │ ├── gfpganv1_clean_arch.py
│ │ │ ├── restoreformer_arch.py
│ │ │ └── stylegan2_clean_arch.py
│ │ ├── gfpgan_plugin.py
│ │ ├── gfpganer.py
│ │ ├── interactive_seg.py
│ │ ├── realesrgan.py
│ │ ├── remove_bg.py
│ │ ├── restoreformer.py
│ │ ├── segment_anything/
│ │ │ ├── __init__.py
│ │ │ ├── build_sam.py
│ │ │ ├── modeling/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── common.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ ├── image_encoder_hq.py
│ │ │ │ ├── mask_decoder.py
│ │ │ │ ├── prompt_encoder.py
│ │ │ │ ├── sam.py
│ │ │ │ ├── sam_hq.py
│ │ │ │ ├── tiny_vit_sam.py
│ │ │ │ └── transformer.py
│ │ │ ├── predictor.py
│ │ │ ├── predictor_hq.py
│ │ │ └── utils/
│ │ │ ├── __init__.py
│ │ │ └── transforms.py
│ │ └── segment_anything2/
│ │ ├── __init__.py
│ │ ├── build_sam.py
│ │ ├── modeling/
│ │ │ ├── __init__.py
│ │ │ ├── backbones/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── hieradet.py
│ │ │ │ ├── image_encoder.py
│ │ │ │ └── utils.py
│ │ │ ├── memory_attention.py
│ │ │ ├── memory_encoder.py
│ │ │ ├── position_encoding.py
│ │ │ ├── sam/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── mask_decoder.py
│ │ │ │ ├── prompt_encoder.py
│ │ │ │ └── transformer.py
│ │ │ ├── sam2_base.py
│ │ │ └── sam2_utils.py
│ │ ├── sam2_image_predictor.py
│ │ └── utils/
│ │ ├── __init__.py
│ │ ├── misc.py
│ │ └── transforms.py
│ ├── runtime.py
│ ├── schema.py
│ ├── tests/
│ │ ├── .gitignore
│ │ ├── __init__.py
│ │ ├── test_adjust_mask.py
│ │ ├── test_anytext.py
│ │ ├── test_brushnet.py
│ │ ├── test_controlnet.py
│ │ ├── test_instruct_pix2pix.py
│ │ ├── test_load_img.py
│ │ ├── test_low_mem.py
│ │ ├── test_match_histograms.py
│ │ ├── test_model.py
│ │ ├── test_model_md5.py
│ │ ├── test_model_switch.py
│ │ ├── test_outpainting.py
│ │ ├── test_paint_by_example.py
│ │ ├── test_plugins.py
│ │ ├── test_save_exif.py
│ │ ├── test_save_quality.py
│ │ ├── test_sd_model.py
│ │ ├── test_sdxl.py
│ │ └── utils.py
│ └── web_config.py
├── main.py
├── publish.sh
├── requirements-dev.txt
├── requirements.txt
├── scripts/
│ ├── .gitignore
│ ├── README.md
│ ├── convert_vae_pt_to_diffusers.py
│ ├── environment.yaml
│ ├── pack.bat
│ ├── pack.sh
│ ├── tool.py
│ └── user_scripts/
│ ├── win_config.bat
│ ├── win_setup.bat
│ ├── win_setup_cn.bat
│ ├── win_start.bat
│ ├── win_start_cn.bat
│ └── win_update.bat
├── setup.py
└── web_app/
├── .eslintrc.cjs
├── .gitignore
├── README.md
├── components.json
├── index.html
├── package.json
├── postcss.config.js
├── src/
│ ├── App.tsx
│ ├── components/
│ │ ├── Coffee.tsx
│ │ ├── Cropper.tsx
│ │ ├── DiffusionProgress.tsx
│ │ ├── Editor.tsx
│ │ ├── Extender.tsx
│ │ ├── FileManager.tsx
│ │ ├── FileSelect.tsx
│ │ ├── Header.tsx
│ │ ├── ImageSize.tsx
│ │ ├── InteractiveSeg.tsx
│ │ ├── Plugins.tsx
│ │ ├── PromptInput.tsx
│ │ ├── Settings.tsx
│ │ ├── Shortcuts.tsx
│ │ ├── SidePanel/
│ │ │ ├── CV2Options.tsx
│ │ │ ├── DiffusionOptions.tsx
│ │ │ ├── LDMOptions.tsx
│ │ │ ├── LabelTitle.tsx
│ │ │ └── index.tsx
│ │ ├── Workspace.tsx
│ │ └── ui/
│ │ ├── accordion.tsx
│ │ ├── alert-dialog.tsx
│ │ ├── button.tsx
│ │ ├── context-menu.tsx
│ │ ├── dialog.tsx
│ │ ├── dropdown-menu.tsx
│ │ ├── form.tsx
│ │ ├── input.tsx
│ │ ├── label.tsx
│ │ ├── popover.tsx
│ │ ├── progress.tsx
│ │ ├── radio-group.tsx
│ │ ├── scroll-area.tsx
│ │ ├── select.tsx
│ │ ├── separator.tsx
│ │ ├── sheet.tsx
│ │ ├── slider.tsx
│ │ ├── switch.tsx
│ │ ├── tabs.tsx
│ │ ├── textarea.tsx
│ │ ├── toast.tsx
│ │ ├── toaster.tsx
│ │ ├── toggle.tsx
│ │ ├── tooltip.tsx
│ │ └── use-toast.ts
│ ├── globals.css
│ ├── hooks/
│ │ ├── useAsyncMemo.tsx
│ │ ├── useHotkey.tsx
│ │ ├── useImage.tsx
│ │ ├── useInputImage.tsx
│ │ └── useResolution.tsx
│ ├── lib/
│ │ ├── api.ts
│ │ ├── const.ts
│ │ ├── states.ts
│ │ ├── types.ts
│ │ └── utils.ts
│ ├── main.tsx
│ └── vite-env.d.ts
├── tailwind.config.js
├── tsconfig.json
├── tsconfig.node.json
└── vite.config.ts
SYMBOL INDEX (2187 symbols across 192 files)
FILE: iopaint/__init__.py
function fix_window_pytorch (line 21) | def fix_window_pytorch():
function entry_point (line 49) | def entry_point():
FILE: iopaint/api.py
function api_middleware (line 70) | def api_middleware(app: FastAPI):
function diffuser_callback (line 138) | def diffuser_callback(pipe, step: int, timestep: int, callback_kwargs: D...
class Api (line 148) | class Api:
method __init__ (line 149) | def __init__(self, app: FastAPI, config: ApiConfig):
method add_api_route (line 183) | def add_api_route(self, path: str, endpoint, **kwargs):
method api_save_image (line 186) | def api_save_image(self, file: UploadFile):
method api_current_model (line 205) | def api_current_model(self) -> ModelInfo:
method api_switch_model (line 208) | def api_switch_model(self, req: SwitchModelRequest) -> ModelInfo:
method api_switch_plugin_model (line 214) | def api_switch_plugin_model(self, req: SwitchPluginModelRequest):
method api_server_config (line 225) | def api_server_config(self) -> ServerConfigResponse:
method api_input_image (line 254) | def api_input_image(self) -> FileResponse:
method api_geninfo (line 262) | def api_geninfo(self, file: UploadFile) -> GenInfoResponse:
method api_inpaint (line 271) | def api_inpaint(self, req: InpaintRequest):
method api_run_plugin_gen_image (line 306) | def api_run_plugin_gen_image(self, req: RunPluginRequest):
method api_run_plugin_gen_mask (line 334) | def api_run_plugin_gen_mask(self, req: RunPluginRequest):
method api_samplers (line 350) | def api_samplers(self) -> List[str]:
method api_adjust_mask (line 353) | def api_adjust_mask(self, req: AdjustMaskRequest):
method launch (line 358) | def launch(self):
method _build_file_manager (line 367) | def _build_file_manager(self) -> Optional[FileManager]:
method _build_plugins (line 381) | def _build_plugins(self) -> Dict[str, BasePlugin]:
method _build_model_manager (line 400) | def _build_model_manager(self):
FILE: iopaint/batch_processing.py
function glob_images (line 26) | def glob_images(path: Path) -> Dict[str, Path]:
function batch_inpaint (line 38) | def batch_inpaint(
FILE: iopaint/benchmark.py
function run_model (line 34) | def run_model(model, size):
function benchmark (line 52) | def benchmark(model, times: int, empty_cache: bool):
function get_args_parser (line 91) | def get_args_parser():
FILE: iopaint/cli.py
function install_plugins_packages (line 20) | def install_plugins_packages():
function download (line 27) | def download(
function list_model (line 44) | def list_model(
function run (line 60) | def run(
function start (line 99) | def start(
function start_web_config (line 234) | def start_web_config(
FILE: iopaint/download.py
function cli_download_model (line 22) | def cli_download_model(model: str):
function folder_name_to_show_name (line 44) | def folder_name_to_show_name(name: str) -> str:
function get_sd_model_type (line 49) | def get_sd_model_type(model_abs_path: str) -> Optional[ModelType]:
function get_sdxl_model_type (line 77) | def get_sdxl_model_type(model_abs_path: str) -> Optional[ModelType]:
function scan_single_file_diffusion_models (line 108) | def scan_single_file_diffusion_models(cache_dir) -> List[ModelInfo]:
function scan_inpaint_models (line 182) | def scan_inpaint_models(model_dir: Path) -> List[ModelInfo]:
function scan_diffusers_models (line 200) | def scan_diffusers_models() -> List[ModelInfo]:
function _scan_converted_diffusers_models (line 253) | def _scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
function scan_converted_diffusers_models (line 297) | def scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
function scan_models (line 307) | def scan_models() -> List[ModelInfo]:
FILE: iopaint/file_manager/file_manager.py
class FileManager (line 18) | class FileManager:
method __init__ (line 19) | def __init__(self, app: FastAPI, input_dir: Path, mask_dir: Path, outp...
method api_medias (line 36) | def api_medias(self, tab: MediaTab) -> List[MediasResponse]:
method api_media_file (line 40) | def api_media_file(self, tab: MediaTab, filename: str) -> FileResponse:
method api_media_thumbnail_file (line 45) | def api_media_thumbnail_file(
method _get_dir (line 62) | def _get_dir(self, tab: MediaTab) -> Path:
method _get_file (line 72) | def _get_file(self, tab: MediaTab, filename: str) -> Path:
method thumbnail_directory (line 79) | def thumbnail_directory(self) -> Path:
method _media_names (line 83) | def _media_names(directory: Path) -> List[MediasResponse]:
method get_thumbnail (line 102) | def get_thumbnail(
method get_raw_data (line 160) | def get_raw_data(self, image, **options):
method colormode (line 171) | def colormode(image, colormode="RGB"):
method background (line 185) | def background(original_image, color=0xFF):
method _get_format (line 195) | def _get_format(self, image, **options):
method _create_thumbnail (line 203) | def _create_thumbnail(self, image, size, crop="fit", background=None):
FILE: iopaint/file_manager/storage_backends.py
class BaseStorageBackend (line 7) | class BaseStorageBackend(ABC):
method __init__ (line 8) | def __init__(self, app=None):
method read (line 12) | def read(self, filepath, mode="rb", **kwargs):
method exists (line 16) | def exists(self, filepath):
method save (line 20) | def save(self, filepath, data):
class FilesystemStorageBackend (line 24) | class FilesystemStorageBackend(BaseStorageBackend):
method read (line 25) | def read(self, filepath, mode="rb", **kwargs):
method exists (line 29) | def exists(self, filepath):
method save (line 32) | def save(self, filepath, data):
FILE: iopaint/file_manager/utils.py
function generate_filename (line 8) | def generate_filename(directory: Path, original_filename, *options) -> str:
function parse_size (line 17) | def parse_size(size):
function aspect_to_string (line 42) | def aspect_to_string(size):
function glob_img (line 52) | def glob_img(p: Union[Path, str], recursive: bool = False):
FILE: iopaint/helper.py
function md5sum (line 19) | def md5sum(filename):
function switch_mps_device (line 27) | def switch_mps_device(model_name, device):
function get_cache_path_by_url (line 34) | def get_cache_path_by_url(url):
function download_model (line 45) | def download_model(url, model_md5: str = None):
function ceil_modulo (line 74) | def ceil_modulo(x, mod):
function handle_error (line 80) | def handle_error(model_path, model_md5, e):
function load_jit_model (line 101) | def load_jit_model(url_or_path, device, model_md5: str):
function load_model (line 116) | def load_model(model: torch.nn.Module, url_or_path, device, model_md5):
function numpy_to_bytes (line 133) | def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
function pil_to_bytes (line 143) | def pil_to_bytes(pil_img, ext: str, quality: int = 95, infos={}) -> bytes:
function load_img (line 158) | def load_img(img_bytes, gray: bool = False, return_info: bool = False):
function norm_img (line 187) | def norm_img(np_img):
function resize_max_size (line 195) | def resize_max_size(
function pad_img_to_modulo (line 209) | def pad_img_to_modulo(
function boxes_from_mask (line 246) | def boxes_from_mask(mask: np.ndarray) -> List[np.ndarray]:
function only_keep_largest_contour (line 270) | def only_keep_largest_contour(mask: np.ndarray) -> List[np.ndarray]:
function is_mac (line 296) | def is_mac():
function get_image_ext (line 300) | def get_image_ext(img_bytes):
function decode_base64_to_image (line 307) | def decode_base64_to_image(
function encode_pil_to_base64 (line 341) | def encode_pil_to_base64(image: Image, quality: int, infos: Dict) -> bytes:
function concat_alpha_channel (line 351) | def concat_alpha_channel(rgb_np_img, alpha_channel) -> np.ndarray:
function adjust_mask (line 363) | def adjust_mask(mask: np.ndarray, kernel_size: int, operate):
function gen_frontend_mask (line 393) | def gen_frontend_mask(bgr_or_gray_mask):
FILE: iopaint/installer.py
function install (line 5) | def install(package):
function install_plugins_package (line 9) | def install_plugins_package():
FILE: iopaint/model/anytext/anytext_model.py
class AnyText (line 11) | class AnyText(DiffusionInpaintModel):
method download (line 17) | def download(local_files_only=False):
method init_model (line 35) | def init_model(self, device, **kwargs):
method forward (line 47) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/anytext/anytext_pipeline.py
function check_limits (line 38) | def check_limits(tensor):
class AnyTextPipeline (line 49) | class AnyTextPipeline:
method __init__ (line 50) | def __init__(self, ckpt_path, font_path, device, use_fp16=True):
method __call__ (line 72) | def __call__(
method modify_prompt (line 338) | def modify_prompt(self, prompt):
method separate_pos_imgs (line 364) | def separate_pos_imgs(self, img, sort_priority, gap=102):
method find_polygon (line 379) | def find_polygon(self, image, min_rect=False):
method arr2tensor (line 397) | def arr2tensor(self, arr, bs):
FILE: iopaint/model/anytext/cldm/cldm.py
function count_parameters (line 30) | def count_parameters(model):
class ControlledUnetModel (line 34) | class ControlledUnetModel(UNetModel):
method forward (line 35) | def forward(self, x, timesteps=None, context=None, control=None, only_...
class ControlNet (line 62) | class ControlNet(nn.Module):
method __init__ (line 63) | def __init__(
method make_zero_conv (line 317) | def make_zero_conv(self, channels):
method forward (line 320) | def forward(self, x, hint, text_info, timesteps, context, **kwargs):
class ControlLDM (line 352) | class ControlLDM(LatentDiffusion):
method __init__ (line 354) | def __init__(self, control_stage_config, control_key, glyph_key, posit...
method get_input (line 390) | def get_input(self, batch, k, bs=None, *args, **kwargs):
method apply_model (line 442) | def apply_model(self, x_noisy, t, cond, *args, **kwargs):
method instantiate_embedding_manager (line 455) | def instantiate_embedding_manager(self, config, embedder):
method get_unconditional_conditioning (line 460) | def get_unconditional_conditioning(self, N):
method get_learned_conditioning (line 463) | def get_learned_conditioning(self, c):
method fill_caption (line 489) | def fill_caption(self, batch, place_holder='*'):
method log_images (line 504) | def log_images(self, batch, N=4, n_row=2, sample=False, ddim_steps=50,...
method sample_log (line 592) | def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
method configure_optimizers (line 599) | def configure_optimizers(self):
method low_vram_shift (line 620) | def low_vram_shift(self, is_diffusing):
FILE: iopaint/model/anytext/cldm/ddim_hacked.py
class DDIMSampler (line 15) | class DDIMSampler(object):
method __init__ (line 16) | def __init__(self, model, device, schedule="linear", **kwargs):
method register_buffer (line 23) | def register_buffer(self, name, attr):
method make_schedule (line 29) | def make_schedule(
method sample (line 90) | def sample(
method ddim_sampling (line 169) | def ddim_sampling(
method p_sample_ddim (line 268) | def p_sample_ddim(
method encode (line 352) | def encode(
method stochastic_encode (line 429) | def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
method decode (line 447) | def decode(
FILE: iopaint/model/anytext/cldm/embedding_manager.py
function get_clip_token_for_string (line 11) | def get_clip_token_for_string(tokenizer, string):
function get_bert_token_for_string (line 19) | def get_bert_token_for_string(tokenizer, string):
function get_clip_vision_emb (line 26) | def get_clip_vision_emb(encoder, processor, img):
function get_recog_emb (line 35) | def get_recog_emb(encoder, img_list):
function pad_H (line 42) | def pad_H(x):
class EncodeNet (line 49) | class EncodeNet(nn.Module):
method __init__ (line 50) | def __init__(self, in_channels, out_channels):
method forward (line 65) | def forward(self, x):
class EmbeddingManager (line 75) | class EmbeddingManager(nn.Module):
method __init__ (line 76) | def __init__(
method encode_text (line 111) | def encode_text(self, text_info):
method forward (line 146) | def forward(
method embedding_parameters (line 164) | def embedding_parameters(self):
FILE: iopaint/model/anytext/cldm/hack.py
function disable_verbosity (line 11) | def disable_verbosity():
function enable_sliced_attention (line 17) | def enable_sliced_attention():
function hack_everything (line 23) | def hack_everything(clip_skip=0):
function _hacked_clip_forward (line 32) | def _hacked_clip_forward(self, text):
function _hacked_sliced_attentin_forward (line 72) | def _hacked_sliced_attentin_forward(self, x, context=None, mask=None):
FILE: iopaint/model/anytext/cldm/model.py
function get_state_dict (line 8) | def get_state_dict(d):
function load_state_dict (line 12) | def load_state_dict(ckpt_path, location="cpu"):
function create_model (line 27) | def create_model(config_path, device, cond_stage_path=None, use_fp16=Fal...
FILE: iopaint/model/anytext/cldm/recognizer.py
function min_bounding_rect (line 16) | def min_bounding_rect(img):
function create_predictor (line 43) | def create_predictor(model_dir=None, model_lang="ch", is_onnx=False):
function _check_image_file (line 93) | def _check_image_file(path):
function get_image_file_list (line 98) | def get_image_file_list(img_file):
class TextRecognizer (line 115) | class TextRecognizer(object):
method __init__ (line 116) | def __init__(self, args, predictor):
method resize_norm_img (line 126) | def resize_norm_img(self, img, max_wh_ratio):
method pred_imglist (line 151) | def pred_imglist(self, img_list, show_debug=False, is_ori=False):
method get_char_dict (line 210) | def get_char_dict(self, character_dict_path):
method get_text (line 221) | def get_text(self, order):
method decode (line 225) | def decode(self, mat):
method get_ctcloss (line 234) | def get_ctcloss(self, preds, gt_text, weight):
function main (line 254) | def main():
FILE: iopaint/model/anytext/ldm/models/autoencoder.py
class AutoencoderKL (line 12) | class AutoencoderKL(torch.nn.Module):
method __init__ (line 13) | def __init__(self,
method init_from_ckpt (line 51) | def init_from_ckpt(self, path, ignore_keys=list()):
method ema_scope (line 63) | def ema_scope(self, context=None):
method on_train_batch_end (line 77) | def on_train_batch_end(self, *args, **kwargs):
method encode (line 81) | def encode(self, x):
method decode (line 87) | def decode(self, z):
method forward (line 92) | def forward(self, input, sample_posterior=True):
method get_input (line 101) | def get_input(self, batch, k):
method training_step (line 108) | def training_step(self, batch, batch_idx, optimizer_idx):
method validation_step (line 129) | def validation_step(self, batch, batch_idx):
method _validation_step (line 135) | def _validation_step(self, batch, batch_idx, postfix=""):
method configure_optimizers (line 149) | def configure_optimizers(self):
method get_last_layer (line 162) | def get_last_layer(self):
method log_images (line 166) | def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs):
method to_rgb (line 191) | def to_rgb(self, x):
class IdentityFirstStage (line 200) | class IdentityFirstStage(torch.nn.Module):
method __init__ (line 201) | def __init__(self, *args, vq_interface=False, **kwargs):
method encode (line 205) | def encode(self, x, *args, **kwargs):
method decode (line 208) | def decode(self, x, *args, **kwargs):
method quantize (line 211) | def quantize(self, x, *args, **kwargs):
method forward (line 216) | def forward(self, x, *args, **kwargs):
FILE: iopaint/model/anytext/ldm/models/diffusion/ddim.py
class DDIMSampler (line 10) | class DDIMSampler(object):
method __init__ (line 11) | def __init__(self, model, schedule="linear", **kwargs):
method register_buffer (line 17) | def register_buffer(self, name, attr):
method make_schedule (line 23) | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddi...
method sample (line 55) | def sample(self,
method ddim_sampling (line 124) | def ddim_sampling(self, cond, shape,
method p_sample_ddim (line 185) | def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_origin...
method encode (line 272) | def encode(self, x0, c, t_enc, use_original_steps=False, return_interm...
method stochastic_encode (line 319) | def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
method decode (line 335) | def decode(self, x_latent, cond, t_start, unconditional_guidance_scale...
FILE: iopaint/model/anytext/ldm/models/diffusion/ddpm.py
function print_grad (line 47) | def print_grad(grad):
function disabled_train (line 64) | def disabled_train(self, mode=True):
function uniform_on_device (line 70) | def uniform_on_device(r1, r2, shape, device):
class DDPM (line 74) | class DDPM(torch.nn.Module):
method __init__ (line 76) | def __init__(
method register_schedule (line 186) | def register_schedule(
method ema_scope (line 291) | def ema_scope(self, context=None):
method init_from_ckpt (line 306) | def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
method q_mean_variance (line 379) | def q_mean_variance(self, x_start, t):
method predict_start_from_noise (line 393) | def predict_start_from_noise(self, x_t, t, noise):
method predict_start_from_z_and_v (line 400) | def predict_start_from_z_and_v(self, x_t, t, v):
method predict_eps_from_z_and_v (line 408) | def predict_eps_from_z_and_v(self, x_t, t, v):
method q_posterior (line 415) | def q_posterior(self, x_start, x_t, t):
method p_mean_variance (line 426) | def p_mean_variance(self, x, t, clip_denoised: bool):
method p_sample (line 441) | def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
method p_sample_loop (line 452) | def p_sample_loop(self, shape, return_intermediates=False):
method sample (line 474) | def sample(self, batch_size=16, return_intermediates=False):
method q_sample (line 482) | def q_sample(self, x_start, t, noise=None):
method get_v (line 490) | def get_v(self, x, noise, t):
method get_loss (line 496) | def get_loss(self, pred, target, mean=True):
method p_losses (line 511) | def p_losses(self, x_start, t, noise=None):
method forward (line 544) | def forward(self, x, *args, **kwargs):
method get_input (line 552) | def get_input(self, batch, k):
method shared_step (line 560) | def shared_step(self, batch):
method training_step (line 565) | def training_step(self, batch, batch_idx):
method validation_step (line 599) | def validation_step(self, batch, batch_idx):
method on_train_batch_end (line 611) | def on_train_batch_end(self, *args, **kwargs):
method _get_rows_from_list (line 615) | def _get_rows_from_list(self, samples):
method log_images (line 623) | def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=Non...
method configure_optimizers (line 662) | def configure_optimizers(self):
class LatentDiffusion (line 671) | class LatentDiffusion(DDPM):
method __init__ (line 674) | def __init__(
method make_cond_schedule (line 741) | def make_cond_schedule(
method on_train_batch_start (line 755) | def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
method register_schedule (line 778) | def register_schedule(
method instantiate_first_stage (line 795) | def instantiate_first_stage(self, config):
method instantiate_cond_stage (line 802) | def instantiate_cond_stage(self, config):
method _get_denoise_row_from_list (line 823) | def _get_denoise_row_from_list(
method get_first_stage_encoding (line 840) | def get_first_stage_encoding(self, encoder_posterior):
method get_learned_conditioning (line 851) | def get_learned_conditioning(self, c):
method meshgrid (line 866) | def meshgrid(self, h, w):
method delta_border (line 873) | def delta_border(self, h, w):
method get_weighting (line 889) | def get_weighting(self, h, w, Ly, Lx, device):
method get_fold_unfold (line 910) | def get_fold_unfold(
method get_input (line 995) | def get_input(
method decode_first_stage (line 1068) | def decode_first_stage(self, z, predict_cids=False, force_not_quantize...
method decode_first_stage_grad (line 1078) | def decode_first_stage_grad(self, z, predict_cids=False, force_not_qua...
method encode_first_stage (line 1089) | def encode_first_stage(self, x):
method shared_step (line 1092) | def shared_step(self, batch, **kwargs):
method forward (line 1097) | def forward(self, x, c, *args, **kwargs):
method apply_model (line 1111) | def apply_model(self, x_noisy, t, cond, return_ids=False):
method _predict_eps_from_xstart (line 1130) | def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
method _prior_bpd (line 1136) | def _prior_bpd(self, x_start):
method p_mean_variance (line 1152) | def p_mean_variance(
method p_sample (line 1198) | def p_sample(
method progressive_denoising (line 1252) | def progressive_denoising(
method p_sample_loop (line 1346) | def p_sample_loop(
method sample (line 1417) | def sample(
method sample_log (line 1460) | def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
method get_unconditional_conditioning (line 1476) | def get_unconditional_conditioning(self, batch_size, null_label=None):
method log_images (line 1503) | def log_images(
method configure_optimizers (line 1694) | def configure_optimizers(self):
method to_rgb (line 1720) | def to_rgb(self, x):
class DiffusionWrapper (line 1729) | class DiffusionWrapper(torch.nn.Module):
method __init__ (line 1730) | def __init__(self, diff_model_config, conditioning_key):
method forward (line 1747) | def forward(
class LatentUpscaleDiffusion (line 1783) | class LatentUpscaleDiffusion(LatentDiffusion):
method __init__ (line 1784) | def __init__(
method instantiate_low_stage (line 1799) | def instantiate_low_stage(self, config):
method get_input (line 1807) | def get_input(self, batch, k, cond_key=None, bs=None, log_mode=False):
method log_images (line 1835) | def log_images(
class LatentFinetuneDiffusion (line 1975) | class LatentFinetuneDiffusion(LatentDiffusion):
method __init__ (line 1981) | def __init__(
method init_from_ckpt (line 2008) | def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
method log_images (line 2046) | def log_images(
class LatentInpaintDiffusion (line 2164) | class LatentInpaintDiffusion(LatentFinetuneDiffusion):
method __init__ (line 2171) | def __init__(
method get_input (line 2183) | def get_input(
method log_images (line 2223) | def log_images(self, *args, **kwargs):
class LatentDepth2ImageDiffusion (line 2233) | class LatentDepth2ImageDiffusion(LatentFinetuneDiffusion):
method __init__ (line 2238) | def __init__(self, depth_stage_config, concat_keys=("midas_in",), *arg...
method get_input (line 2244) | def get_input(
method log_images (line 2288) | def log_images(self, *args, **kwargs):
class LatentUpscaleFinetuneDiffusion (line 2298) | class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion):
method __init__ (line 2303) | def __init__(
method instantiate_low_stage (line 2321) | def instantiate_low_stage(self, config):
method get_input (line 2329) | def get_input(
method log_images (line 2377) | def log_images(self, *args, **kwargs):
FILE: iopaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py
class NoiseScheduleVP (line 7) | class NoiseScheduleVP:
method __init__ (line 8) | def __init__(
method marginal_log_mean_coeff (line 106) | def marginal_log_mean_coeff(self, t):
method marginal_alpha (line 120) | def marginal_alpha(self, t):
method marginal_std (line 126) | def marginal_std(self, t):
method marginal_lambda (line 132) | def marginal_lambda(self, t):
method inverse_lambda (line 140) | def inverse_lambda(self, lamb):
function model_wrapper (line 161) | def model_wrapper(
class DPM_Solver (line 319) | class DPM_Solver:
method __init__ (line 320) | def __init__(self, model_fn, noise_schedule, predict_x0=False, thresho...
method noise_prediction_fn (line 346) | def noise_prediction_fn(self, x, t):
method data_prediction_fn (line 352) | def data_prediction_fn(self, x, t):
method model_fn (line 367) | def model_fn(self, x, t):
method get_time_steps (line 376) | def get_time_steps(self, skip_type, t_T, t_0, N, device):
method get_orders_and_timesteps_for_singlestep_solver (line 405) | def get_orders_and_timesteps_for_singlestep_solver(self, steps, order,...
method denoise_to_zero_fn (line 463) | def denoise_to_zero_fn(self, x, s):
method dpm_solver_first_update (line 469) | def dpm_solver_first_update(self, x, s, t, model_s=None, return_interm...
method singlestep_dpm_solver_second_update (line 515) | def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s...
method singlestep_dpm_solver_third_update (line 599) | def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2...
method multistep_dpm_solver_second_update (line 723) | def multistep_dpm_solver_second_update(self, x, model_prev_list, t_pre...
method multistep_dpm_solver_third_update (line 780) | def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev...
method singlestep_dpm_solver_update (line 827) | def singlestep_dpm_solver_update(self, x, s, t, order, return_intermed...
method multistep_dpm_solver_update (line 855) | def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list,...
method dpm_solver_adaptive (line 878) | def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0....
method sample (line 939) | def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_...
function interpolate_fn (line 1104) | def interpolate_fn(x, xp, yp):
function expand_dims (line 1145) | def expand_dims(v, dims):
FILE: iopaint/model/anytext/ldm/models/diffusion/dpm_solver/sampler.py
class DPMSolverSampler (line 13) | class DPMSolverSampler(object):
method __init__ (line 14) | def __init__(self, model, **kwargs):
method register_buffer (line 20) | def register_buffer(self, name, attr):
method sample (line 27) | def sample(self,
FILE: iopaint/model/anytext/ldm/models/diffusion/plms.py
class PLMSSampler (line 12) | class PLMSSampler(object):
method __init__ (line 13) | def __init__(self, model, schedule="linear", **kwargs):
method register_buffer (line 19) | def register_buffer(self, name, attr):
method make_schedule (line 25) | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddi...
method sample (line 59) | def sample(self,
method plms_sampling (line 118) | def plms_sampling(self, cond, shape,
method p_sample_plms (line 178) | def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_origin...
FILE: iopaint/model/anytext/ldm/models/diffusion/sampling_util.py
function append_dims (line 5) | def append_dims(x, target_dims):
function norm_thresholding (line 14) | def norm_thresholding(x0, value):
function spatial_norm_thresholding (line 19) | def spatial_norm_thresholding(x0, value):
FILE: iopaint/model/anytext/ldm/modules/attention.py
function exists (line 18) | def exists(val):
function uniq (line 22) | def uniq(arr):
function default (line 26) | def default(val, d):
function max_neg_value (line 32) | def max_neg_value(t):
function init_ (line 36) | def init_(tensor):
class GEGLU (line 44) | class GEGLU(nn.Module):
method __init__ (line 45) | def __init__(self, dim_in, dim_out):
method forward (line 49) | def forward(self, x):
class FeedForward (line 54) | class FeedForward(nn.Module):
method __init__ (line 55) | def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
method forward (line 69) | def forward(self, x):
function zero_module (line 73) | def zero_module(module):
function Normalize (line 82) | def Normalize(in_channels):
class SpatialSelfAttention (line 88) | class SpatialSelfAttention(nn.Module):
method __init__ (line 89) | def __init__(self, in_channels):
method forward (line 107) | def forward(self, x):
class CrossAttention (line 133) | class CrossAttention(nn.Module):
method __init__ (line 134) | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, ...
method forward (line 150) | def forward(self, x, context=None, mask=None):
class SDPACrossAttention (line 184) | class SDPACrossAttention(CrossAttention):
method forward (line 185) | def forward(self, x, context=None, mask=None):
class BasicTransformerBlock (line 227) | class BasicTransformerBlock(nn.Module):
method __init__ (line 228) | def __init__(
method forward (line 267) | def forward(self, x, context=None):
method _forward (line 272) | def _forward(self, x, context=None):
class SpatialTransformer (line 284) | class SpatialTransformer(nn.Module):
method __init__ (line 294) | def __init__(
method forward (line 341) | def forward(self, x, context=None):
FILE: iopaint/model/anytext/ldm/modules/diffusionmodules/model.py
function get_timestep_embedding (line 9) | def get_timestep_embedding(timesteps, embedding_dim):
function nonlinearity (line 30) | def nonlinearity(x):
function Normalize (line 35) | def Normalize(in_channels, num_groups=32):
class Upsample (line 41) | class Upsample(nn.Module):
method __init__ (line 42) | def __init__(self, in_channels, with_conv):
method forward (line 50) | def forward(self, x):
class Downsample (line 57) | class Downsample(nn.Module):
method __init__ (line 58) | def __init__(self, in_channels, with_conv):
method forward (line 67) | def forward(self, x):
class ResnetBlock (line 77) | class ResnetBlock(nn.Module):
method __init__ (line 78) | def __init__(
method forward (line 114) | def forward(self, x, temb):
class AttnBlock (line 137) | class AttnBlock(nn.Module):
method __init__ (line 138) | def __init__(self, in_channels):
method forward (line 156) | def forward(self, x):
class AttnBlock2_0 (line 183) | class AttnBlock2_0(nn.Module):
method __init__ (line 184) | def __init__(self, in_channels):
method forward (line 202) | def forward(self, x):
function make_attn (line 231) | def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
class Model (line 246) | class Model(nn.Module):
method __init__ (line 247) | def __init__(
method forward (line 368) | def forward(self, x, t=None, context=None):
method get_last_layer (line 417) | def get_last_layer(self):
class Encoder (line 421) | class Encoder(nn.Module):
method __init__ (line 422) | def __init__(
method forward (line 510) | def forward(self, x):
class Decoder (line 538) | class Decoder(nn.Module):
method __init__ (line 539) | def __init__(
method forward (line 634) | def forward(self, z):
class SimpleDecoder (line 670) | class SimpleDecoder(nn.Module):
method __init__ (line 671) | def __init__(self, in_channels, out_channels, *args, **kwargs):
method forward (line 704) | def forward(self, x):
class UpsampleDecoder (line 717) | class UpsampleDecoder(nn.Module):
method __init__ (line 718) | def __init__(
method forward (line 761) | def forward(self, x):
class LatentRescaler (line 775) | class LatentRescaler(nn.Module):
method __init__ (line 776) | def __init__(self, factor, in_channels, mid_channels, out_channels, de...
method forward (line 813) | def forward(self, x):
class MergedRescaleEncoder (line 831) | class MergedRescaleEncoder(nn.Module):
method __init__ (line 832) | def __init__(
method forward (line 869) | def forward(self, x):
class MergedRescaleDecoder (line 875) | class MergedRescaleDecoder(nn.Module):
method __init__ (line 876) | def __init__(
method forward (line 912) | def forward(self, x):
class Upsampler (line 918) | class Upsampler(nn.Module):
method __init__ (line 919) | def __init__(self, in_size, out_size, in_channels, out_channels, ch_mu...
method forward (line 944) | def forward(self, x):
class Resize (line 950) | class Resize(nn.Module):
method __init__ (line 951) | def __init__(self, in_channels=None, learned=False, mode="bilinear"):
method forward (line 966) | def forward(self, x, scale_factor=1.0):
FILE: iopaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py
function convert_module_to_f16 (line 23) | def convert_module_to_f16(x):
function convert_module_to_f32 (line 26) | def convert_module_to_f32(x):
class AttentionPool2d (line 31) | class AttentionPool2d(nn.Module):
method __init__ (line 36) | def __init__(
method forward (line 50) | def forward(self, x):
class TimestepBlock (line 61) | class TimestepBlock(nn.Module):
method forward (line 67) | def forward(self, x, emb):
class TimestepEmbedSequential (line 73) | class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
method forward (line 79) | def forward(self, x, emb, context=None):
class Upsample (line 90) | class Upsample(nn.Module):
method __init__ (line 99) | def __init__(self, channels, use_conv, dims=2, out_channels=None, padd...
method forward (line 108) | def forward(self, x):
class TransposedUpsample (line 120) | class TransposedUpsample(nn.Module):
method __init__ (line 122) | def __init__(self, channels, out_channels=None, ks=5):
method forward (line 129) | def forward(self,x):
class Downsample (line 133) | class Downsample(nn.Module):
method __init__ (line 142) | def __init__(self, channels, use_conv, dims=2, out_channels=None,paddi...
method forward (line 157) | def forward(self, x):
class ResBlock (line 162) | class ResBlock(TimestepBlock):
method __init__ (line 178) | def __init__(
method forward (line 242) | def forward(self, x, emb):
method _forward (line 254) | def _forward(self, x, emb):
class AttentionBlock (line 277) | class AttentionBlock(nn.Module):
method __init__ (line 284) | def __init__(
method forward (line 313) | def forward(self, x):
method _forward (line 317) | def _forward(self, x):
function count_flops_attn (line 326) | def count_flops_attn(model, _x, y):
class QKVAttentionLegacy (line 346) | class QKVAttentionLegacy(nn.Module):
method __init__ (line 351) | def __init__(self, n_heads):
method forward (line 355) | def forward(self, qkv):
method count_flops (line 374) | def count_flops(model, _x, y):
class QKVAttention (line 378) | class QKVAttention(nn.Module):
method __init__ (line 383) | def __init__(self, n_heads):
method forward (line 387) | def forward(self, qkv):
method count_flops (line 408) | def count_flops(model, _x, y):
class UNetModel (line 412) | class UNetModel(nn.Module):
method __init__ (line 442) | def __init__(
method convert_to_fp16 (line 738) | def convert_to_fp16(self):
method convert_to_fp32 (line 746) | def convert_to_fp32(self):
method forward (line 754) | def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
FILE: iopaint/model/anytext/ldm/modules/diffusionmodules/upscaling.py
class AbstractLowScaleModel (line 10) | class AbstractLowScaleModel(nn.Module):
method __init__ (line 12) | def __init__(self, noise_schedule_config=None):
method register_schedule (line 17) | def register_schedule(self, beta_schedule="linear", timesteps=1000,
method q_sample (line 44) | def q_sample(self, x_start, t, noise=None):
method forward (line 49) | def forward(self, x):
method decode (line 52) | def decode(self, x):
class SimpleImageConcat (line 56) | class SimpleImageConcat(AbstractLowScaleModel):
method __init__ (line 58) | def __init__(self):
method forward (line 62) | def forward(self, x):
class ImageConcatWithNoiseAugmentation (line 67) | class ImageConcatWithNoiseAugmentation(AbstractLowScaleModel):
method __init__ (line 68) | def __init__(self, noise_schedule_config, max_noise_level=1000, to_cud...
method forward (line 72) | def forward(self, x, noise_level=None):
FILE: iopaint/model/anytext/ldm/modules/diffusionmodules/util.py
function make_beta_schedule (line 21) | def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_e...
function make_ddim_timesteps (line 46) | def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_...
function make_ddim_sampling_parameters (line 63) | def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbos...
function betas_for_alpha_bar (line 77) | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.9...
function extract_into_tensor (line 96) | def extract_into_tensor(a, t, x_shape):
function checkpoint (line 102) | def checkpoint(func, inputs, params, flag):
class CheckpointFunction (line 119) | class CheckpointFunction(torch.autograd.Function):
method forward (line 121) | def forward(ctx, run_function, length, *args):
method backward (line 133) | def backward(ctx, *output_grads):
function timestep_embedding (line 154) | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=Fal...
function zero_module (line 177) | def zero_module(module):
function scale_module (line 186) | def scale_module(module, scale):
function mean_flat (line 195) | def mean_flat(tensor):
function normalization (line 202) | def normalization(channels):
class SiLU (line 212) | class SiLU(nn.Module):
method forward (line 213) | def forward(self, x):
class GroupNorm32 (line 217) | class GroupNorm32(nn.GroupNorm):
method forward (line 218) | def forward(self, x):
function conv_nd (line 222) | def conv_nd(dims, *args, **kwargs):
function linear (line 235) | def linear(*args, **kwargs):
function avg_pool_nd (line 242) | def avg_pool_nd(dims, *args, **kwargs):
class HybridConditioner (line 255) | class HybridConditioner(nn.Module):
method __init__ (line 257) | def __init__(self, c_concat_config, c_crossattn_config):
method forward (line 262) | def forward(self, c_concat, c_crossattn):
function noise_like (line 268) | def noise_like(shape, device, repeat=False):
FILE: iopaint/model/anytext/ldm/modules/distributions/distributions.py
class AbstractDistribution (line 5) | class AbstractDistribution:
method sample (line 6) | def sample(self):
method mode (line 9) | def mode(self):
class DiracDistribution (line 13) | class DiracDistribution(AbstractDistribution):
method __init__ (line 14) | def __init__(self, value):
method sample (line 17) | def sample(self):
method mode (line 20) | def mode(self):
class DiagonalGaussianDistribution (line 24) | class DiagonalGaussianDistribution(object):
method __init__ (line 25) | def __init__(self, parameters, deterministic=False):
method sample (line 35) | def sample(self):
method kl (line 39) | def kl(self, other=None):
method nll (line 53) | def nll(self, sample, dims=[1,2,3]):
method mode (line 61) | def mode(self):
function normal_kl (line 65) | def normal_kl(mean1, logvar1, mean2, logvar2):
FILE: iopaint/model/anytext/ldm/modules/ema.py
class LitEma (line 5) | class LitEma(nn.Module):
method __init__ (line 6) | def __init__(self, model, decay=0.9999, use_num_upates=True):
method reset_num_updates (line 25) | def reset_num_updates(self):
method forward (line 29) | def forward(self, model):
method copy_to (line 50) | def copy_to(self, model):
method store (line 59) | def store(self, parameters):
method restore (line 68) | def restore(self, parameters):
FILE: iopaint/model/anytext/ldm/modules/encoders/modules.py
function _expand_mask (line 17) | def _expand_mask(mask, dtype, tgt_len=None):
function _build_causal_attention_mask (line 33) | def _build_causal_attention_mask(bsz, seq_len, dtype):
class AbstractEncoder (line 43) | class AbstractEncoder(nn.Module):
method __init__ (line 44) | def __init__(self):
method encode (line 47) | def encode(self, *args, **kwargs):
class IdentityEncoder (line 51) | class IdentityEncoder(AbstractEncoder):
method encode (line 52) | def encode(self, x):
class ClassEmbedder (line 56) | class ClassEmbedder(nn.Module):
method __init__ (line 57) | def __init__(self, embed_dim, n_classes=1000, key="class", ucg_rate=0.1):
method forward (line 64) | def forward(self, batch, key=None, disable_dropout=False):
method get_unconditional_conditioning (line 76) | def get_unconditional_conditioning(self, bs, device="cuda"):
function disabled_train (line 85) | def disabled_train(self, mode=True):
class FrozenT5Embedder (line 91) | class FrozenT5Embedder(AbstractEncoder):
method __init__ (line 94) | def __init__(
method freeze (line 105) | def freeze(self):
method forward (line 111) | def forward(self, text):
method encode (line 127) | def encode(self, text):
class FrozenCLIPEmbedder (line 131) | class FrozenCLIPEmbedder(AbstractEncoder):
method __init__ (line 136) | def __init__(
method freeze (line 159) | def freeze(self):
method forward (line 165) | def forward(self, text):
method encode (line 187) | def encode(self, text):
class FrozenCLIPT5Encoder (line 191) | class FrozenCLIPT5Encoder(AbstractEncoder):
method __init__ (line 192) | def __init__(
method encode (line 210) | def encode(self, text):
method forward (line 213) | def forward(self, text):
class FrozenCLIPEmbedderT3 (line 219) | class FrozenCLIPEmbedderT3(AbstractEncoder):
method __init__ (line 222) | def __init__(
method freeze (line 391) | def freeze(self):
method forward (line 396) | def forward(self, text, **kwargs):
method encode (line 410) | def encode(self, text, **kwargs):
FILE: iopaint/model/anytext/ldm/util.py
function log_txt_as_img (line 11) | def log_txt_as_img(wh, xc, size=10):
function ismap (line 35) | def ismap(x):
function isimage (line 41) | def isimage(x):
function exists (line 47) | def exists(x):
function default (line 51) | def default(val, d):
function mean_flat (line 57) | def mean_flat(tensor):
function count_params (line 65) | def count_params(model, verbose=False):
function instantiate_from_config (line 72) | def instantiate_from_config(config, **kwargs):
function get_obj_from_str (line 82) | def get_obj_from_str(string, reload=False):
class AdamWwithEMAandWings (line 90) | class AdamWwithEMAandWings(optim.Optimizer):
method __init__ (line 92) | def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, #...
method __setstate__ (line 113) | def __setstate__(self, state):
method step (line 119) | def step(self, closure=None):
FILE: iopaint/model/anytext/ocr_recog/RNN.py
class Swish (line 5) | class Swish(nn.Module):
method __int__ (line 6) | def __int__(self):
method forward (line 9) | def forward(self,x):
class Im2Im (line 12) | class Im2Im(nn.Module):
method __init__ (line 13) | def __init__(self, in_channels, **kwargs):
method forward (line 17) | def forward(self, x):
class Im2Seq (line 20) | class Im2Seq(nn.Module):
method __init__ (line 21) | def __init__(self, in_channels, **kwargs):
method forward (line 25) | def forward(self, x):
class EncoderWithRNN (line 32) | class EncoderWithRNN(nn.Module):
method __init__ (line 33) | def __init__(self, in_channels,**kwargs):
method forward (line 39) | def forward(self, x):
class SequenceEncoder (line 44) | class SequenceEncoder(nn.Module):
method __init__ (line 45) | def __init__(self, in_channels, encoder_type='rnn', **kwargs):
method forward (line 66) | def forward(self, x):
class ConvBNLayer (line 77) | class ConvBNLayer(nn.Module):
method __init__ (line 78) | def __init__(self,
method forward (line 100) | def forward(self, inputs):
class EncoderWithSVTR (line 107) | class EncoderWithSVTR(nn.Module):
method __init__ (line 108) | def __init__(
method _init_weights (line 159) | def _init_weights(self, m):
method forward (line 180) | def forward(self, x):
FILE: iopaint/model/anytext/ocr_recog/RecCTCHead.py
class CTCHead (line 4) | class CTCHead(nn.Module):
method __init__ (line 5) | def __init__(self,
method forward (line 34) | def forward(self, x, labels=None):
FILE: iopaint/model/anytext/ocr_recog/RecModel.py
class RecModel (line 12) | class RecModel(nn.Module):
method __init__ (line 13) | def __init__(self, config):
method load_3rd_state_dict (line 30) | def load_3rd_state_dict(self, _3rd_name, _state):
method forward (line 35) | def forward(self, x):
method encode (line 41) | def encode(self, x):
FILE: iopaint/model/anytext/ocr_recog/RecMv1_enhance.py
class ConvBNLayer (line 7) | class ConvBNLayer(nn.Module):
method __init__ (line 8) | def __init__(self,
method forward (line 34) | def forward(self, inputs):
class DepthwiseSeparable (line 42) | class DepthwiseSeparable(nn.Module):
method __init__ (line 43) | def __init__(self,
method forward (line 71) | def forward(self, inputs):
class MobileNetV1Enhance (line 79) | class MobileNetV1Enhance(nn.Module):
method __init__ (line 80) | def __init__(self,
method forward (line 196) | def forward(self, inputs):
function hardsigmoid (line 202) | def hardsigmoid(x):
class SEModule (line 205) | class SEModule(nn.Module):
method __init__ (line 206) | def __init__(self, channel, reduction=4):
method forward (line 224) | def forward(self, inputs):
FILE: iopaint/model/anytext/ocr_recog/RecSVTR.py
function drop_path (line 8) | def drop_path(x, drop_prob=0., training=False):
class Swish (line 23) | class Swish(nn.Module):
method __int__ (line 24) | def __int__(self):
method forward (line 27) | def forward(self,x):
class ConvBNLayer (line 31) | class ConvBNLayer(nn.Module):
method __init__ (line 32) | def __init__(self,
method forward (line 54) | def forward(self, inputs):
class DropPath (line 61) | class DropPath(nn.Module):
method __init__ (line 65) | def __init__(self, drop_prob=None):
method forward (line 69) | def forward(self, x):
class Identity (line 73) | class Identity(nn.Module):
method __init__ (line 74) | def __init__(self):
method forward (line 77) | def forward(self, input):
class Mlp (line 81) | class Mlp(nn.Module):
method __init__ (line 82) | def __init__(self,
method forward (line 99) | def forward(self, x):
class ConvMixer (line 108) | class ConvMixer(nn.Module):
method __init__ (line 109) | def __init__(
method forward (line 127) | def forward(self, x):
class Attention (line 136) | class Attention(nn.Module):
method __init__ (line 137) | def __init__(self,
method forward (line 177) | def forward(self, x):
class Block (line 198) | class Block(nn.Module):
method __init__ (line 199) | def __init__(self,
method forward (line 251) | def forward(self, x):
class PatchEmbed (line 261) | class PatchEmbed(nn.Module):
method __init__ (line 265) | def __init__(self,
method forward (line 322) | def forward(self, x):
class SubSample (line 330) | class SubSample(nn.Module):
method __init__ (line 331) | def __init__(self,
method forward (line 361) | def forward(self, x):
class SVTRNet (line 378) | class SVTRNet(nn.Module):
method __init__ (line 379) | def __init__(
method _init_weights (line 533) | def _init_weights(self, m):
method forward_features (line 542) | def forward_features(self, x):
method forward (line 564) | def forward(self, x):
FILE: iopaint/model/anytext/ocr_recog/common.py
class Hswish (line 8) | class Hswish(nn.Module):
method __init__ (line 9) | def __init__(self, inplace=True):
method forward (line 13) | def forward(self, x):
class Hsigmoid (line 18) | class Hsigmoid(nn.Module):
method __init__ (line 19) | def __init__(self, inplace=True):
method forward (line 23) | def forward(self, x):
class GELU (line 28) | class GELU(nn.Module):
method __init__ (line 29) | def __init__(self, inplace=True):
method forward (line 33) | def forward(self, x):
class Swish (line 37) | class Swish(nn.Module):
method __init__ (line 38) | def __init__(self, inplace=True):
method forward (line 42) | def forward(self, x):
class Activation (line 50) | class Activation(nn.Module):
method __init__ (line 51) | def __init__(self, act_type, inplace=True):
method forward (line 73) | def forward(self, inputs):
FILE: iopaint/model/anytext/utils.py
function save_images (line 8) | def save_images(img_list, folder):
function check_channels (line 24) | def check_channels(image):
function resize_image (line 33) | def resize_image(img, max_length=768):
function insert_spaces (line 48) | def insert_spaces(string, nSpace):
function draw_glyph (line 57) | def draw_glyph(font, text):
function draw_glyph2 (line 78) | def draw_glyph2(
FILE: iopaint/model/base.py
class InpaintModel (line 20) | class InpaintModel:
method __init__ (line 27) | def __init__(self, device, **kwargs):
method init_model (line 38) | def init_model(self, device, **kwargs): ...
method is_downloaded (line 42) | def is_downloaded() -> bool:
method forward (line 46) | def forward(self, image, mask, config: InpaintRequest):
method download (line 55) | def download(): ...
method _pad_forward (line 57) | def _pad_forward(self, image, mask, config: InpaintRequest):
method forward_pre_process (line 80) | def forward_pre_process(self, image, mask, config):
method forward_post_process (line 83) | def forward_post_process(self, result, image, mask, config):
method __call__ (line 87) | def __call__(self, image, mask, config: InpaintRequest):
method _crop_box (line 142) | def _crop_box(self, image, mask, box, config: InpaintRequest):
method _calculate_cdf (line 194) | def _calculate_cdf(self, histogram):
method _calculate_lookup (line 199) | def _calculate_lookup(self, source_cdf, reference_cdf):
method _match_histograms (line 210) | def _match_histograms(self, source, reference, mask):
method _apply_cropper (line 237) | def _apply_cropper(self, image, mask, config: InpaintRequest):
method _run_box (line 257) | def _run_box(self, image, mask, box, config: InpaintRequest):
class DiffusionInpaintModel (line 273) | class DiffusionInpaintModel(InpaintModel):
method __init__ (line 274) | def __init__(self, device, **kwargs):
method __call__ (line 280) | def __call__(self, image, mask, config: InpaintRequest):
method _do_outpainting (line 299) | def _do_outpainting(self, image, config: InpaintRequest):
method _scaled_pad_forward (line 363) | def _scaled_pad_forward(self, image, mask, config: InpaintRequest):
method set_scheduler (line 382) | def set_scheduler(self, config: InpaintRequest):
method forward_pre_process (line 391) | def forward_pre_process(self, image, mask, config):
method forward_post_process (line 398) | def forward_post_process(self, result, image, mask, config):
FILE: iopaint/model/brushnet/brushnet.py
class BrushNetOutput (line 31) | class BrushNetOutput(BaseOutput):
class BrushNetModel (line 55) | class BrushNetModel(ModelMixin, ConfigMixin):
method __init__ (line 132) | def __init__(
method from_unet (line 445) | def from_unet(
method attn_processors (line 526) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 550) | def set_attn_processor(self, processor: Union[AttentionProcessor, Dict...
method set_default_attn_processor (line 585) | def set_default_attn_processor(self):
method set_attention_slice (line 601) | def set_attention_slice(self, slice_size: Union[str, int, List[int]]) ...
method _set_gradient_checkpointing (line 666) | def _set_gradient_checkpointing(self, module, value: bool = False) -> ...
method forward (line 670) | def forward(
function zero_module (line 923) | def zero_module(module):
FILE: iopaint/model/brushnet/brushnet_unet_forward.py
function brushnet_unet_forward (line 13) | def brushnet_unet_forward(
FILE: iopaint/model/brushnet/brushnet_wrapper.py
class BrushNetWrapper (line 27) | class BrushNetWrapper(DiffusionInpaintModel):
method init_model (line 31) | def init_model(self, device: torch.device, **kwargs):
method switch_brushnet_method (line 136) | def switch_brushnet_method(self, new_method: str):
method forward (line 145) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/brushnet/brushnet_xl_wrapper.py
class BrushNetXLWrapper (line 28) | class BrushNetXLWrapper(DiffusionInpaintModel):
method init_model (line 34) | def init_model(self, device: torch.device, **kwargs):
method switch_brushnet_method (line 140) | def switch_brushnet_method(self, new_method: str):
method forward (line 149) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/brushnet/pipeline_brushnet.py
function retrieve_timesteps (line 85) | def retrieve_timesteps(
class StableDiffusionBrushNetPipeline (line 129) | class StableDiffusionBrushNetPipeline(
method __init__ (line 177) | def __init__(
method _encode_prompt (line 224) | def _encode_prompt(
method encode_prompt (line 257) | def encode_prompt(
method encode_image (line 439) | def encode_image(self, image, device, num_images_per_prompt, output_hi...
method prepare_ip_adapter_image_embeds (line 464) | def prepare_ip_adapter_image_embeds(
method run_safety_checker (line 516) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 531) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 543) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 560) | def check_inputs(
method check_image (line 675) | def check_image(self, image, mask, prompt, prompt_embeds):
method prepare_image (line 731) | def prepare_image(
method prepare_latents (line 762) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method get_guidance_scale_embedding (line 780) | def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=tor...
method guidance_scale (line 809) | def guidance_scale(self):
method clip_skip (line 813) | def clip_skip(self):
method do_classifier_free_guidance (line 820) | def do_classifier_free_guidance(self):
method cross_attention_kwargs (line 824) | def cross_attention_kwargs(self):
method num_timesteps (line 828) | def num_timesteps(self):
method __call__ (line 833) | def __call__(
FILE: iopaint/model/brushnet/pipeline_brushnet_sd_xl.py
class StableDiffusionXLBrushNetPipeline (line 103) | class StableDiffusionXLBrushNetPipeline(
method __init__ (line 166) | def __init__(
method encode_prompt (line 213) | def encode_prompt(
method encode_image (line 448) | def encode_image(self, image, device, num_images_per_prompt, output_hi...
method prepare_ip_adapter_image_embeds (line 473) | def prepare_ip_adapter_image_embeds(
method prepare_extra_step_kwargs (line 525) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 542) | def check_inputs(
method check_image (line 733) | def check_image(self, image, mask, prompt, prompt_embeds):
method prepare_image (line 791) | def prepare_image(
method prepare_latents (line 822) | def prepare_latents(self, batch_size, num_channels_latents, height, wi...
method _get_add_time_ids (line 840) | def _get_add_time_ids(
method upcast_vae (line 859) | def upcast_vae(self):
method get_guidance_scale_embedding (line 879) | def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=tor...
method guidance_scale (line 908) | def guidance_scale(self):
method clip_skip (line 912) | def clip_skip(self):
method do_classifier_free_guidance (line 919) | def do_classifier_free_guidance(self):
method cross_attention_kwargs (line 923) | def cross_attention_kwargs(self):
method denoising_end (line 927) | def denoising_end(self):
method num_timesteps (line 931) | def num_timesteps(self):
method __call__ (line 936) | def __call__(
FILE: iopaint/model/brushnet/unet_2d_blocks.py
class MidBlock2D (line 10) | class MidBlock2D(nn.Module):
method __init__ (line 11) | def __init__(
method forward (line 66) | def forward(
function DownBlock2D_forward (line 98) | def DownBlock2D_forward(
function CrossAttnDownBlock2D_forward (line 141) | def CrossAttnDownBlock2D_forward(
function CrossAttnUpBlock2D_forward (line 217) | def CrossAttnUpBlock2D_forward(
function UpBlock2D_forward (line 314) | def UpBlock2D_forward(
FILE: iopaint/model/controlnet.py
class ControlNet (line 26) | class ControlNet(DiffusionInpaintModel):
method lcm_lora_id (line 32) | def lcm_lora_id(self):
method init_model (line 45) | def init_model(self, device: torch.device, **kwargs):
method switch_controlnet_method (line 139) | def switch_controlnet_method(self, new_method: str):
method _get_control_image (line 148) | def _get_control_image(self, image, mask):
method forward (line 161) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/ddim_sampler.py
class DDIMSampler (line 10) | class DDIMSampler(object):
method __init__ (line 11) | def __init__(self, model, schedule="linear"):
method register_buffer (line 17) | def register_buffer(self, name, attr):
method make_schedule (line 20) | def make_schedule(
method sample (line 82) | def sample(self, steps, conditioning, batch_size, shape):
method ddim_sampling (line 99) | def ddim_sampling(
method p_sample_ddim (line 144) | def p_sample_ddim(
FILE: iopaint/model/fcf.py
function upfirdn2d (line 39) | def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, ...
function _upfirdn2d_ref (line 46) | def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gai...
class EncoderEpilogue (line 94) | class EncoderEpilogue(torch.nn.Module):
method __init__ (line 95) | def __init__(
method forward (line 139) | def forward(self, x, cmap, force_fp32=False):
class EncoderBlock (line 162) | class EncoderBlock(torch.nn.Module):
method __init__ (line 163) | def __init__(
method forward (line 252) | def forward(self, x, img, force_fp32=False):
class EncoderNetwork (line 292) | class EncoderNetwork(torch.nn.Module):
method __init__ (line 293) | def __init__(
method forward (line 369) | def forward(self, img, c, **block_kwargs):
function fma (line 390) | def fma(a, b, c): # => a * b + c
class _FusedMultiplyAdd (line 394) | class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c
method forward (line 396) | def forward(ctx, a, b, c): # pylint: disable=arguments-differ
method backward (line 403) | def backward(ctx, dout): # pylint: disable=arguments-differ
function _unbroadcast (line 422) | def _unbroadcast(x, shape):
function modulated_conv2d (line 438) | def modulated_conv2d(
class SynthesisLayer (line 516) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 517) | def __init__(
method forward (line 560) | def forward(self, x, w, noise_mode="none", fused_modconv=True, gain=1):
class ToRGBLayer (line 599) | class ToRGBLayer(torch.nn.Module):
method __init__ (line 600) | def __init__(
method forward (line 623) | def forward(self, x, w, fused_modconv=True):
class SynthesisForeword (line 636) | class SynthesisForeword(torch.nn.Module):
method __init__ (line 637) | def __init__(
method forward (line 668) | def forward(self, x, ws, feats, img, force_fp32=False):
class SELayer (line 703) | class SELayer(nn.Module):
method __init__ (line 704) | def __init__(self, channel, reduction=16):
method forward (line 714) | def forward(self, x):
class FourierUnit (line 722) | class FourierUnit(nn.Module):
method __init__ (line 723) | def __init__(
method forward (line 764) | def forward(self, x):
class SpectralTransform (line 840) | class SpectralTransform(nn.Module):
method __init__ (line 841) | def __init__(
method forward (line 873) | def forward(self, x):
class FFC (line 896) | class FFC(nn.Module):
method __init__ (line 897) | def __init__(
method forward (line 982) | def forward(self, x, fname=None):
class FFC_BN_ACT (line 1007) | class FFC_BN_ACT(nn.Module):
method __init__ (line 1008) | def __init__(
method forward (line 1053) | def forward(self, x, fname=None):
class FFCResnetBlock (line 1063) | class FFCResnetBlock(nn.Module):
method __init__ (line 1064) | def __init__(
method forward (line 1103) | def forward(self, x, fname=None):
class ConcatTupleLayer (line 1124) | class ConcatTupleLayer(nn.Module):
method forward (line 1125) | def forward(self, x):
class FFCBlock (line 1134) | class FFCBlock(torch.nn.Module):
method __init__ (line 1135) | def __init__(
method forward (line 1163) | def forward(self, gen_ft, mask, fname=None):
class FFCSkipLayer (line 1179) | class FFCSkipLayer(torch.nn.Module):
method __init__ (line 1180) | def __init__(
method forward (line 1199) | def forward(self, gen_ft, mask, fname=None):
class SynthesisBlock (line 1204) | class SynthesisBlock(torch.nn.Module):
method __init__ (line 1205) | def __init__(
method forward (line 1296) | def forward(
class SynthesisNetwork (line 1384) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 1385) | def __init__(
method forward (line 1439) | def forward(self, x_global, mask, feats, ws, fname=None, **block_kwargs):
class MappingNetwork (line 1472) | class MappingNetwork(torch.nn.Module):
method __init__ (line 1473) | def __init__(
method forward (line 1520) | def forward(
class Generator (line 1562) | class Generator(torch.nn.Module):
method __init__ (line 1563) | def __init__(
method forward (line 1599) | def forward(
class FcF (line 1624) | class FcF(InpaintModel):
method init_model (line 1631) | def init_model(self, device, **kwargs):
method download (line 1660) | def download():
method is_downloaded (line 1664) | def is_downloaded() -> bool:
method __call__ (line 1668) | def __call__(self, image, mask, config: InpaintRequest):
method forward (line 1708) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/helper/controlnet_preprocess.py
function make_canny_control_image (line 10) | def make_canny_control_image(image: np.ndarray) -> Image:
function make_openpose_control_image (line 19) | def make_openpose_control_image(image: np.ndarray) -> Image:
function resize_image (line 27) | def resize_image(input_image, resolution):
function make_depth_control_image (line 44) | def make_depth_control_image(image: np.ndarray) -> Image:
function make_inpaint_control_image (line 59) | def make_inpaint_control_image(image: np.ndarray, mask: np.ndarray) -> t...
FILE: iopaint/model/helper/cpu_text_encoder.py
class CPUTextEncoderWrapper (line 7) | class CPUTextEncoderWrapper(PreTrainedModel):
method __init__ (line 8) | def __init__(self, text_encoder, torch_dtype):
method __call__ (line 19) | def __call__(self, x, **kwargs):
method dtype (line 32) | def dtype(self):
method device (line 36) | def device(self) -> torch.device:
FILE: iopaint/model/helper/g_diffuser_bot.py
function expand_image (line 5) | def expand_image(cv2_img, top: int, right: int, bottom: int, left: int):
FILE: iopaint/model/instruct_pix2pix.py
class InstructPix2Pix (line 12) | class InstructPix2Pix(DiffusionInpaintModel):
method init_model (line 17) | def init_model(self, device: torch.device, **kwargs):
method forward (line 44) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/kandinsky.py
class Kandinsky (line 12) | class Kandinsky(DiffusionInpaintModel):
method init_model (line 16) | def init_model(self, device: torch.device, **kwargs):
method forward (line 32) | def forward(self, image, mask, config: InpaintRequest):
class Kandinsky22 (line 64) | class Kandinsky22(Kandinsky):
FILE: iopaint/model/lama.py
class LaMa (line 31) | class LaMa(InpaintModel):
method download (line 37) | def download():
method init_model (line 40) | def init_model(self, device, **kwargs):
method is_downloaded (line 44) | def is_downloaded() -> bool:
method forward (line 47) | def forward(self, image, mask, config: InpaintRequest):
class AnimeLaMa (line 68) | class AnimeLaMa(LaMa):
method download (line 72) | def download():
method init_model (line 75) | def init_model(self, device, **kwargs):
method is_downloaded (line 81) | def is_downloaded() -> bool:
FILE: iopaint/model/ldm.py
class DDPM (line 51) | class DDPM(nn.Module):
method __init__ (line 53) | def __init__(
method register_schedule (line 84) | def register_schedule(
class LatentDiffusion (line 177) | class LatentDiffusion(DDPM):
method __init__ (line 178) | def __init__(
method make_cond_schedule (line 200) | def make_cond_schedule(
method register_schedule (line 213) | def register_schedule(
method apply_model (line 230) | def apply_model(self, x_noisy, t, cond):
class LDM (line 237) | class LDM(InpaintModel):
method __init__ (line 242) | def __init__(self, device, fp16: bool = True, **kwargs):
method init_model (line 247) | def init_model(self, device, **kwargs):
method download (line 265) | def download():
method is_downloaded (line 271) | def is_downloaded() -> bool:
method forward (line 280) | def forward(self, image, mask, config: InpaintRequest):
method _norm (line 335) | def _norm(self, tensor):
FILE: iopaint/model/manga.py
class Manga (line 32) | class Manga(InpaintModel):
method init_model (line 37) | def init_model(self, device, **kwargs):
method download (line 47) | def download():
method is_downloaded (line 52) | def is_downloaded() -> bool:
method forward (line 59) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/mat.py
class ModulatedConv2d (line 34) | class ModulatedConv2d(nn.Module):
method __init__ (line 35) | def __init__(
method forward (line 69) | def forward(self, x, style):
class StyleConv (line 96) | class StyleConv(torch.nn.Module):
method __init__ (line 97) | def __init__(
method forward (line 140) | def forward(self, x, style, noise_mode="random", gain=1):
class ToRGB (line 165) | class ToRGB(torch.nn.Module):
method __init__ (line 166) | def __init__(
method forward (line 191) | def forward(self, x, style, skip=None):
function get_style_code (line 203) | def get_style_code(a, b):
class DecBlockFirst (line 207) | class DecBlockFirst(nn.Module):
method __init__ (line 208) | def __init__(
method forward (line 242) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
class DecBlockFirstV2 (line 253) | class DecBlockFirstV2(nn.Module):
method __init__ (line 254) | def __init__(
method forward (line 289) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
method __init__ (line 1280) | def __init__(
method forward (line 1318) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
class DecBlock (line 301) | class DecBlock(nn.Module):
method __init__ (line 302) | def __init__(
method forward (line 345) | def forward(self, x, img, ws, gs, E_features, noise_mode="random"):
method __init__ (line 1331) | def __init__(
method forward (line 1374) | def forward(self, x, img, ws, gs, E_features, noise_mode="random"):
class MappingNet (line 357) | class MappingNet(torch.nn.Module):
method __init__ (line 358) | def __init__(
method forward (line 407) | def forward(
class DisFromRGB (line 444) | class DisFromRGB(nn.Module):
method __init__ (line 445) | def __init__(
method forward (line 456) | def forward(self, x):
class DisBlock (line 460) | class DisBlock(nn.Module):
method __init__ (line 461) | def __init__(
method forward (line 486) | def forward(self, x):
class Discriminator (line 495) | class Discriminator(torch.nn.Module):
method __init__ (line 496) | def __init__(
method forward (line 554) | def forward(self, images_in, masks_in, c):
method __init__ (line 1770) | def __init__(
method forward (line 1851) | def forward(self, images_in, masks_in, images_stg1, c):
function nf (line 568) | def nf(stage, channel_base=32768, channel_decay=1.0, channel_max=512):
class Mlp (line 573) | class Mlp(nn.Module):
method __init__ (line 574) | def __init__(
method forward (line 592) | def forward(self, x):
function window_partition (line 598) | def window_partition(x, window_size):
function window_reverse (line 614) | def window_reverse(windows, window_size: int, H: int, W: int):
class Conv2dLayerPartial (line 633) | class Conv2dLayerPartial(nn.Module):
method __init__ (line 634) | def __init__(
method forward (line 671) | def forward(self, x, mask=None):
class WindowAttention (line 694) | class WindowAttention(nn.Module):
method __init__ (line 707) | def __init__(
method forward (line 732) | def forward(self, x, mask_windows=None, mask=None):
class SwinTransformerBlock (line 781) | class SwinTransformerBlock(nn.Module):
method __init__ (line 799) | def __init__(
method calculate_mask (line 863) | def calculate_mask(self, x_size):
method forward (line 894) | def forward(self, x, x_size, mask=None):
class PatchMerging (line 975) | class PatchMerging(nn.Module):
method __init__ (line 976) | def __init__(self, in_channels, out_channels, down=2):
method forward (line 987) | def forward(self, x, x_size, mask=None):
class PatchUpsampling (line 1001) | class PatchUpsampling(nn.Module):
method __init__ (line 1002) | def __init__(self, in_channels, out_channels, up=2):
method forward (line 1013) | def forward(self, x, x_size, mask=None):
class BasicLayer (line 1026) | class BasicLayer(nn.Module):
method __init__ (line 1045) | def __init__(
method forward (line 1104) | def forward(self, x, x_size, mask=None):
class ToToken (line 1122) | class ToToken(nn.Module):
method __init__ (line 1123) | def __init__(self, in_channels=3, dim=128, kernel_size=5, stride=1):
method forward (line 1133) | def forward(self, x, mask):
class EncFromRGB (line 1139) | class EncFromRGB(nn.Module):
method __init__ (line 1140) | def __init__(
method forward (line 1157) | def forward(self, x):
class ConvBlockDown (line 1164) | class ConvBlockDown(nn.Module):
method __init__ (line 1165) | def __init__(
method forward (line 1184) | def forward(self, x):
function token2feature (line 1191) | def token2feature(x, x_size):
function feature2token (line 1198) | def feature2token(x):
class Encoder (line 1204) | class Encoder(nn.Module):
method __init__ (line 1205) | def __init__(
method forward (line 1227) | def forward(self, x):
class ToStyle (line 1237) | class ToStyle(nn.Module):
method __init__ (line 1238) | def __init__(self, in_channels, out_channels, activation, drop_rate):
method forward (line 1270) | def forward(self, x):
class DecBlockFirstV2 (line 1279) | class DecBlockFirstV2(nn.Module):
method __init__ (line 254) | def __init__(
method forward (line 289) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
method __init__ (line 1280) | def __init__(
method forward (line 1318) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
class DecBlock (line 1330) | class DecBlock(nn.Module):
method __init__ (line 302) | def __init__(
method forward (line 345) | def forward(self, x, img, ws, gs, E_features, noise_mode="random"):
method __init__ (line 1331) | def __init__(
method forward (line 1374) | def forward(self, x, img, ws, gs, E_features, noise_mode="random"):
class Decoder (line 1386) | class Decoder(nn.Module):
method __init__ (line 1387) | def __init__(
method forward (line 1411) | def forward(self, x, ws, gs, E_features, noise_mode="random"):
class DecStyleBlock (line 1420) | class DecStyleBlock(nn.Module):
method __init__ (line 1421) | def __init__(
method forward (line 1464) | def forward(self, x, img, style, skip, noise_mode="random"):
class FirstStage (line 1473) | class FirstStage(nn.Module):
method __init__ (line 1474) | def __init__(
method forward (line 1577) | def forward(self, images_in, masks_in, ws, noise_mode="random"):
class SynthesisNet (line 1632) | class SynthesisNet(nn.Module):
method __init__ (line 1633) | def __init__(
method forward (line 1681) | def forward(self, images_in, masks_in, ws, noise_mode="random", return...
class Generator (line 1714) | class Generator(nn.Module):
method __init__ (line 1715) | def __init__(
method forward (line 1746) | def forward(
class Discriminator (line 1769) | class Discriminator(torch.nn.Module):
method __init__ (line 496) | def __init__(
method forward (line 554) | def forward(self, images_in, masks_in, c):
method __init__ (line 1770) | def __init__(
method forward (line 1851) | def forward(self, images_in, masks_in, images_stg1, c):
class MAT (line 1878) | class MAT(InpaintModel):
method init_model (line 1885) | def init_model(self, device, **kwargs):
method download (line 1908) | def download():
method is_downloaded (line 1912) | def is_downloaded() -> bool:
method forward (line 1915) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/mi_gan.py
class MIGAN (line 24) | class MIGAN(InpaintModel):
method init_model (line 31) | def init_model(self, device, **kwargs):
method download (line 35) | def download():
method is_downloaded (line 39) | def is_downloaded() -> bool:
method __call__ (line 43) | def __call__(self, image, mask, config: InpaintRequest):
method forward (line 83) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/opencv2.py
class OpenCV2 (line 8) | class OpenCV2(InpaintModel):
method is_downloaded (line 14) | def is_downloaded() -> bool:
method forward (line 17) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/original_sd_configs/__init__.py
function get_config_files (line 7) | def get_config_files() -> Dict[str, Path]:
FILE: iopaint/model/paint_by_example.py
class PaintByExample (line 13) | class PaintByExample(DiffusionInpaintModel):
method init_model (line 18) | def init_model(self, device: torch.device, **kwargs):
method forward (line 44) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/plms_sampler.py
class PLMSSampler (line 8) | class PLMSSampler(object):
method __init__ (line 9) | def __init__(self, model, schedule="linear", **kwargs):
method register_buffer (line 15) | def register_buffer(self, name, attr):
method make_schedule (line 18) | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddi...
method sample (line 52) | def sample(self,
method plms_sampling (line 109) | def plms_sampling(self, cond, shape,
method p_sample_plms (line 162) | def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_origin...
FILE: iopaint/model/power_paint/pipeline_powerpaint.py
function prepare_mask_and_masked_image (line 53) | def prepare_mask_and_masked_image(
class StableDiffusionInpaintPipeline (line 188) | class StableDiffusionInpaintPipeline(
method __init__ (line 223) | def __init__(
method enable_model_cpu_offload (line 342) | def enable_model_cpu_offload(self, gpu_id=0):
method _encode_prompt (line 377) | def _encode_prompt(
method run_safety_checker (line 608) | def run_safety_checker(self, image, device, dtype):
method prepare_extra_step_kwargs (line 627) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 648) | def check_inputs(
method prepare_latents (line 708) | def prepare_latents(
method _encode_vae_image (line 774) | def _encode_vae_image(self, image: torch.Tensor, generator: torch.Gene...
method prepare_mask_latents (line 792) | def prepare_mask_latents(
method get_timesteps (line 847) | def get_timesteps(self, num_inference_steps, strength, device):
method __call__ (line 857) | def __call__(
FILE: iopaint/model/power_paint/power_paint.py
class PowerPaint (line 20) | class PowerPaint(DiffusionInpaintModel):
method init_model (line 26) | def init_model(self, device: torch.device, **kwargs):
method forward (line 66) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/power_paint/power_paint_v2.py
class PowerPaintV2 (line 31) | class PowerPaintV2(DiffusionInpaintModel):
method init_model (line 37) | def init_model(self, device: torch.device, **kwargs):
method forward (line 144) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/power_paint/powerpaint_tokenizer.py
function add_task_to_prompt (line 9) | def add_task_to_prompt(prompt, negative_prompt, task: PowerPaintTask):
function task_to_prompt (line 39) | def task_to_prompt(task: PowerPaintTask):
class PowerPaintTokenizer (line 51) | class PowerPaintTokenizer:
method __init__ (line 52) | def __init__(self, tokenizer: CLIPTokenizer):
method __getattr__ (line 64) | def __getattr__(self, name: str) -> Any:
method try_adding_tokens (line 80) | def try_adding_tokens(self, tokens: Union[str, List[str]], *args, **kw...
method get_token_info (line 93) | def get_token_info(self, token: str) -> dict:
method add_placeholder_token (line 108) | def add_placeholder_token(
method replace_placeholder_tokens_in_text (line 139) | def replace_placeholder_tokens_in_text(
method replace_text_with_placeholder_tokens (line 178) | def replace_text_with_placeholder_tokens(
method __call__ (line 202) | def __call__(
method encode (line 226) | def encode(self, text: Union[str, List[str]], *args, **kwargs):
method decode (line 236) | def decode(
FILE: iopaint/model/power_paint/v2/BrushNet_CA.py
class BrushNetOutput (line 37) | class BrushNetOutput(BaseOutput):
class BrushNetModel (line 61) | class BrushNetModel(ModelMixin, ConfigMixin):
method __init__ (line 138) | def __init__(
method from_unet (line 496) | def from_unet(
method attn_processors (line 621) | def attn_processors(self) -> Dict[str, AttentionProcessor]:
method set_attn_processor (line 651) | def set_attn_processor(
method set_default_attn_processor (line 688) | def set_default_attn_processor(self):
method set_attention_slice (line 710) | def set_attention_slice(self, slice_size: Union[str, int, List[int]]) ...
method _set_gradient_checkpointing (line 781) | def _set_gradient_checkpointing(self, module, value: bool = False) -> ...
method forward (line 785) | def forward(
function zero_module (line 1091) | def zero_module(module):
FILE: iopaint/model/power_paint/v2/pipeline_PowerPaint_Brushnet_CA.py
function retrieve_timesteps (line 103) | def retrieve_timesteps(
class StableDiffusionPowerPaintBrushNetPipeline (line 149) | class StableDiffusionPowerPaintBrushNetPipeline(
method __init__ (line 197) | def __init__(
method _encode_prompt (line 248) | def _encode_prompt(
method encode_prompt (line 479) | def encode_prompt(
method encode_image (line 691) | def encode_image(
method prepare_ip_adapter_image_embeds (line 724) | def prepare_ip_adapter_image_embeds(
method run_safety_checker (line 794) | def run_safety_checker(self, image, device, dtype):
method decode_latents (line 813) | def decode_latents(self, latents):
method prepare_extra_step_kwargs (line 825) | def prepare_extra_step_kwargs(self, generator, eta):
method check_inputs (line 846) | def check_inputs(
method check_image (line 974) | def check_image(self, image, mask, prompt, prompt_embeds):
method prepare_image (line 1038) | def prepare_image(
method prepare_latents (line 1071) | def prepare_latents(
method get_guidance_scale_embedding (line 1104) | def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=tor...
method guidance_scale (line 1133) | def guidance_scale(self):
method clip_skip (line 1137) | def clip_skip(self):
method do_classifier_free_guidance (line 1144) | def do_classifier_free_guidance(self):
method cross_attention_kwargs (line 1148) | def cross_attention_kwargs(self):
method num_timesteps (line 1152) | def num_timesteps(self):
method __call__ (line 1157) | def __call__(
FILE: iopaint/model/power_paint/v2/unet_2d_blocks.py
function CrossAttnDownBlock2D_forward (line 23) | def CrossAttnDownBlock2D_forward(
function DownBlock2D_forward (line 107) | def DownBlock2D_forward(
function CrossAttnUpBlock2D_forward (line 158) | def CrossAttnUpBlock2D_forward(
function UpBlock2D_forward (line 261) | def UpBlock2D_forward(
FILE: iopaint/model/power_paint/v2/unet_2d_condition.py
function UNet2DConditionModel_forward (line 30) | def UNet2DConditionModel_forward(
FILE: iopaint/model/sd.py
class SD (line 18) | class SD(DiffusionInpaintModel):
method init_model (line 23) | def init_model(self, device: torch.device, **kwargs):
method forward (line 82) | def forward(self, image, mask, config: InpaintRequest):
class SD15 (line 112) | class SD15(SD):
class Anything4 (line 117) | class Anything4(SD):
class RealisticVision14 (line 122) | class RealisticVision14(SD):
class SD2 (line 127) | class SD2(SD):
FILE: iopaint/model/sdxl.py
class SDXL (line 22) | class SDXL(DiffusionInpaintModel):
method init_model (line 29) | def init_model(self, device: torch.device, **kwargs):
method forward (line 83) | def forward(self, image, mask, config: InpaintRequest):
FILE: iopaint/model/utils.py
function make_beta_schedule (line 32) | def make_beta_schedule(
function make_ddim_sampling_parameters (line 67) | def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbos...
function make_ddim_timesteps (line 87) | def make_ddim_timesteps(
function noise_like (line 110) | def noise_like(shape, device, repeat=False):
function timestep_embedding (line 118) | def timestep_embedding(device, timesteps, dim, max_period=10000, repeat_...
function normalize_2nd_moment (line 145) | def normalize_2nd_moment(x, dim=1):
class EasyDict (line 151) | class EasyDict(dict):
method __getattr__ (line 154) | def __getattr__(self, name: str) -> Any:
method __setattr__ (line 160) | def __setattr__(self, name: str, value: Any) -> None:
method __delattr__ (line 163) | def __delattr__(self, name: str) -> None:
function _bias_act_ref (line 167) | def _bias_act_ref(x, b=None, dim=1, act="linear", alpha=None, gain=None,...
function bias_act (line 198) | def bias_act(
function _get_filter_size (line 237) | def _get_filter_size(f):
function _get_weight_shape (line 251) | def _get_weight_shape(w):
function _parse_scaling (line 256) | def _parse_scaling(scaling):
function _parse_padding (line 266) | def _parse_padding(padding):
function setup_filter (line 278) | def setup_filter(
function _ntuple (line 332) | def _ntuple(n):
function upfirdn2d (line 419) | def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, ...
function _upfirdn2d_ref (line 466) | def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gai...
function downsample2d (line 519) | def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, imp...
function upsample2d (line 561) | def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl="c...
class MinibatchStdLayer (line 609) | class MinibatchStdLayer(torch.nn.Module):
method __init__ (line 610) | def __init__(self, group_size, num_channels=1):
method forward (line 615) | def forward(self, x):
class FullyConnectedLayer (line 638) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 639) | def __init__(
method forward (line 662) | def forward(self, x):
function _conv2d_wrapper (line 678) | def _conv2d_wrapper(
function conv2d_resample (line 717) | def conv2d_resample(
class Conv2dLayer (line 854) | class Conv2dLayer(torch.nn.Module):
method __init__ (line 855) | def __init__(
method forward (line 901) | def forward(self, x, gain=1):
function torch_gc (line 920) | def torch_gc():
function set_seed (line 927) | def set_seed(seed: int):
function get_scheduler (line 934) | def get_scheduler(sd_sampler, scheduler_config):
function is_local_files_only (line 974) | def is_local_files_only(**kwargs) -> bool:
function handle_from_pretrained_exceptions (line 980) | def handle_from_pretrained_exceptions(func, **kwargs):
function get_torch_dtype (line 1009) | def get_torch_dtype(device, no_half: bool):
function enable_low_mem (line 1021) | def enable_low_mem(pipe, enable: bool):
FILE: iopaint/model/zits.py
function resize (line 47) | def resize(img, height, width, center_crop=False):
function to_tensor (line 66) | def to_tensor(img, scale=True, norm=False):
function load_masked_position_encoding (line 83) | def load_masked_position_encoding(mask):
function load_image (line 146) | def load_image(img, mask, device, sigma256=3.0):
function to_device (line 212) | def to_device(data, device):
class ZITS (line 224) | class ZITS(InpaintModel):
method __init__ (line 231) | def __init__(self, device, **kwargs):
method init_model (line 241) | def init_model(self, device, **kwargs):
method download (line 256) | def download():
method is_downloaded (line 265) | def is_downloaded() -> bool:
method wireframe_edge_and_line (line 274) | def wireframe_edge_and_line(self, items, enable: bool):
method forward (line 346) | def forward(self, image, mask, config: InpaintRequest):
method wireframe_forward (line 377) | def wireframe_forward(self, images, h, w, masks, mask_th=0.925):
method sample_edge_line_logits (line 430) | def sample_edge_line_logits(
FILE: iopaint/model_manager.py
class ModelManager (line 17) | class ModelManager:
method __init__ (line 18) | def __init__(self, name: str, device: torch.device, **kwargs):
method current_model (line 43) | def current_model(self) -> ModelInfo:
method init_model (line 46) | def init_model(self, name: str, device, **kwargs):
method __call__ (line 93) | def __call__(self, image, mask, config: InpaintRequest):
method scan_models (line 113) | def scan_models(self) -> List[ModelInfo]:
method switch (line 118) | def switch(self, new_name: str):
method switch_brushnet_method (line 149) | def switch_brushnet_method(self, config):
method switch_controlnet_method (line 193) | def switch_controlnet_method(self, config):
method enable_disable_powerpaint_v2 (line 231) | def enable_disable_powerpaint_v2(self, config: InpaintRequest):
method enable_disable_lcm_lora (line 250) | def enable_disable_lcm_lora(self, config: InpaintRequest):
FILE: iopaint/plugins/__init__.py
function build_plugins (line 14) | def build_plugins(
FILE: iopaint/plugins/anime_seg.py
class REBNCONV (line 13) | class REBNCONV(nn.Module):
method __init__ (line 14) | def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
method forward (line 23) | def forward(self, x):
function _upsample_like (line 31) | def _upsample_like(src, tar):
class RSU7 (line 38) | class RSU7(nn.Module):
method __init__ (line 39) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
method forward (line 74) | def forward(self, x):
class RSU6 (line 120) | class RSU6(nn.Module):
method __init__ (line 121) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 148) | def forward(self, x):
class RSU5 (line 187) | class RSU5(nn.Module):
method __init__ (line 188) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 211) | def forward(self, x):
class RSU4 (line 244) | class RSU4(nn.Module):
method __init__ (line 245) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 264) | def forward(self, x):
class RSU4F (line 291) | class RSU4F(nn.Module):
method __init__ (line 292) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 307) | def forward(self, x):
class ISNetDIS (line 325) | class ISNetDIS(nn.Module):
method __init__ (line 326) | def __init__(self, in_ch=3, out_ch=1):
method forward (line 358) | def forward(self, x):
class AnimeSeg (line 416) | class AnimeSeg(BasePlugin):
method __init__ (line 422) | def __init__(self):
method gen_image (line 431) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method gen_mask (line 440) | def gen_mask(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method forward (line 444) | def forward(self, rgb_np_img):
FILE: iopaint/plugins/base_plugin.py
class BasePlugin (line 7) | class BasePlugin:
method __init__ (line 12) | def __init__(self):
method gen_image (line 18) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method gen_mask (line 22) | def gen_mask(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method check_dep (line 26) | def check_dep(self):
method switch_model (line 29) | def switch_model(self, new_model_name: str):
FILE: iopaint/plugins/basicsr/arch_util.py
function default_init_weights (line 10) | def default_init_weights(
function make_layer (line 45) | def make_layer(
function pixel_unshuffle (line 64) | def pixel_unshuffle(x: torch.Tensor, scale: int) -> torch.Tensor:
FILE: iopaint/plugins/basicsr/img_util.py
function img2tensor (line 9) | def img2tensor(imgs, bgr2rgb=True, float32=True):
function tensor2img (line 38) | def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
function tensor2img_fast (line 97) | def tensor2img_fast(tensor, rgb2bgr=True, min_max=(0, 1)):
function imfrombytes (line 114) | def imfrombytes(content, flag='color', float32=False):
function imwrite (line 135) | def imwrite(img, file_path, params=None, auto_mkdir=True):
function crop_border (line 156) | def crop_border(imgs, crop_border):
FILE: iopaint/plugins/basicsr/rrdbnet_arch.py
class ResidualDenseBlock (line 8) | class ResidualDenseBlock(nn.Module):
method __init__ (line 18) | def __init__(self, num_feat: int = 64, num_grow_ch: int = 32) -> None:
method forward (line 33) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class RRDB (line 43) | class RRDB(nn.Module):
method __init__ (line 53) | def __init__(self, num_feat: int, num_grow_ch: int = 32) -> None:
method forward (line 59) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class RRDBNet (line 67) | class RRDBNet(nn.Module):
method __init__ (line 87) | def __init__(
method forward (line 115) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/briarmbg.py
class REBNCONV (line 11) | class REBNCONV(nn.Module):
method __init__ (line 12) | def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
method forward (line 21) | def forward(self, x):
function _upsample_like (line 29) | def _upsample_like(src, tar):
class RSU7 (line 36) | class RSU7(nn.Module):
method __init__ (line 37) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
method forward (line 72) | def forward(self, x):
class RSU6 (line 118) | class RSU6(nn.Module):
method __init__ (line 119) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 146) | def forward(self, x):
class RSU5 (line 185) | class RSU5(nn.Module):
method __init__ (line 186) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 209) | def forward(self, x):
class RSU4 (line 242) | class RSU4(nn.Module):
method __init__ (line 243) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 262) | def forward(self, x):
class RSU4F (line 289) | class RSU4F(nn.Module):
method __init__ (line 290) | def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
method forward (line 305) | def forward(self, x):
class myrebnconv (line 323) | class myrebnconv(nn.Module):
method __init__ (line 324) | def __init__(
method forward (line 348) | def forward(self, x):
class BriaRMBG (line 352) | class BriaRMBG(nn.Module):
method __init__ (line 353) | def __init__(self, in_ch=3, out_ch=1):
method forward (line 392) | def forward(self, x):
function resize_image (line 466) | def resize_image(image):
function create_briarmbg_session (line 473) | def create_briarmbg_session():
function briarmbg_process (line 483) | def briarmbg_process(device, bgr_np_image, session, only_mask=False):
FILE: iopaint/plugins/briarmbg2.py
class Config (line 29) | class Config:
method __init__ (line 30) | def __init__(self) -> None:
method print_task (line 234) | def print_task(self) -> None:
class Mlp (line 239) | class Mlp(nn.Module):
method __init__ (line 240) | def __init__(
method _init_weights (line 259) | def _init_weights(self, m):
method forward (line 274) | def forward(self, x, H, W):
method __init__ (line 871) | def __init__(
method forward (line 887) | def forward(self, x):
class Attention (line 284) | class Attention(nn.Module):
method __init__ (line 285) | def __init__(
method _init_weights (line 319) | def _init_weights(self, m):
method forward (line 334) | def forward(self, x, H, W):
class Block (line 384) | class Block(nn.Module):
method __init__ (line 385) | def __init__(
method _init_weights (line 423) | def _init_weights(self, m):
method forward (line 438) | def forward(self, x, H, W):
class OverlapPatchEmbed (line 445) | class OverlapPatchEmbed(nn.Module):
method __init__ (line 448) | def __init__(
method _init_weights (line 470) | def _init_weights(self, m):
method forward (line 485) | def forward(self, x):
class PyramidVisionTransformerImpr (line 494) | class PyramidVisionTransformerImpr(nn.Module):
method __init__ (line 495) | def __init__(
method _init_weights (line 636) | def _init_weights(self, m):
method init_weights (line 651) | def init_weights(self, pretrained=None):
method reset_drop_path (line 656) | def reset_drop_path(self, drop_path_rate):
method freeze_patch_emb (line 674) | def freeze_patch_emb(self):
method no_weight_decay (line 678) | def no_weight_decay(self):
method get_classifier (line 687) | def get_classifier(self):
method reset_classifier (line 690) | def reset_classifier(self, num_classes, global_pool=""):
method forward_features (line 696) | def forward_features(self, x):
method forward (line 736) | def forward(self, x):
class DWConv (line 743) | class DWConv(nn.Module):
method __init__ (line 744) | def __init__(self, dim=768):
method forward (line 748) | def forward(self, x, H, W):
class pvt_v2_b0 (line 757) | class pvt_v2_b0(PyramidVisionTransformerImpr):
method __init__ (line 758) | def __init__(self, **kwargs):
class pvt_v2_b1 (line 773) | class pvt_v2_b1(PyramidVisionTransformerImpr):
method __init__ (line 774) | def __init__(self, **kwargs):
class pvt_v2_b2 (line 790) | class pvt_v2_b2(PyramidVisionTransformerImpr):
method __init__ (line 791) | def __init__(self, in_channels=3, **kwargs):
class pvt_v2_b3 (line 808) | class pvt_v2_b3(PyramidVisionTransformerImpr):
method __init__ (line 809) | def __init__(self, **kwargs):
class pvt_v2_b4 (line 825) | class pvt_v2_b4(PyramidVisionTransformerImpr):
method __init__ (line 826) | def __init__(self, **kwargs):
class pvt_v2_b5 (line 842) | class pvt_v2_b5(PyramidVisionTransformerImpr):
method __init__ (line 843) | def __init__(self, **kwargs):
class Mlp (line 868) | class Mlp(nn.Module):
method __init__ (line 240) | def __init__(
method _init_weights (line 259) | def _init_weights(self, m):
method forward (line 274) | def forward(self, x, H, W):
method __init__ (line 871) | def __init__(
method forward (line 887) | def forward(self, x):
function window_partition (line 896) | def window_partition(x, window_size):
function window_reverse (line 913) | def window_reverse(windows, window_size, H, W):
class WindowAttention (line 932) | class WindowAttention(nn.Module):
method __init__ (line 946) | def __init__(
method forward (line 996) | def forward(self, x, mask=None):
class SwinTransformerBlock (line 1063) | class SwinTransformerBlock(nn.Module):
method __init__ (line 1081) | def __init__(
method forward (line 1130) | def forward(self, x, mask_matrix):
class PatchMerging (line 1200) | class PatchMerging(nn.Module):
method __init__ (line 1208) | def __init__(self, dim, norm_layer=nn.LayerNorm):
method forward (line 1214) | def forward(self, x, H, W):
class BasicLayer (line 1244) | class BasicLayer(nn.Module):
method __init__ (line 1263) | def __init__(
method forward (line 1313) | def forward(self, x, H, W):
class PatchEmbed (line 1364) | class PatchEmbed(nn.Module):
method __init__ (line 1374) | def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_lay...
method forward (line 1390) | def forward(self, x):
class SwinTransformer (line 1409) | class SwinTransformer(nn.Module):
method __init__ (line 1438) | def __init__(
method _freeze_stages (line 1530) | def _freeze_stages(self):
method forward (line 1547) | def forward(self, x):
method train (line 1579) | def train(self, mode=True):
function swin_v1_t (line 1585) | def swin_v1_t():
function swin_v1_s (line 1592) | def swin_v1_s():
function swin_v1_b (line 1599) | def swin_v1_b():
function swin_v1_l (line 1606) | def swin_v1_l():
class DeformableConv2d (line 1620) | class DeformableConv2d(nn.Module):
method __init__ (line 1621) | def __init__(
method forward (line 1667) | def forward(self, x):
function build_act_layer (line 1691) | def build_act_layer(act_layer):
function build_norm_layer (line 1702) | def build_norm_layer(
class to_channels_first (line 1723) | class to_channels_first(nn.Module):
method __init__ (line 1724) | def __init__(self):
method forward (line 1727) | def forward(self, x):
class to_channels_last (line 1731) | class to_channels_last(nn.Module):
method __init__ (line 1732) | def __init__(self):
method forward (line 1735) | def forward(self, x):
function build_backbone (line 1765) | def build_backbone(bb_name, pretrained=True, params_settings=""):
function load_weights (line 1819) | def load_weights(model, model_name):
class BasicDecBlk (line 1863) | class BasicDecBlk(nn.Module):
method __init__ (line 1864) | def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
method forward (line 1881) | def forward(self, x):
class ResBlk (line 1892) | class ResBlk(nn.Module):
method __init__ (line 1893) | def __init__(self, in_channels=64, out_channels=None, inter_channels=64):
method forward (line 1917) | def forward(self, x):
class BasicLatBlk (line 1943) | class BasicLatBlk(nn.Module):
method __init__ (line 1944) | def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
method forward (line 1949) | def forward(self, x):
class _ASPPModule (line 1966) | class _ASPPModule(nn.Module):
method __init__ (line 1967) | def __init__(self, in_channels, planes, kernel_size, padding, dilation):
method forward (line 1981) | def forward(self, x):
class ASPP (line 1988) | class ASPP(nn.Module):
method __init__ (line 1989) | def __init__(self, in_channels=64, out_channels=None, output_stride=16):
method forward (line 2042) | def forward(self, x):
class _ASPPModuleDeformable (line 2059) | class _ASPPModuleDeformable(nn.Module):
method __init__ (line 2060) | def __init__(self, in_channels, planes, kernel_size, padding):
method forward (line 2073) | def forward(self, x):
class ASPPDeformable (line 2080) | class ASPPDeformable(nn.Module):
method __init__ (line 2081) | def __init__(self, in_channels, out_channels=None, parallel_block_size...
method forward (line 2123) | def forward(self, x):
class RefinerPVTInChannels4 (line 2140) | class RefinerPVTInChannels4(nn.Module):
method __init__ (line 2141) | def __init__(self, in_channels=3 + 1):
method forward (line 2166) | def forward(self, x):
class Refiner (line 2188) | class Refiner(nn.Module):
method __init__ (line 2189) | def __init__(self, in_channels=3 + 1):
method forward (line 2220) | def forward(self, x):
class Decoder (line 2243) | class Decoder(nn.Module):
method __init__ (line 2244) | def __init__(self, channels):
method forward (line 2265) | def forward(self, features):
method __init__ (line 2608) | def __init__(self, channels):
method get_patches_batch (line 2726) | def get_patches_batch(self, x, p):
method forward (line 2743) | def forward(self, features):
class RefUNet (line 2292) | class RefUNet(nn.Module):
method __init__ (line 2294) | def __init__(self, in_channels=3 + 1):
method forward (line 2350) | def forward(self, x):
class StemLayer (line 2383) | class StemLayer(nn.Module):
method __init__ (line 2392) | def __init__(
method forward (line 2415) | def forward(self, x):
class BiRefNetConfig (line 2427) | class BiRefNetConfig(PretrainedConfig):
method __init__ (line 2430) | def __init__(self, bb_pretrained=False, **kwargs):
class BiRefNet (line 2435) | class BiRefNet(PreTrainedModel):
method __init__ (line 2438) | def __init__(self, bb_pretrained=True, config=BiRefNetConfig()):
method forward_enc (line 2494) | def forward_enc(self, x):
method forward_ori (line 2589) | def forward_ori(self, x):
method forward (line 2601) | def forward(self, x):
class Decoder (line 2607) | class Decoder(nn.Module):
method __init__ (line 2244) | def __init__(self, channels):
method forward (line 2265) | def forward(self, features):
method __init__ (line 2608) | def __init__(self, channels):
method get_patches_batch (line 2726) | def get_patches_batch(self, x, p):
method forward (line 2743) | def forward(self, features):
class SimpleConvs (line 2914) | class SimpleConvs(nn.Module):
method __init__ (line 2915) | def __init__(self, in_channels: int, out_channels: int, inter_channels...
method forward (line 2920) | def forward(self, x):
function create_briarmbg2_session (line 2924) | def create_briarmbg2_session():
function briarmbg2_process (line 2929) | def briarmbg2_process(device, bgr_np_image, session, only_mask=False):
FILE: iopaint/plugins/facexlib/detection/__init__.py
function init_detection_model (line 8) | def init_detection_model(model_name, half=False, device='cuda', model_ro...
FILE: iopaint/plugins/facexlib/detection/align_trans.py
class FaceWarpException (line 13) | class FaceWarpException(Exception):
method __str__ (line 15) | def __str__(self):
function get_reference_facial_points (line 19) | def get_reference_facial_points(output_size=None, inner_padding_factor=0...
function get_affine_transform_matrix (line 112) | def get_affine_transform_matrix(src_pts, dst_pts):
function warp_and_crop_face (line 145) | def warp_and_crop_face(src_img, facial_pts, reference_pts=None, crop_siz...
FILE: iopaint/plugins/facexlib/detection/matlab_cp2tform.py
class MatlabCp2tormException (line 7) | class MatlabCp2tormException(Exception):
method __str__ (line 9) | def __str__(self):
function tformfwd (line 13) | def tformfwd(trans, uv):
function tforminv (line 37) | def tforminv(trans, uv):
function findNonreflectiveSimilarity (line 60) | def findNonreflectiveSimilarity(uv, xy, options=None):
function findSimilarity (line 94) | def findSimilarity(uv, xy, options=None):
function get_similarity_transform (line 130) | def get_similarity_transform(src_pts, dst_pts, reflective=True):
function cvt_tform_mat_for_cv2 (line 170) | def cvt_tform_mat_for_cv2(trans):
function get_similarity_transform_for_cv2 (line 198) | def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
FILE: iopaint/plugins/facexlib/detection/retinaface.py
function generate_config (line 28) | def generate_config(network_name):
class RetinaFace (line 75) | class RetinaFace(nn.Module):
method __init__ (line 76) | def __init__(self, network_name="resnet50", half=False, phase="test", ...
method forward (line 130) | def forward(self, inputs):
method __detect_faces (line 163) | def __detect_faces(self, inputs):
method transform (line 196) | def transform(self, image, use_origin_size):
method detect_faces (line 225) | def detect_faces(
method __align_multi (line 270) | def __align_multi(self, image, boxes, landmarks, limit=None):
method align_multi (line 289) | def align_multi(self, img, conf_threshold=0.8, limit=None):
method batched_transform (line 296) | def batched_transform(self, frames, use_origin_size):
method batched_detect_faces (line 348) | def batched_detect_faces(
FILE: iopaint/plugins/facexlib/detection/retinaface_net.py
function conv_bn (line 6) | def conv_bn(inp, oup, stride=1, leaky=0):
function conv_bn_no_relu (line 12) | def conv_bn_no_relu(inp, oup, stride):
function conv_bn1X1 (line 19) | def conv_bn1X1(inp, oup, stride, leaky=0):
function conv_dw (line 25) | def conv_dw(inp, oup, stride, leaky=0.1):
class SSH (line 36) | class SSH(nn.Module):
method __init__ (line 38) | def __init__(self, in_channel, out_channel):
method forward (line 52) | def forward(self, input):
class FPN (line 66) | class FPN(nn.Module):
method __init__ (line 68) | def __init__(self, in_channels_list, out_channels):
method forward (line 80) | def forward(self, input):
class MobileNetV1 (line 100) | class MobileNetV1(nn.Module):
method __init__ (line 102) | def __init__(self):
method forward (line 127) | def forward(self, x):
class ClassHead (line 138) | class ClassHead(nn.Module):
method __init__ (line 140) | def __init__(self, inchannels=512, num_anchors=3):
method forward (line 145) | def forward(self, x):
class BboxHead (line 152) | class BboxHead(nn.Module):
method __init__ (line 154) | def __init__(self, inchannels=512, num_anchors=3):
method forward (line 158) | def forward(self, x):
class LandmarkHead (line 165) | class LandmarkHead(nn.Module):
method __init__ (line 167) | def __init__(self, inchannels=512, num_anchors=3):
method forward (line 171) | def forward(self, x):
function make_class_head (line 178) | def make_class_head(fpn_num=3, inchannels=64, anchor_num=2):
function make_bbox_head (line 185) | def make_bbox_head(fpn_num=3, inchannels=64, anchor_num=2):
function make_landmark_head (line 192) | def make_landmark_head(fpn_num=3, inchannels=64, anchor_num=2):
FILE: iopaint/plugins/facexlib/detection/retinaface_utils.py
class PriorBox (line 8) | class PriorBox(object):
method __init__ (line 10) | def __init__(self, cfg, image_size=None, phase='train'):
method forward (line 19) | def forward(self):
function py_cpu_nms (line 39) | def py_cpu_nms(dets, thresh):
function point_form (line 50) | def point_form(boxes):
function center_size (line 65) | def center_size(boxes):
function intersect (line 79) | def intersect(box_a, box_b):
function jaccard (line 98) | def jaccard(box_a, box_b):
function matrix_iou (line 117) | def matrix_iou(a, b):
function matrix_iof (line 130) | def matrix_iof(a, b):
function match (line 142) | def match(threshold, truths, priors, variances, labels, landms, loc_t, c...
function encode (line 200) | def encode(matched, priors, variances):
function encode_landm (line 224) | def encode_landm(matched, priors, variances):
function decode (line 254) | def decode(loc, priors, variances):
function decode_landm (line 274) | def decode_landm(pre, priors, variances):
function batched_decode (line 297) | def batched_decode(b_loc, priors, variances):
function batched_decode_landm (line 320) | def batched_decode_landm(pre, priors, variances):
function log_sum_exp (line 343) | def log_sum_exp(x):
function nms (line 357) | def nms(boxes, scores, overlap=0.5, top_k=200):
FILE: iopaint/plugins/facexlib/parsing/__init__.py
function init_parsing_model (line 8) | def init_parsing_model(model_name='bisenet', half=False, device='cuda', ...
FILE: iopaint/plugins/facexlib/parsing/bisenet.py
class ConvBNReLU (line 8) | class ConvBNReLU(nn.Module):
method __init__ (line 10) | def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1):
method forward (line 15) | def forward(self, x):
class BiSeNetOutput (line 21) | class BiSeNetOutput(nn.Module):
method __init__ (line 23) | def __init__(self, in_chan, mid_chan, num_class):
method forward (line 28) | def forward(self, x):
class AttentionRefinementModule (line 34) | class AttentionRefinementModule(nn.Module):
method __init__ (line 36) | def __init__(self, in_chan, out_chan):
method forward (line 43) | def forward(self, x):
class ContextPath (line 53) | class ContextPath(nn.Module):
method __init__ (line 55) | def __init__(self):
method forward (line 64) | def forward(self, x):
class FeatureFusionModule (line 87) | class FeatureFusionModule(nn.Module):
method __init__ (line 89) | def __init__(self, in_chan, out_chan):
method forward (line 97) | def forward(self, fsp, fcp):
class BiSeNet (line 110) | class BiSeNet(nn.Module):
method __init__ (line 112) | def __init__(self, num_class):
method forward (line 120) | def forward(self, x, return_feat=False):
FILE: iopaint/plugins/facexlib/parsing/parsenet.py
class NormLayer (line 8) | class NormLayer(nn.Module):
method __init__ (line 16) | def __init__(self, channels, normalize_shape=None, norm_type='bn'):
method forward (line 35) | def forward(self, x, ref=None):
class ReluLayer (line 42) | class ReluLayer(nn.Module):
method __init__ (line 54) | def __init__(self, channels, relu_type='relu'):
method forward (line 70) | def forward(self, x):
class ConvLayer (line 74) | class ConvLayer(nn.Module):
method __init__ (line 76) | def __init__(self,
method forward (line 103) | def forward(self, x):
class ResidualBlock (line 113) | class ResidualBlock(nn.Module):
method __init__ (line 118) | def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', sca...
method forward (line 132) | def forward(self, x):
class ParseNet (line 140) | class ParseNet(nn.Module):
method __init__ (line 142) | def __init__(self,
method forward (line 188) | def forward(self, x):
FILE: iopaint/plugins/facexlib/parsing/resnet.py
function conv3x3 (line 5) | def conv3x3(in_planes, out_planes, stride=1):
class BasicBlock (line 10) | class BasicBlock(nn.Module):
method __init__ (line 12) | def __init__(self, in_chan, out_chan, stride=1):
method forward (line 26) | def forward(self, x):
function create_layer_basic (line 41) | def create_layer_basic(in_chan, out_chan, bnum, stride=1):
class ResNet18 (line 48) | class ResNet18(nn.Module):
method __init__ (line 50) | def __init__(self):
method forward (line 60) | def forward(self, x):
FILE: iopaint/plugins/facexlib/utils/face_restoration_helper.py
function get_largest_face (line 12) | def get_largest_face(det_faces, h, w):
function get_center_face (line 33) | def get_center_face(det_faces, h=0, w=0, center=None):
class FaceRestoreHelper (line 49) | class FaceRestoreHelper(object):
method __init__ (line 52) | def __init__(
method set_upscale_factor (line 124) | def set_upscale_factor(self, upscale_factor):
method read_image (line 127) | def read_image(self, img):
method get_face_landmarks_5 (line 142) | def get_face_landmarks_5(
method align_warp_face (line 278) | def align_warp_face(self, save_cropped_path=None, border_mode="constan...
method get_inverse_affine (line 317) | def get_inverse_affine(self, save_inverse_affine_path=None):
method add_restored_face (line 329) | def add_restored_face(self, face):
method paste_faces_to_input_image (line 332) | def paste_faces_to_input_image(self, save_path=None, upsample_img=None):
method clean_all (line 466) | def clean_all(self):
FILE: iopaint/plugins/facexlib/utils/face_utils.py
function compute_increased_bbox (line 6) | def compute_increased_bbox(bbox, increase_area, preserve_aspect=True):
function get_valid_bboxes (line 23) | def get_valid_bboxes(bboxes, h, w):
function align_crop_face_landmarks (line 31) | def align_crop_face_landmarks(img,
function paste_face_back (line 190) | def paste_face_back(img, face, inverse_affine):
FILE: iopaint/plugins/facexlib/utils/misc.py
function imwrite (line 11) | def imwrite(img, file_path, params=None, auto_mkdir=True):
function img2tensor (line 30) | def img2tensor(imgs, bgr2rgb=True, float32=True):
function load_file_from_url (line 59) | def load_file_from_url(url, model_dir=None, progress=True, file_name=Non...
function scandir (line 81) | def scandir(dir_path, suffix=None, recursive=False, full_path=False):
FILE: iopaint/plugins/gfpgan/archs/gfpganv1_clean_arch.py
class StyleGAN2GeneratorCSFT (line 10) | class StyleGAN2GeneratorCSFT(StyleGAN2GeneratorClean):
method __init__ (line 24) | def __init__(self, out_size, num_style_feat=512, num_mlp=8, channel_mu...
method forward (line 33) | def forward(self,
class ResBlock (line 119) | class ResBlock(nn.Module):
method __init__ (line 128) | def __init__(self, in_channels, out_channels, mode='down'):
method forward (line 139) | def forward(self, x):
class GFPGANv1Clean (line 151) | class GFPGANv1Clean(nn.Module):
method __init__ (line 172) | def __init__(
method forward (line 275) | def forward(self, x, return_latents=False, return_rgb=True, randomize_...
FILE: iopaint/plugins/gfpgan/archs/restoreformer_arch.py
class VectorQuantizer (line 9) | class VectorQuantizer(nn.Module):
method __init__ (line 21) | def __init__(self, n_e, e_dim, beta):
method forward (line 30) | def forward(self, z):
method get_codebook_entry (line 94) | def get_codebook_entry(self, indices, shape):
function nonlinearity (line 113) | def nonlinearity(x):
function Normalize (line 118) | def Normalize(in_channels):
class Upsample (line 124) | class Upsample(nn.Module):
method __init__ (line 125) | def __init__(self, in_channels, with_conv):
method forward (line 133) | def forward(self, x):
class Downsample (line 140) | class Downsample(nn.Module):
method __init__ (line 141) | def __init__(self, in_channels, with_conv):
method forward (line 150) | def forward(self, x):
class ResnetBlock (line 160) | class ResnetBlock(nn.Module):
method __init__ (line 161) | def __init__(
method forward (line 197) | def forward(self, x, temb):
class MultiHeadAttnBlock (line 220) | class MultiHeadAttnBlock(nn.Module):
method __init__ (line 221) | def __init__(self, in_channels, head_size=1):
method forward (line 247) | def forward(self, x, y=None):
class MultiHeadEncoder (line 290) | class MultiHeadEncoder(nn.Module):
method __init__ (line 291) | def __init__(
method forward (line 377) | def forward(self, x):
class MultiHeadDecoder (line 415) | class MultiHeadDecoder(nn.Module):
method __init__ (line 416) | def __init__(
method forward (line 507) | def forward(self, z):
class MultiHeadDecoderTransformer (line 542) | class MultiHeadDecoderTransformer(nn.Module):
method __init__ (line 543) | def __init__(
method forward (line 634) | def forward(self, z, hs):
class RestoreFormer (line 672) | class RestoreFormer(nn.Module):
method __init__ (line 673) | def __init__(
method encode (line 743) | def encode(self, x):
method decode (line 749) | def decode(self, quant, hs):
method forward (line 755) | def forward(self, input, **kwargs):
FILE: iopaint/plugins/gfpgan/archs/stylegan2_clean_arch.py
class NormStyleCode (line 10) | class NormStyleCode(nn.Module):
method forward (line 11) | def forward(self, x):
class ModulatedConv2d (line 23) | class ModulatedConv2d(nn.Module):
method __init__ (line 38) | def __init__(
method forward (line 74) | def forward(self, x, style):
method __repr__ (line 112) | def __repr__(self):
class StyleConv (line 119) | class StyleConv(nn.Module):
method __init__ (line 131) | def __init__(
method forward (line 153) | def forward(self, x, style, noise=None):
class ToRGB (line 168) | class ToRGB(nn.Module):
method __init__ (line 177) | def __init__(self, in_channels, num_style_feat, upsample=True):
method forward (line 190) | def forward(self, x, style, skip=None):
class ConstantInput (line 212) | class ConstantInput(nn.Module):
method __init__ (line 220) | def __init__(self, num_channel, size):
method forward (line 224) | def forward(self, batch):
class StyleGAN2GeneratorClean (line 229) | class StyleGAN2GeneratorClean(nn.Module):
method __init__ (line 240) | def __init__(
method make_noise (line 330) | def make_noise(self):
method get_latent (line 341) | def get_latent(self, x):
method mean_latent (line 344) | def mean_latent(self, num_latent):
method forward (line 351) | def forward(
FILE: iopaint/plugins/gfpgan_plugin.py
class GFPGANPlugin (line 10) | class GFPGANPlugin(BasePlugin):
method __init__ (line 14) | def __init__(self, device, upscaler=None):
method gen_image (line 37) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
FILE: iopaint/plugins/gfpganer.py
class MyGFPGANer (line 13) | class MyGFPGANer:
method __init__ (line 29) | def __init__(
method enhance (line 91) | def enhance(
FILE: iopaint/plugins/interactive_seg.py
class InteractiveSeg (line 81) | class InteractiveSeg(BasePlugin):
method __init__ (line 85) | def __init__(self, model_name, device):
method _init_session (line 91) | def _init_session(self, model_name: str):
method switch_model (line 112) | def switch_model(self, new_model_name):
method gen_mask (line 122) | def gen_mask(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method forward (line 127) | def forward(self, rgb_np_img, clicks: List[List], img_md5: str):
FILE: iopaint/plugins/realesrgan.py
class RealESRGANer (line 15) | class RealESRGANer:
method __init__ (line 30) | def __init__(
method dni (line 86) | def dni(self, net_a, net_b, dni_weight, key="params", loc="cpu"):
method pre_process (line 97) | def pre_process(self, img):
method process (line 123) | def process(self):
method tile_process (line 127) | def tile_process(self):
method post_process (line 202) | def post_process(self):
method enhance (line 224) | def enhance(self, img, outscale=None, alpha_upsampler="realesrgan"):
class SRVGGNetCompact (line 305) | class SRVGGNetCompact(nn.Module):
method __init__ (line 320) | def __init__(
method forward (line 366) | def forward(self, x):
class RealESRGANUpscaler (line 378) | class RealESRGANUpscaler(BasePlugin):
method __init__ (line 382) | def __init__(self, name, device, no_half=False):
method _init_model (line 389) | def _init_model(self, name):
method switch_model (line 451) | def switch_model(self, new_model_name: str):
method gen_image (line 457) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method forward (line 465) | def forward(self, bgr_np_img, scale: float):
FILE: iopaint/plugins/remove_bg.py
function _rmbg_remove (line 12) | def _rmbg_remove(device, *args, **kwargs):
class RemoveBG (line 18) | class RemoveBG(BasePlugin):
method __init__ (line 23) | def __init__(self, model_name, device):
method _init_session (line 42) | def _init_session(self, model_name: str):
method switch_model (line 67) | def switch_model(self, new_model_name):
method gen_image (line 78) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method gen_mask (line 86) | def gen_mask(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
method check_dep (line 95) | def check_dep(self):
method device_warning (line 104) | def device_warning(self):
FILE: iopaint/plugins/restoreformer.py
class RestoreFormerPlugin (line 10) | class RestoreFormerPlugin(BasePlugin):
method __init__ (line 14) | def __init__(self, device, upscaler=None):
method gen_image (line 32) | def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
FILE: iopaint/plugins/segment_anything/build_sam.py
function build_sam_vit_h (line 25) | def build_sam_vit_h(checkpoint=None):
function build_sam_vit_l (line 35) | def build_sam_vit_l(checkpoint=None):
function build_sam_vit_b (line 45) | def build_sam_vit_b(checkpoint=None):
function build_sam_vit_t (line 55) | def build_sam_vit_t(checkpoint=None):
function build_sam_vit_h_hq (line 107) | def build_sam_vit_h_hq(checkpoint=None):
function build_sam_vit_l_hq (line 117) | def build_sam_vit_l_hq(checkpoint=None):
function build_sam_vit_b_hq (line 127) | def build_sam_vit_b_hq(checkpoint=None):
function _build_sam (line 149) | def _build_sam(
function _build_sam_hq (line 204) | def _build_sam_hq(
FILE: iopaint/plugins/segment_anything/modeling/common.py
class MLPBlock (line 13) | class MLPBlock(nn.Module):
method __init__ (line 14) | def __init__(
method forward (line 25) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class LayerNorm2d (line 31) | class LayerNorm2d(nn.Module):
method __init__ (line 32) | def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
method forward (line 38) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything/modeling/image_encoder.py
class ImageEncoderViT (line 17) | class ImageEncoderViT(nn.Module):
method __init__ (line 18) | def __init__(
method forward (line 106) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Block (line 119) | class Block(nn.Module):
method __init__ (line 122) | def __init__(
method forward (line 166) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Attention (line 185) | class Attention(nn.Module):
method __init__ (line 188) | def __init__(
method forward (line 224) | def forward(self, x: torch.Tensor) -> torch.Tensor:
function window_partition (line 243) | def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.T...
function window_unpartition (line 267) | def window_unpartition(
function get_rel_pos (line 292) | def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torc...
function add_decomposed_rel_pos (line 325) | def add_decomposed_rel_pos(
class PatchEmbed (line 364) | class PatchEmbed(nn.Module):
method __init__ (line 369) | def __init__(
method forward (line 391) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything/modeling/image_encoder_hq.py
class ImageEncoderViTHQ (line 17) | class ImageEncoderViTHQ(nn.Module):
method __init__ (line 18) | def __init__(
method forward (line 108) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Block (line 124) | class Block(nn.Module):
method __init__ (line 127) | def __init__(
method forward (line 173) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Attention (line 192) | class Attention(nn.Module):
method __init__ (line 195) | def __init__(
method forward (line 231) | def forward(self, x: torch.Tensor) -> torch.Tensor:
function window_partition (line 259) | def window_partition(
function window_unpartition (line 287) | def window_unpartition(
function get_rel_pos (line 317) | def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torc...
function add_decomposed_rel_pos (line 350) | def add_decomposed_rel_pos(
class PatchEmbed (line 391) | class PatchEmbed(nn.Module):
method __init__ (line 396) | def __init__(
method forward (line 418) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything/modeling/mask_decoder.py
class MaskDecoder (line 16) | class MaskDecoder(nn.Module):
method __init__ (line 17) | def __init__(
method forward (line 75) | def forward(
method predict_masks (line 116) | def predict_masks(
class MaskDecoderHQ (line 162) | class MaskDecoderHQ(nn.Module):
method __init__ (line 163) | def __init__(
method forward (line 256) | def forward(
method predict_masks (line 321) | def predict_masks(
class MLP (line 388) | class MLP(nn.Module):
method __init__ (line 389) | def __init__(
method forward (line 405) | def forward(self, x):
FILE: iopaint/plugins/segment_anything/modeling/prompt_encoder.py
class PromptEncoder (line 16) | class PromptEncoder(nn.Module):
method __init__ (line 17) | def __init__(
method get_dense_pe (line 62) | def get_dense_pe(self) -> torch.Tensor:
method _embed_points (line 73) | def _embed_points(
method _embed_boxes (line 93) | def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
method _embed_masks (line 102) | def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
method _get_batch_size (line 107) | def _get_batch_size(
method _get_device (line 125) | def _get_device(self) -> torch.device:
method forward (line 128) | def forward(
class PositionEmbeddingRandom (line 171) | class PositionEmbeddingRandom(nn.Module):
method __init__ (line 176) | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = N...
method _pe_encoding (line 185) | def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
method forward (line 194) | def forward(self, size: Tuple[int, int]) -> torch.Tensor:
method forward_with_coords (line 207) | def forward_with_coords(
FILE: iopaint/plugins/segment_anything/modeling/sam.py
class Sam (line 18) | class Sam(nn.Module):
method __init__ (line 22) | def __init__(
method device (line 50) | def device(self) -> Any:
method forward (line 54) | def forward(
method postprocess_masks (line 133) | def postprocess_masks(
method preprocess (line 164) | def preprocess(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything/modeling/sam_hq.py
class SamHQ (line 18) | class SamHQ(nn.Module):
method __init__ (line 22) | def __init__(
method device (line 50) | def device(self) -> Any:
method forward (line 53) | def forward(
method postprocess_masks (line 136) | def postprocess_masks(
method preprocess (line 167) | def preprocess(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything/modeling/tiny_vit_sam.py
function _ntuple (line 21) | def _ntuple(n):
function _trunc_normal_ (line 33) | def _trunc_normal_(tensor, mean, std, a, b):
function trunc_normal_ (line 70) | def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
function drop_path (line 97) | def drop_path(
class TimmDropPath (line 121) | class TimmDropPath(nn.Module):
method __init__ (line 124) | def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
method forward (line 129) | def forward(self, x):
method extra_repr (line 132) | def extra_repr(self):
class Conv2d_BN (line 136) | class Conv2d_BN(torch.nn.Sequential):
method __init__ (line 137) | def __init__(
method fuse (line 150) | def fuse(self):
class DropPath (line 169) | class DropPath(TimmDropPath):
method __init__ (line 170) | def __init__(self, drop_prob=None):
method __repr__ (line 174) | def __repr__(self):
class PatchEmbed (line 180) | class PatchEmbed(nn.Module):
method __init__ (line 181) | def __init__(self, in_chans, embed_dim, resolution, activation):
method forward (line 195) | def forward(self, x):
class MBConv (line 199) | class MBConv(nn.Module):
method __init__ (line 200) | def __init__(self, in_chans, out_chans, expand_ratio, activation, drop...
method forward (line 224) | def forward(self, x):
class PatchMerging (line 243) | class PatchMerging(nn.Module):
method __init__ (line 244) | def __init__(self, input_resolution, dim, out_dim, activation):
method forward (line 258) | def forward(self, x):
class ConvLayer (line 275) | class ConvLayer(nn.Module):
method __init__ (line 276) | def __init__(
method forward (line 316) | def forward(self, x):
class Mlp (line 327) | class Mlp(nn.Module):
method __init__ (line 328) | def __init__(
method forward (line 345) | def forward(self, x):
class Attention (line 356) | class Attention(torch.nn.Module):
method __init__ (line 357) | def __init__(
method train (line 399) | def train(self, mode=True):
method forward (line 410) | def forward(self, x): # x (B,N,C)
class TinyViTBlock (line 437) | class TinyViTBlock(nn.Module):
method __init__ (line 453) | def __init__(
method forward (line 497) | def forward(self, x):
method extra_repr (line 544) | def extra_repr(self) -> str:
class BasicLayer (line 551) | class BasicLayer(nn.Module):
method __init__ (line 570) | def __init__(
method forward (line 620) | def forward(self, x):
method extra_repr (line 630) | def extra_repr(self) -> str:
class LayerNorm2d (line 634) | class LayerNorm2d(nn.Module):
method __init__ (line 635) | def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
method forward (line 641) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class TinyViT (line 649) | class TinyViT(nn.Module):
method __init__ (line 650) | def __init__(
method set_layer_lr_decay (line 756) | def set_layer_lr_decay(self, layer_lr_decay):
method _init_weights (line 789) | def _init_weights(self, m):
method no_weight_decay_keywords (line 799) | def no_weight_decay_keywords(self):
method forward_features (line 802) | def forward_features(self, x):
method forward (line 818) | def forward(self, x):
FILE: iopaint/plugins/segment_anything/modeling/transformer.py
class TwoWayTransformer (line 16) | class TwoWayTransformer(nn.Module):
method __init__ (line 17) | def __init__(
method forward (line 62) | def forward(
class TwoWayAttentionBlock (line 109) | class TwoWayAttentionBlock(nn.Module):
method __init__ (line 110) | def __init__(
method forward (line 151) | def forward(
class Attention (line 185) | class Attention(nn.Module):
method __init__ (line 191) | def __init__(
method _separate_heads (line 208) | def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
method _recombine_heads (line 213) | def _recombine_heads(self, x: Tensor) -> Tensor:
method forward (line 218) | def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
FILE: iopaint/plugins/segment_anything/predictor.py
class SamPredictor (line 15) | class SamPredictor:
method __init__ (line 16) | def __init__(
method set_image (line 34) | def set_image(
method set_torch_image (line 65) | def set_torch_image(
method predict (line 94) | def predict(
method predict_torch (line 179) | def predict_torch(
method get_image_embedding (line 259) | def get_image_embedding(self) -> torch.Tensor:
method device (line 275) | def device(self) -> torch.device:
method reset_image (line 278) | def reset_image(self) -> None:
FILE: iopaint/plugins/segment_anything/predictor_hq.py
class SamHQPredictor (line 17) | class SamHQPredictor:
method __init__ (line 18) | def __init__(
method set_image (line 34) | def set_image(
method set_torch_image (line 67) | def set_torch_image(
method predict (line 96) | def predict(
method predict_torch (line 183) | def predict_torch(
method get_image_embedding (line 266) | def get_image_embedding(self) -> torch.Tensor:
method device (line 282) | def device(self) -> torch.device:
method reset_image (line 285) | def reset_image(self) -> None:
FILE: iopaint/plugins/segment_anything/utils/transforms.py
class ResizeLongestSide (line 16) | class ResizeLongestSide:
method __init__ (line 23) | def __init__(self, target_length: int) -> None:
method apply_image (line 26) | def apply_image(self, image: np.ndarray) -> np.ndarray:
method apply_coords (line 35) | def apply_coords(
method apply_boxes (line 51) | def apply_boxes(
method apply_image_torch (line 61) | def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
method apply_coords_torch (line 75) | def apply_coords_torch(
method apply_boxes_torch (line 91) | def apply_boxes_torch(
method get_preprocess_shape (line 102) | def get_preprocess_shape(
FILE: iopaint/plugins/segment_anything2/build_sam.py
function build_memory_attention (line 72) | def build_memory_attention():
function build_memory_encoder (line 106) | def build_memory_encoder():
function build_image_encoder_tiny (line 130) | def build_image_encoder_tiny():
function build_image_encoder_small (line 156) | def build_image_encoder_small():
function build_image_encoder_base (line 182) | def build_image_encoder_base():
function build_image_encoder_large (line 208) | def build_image_encoder_large():
function build_sam2_tiny (line 234) | def build_sam2_tiny():
function build_sam2_small (line 243) | def build_sam2_small():
function build_sam2_base (line 252) | def build_sam2_base():
function build_sam2_large (line 261) | def build_sam2_large():
function build_sam2_1_tiny (line 270) | def build_sam2_1_tiny():
function build_sam2_1_small (line 279) | def build_sam2_1_small():
function build_sam2_1_base (line 288) | def build_sam2_1_base():
function build_sam2_1_large (line 297) | def build_sam2_1_large():
function build_sam2 (line 318) | def build_sam2(
function _load_checkpoint (line 332) | def _load_checkpoint(model, ckpt_path):
FILE: iopaint/plugins/segment_anything2/modeling/backbones/hieradet.py
function do_pool (line 24) | def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) ->...
class MultiScaleAttention (line 38) | class MultiScaleAttention(nn.Module):
method __init__ (line 39) | def __init__(
method forward (line 55) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class MultiScaleBlock (line 83) | class MultiScaleBlock(nn.Module):
method __init__ (line 84) | def __init__(
method forward (line 133) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Hiera (line 168) | class Hiera(nn.Module):
method __init__ (line 173) | def __init__(
method _get_pos_embed (line 271) | def _get_pos_embed(self, hw: Tuple[int, int]) -> torch.Tensor:
method forward (line 281) | def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
method get_layer_id (line 299) | def get_layer_id(self, layer_name):
method get_num_layers (line 314) | def get_num_layers(self) -> int:
FILE: iopaint/plugins/segment_anything2/modeling/backbones/image_encoder.py
class ImageEncoder (line 14) | class ImageEncoder(nn.Module):
method __init__ (line 15) | def __init__(
method forward (line 29) | def forward(self, sample: torch.Tensor):
class FpnNeck (line 45) | class FpnNeck(nn.Module):
method __init__ (line 52) | def __init__(
method forward (line 102) | def forward(self, xs: List[torch.Tensor]):
FILE: iopaint/plugins/segment_anything2/modeling/backbones/utils.py
function window_partition (line 16) | def window_partition(x, window_size):
function window_unpartition (line 41) | def window_unpartition(windows, window_size, pad_hw, hw):
class PatchEmbed (line 65) | class PatchEmbed(nn.Module):
method __init__ (line 70) | def __init__(
method forward (line 91) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: iopaint/plugins/segment_anything2/modeling/memory_attention.py
class MemoryAttentionLayer (line 17) | class MemoryAttentionLayer(nn.Module):
method __init__ (line 19) | def __init__(
method _forward_sa (line 58) | def _forward_sa(self, tgt, query_pos):
method _forward_ca (line 66) | def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0):
method forward (line 83) | def forward(
class MemoryAttention (line 102) | class MemoryAttention(nn.Module):
method __init__ (line 103) | def __init__(
method forward (line 119) | def forward(
FILE: iopaint/plugins/segment_anything2/modeling/memory_encoder.py
class MaskDownSampler (line 17) | class MaskDownSampler(nn.Module):
method __init__ (line 26) | def __init__(
method forward (line 57) | def forward(self, x):
class CXBlock (line 62) | class CXBlock(nn.Module):
method __init__ (line 74) | def __init__(
method forward (line 104) | def forward(self, x):
class Fuser (line 120) | class Fuser(nn.Module):
method __init__ (line 121) | def __init__(self, layer, num_layers, dim=None, input_projection=False):
method forward (line 130) | def forward(self, x):
class MemoryEncoder (line 138) | class MemoryEncoder(nn.Module):
method __init__ (line 139) | def __init__(
method forward (line 158) | def forward(
FILE: iopaint/plugins/segment_anything2/modeling/position_encoding.py
class PositionEmbeddingSine (line 16) | class PositionEmbeddingSine(nn.Module):
method __init__ (line 22) | def __init__(
method _encode_xy (line 42) | def _encode_xy(self, x, y):
method encode_boxes (line 62) | def encode_boxes(self, x, y, w, h):
method encode_points (line 70) | def encode_points(self, x, y, labels):
method forward (line 79) | def forward(self, x: torch.Tensor):
class PositionEmbeddingRandom (line 115) | class PositionEmbeddingRandom(nn.Module):
method __init__ (line 120) | def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = N...
method _pe_encoding (line 129) | def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
method forward (line 138) | def forward(self, size: Tuple[int, int]) -> torch.Tensor:
method forward_with_coords (line 151) | def forward_with_coords(
function init_t_xy (line 167) | def init_t_xy(end_x: int, end_y: int):
function compute_axial_cis (line 174) | def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 1...
function reshape_for_broadcast (line 186) | def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
function apply_rotary_enc (line 194) | def apply_rotary_enc(
FILE: iopaint/plugins/segment_anything2/modeling/sam/mask_decoder.py
class MaskDecoder (line 15) | class MaskDecoder(nn.Module):
method __init__ (line 16) | def __init__(
method forward (line 110) | def forward(
method predict_masks (line 168) | def predict_masks(
method _get_stability_scores (line 247) | def _get_stability_scores(self, mask_logits):
method _dynamic_multimask_via_stability (line 259) | def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_sc...
FILE: iopaint/plugins/segment_anything2/modeling/sam/prompt_encoder.py
class PromptEncoder (line 17) | class PromptEncoder(nn.Module):
method __init__ (line 18) | def __init__(
method get_dense_pe (line 68) | def get_dense_pe(self) -> torch.Tensor:
method _embed_points (line 79) | def _embed_points(
method _embed_boxes (line 103) | def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
method _embed_masks (line 114) | def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
method _get_batch_size (line 119) | def _get_batch_size(
method _get_device (line 137) | def _get_device(self) -> torch.device:
method forward (line 140) | def forward(
FILE: iopaint/plugins/segment_anything2/modeling/sam/transformer.py
class TwoWayTransformer (line 25) | class TwoWayTransformer(nn.Module):
method __init__ (line 26) | def __init__(
method forward (line 71) | def forward(
class TwoWayAttentionBlock (line 118) | class TwoWayAttentionBlock(nn.Module):
method __init__ (line 119) | def __init__(
method forward (line 162) | def forward(
class Attention (line 196) | class Attention(nn.Module):
method __init__ (line 202) | def __init__(
method _separate_heads (line 226) | def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
method _recombine_heads (line 231) | def _recombine_heads(self, x: Tensor) -> Tensor:
method forward (line 236) | def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
class RoPEAttention (line 263) | class RoPEAttention(Attention):
method __init__ (line 266) | def __init__(
method forward (line 285) | def forward(
FILE: iopaint/plugins/segment_anything2/modeling/sam2_base.py
class SAM2Base (line 22) | class SAM2Base(torch.nn.Module):
method __init__ (line 23) | def __init__(
method device (line 199) | def device(self):
method forward (line 202) | def forward(self, *args, **kwargs):
method _build_sam_heads (line 208) | def _build_sam_heads(self):
method _forward_sam_heads (line 258) | def _forward_sam_heads(
method _use_mask_as_output (line 416) | def _use_mask_as_output(self, backbone_features, high_res_features, ma...
method forward_image (line 468) | def forward_image(self, img_batch: torch.Tensor):
method _prepare_backbone_features (line 482) | def _prepare_backbone_features(self, backbone_out):
method _prepare_memory_conditioned_features (line 498) | def _prepare_memory_conditioned_features(
method _encode_new_memory (line 677) | def _encode_new_memory(
method _track_step (line 729) | def _track_step(
method _encode_memory_in_output (line 790) | def _encode_memory_in_output(
method track_step (line 815) | def track_step(
method _use_multimask (line 882) | def _use_multimask(self, is_init_cond_frame, point_inputs):
method _apply_non_overlapping_constraints (line 892) | def _apply_non_overlapping_constraints(self, pred_masks):
FILE: iopaint/plugins/segment_anything2/modeling/sam2_utils.py
function select_closest_cond_frames (line 19) | def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_f...
function get_1d_sine_pe (line 64) | def get_1d_sine_pe(pos_inds, dim, temperature=10000):
function get_activation_fn (line 77) | def get_activation_fn(activation):
function get_clones (line 88) | def get_clones(module, N):
class DropPath (line 92) | class DropPath(nn.Module):
method __init__ (line 94) | def __init__(self, drop_prob=0.0, scale_by_keep=True):
method forward (line 99) | def forward(self, x):
class MLP (line 112) | class MLP(nn.Module):
method __init__ (line 113) | def __init__(
method forward (line 131) | def forward(self, x):
class LayerNorm2d (line 141) | class LayerNorm2d(nn.Module):
method __init__ (line 142) | def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
method forward (line 148) | def forward(self, x: torch.Tensor) -> torch.Tensor:
function sample_box_points (line 156) | def sample_box_points(
function sample_random_points_from_errors (line 202) | def sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1):
function sample_one_point_from_error_center (line 252) | def sample_one_point_from_error_center(gt_masks, pred_masks, padding=True):
function get_next_point (line 317) | def get_next_point(gt_masks, pred_masks, method):
FILE: iopaint/plugins/segment_anything2/sam2_image_predictor.py
class SAM2ImagePredictor (line 20) | class SAM2ImagePredictor:
method __init__ (line 21) | def __init__(
method set_image (line 66) | def set_image(
method set_image_batch (line 112) | def set_image_batch(
method predict_batch (line 155) | def predict_batch(
method predict (line 217) | def predict(
method _prep_prompts (line 285) | def _prep_prompts(
method _predict (line 316) | def _predict(
method get_image_embedding (line 419) | def get_image_embedding(self) -> torch.Tensor:
method device (line 435) | def device(self) -> torch.device:
method reset_predictor (line 438) | def reset_predictor(self) -> None:
FILE: iopaint/plugins/segment_anything2/utils/misc.py
function get_sdpa_settings (line 14) | def get_sdpa_settings():
function mask_to_box (line 44) | def mask_to_box(masks: torch.Tensor):
function _load_img_as_tensor (line 70) | def _load_img_as_tensor(img_path, image_size):
function concat_points (line 82) | def concat_points(old_point_inputs, new_points, new_labels):
FILE: iopaint/plugins/segment_anything2/utils/transforms.py
class SAM2Transforms (line 12) | class SAM2Transforms(nn.Module):
method __init__ (line 13) | def __init__(
method __call__ (line 34) | def __call__(self, x):
method forward_batch (line 38) | def forward_batch(self, img_list):
method transform_coords (line 43) | def transform_coords(
method transform_boxes (line 63) | def transform_boxes(
method postprocess_masks (line 73) | def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Ten...
FILE: iopaint/runtime.py
function dump_environment_info (line 44) | def dump_environment_info() -> Dict[str, str]:
function check_device (line 57) | def check_device(device: Device) -> Device:
function setup_model_dir (line 79) | def setup_model_dir(model_dir: Path):
FILE: iopaint/schema.py
class ModelType (line 22) | class ModelType(str, Enum):
class ModelInfo (line 31) | class ModelInfo(BaseModel):
method need_prompt (line 39) | def need_prompt(self) -> bool:
method controlnets (line 54) | def controlnets(self) -> List[str]:
method brushnets (line 71) | def brushnets(self) -> List[str]:
method support_strength (line 80) | def support_strength(self) -> bool:
method support_outpainting (line 90) | def support_outpainting(self) -> bool:
method support_lcm_lora (line 100) | def support_lcm_lora(self) -> bool:
method support_controlnet (line 110) | def support_controlnet(self) -> bool:
method support_brushnet (line 120) | def support_brushnet(self) -> bool:
method support_powerpaint_v2 (line 128) | def support_powerpaint_v2(self) -> bool:
class Choices (line 138) | class Choices(str, Enum):
method values (line 140) | def values(cls):
class RealESRGANModel (line 144) | class RealESRGANModel(Choices):
class RemoveBGModel (line 150) | class RemoveBGModel(Choices):
class Device (line 169) | class Device(Choices):
class InteractiveSegModel (line 175) | class InteractiveSegModel(Choices):
class PluginInfo (line 195) | class PluginInfo(BaseModel):
class CV2Flag (line 201) | class CV2Flag(str, Enum):
class HDStrategy (line 206) | class HDStrategy(str, Enum):
class LDMSampler (line 217) | class LDMSampler(str, Enum):
class SDSampler (line 222) | class SDSampler(str, Enum):
class PowerPaintTask (line 245) | class PowerPaintTask(Choices):
class ApiConfig (line 253) | class ApiConfig(BaseModel):
class InpaintRequest (line 285) | class InpaintRequest(BaseModel):
method validate_field (line 422) | def validate_field(cls, values: "InpaintRequest"):
class RunPluginRequest (line 451) | class RunPluginRequest(BaseModel):
class MediasResponse (line 463) | class MediasResponse(BaseModel):
class GenInfoResponse (line 471) | class GenInfoResponse(BaseModel):
class ServerConfigResponse (line 476) | class ServerConfigResponse(BaseModel):
class SwitchModelRequest (line 494) | class SwitchModelRequest(BaseModel):
class SwitchPluginModelRequest (line 498) | class SwitchPluginModelRequest(BaseModel):
class AdjustMaskRequest (line 506) | class AdjustMaskRequest(BaseModel):
FILE: iopaint/tests/test_adjust_mask.py
function test_adjust_mask (line 8) | def test_adjust_mask():
FILE: iopaint/tests/test_anytext.py
function test_anytext (line 20) | def test_anytext(device):
FILE: iopaint/tests/test_brushnet.py
function test_runway_brushnet (line 22) | def test_runway_brushnet(device, sampler):
function test_runway_powerpaint_v2 (line 51) | def test_runway_powerpaint_v2(device, sampler):
FILE: iopaint/tests/test_controlnet.py
function convert_controlnet_method_name (line 19) | def convert_controlnet_method_name(name):
function test_runway_sd_1_5 (line 25) | def test_runway_sd_1_5(device, controlnet_method):
function test_controlnet_switch (line 56) | def test_controlnet_switch(device):
function test_local_file_path (line 88) | def test_local_file_path(device, local_file):
FILE: iopaint/tests/test_instruct_pix2pix.py
function test_instruct_pix2pix (line 16) | def test_instruct_pix2pix(device, disable_nsfw, cpu_offload):
FILE: iopaint/tests/test_load_img.py
function test_load_png_image (line 8) | def test_load_png_image():
function test_load_jpg_image (line 15) | def test_load_jpg_image():
FILE: iopaint/tests/test_low_mem.py
function test_runway_sd_1_5_low_mem (line 17) | def test_runway_sd_1_5_low_mem(device):
function test_runway_sd_lcm_lora_low_mem (line 49) | def test_runway_sd_lcm_lora_low_mem(device, sampler):
function test_runway_norm_sd_model (line 82) | def test_runway_norm_sd_model(device, strategy, sampler):
FILE: iopaint/tests/test_match_histograms.py
function test_sd_match_histograms (line 11) | def test_sd_match_histograms(device, sampler):
FILE: iopaint/tests/test_model.py
function test_lama (line 13) | def test_lama(device, strategy):
function test_ldm (line 36) | def test_ldm(device, strategy, ldm_sampler):
function test_zits (line 58) | def test_zits(device, strategy, zits_wireframe):
function test_mat (line 80) | def test_mat(device, strategy, no_half):
function test_fcf (line 94) | def test_fcf(device, strategy):
function test_cv2 (line 108) | def test_cv2(strategy, cv2_flag, cv2_radius):
function test_manga (line 127) | def test_manga(device, strategy):
function test_mi_gan (line 145) | def test_mi_gan(device, strategy):
FILE: iopaint/tests/test_model_md5.py
function test_load_model (line 1) | def test_load_model():
FILE: iopaint/tests/test_model_switch.py
function test_model_switch (line 12) | def test_model_switch():
function test_controlnet_switch_onoff (line 26) | def test_controlnet_switch_onoff(caplog):
function test_switch_controlnet_method (line 48) | def test_switch_controlnet_method(caplog):
FILE: iopaint/tests/test_outpainting.py
function test_outpainting (line 29) | def test_outpainting(name, device, rect):
function test_kandinsky_outpainting (line 67) | def test_kandinsky_outpainting(name, device, rect):
function test_powerpaint_outpainting (line 108) | def test_powerpaint_outpainting(name, device, rect):
FILE: iopaint/tests/test_paint_by_example.py
function assert_equal (line 19) | def assert_equal(
function test_paint_by_example (line 46) | def test_paint_by_example(device):
FILE: iopaint/tests/test_plugins.py
function _save (line 35) | def _save(img, name):
function test_remove_bg (line 42) | def test_remove_bg(model_name, device):
function test_anime_seg (line 63) | def test_anime_seg():
function test_upscale (line 78) | def test_upscale(device):
function test_gfpgan (line 95) | def test_gfpgan(device):
function test_restoreformer (line 105) | def test_restoreformer(device):
function test_segment_anything (line 117) | def test_segment_anything(name, device):
FILE: iopaint/tests/test_save_exif.py
function print_exif (line 13) | def print_exif(exif):
function extra_info (line 18) | def extra_info(img_p: Path):
function assert_keys (line 27) | def assert_keys(keys: List[str], infos, res_infos):
function run_test (line 34) | def run_test(file_path, keys):
function test_png_icc_profile_png (line 44) | def test_png_icc_profile_png():
function test_png_icc_profile_jpeg (line 48) | def test_png_icc_profile_jpeg():
function test_jpeg (line 52) | def test_jpeg():
function test_png_parameter (line 57) | def test_png_parameter():
FILE: iopaint/tests/test_save_quality.py
function test_jpeg_quality (line 9) | def test_jpeg_quality():
function test_png_parameters (line 30) | def test_png_parameters():
function test_format_conversion (line 51) | def test_format_conversion():
FILE: iopaint/tests/test_sd_model.py
function test_runway_sd_1_5_all_samplers (line 22) | def test_runway_sd_1_5_all_samplers(device):
function test_runway_sd_lcm_lora (line 65) | def test_runway_sd_lcm_lora(device, sampler):
function test_runway_sd_sd_strength (line 96) | def test_runway_sd_sd_strength(device, strategy, sampler):
function test_runway_sd_cpu_textencoder (line 124) | def test_runway_sd_cpu_textencoder(device, strategy, sampler):
function test_runway_norm_sd_model (line 151) | def test_runway_norm_sd_model(device, strategy, sampler):
function test_runway_sd_1_5_cpu_offload (line 176) | def test_runway_sd_1_5_cpu_offload(device, strategy, sampler):
function test_local_file_path (line 212) | def test_local_file_path(device, sampler, name):
FILE: iopaint/tests/test_sdxl.py
function test_sdxl (line 18) | def test_sdxl(device, strategy, sampler):
function test_sdxl_cpu_text_encoder (line 50) | def test_sdxl_cpu_text_encoder(device, strategy, sampler):
function test_sdxl_outpainting (line 86) | def test_sdxl_outpainting(device, rect):
FILE: iopaint/tests/utils.py
function check_device (line 14) | def check_device(device: str) -> int:
function assert_equal (line 23) | def assert_equal(
function get_data (line 52) | def get_data(
function get_config (line 66) | def get_config(**kwargs):
FILE: iopaint/web_config.py
class WebConfig (line 67) | class WebConfig(ApiConfig):
function load_config (line 71) | def load_config(p: Path) -> WebConfig:
function save_config (line 83) | def save_config(
function change_current_model (line 138) | def change_current_model(new_model):
function main (line 142) | def main(config_file: Path):
FILE: scripts/convert_vae_pt_to_diffusers.py
function custom_convert_ldm_vae_checkpoint (line 18) | def custom_convert_ldm_vae_checkpoint(checkpoint, config):
function vae_pt_to_vae_diffuser (line 187) | def vae_pt_to_vae_diffuser(
FILE: scripts/tool.py
class CheckpointMergerPipeline (line 24) | class CheckpointMergerPipeline(DiffusionPipeline):
method __init__ (line 45) | def __init__(self):
method _compare_model_configs (line 49) | def _compare_model_configs(self, dict0, dict1):
method _remove_meta_keys (line 60) | def _remove_meta_keys(self, config_dict: Dict):
method merge (line 70) | def merge(
method weighted_sum (line 314) | def weighted_sum(theta0, theta1, theta2, alpha):
method sigmoid (line 319) | def sigmoid(theta0, theta1, theta2, alpha):
method inv_sigmoid (line 325) | def inv_sigmoid(theta0, theta1, theta2, alpha):
method add_difference (line 332) | def add_difference(theta0, theta1, theta2, alpha):
FILE: setup.py
function load_requirements (line 18) | def load_requirements():
FILE: web_app/src/App.tsx
constant SUPPORTED_FILE_TYPE (line 13) | const SUPPORTED_FILE_TYPE = [
function Home (line 20) | function Home() {
FILE: web_app/src/components/Coffee.tsx
function Coffee (line 7) | function Coffee() {
FILE: web_app/src/components/Cropper.tsx
constant DOC_MOVE_OPTS (line 6) | const DOC_MOVE_OPTS = { capture: true, passive: false }
constant DRAG_HANDLE_BORDER (line 8) | const DRAG_HANDLE_BORDER = 2
type EVData (line 10) | interface EVData {
type Props (line 20) | interface Props {
FILE: web_app/src/components/DiffusionProgress.tsx
constant API_ENDPOINT (line 6) | const API_ENDPOINT = import.meta.env.DEV
FILE: web_app/src/components/Editor.tsx
constant TOOLBAR_HEIGHT (line 39) | const TOOLBAR_HEIGHT = 200
constant COMPARE_SLIDER_DURATION_MS (line 40) | const COMPARE_SLIDER_DURATION_MS = 300
type EditorProps (line 42) | interface EditorProps {
function Editor (line 46) | function Editor(props: EditorProps) {
FILE: web_app/src/components/Extender.tsx
constant DOC_MOVE_OPTS (line 7) | const DOC_MOVE_OPTS = { capture: true, passive: false }
constant DRAG_HANDLE_BORDER (line 9) | const DRAG_HANDLE_BORDER = 2
type EVData (line 11) | interface EVData {
type Props (line 21) | interface Props {
FILE: web_app/src/components/FileManager.tsx
type Photo (line 40) | interface Photo {
constant SORT_BY_NAME (line 47) | const SORT_BY_NAME = "Name"
constant SORT_BY_CREATED_TIME (line 48) | const SORT_BY_CREATED_TIME = "Created time"
constant SORT_BY_MODIFIED_TIME (line 49) | const SORT_BY_MODIFIED_TIME = "Modified time"
constant IMAGE_TAB (line 51) | const IMAGE_TAB = "input"
constant OUTPUT_TAB (line 52) | const OUTPUT_TAB = "output"
constant MASK_TAB (line 53) | const MASK_TAB = "mask"
type Props (line 61) | interface Props {
function FileManager (line 66) | function FileManager(props: Props) {
FILE: web_app/src/components/FileSelect.tsx
type FileSelectProps (line 4) | type FileSelectProps = {
function FileSelect (line 8) | function FileSelect(props: FileSelectProps) {
FILE: web_app/src/components/InteractiveSeg.tsx
type InteractiveSegReplaceModal (line 5) | interface InteractiveSegReplaceModal {
type ItemProps (line 79) | interface ItemProps {
FILE: web_app/src/components/Plugins.tsx
type PluginName (line 21) | enum PluginName {
FILE: web_app/src/components/Settings.tsx
constant TAB_GENERAL (line 64) | const TAB_GENERAL = "General"
constant TAB_MODEL (line 65) | const TAB_MODEL = "Model"
constant TAB_PLUGINS (line 66) | const TAB_PLUGINS = "Plugins"
constant TAB_NAMES (line 69) | const TAB_NAMES = [TAB_MODEL, TAB_GENERAL, TAB_PLUGINS]
function SettingsDialog (line 71) | function SettingsDialog() {
FILE: web_app/src/components/Shortcuts.tsx
type ShortcutProps (line 13) | interface ShortcutProps {
function ShortCut (line 18) | function ShortCut(props: ShortcutProps) {
function Shortcuts (line 44) | function Shortcuts() {
FILE: web_app/src/components/ui/button.tsx
type ButtonProps (line 39) | interface ButtonProps
type IconButtonProps (line 63) | interface IconButtonProps extends ButtonProps {
type UploadButtonProps (line 91) | interface UploadButtonProps extends IconButtonProps {
FILE: web_app/src/components/ui/form.tsx
type FormFieldContextValue (line 18) | type FormFieldContextValue<
type FormItemContextValue (line 65) | type FormItemContextValue = {
FILE: web_app/src/components/ui/input.tsx
type InputProps (line 6) | interface InputProps
type NumberInputProps (line 40) | interface NumberInputProps extends InputProps {
FILE: web_app/src/components/ui/label.tsx
type LabelProps (line 11) | interface LabelProps {
FILE: web_app/src/components/ui/sheet.tsx
type SheetContentProps (line 49) | interface SheetContentProps
FILE: web_app/src/components/ui/textarea.tsx
type TextareaProps (line 6) | interface TextareaProps
FILE: web_app/src/components/ui/toast.tsx
type ToastProps (line 119) | type ToastProps = React.ComponentPropsWithoutRef<typeof Toast>
type ToastActionElement (line 121) | type ToastActionElement = React.ReactElement<typeof ToastAction>
FILE: web_app/src/components/ui/toaster.tsx
function Toaster (line 11) | function Toaster() {
FILE: web_app/src/components/ui/use-toast.ts
constant TOAST_LIMIT (line 9) | const TOAST_LIMIT = 1
constant TOAST_REMOVE_DELAY (line 10) | const TOAST_REMOVE_DELAY = 1000000
type ToasterToast (line 12) | type ToasterToast = ToastProps & {
function genId (line 28) | function genId() {
type ActionType (line 33) | type ActionType = typeof actionTypes
type Action (line 35) | type Action =
type State (line 53) | interface State {
function dispatch (line 134) | function dispatch(action: Action) {
type Toast (line 141) | type Toast = Omit<ToasterToast, "id">
function toast (line 143) | function toast({ ...props }: Toast) {
function useToast (line 172) | function useToast() {
FILE: web_app/src/hooks/useAsyncMemo.tsx
function useAsyncMemo (line 12) | function useAsyncMemo<T>(
FILE: web_app/src/hooks/useImage.tsx
function useImage (line 3) | function useImage(file: File | null): [HTMLImageElement, boolean] {
FILE: web_app/src/hooks/useInputImage.tsx
function useInputImage (line 4) | function useInputImage() {
FILE: web_app/src/lib/api.ts
constant API_ENDPOINT (line 13) | const API_ENDPOINT = import.meta.env.DEV
function inpaint (line 28) | async function inpaint(
function getServerConfig (line 105) | async function getServerConfig(): Promise<ServerConfig> {
function switchModel (line 110) | async function switchModel(name: string): Promise<ModelInfo> {
function switchPluginModel (line 115) | async function switchPluginModel(
function currentModel (line 122) | async function currentModel(): Promise<ModelInfo> {
function runPlugin (line 127) | async function runPlugin(
function getMediaFile (line 155) | async function getMediaFile(tab: string, filename: string) {
function getMediaBlob (line 174) | async function getMediaBlob(tab: string, filename: string) {
function getMedias (line 190) | async function getMedias(tab: string): Promise<Filename[]> {
function downloadToOutput (line 195) | async function downloadToOutput(
function getGenInfo (line 217) | async function getGenInfo(file: File): Promise<GenInfo> {
function getSamplers (line 224) | async function getSamplers(): Promise<string[]> {
function postAdjustMask (line 229) | async function postAdjustMask(
FILE: web_app/src/lib/const.ts
constant ACCENT_COLOR (line 1) | const ACCENT_COLOR = "#ffcc00bb"
constant DEFAULT_BRUSH_SIZE (line 2) | const DEFAULT_BRUSH_SIZE = 40
constant MIN_BRUSH_SIZE (line 3) | const MIN_BRUSH_SIZE = 1
constant MAX_BRUSH_SIZE (line 4) | const MAX_BRUSH_SIZE = 200
constant MODEL_TYPE_INPAINT (line 5) | const MODEL_TYPE_INPAINT = "inpaint"
constant MODEL_TYPE_DIFFUSERS_SD (line 6) | const MODEL_TYPE_DIFFUSERS_SD = "diffusers_sd"
constant MODEL_TYPE_DIFFUSERS_SDXL (line 7) | const MODEL_TYPE_DIFFUSERS_SDXL = "diffusers_sdxl"
constant MODEL_TYPE_DIFFUSERS_SD_INPAINT (line 8) | const MODEL_TYPE_DIFFUSERS_SD_INPAINT = "diffusers_sd_inpaint"
constant MODEL_TYPE_DIFFUSERS_SDXL_INPAINT (line 9) | const MODEL_TYPE_DIFFUSERS_SDXL_INPAINT = "diffusers_sdxl_inpaint"
constant MODEL_TYPE_OTHER (line 10) | const MODEL_TYPE_OTHER = "diffusers_other"
constant BRUSH_COLOR (line 11) | const BRUSH_COLOR = "#ffcc00bb"
constant LDM (line 13) | const LDM = "ldm"
constant CV2 (line 14) | const CV2 = "cv2"
constant PAINT_BY_EXAMPLE (line 16) | const PAINT_BY_EXAMPLE = "Fantasy-Studio/Paint-by-Example"
constant INSTRUCT_PIX2PIX (line 17) | const INSTRUCT_PIX2PIX = "timbrooks/instruct-pix2pix"
constant KANDINSKY_2_2 (line 18) | const KANDINSKY_2_2 = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
constant POWERPAINT (line 19) | const POWERPAINT = "Sanster/PowerPaint-V1-stable-diffusion-inpainting"
constant ANYTEXT (line 20) | const ANYTEXT = "Sanster/AnyText"
constant DEFAULT_NEGATIVE_PROMPT (line 22) | const DEFAULT_NEGATIVE_PROMPT =
constant SHORTCUT_KEY_CHANGE_BRUSH_SIZE (line 25) | const SHORTCUT_KEY_CHANGE_BRUSH_SIZE = "Alt"
FILE: web_app/src/lib/states.ts
type FileManagerState (line 41) | type FileManagerState = {
type CropperState (line 50) | type CropperState = {
type Settings (line 57) | type Settings = {
type InteractiveSegState (line 116) | type InteractiveSegState = {
type EditorState (line 122) | type EditorState = {
type AppState (line 141) | type AppState = {
type AppAction (line 166) | type AppAction = {
method updateEnableControlnet (line 836) | updateEnableControlnet(newValue) {
method updateLCMLora (line 846) | updateLCMLora(newValue) {
FILE: web_app/src/lib/types.ts
type Filename (line 1) | interface Filename {
type PluginInfo (line 9) | interface PluginInfo {
type ServerConfig (line 15) | interface ServerConfig {
type GenInfo (line 33) | interface GenInfo {
type ModelInfo (line 38) | interface ModelInfo {
type PluginName (line 60) | enum PluginName {
type PluginParams (line 69) | interface PluginParams {
type SortBy (line 73) | enum SortBy {
type SortOrder (line 79) | enum SortOrder {
type LDMSampler (line 84) | enum LDMSampler {
type CV2Flag (line 89) | enum CV2Flag {
type Rect (line 94) | interface Rect {
type Point (line 101) | interface Point {
type Line (line 106) | interface Line {
type LineGroup (line 111) | type LineGroup = Array<Line>
type Size (line 113) | interface Size {
type ExtenderDirection (line 118) | enum ExtenderDirection {
type PowerPaintTask (line 124) | enum PowerPaintTask {
type AdjustMaskOperate (line 132) | type AdjustMaskOperate = "expand" | "shrink" | "reverse"
FILE: web_app/src/lib/utils.ts
function cn (line 7) | function cn(...inputs: ClassValue[]) {
function keepGUIAlive (line 11) | function keepGUIAlive() {
function dataURItoBlob (line 33) | function dataURItoBlob(dataURI: string) {
function loadImage (line 43) | function loadImage(image: HTMLImageElement, src: string) {
function blobToImage (line 56) | async function blobToImage(blob: Blob) {
function canvasToImage (line 63) | function canvasToImage(
function fileToImage (line 81) | function fileToImage(file: File): Promise<HTMLImageElement> {
function srcToFile (line 101) | function srcToFile(src: string, fileName: string, mimeType: string) {
function askWritePermission (line 111) | async function askWritePermission() {
function canvasToBlob (line 125) | function canvasToBlob(canvas: HTMLCanvasElement, mime: string): Promise<...
function isRightClick (line 142) | function isRightClick(ev: SyntheticEvent) {
function isMidClick (line 147) | function isMidClick(ev: SyntheticEvent) {
function copyCanvasImage (line 152) | async function copyCanvasImage(canvas: HTMLCanvasElement) {
function downloadImage (line 161) | function downloadImage(uri: string, name: string) {
function mouseXY (line 182) | function mouseXY(ev: SyntheticEvent) {
function drawLines (line 197) | function drawLines(
Condensed preview — 302 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (2,450K chars).
[
{
"path": ".github/FUNDING.yml",
"chars": 64,
"preview": "# These are supported funding model platforms\n\nko_fi: Z8Z1CZJGY\n"
},
{
"path": ".github/ISSUE_TEMPLATE/config.yml",
"chars": 114,
"preview": "contact_links:\n - name: Blank issue\n url: https://github.com/Sanster/lama-cleaner/issues/new\n about: Other\n"
},
{
"path": ".github/ISSUE_TEMPLATE/🐛-bug-report.md",
"chars": 383,
"preview": "---\nname: \"\\U0001F41B Bug Report\"\nabout: Create a report to help us improve\ntitle: \"[BUG]\"\nlabels: bug\nassignees: ''\n\n--"
},
{
"path": ".github/ISSUE_TEMPLATE/🚀-feature-request.md",
"chars": 488,
"preview": "---\nname: \"\\U0001F680 Feature request\"\nabout: Suggest an idea for this project\ntitle: \"[Feature Request]\"\nlabels: ''\nass"
},
{
"path": ".gitignore",
"chars": 128,
"preview": ".DS_Store\n**/__pycache__\nexamples/\n.idea/\n.vscode/\nbuild\n!iopaint/app/build\ndist/\nIOPaint.egg-info/\nvenv/\ntmp/\niopaint/w"
},
{
"path": "LICENSE",
"chars": 11357,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "README.md",
"chars": 6748,
"preview": "<h1 align=\"center\">IOPaint</h1>\n<p align=\"center\">A free and open-source inpainting & outpainting tool powered by SOTA A"
},
{
"path": "build_docker.sh",
"chars": 1077,
"preview": "#!/usr/bin/env bash\nset -e\n\nGIT_TAG=$1\nIMAGE_DESC=\"Image inpainting tool powered by SOTA AI Model\" \nGIT_REPO=\"https://gi"
},
{
"path": "docker/CPUDockerfile",
"chars": 619,
"preview": "FROM python:3.10.11-slim-buster\n\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n software-propert"
},
{
"path": "docker/GPUDockerfile",
"chars": 587,
"preview": "FROM nvidia/cuda:11.7.1-runtime-ubuntu20.04\n\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n soft"
},
{
"path": "iopaint/__init__.py",
"chars": 2000,
"preview": "import os\nimport importlib.util\nimport shutil\nimport ctypes\nimport logging\n\nos.environ[\"PYTORCH_ENABLE_MPS_FALLBACK\"] = "
},
{
"path": "iopaint/__main__.py",
"chars": 78,
"preview": "from iopaint import entry_point\n\nif __name__ == \"__main__\":\n entry_point()\n"
},
{
"path": "iopaint/api.py",
"chars": 15794,
"preview": "import asyncio\nimport os\nimport threading\nimport time\nimport traceback\nfrom pathlib import Path\nfrom typing import Optio"
},
{
"path": "iopaint/batch_processing.py",
"chars": 4124,
"preview": "import json\nfrom pathlib import Path\nfrom typing import Dict, Optional\n\nimport cv2\nimport numpy as np\nfrom PIL import Im"
},
{
"path": "iopaint/benchmark.py",
"chars": 3195,
"preview": "#!/usr/bin/env python3\n\nimport argparse\nimport os\nimport time\n\nimport numpy as np\nimport nvidia_smi\nimport psutil\nimport"
},
{
"path": "iopaint/cli.py",
"chars": 8450,
"preview": "import webbrowser\nfrom contextlib import asynccontextmanager\nfrom pathlib import Path\nfrom typing import Optional\n\nimpor"
},
{
"path": "iopaint/const.py",
"chars": 4468,
"preview": "import os\nfrom typing import List\n\nINSTRUCT_PIX2PIX_NAME = \"timbrooks/instruct-pix2pix\"\nKANDINSKY22_NAME = \"kandinsky-co"
},
{
"path": "iopaint/download.py",
"chars": 10961,
"preview": "import glob\nimport json\nimport os\nfrom functools import lru_cache\nfrom typing import List, Optional\n\nfrom iopaint.schema"
},
{
"path": "iopaint/file_manager/__init__.py",
"chars": 38,
"preview": "from .file_manager import FileManager\n"
},
{
"path": "iopaint/file_manager/file_manager.py",
"chars": 7299,
"preview": "import os\nfrom io import BytesIO\nfrom pathlib import Path\nfrom typing import List\n\nfrom PIL import Image, ImageOps, PngI"
},
{
"path": "iopaint/file_manager/storage_backends.py",
"chars": 1293,
"preview": "# Copy from https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/storage_backends.py\nimport er"
},
{
"path": "iopaint/file_manager/utils.py",
"chars": 1835,
"preview": "# Copy from: https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/utils.py\nimport hashlib\nfrom"
},
{
"path": "iopaint/helper.py",
"chars": 12080,
"preview": "import base64\nimport imghdr\nimport io\nimport os\nimport sys\nfrom typing import List, Optional, Dict, Tuple\n\nfrom urllib.p"
},
{
"path": "iopaint/installer.py",
"chars": 224,
"preview": "import subprocess\nimport sys\n\n\ndef install(package):\n subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", "
},
{
"path": "iopaint/model/__init__.py",
"chars": 1070,
"preview": "from .anytext.anytext_model import AnyText\nfrom .controlnet import ControlNet\nfrom .fcf import FcF\nfrom .instruct_pix2pi"
},
{
"path": "iopaint/model/anytext/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/anytext_model.py",
"chars": 2536,
"preview": "import torch\nfrom huggingface_hub import hf_hub_download\n\nfrom iopaint.const import ANYTEXT_NAME\nfrom iopaint.model.anyt"
},
{
"path": "iopaint/model/anytext/anytext_pipeline.py",
"chars": 15143,
"preview": "\"\"\"\nAnyText: Multilingual Visual Text Generation And Editing\nPaper: https://arxiv.org/abs/2311.03054\nCode: https://githu"
},
{
"path": "iopaint/model/anytext/anytext_sd15.yaml",
"chars": 2777,
"preview": "model:\n target: iopaint.model.anytext.cldm.cldm.ControlLDM\n params:\n linear_start: 0.00085\n linear_end: 0.0120\n "
},
{
"path": "iopaint/model/anytext/cldm/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/cldm/cldm.py",
"chars": 28531,
"preview": "import os\nfrom pathlib import Path\n\nimport einops\nimport torch\nimport torch as th\nimport torch.nn as nn\nimport copy\nfrom"
},
{
"path": "iopaint/model/anytext/cldm/ddim_hacked.py",
"chars": 17275,
"preview": "\"\"\"SAMPLING ONLY.\"\"\"\n\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\n\nfrom iopaint.model.anytext.ldm.modules.diff"
},
{
"path": "iopaint/model/anytext/cldm/embedding_manager.py",
"chars": 6083,
"preview": "'''\nCopyright (c) Alibaba, Inc. and its affiliates.\n'''\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as"
},
{
"path": "iopaint/model/anytext/cldm/hack.py",
"chars": 3699,
"preview": "import torch\nimport einops\n\nimport iopaint.model.anytext.ldm.modules.encoders.modules\nimport iopaint.model.anytext.ldm.m"
},
{
"path": "iopaint/model/anytext/cldm/model.py",
"chars": 1373,
"preview": "import os\nimport torch\n\nfrom omegaconf import OmegaConf\nfrom iopaint.model.anytext.ldm.util import instantiate_from_conf"
},
{
"path": "iopaint/model/anytext/cldm/recognizer.py",
"chars": 11439,
"preview": "\"\"\"\nCopyright (c) Alibaba, Inc. and its affiliates.\n\"\"\"\nimport os\nimport cv2\nimport numpy as np\nimport math\nimport trace"
},
{
"path": "iopaint/model/anytext/ldm/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/models/autoencoder.py",
"chars": 8614,
"preview": "import torch\nimport torch.nn.functional as F\nfrom contextlib import contextmanager\n\nfrom iopaint.model.anytext.ldm.modul"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/ddim.py",
"chars": 18214,
"preview": "\"\"\"SAMPLING ONLY.\"\"\"\n\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\n\nfrom iopaint.model.anytext.ldm.modules.diff"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/ddpm.py",
"chars": 89020,
"preview": "\"\"\"\nPart of the implementation is borrowed and modified from ControlNet, publicly available at https://github.com/lllyas"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/dpm_solver/__init__.py",
"chars": 37,
"preview": "from .sampler import DPMSolverSampler"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py",
"chars": 65968,
"preview": "import torch\nimport torch.nn.functional as F\nimport math\nfrom tqdm import tqdm\n\n\nclass NoiseScheduleVP:\n def __init__"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/dpm_solver/sampler.py",
"chars": 2990,
"preview": "\"\"\"SAMPLING ONLY.\"\"\"\nimport torch\n\nfrom .dpm_solver import NoiseScheduleVP, model_wrapper, DPM_Solver\n\n\nMODEL_TYPES = {\n"
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/plms.py",
"chars": 12971,
"preview": "\"\"\"SAMPLING ONLY.\"\"\"\n\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\nfrom functools import partial\n\nfrom iopaint."
},
{
"path": "iopaint/model/anytext/ldm/models/diffusion/sampling_util.py",
"chars": 753,
"preview": "import torch\nimport numpy as np\n\n\ndef append_dims(x, target_dims):\n \"\"\"Appends dimensions to the end of a tensor unti"
},
{
"path": "iopaint/model/anytext/ldm/modules/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/modules/attention.py",
"chars": 10738,
"preview": "from inspect import isfunction\nimport math\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn, einsum\nfro"
},
{
"path": "iopaint/model/anytext/ldm/modules/diffusionmodules/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/modules/diffusionmodules/model.py",
"chars": 30615,
"preview": "# pytorch_diffusion + derived encoder decoder\nimport math\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\n\ndef g"
},
{
"path": "iopaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py",
"chars": 30462,
"preview": "from abc import abstractmethod\nimport math\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nimport torch.nn."
},
{
"path": "iopaint/model/anytext/ldm/modules/diffusionmodules/upscaling.py",
"chars": 3468,
"preview": "import torch\nimport torch.nn as nn\nimport numpy as np\nfrom functools import partial\n\nfrom iopaint.model.anytext.ldm.modu"
},
{
"path": "iopaint/model/anytext/ldm/modules/diffusionmodules/util.py",
"chars": 9995,
"preview": "# adopted from\n# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py\n# and\n#"
},
{
"path": "iopaint/model/anytext/ldm/modules/distributions/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/modules/distributions/distributions.py",
"chars": 2970,
"preview": "import torch\nimport numpy as np\n\n\nclass AbstractDistribution:\n def sample(self):\n raise NotImplementedError()\n"
},
{
"path": "iopaint/model/anytext/ldm/modules/ema.py",
"chars": 3110,
"preview": "import torch\nfrom torch import nn\n\n\nclass LitEma(nn.Module):\n def __init__(self, model, decay=0.9999, use_num_upates="
},
{
"path": "iopaint/model/anytext/ldm/modules/encoders/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ldm/modules/encoders/modules.py",
"chars": 14135,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.utils.checkpoint import checkpoint\n\nfrom transformers import (\n T5Token"
},
{
"path": "iopaint/model/anytext/ldm/util.py",
"chars": 7250,
"preview": "import importlib\n\nimport torch\nfrom torch import optim\nimport numpy as np\n\nfrom inspect import isfunction\nfrom PIL impor"
},
{
"path": "iopaint/model/anytext/main.py",
"chars": 1554,
"preview": "import cv2\nimport os\n\nfrom anytext_pipeline import AnyTextPipeline\nfrom utils import save_images\n\nseed = 66273235\n# seed"
},
{
"path": "iopaint/model/anytext/ocr_recog/RNN.py",
"chars": 6607,
"preview": "from torch import nn\nimport torch\nfrom .RecSVTR import Block\n\nclass Swish(nn.Module):\n def __int__(self):\n sup"
},
{
"path": "iopaint/model/anytext/ocr_recog/RecCTCHead.py",
"chars": 1260,
"preview": "from torch import nn\n\n\nclass CTCHead(nn.Module):\n def __init__(self,\n in_channels,\n o"
},
{
"path": "iopaint/model/anytext/ocr_recog/RecModel.py",
"chars": 1673,
"preview": "from torch import nn\nfrom .RNN import SequenceEncoder, Im2Seq, Im2Im\nfrom .RecMv1_enhance import MobileNetV1Enhance\n\nfro"
},
{
"path": "iopaint/model/anytext/ocr_recog/RecMv1_enhance.py",
"chars": 6709,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom .common import Activation\n\n\nclass ConvBNLayer(nn"
},
{
"path": "iopaint/model/anytext/ocr_recog/RecSVTR.py",
"chars": 19468,
"preview": "import torch\nimport torch.nn as nn\nimport numpy as np\nfrom torch.nn.init import trunc_normal_, zeros_, ones_\nfrom torch."
},
{
"path": "iopaint/model/anytext/ocr_recog/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/anytext/ocr_recog/common.py",
"chars": 2180,
"preview": "\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass Hswish(nn.Module):\n def __init__(self, i"
},
{
"path": "iopaint/model/anytext/ocr_recog/en_dict.txt",
"chars": 190,
"preview": "0\n1\n2\n3\n4\n5\n6\n7\n8\n9\n:\n;\n<\n=\n>\n?\n@\nA\nB\nC\nD\nE\nF\nG\nH\nI\nJ\nK\nL\nM\nN\nO\nP\nQ\nR\nS\nT\nU\nV\nW\nX\nY\nZ\n[\n\\\n]\n^\n_\n`\na\nb\nc\nd\ne\nf\ng\nh\ni\nj\nk\n"
},
{
"path": "iopaint/model/anytext/ocr_recog/ppocr_keys_v1.txt",
"chars": 13245,
"preview": "'\n疗\n绚\n诚\n娇\n溜\n题\n贿\n者\n廖\n更\n纳\n加\n奉\n公\n一\n就\n汴\n计\n与\n路\n房\n原\n妇\n2\n0\n8\n-\n7\n其\n>\n:\n]\n,\n，\n骑\n刈\n全\n消\n昏\n傈\n安\n久\n钟\n嗅\n不\n影\n处\n驽\n蜿\n资\n关\n椤\n地\n瘸\n专\n问\n忖\n票\n嫉\n"
},
{
"path": "iopaint/model/anytext/utils.py",
"chars": 4989,
"preview": "import os\nimport datetime\nimport cv2\nimport numpy as np\nfrom PIL import Image, ImageDraw\n\n\ndef save_images(img_list, fol"
},
{
"path": "iopaint/model/base.py",
"chars": 13377,
"preview": "import abc\nfrom typing import Optional\n\nimport cv2\nimport torch\nimport numpy as np\nfrom loguru import logger\n\nfrom iopai"
},
{
"path": "iopaint/model/brushnet/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/brushnet/brushnet.py",
"chars": 47487,
"preview": "from dataclasses import dataclass\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom torch im"
},
{
"path": "iopaint/model/brushnet/brushnet_unet_forward.py",
"chars": 16782,
"preview": "from typing import Union, Optional, Dict, Any, Tuple\n\nimport torch\nfrom diffusers.models.unets.unet_2d_condition import "
},
{
"path": "iopaint/model/brushnet/brushnet_wrapper.py",
"chars": 6587,
"preview": "import PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\nimport numpy as np\n\nfrom ..base import DiffusionInpai"
},
{
"path": "iopaint/model/brushnet/brushnet_xl_wrapper.py",
"chars": 6781,
"preview": "import PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\nimport numpy as np\n\nfrom ..base import DiffusionInpai"
},
{
"path": "iopaint/model/brushnet/pipeline_brushnet.py",
"chars": 65579,
"preview": "# https://github.com/TencentARC/BrushNet\nimport inspect\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\ni"
},
{
"path": "iopaint/model/brushnet/pipeline_brushnet_sd_xl.py",
"chars": 79022,
"preview": "# https://github.com/TencentARC/BrushNet\nimport inspect\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, U"
},
{
"path": "iopaint/model/brushnet/unet_2d_blocks.py",
"chars": 14794,
"preview": "from typing import Dict, Any, Optional, Tuple\n\nimport torch\nfrom diffusers.models.resnet import ResnetBlock2D\nfrom diffu"
},
{
"path": "iopaint/model/controlnet.py",
"chars": 7011,
"preview": "import PIL.Image\nimport cv2\nimport torch\nfrom diffusers import ControlNetModel\nfrom loguru import logger\nfrom iopaint.sc"
},
{
"path": "iopaint/model/ddim_sampler.py",
"chars": 6855,
"preview": "import torch\nimport numpy as np\nfrom tqdm import tqdm\n\nfrom .utils import make_ddim_timesteps, make_ddim_sampling_parame"
},
{
"path": "iopaint/model/fcf.py",
"chars": 57219,
"preview": "import os\nimport random\n\nimport cv2\nimport torch\nimport numpy as np\nimport torch.fft as fft\n\nfrom iopaint.schema import "
},
{
"path": "iopaint/model/helper/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/helper/controlnet_preprocess.py",
"chars": 2117,
"preview": "import torch\nimport PIL\nimport cv2\nfrom PIL import Image\nimport numpy as np\n\nfrom iopaint.helper import pad_img_to_modul"
},
{
"path": "iopaint/model/helper/cpu_text_encoder.py",
"chars": 1409,
"preview": "import torch\nfrom transformers import PreTrainedModel\n\nfrom ..utils import torch_gc\n\n\nclass CPUTextEncoderWrapper(PreTra"
},
{
"path": "iopaint/model/helper/g_diffuser_bot.py",
"chars": 1913,
"preview": "import cv2\nimport numpy as np\n\n\ndef expand_image(cv2_img, top: int, right: int, bottom: int, left: int):\n assert cv2_"
},
{
"path": "iopaint/model/instruct_pix2pix.py",
"chars": 2436,
"preview": "import PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\n\nfrom iopaint.const import INSTRUCT_PIX2PIX_NAME\nfrom"
},
{
"path": "iopaint/model/kandinsky.py",
"chars": 2037,
"preview": "import PIL.Image\nimport cv2\nimport numpy as np\nimport torch\n\nfrom iopaint.const import KANDINSKY22_NAME\nfrom .base impor"
},
{
"path": "iopaint/model/lama.py",
"chars": 2330,
"preview": "import os\n\nimport cv2\nimport numpy as np\nimport torch\n\nfrom iopaint.helper import (\n norm_img,\n get_cache_path_by_"
},
{
"path": "iopaint/model/ldm.py",
"chars": 11481,
"preview": "import os\n\nimport numpy as np\nimport torch\nfrom loguru import logger\n\nfrom .base import InpaintModel\nfrom .ddim_sampler "
},
{
"path": "iopaint/model/manga.py",
"chars": 3097,
"preview": "import os\nimport random\n\nimport cv2\nimport numpy as np\nimport torch\nimport time\nfrom loguru import logger\n\nfrom iopaint."
},
{
"path": "iopaint/model/mat.py",
"chars": 62682,
"preview": "import os\nimport random\n\nimport cv2\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as "
},
{
"path": "iopaint/model/mi_gan.py",
"chars": 3438,
"preview": "import os\n\nimport cv2\nimport torch\n\nfrom iopaint.helper import (\n load_jit_model,\n download_model,\n get_cache_p"
},
{
"path": "iopaint/model/opencv2.py",
"chars": 735,
"preview": "import cv2\nfrom .base import InpaintModel\nfrom iopaint.schema import InpaintRequest\n\nflag_map = {\"INPAINT_NS\": cv2.INPAI"
},
{
"path": "iopaint/model/original_sd_configs/__init__.py",
"chars": 590,
"preview": "from pathlib import Path\nfrom typing import Dict\n\nCURRENT_DIR = Path(__file__).parent.absolute()\n\n\ndef get_config_files("
},
{
"path": "iopaint/model/original_sd_configs/sd_xl_base.yaml",
"chars": 2784,
"preview": "model:\n target: sgm.models.diffusion.DiffusionEngine\n params:\n scale_factor: 0.13025\n disable_first_stage_autoca"
},
{
"path": "iopaint/model/original_sd_configs/sd_xl_refiner.yaml",
"chars": 2589,
"preview": "model:\n target: sgm.models.diffusion.DiffusionEngine\n params:\n scale_factor: 0.13025\n disable_first_stage_autoca"
},
{
"path": "iopaint/model/original_sd_configs/v1-inference.yaml",
"chars": 1873,
"preview": "model:\n base_learning_rate: 1.0e-04\n target: ldm.models.diffusion.ddpm.LatentDiffusion\n params:\n linear_start: 0.0"
},
{
"path": "iopaint/model/original_sd_configs/v2-inference-v.yaml",
"chars": 1815,
"preview": "model:\n base_learning_rate: 1.0e-4\n target: ldm.models.diffusion.ddpm.LatentDiffusion\n params:\n parameterization: "
},
{
"path": "iopaint/model/paint_by_example.py",
"chars": 2795,
"preview": "import PIL\nimport PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\n\nfrom iopaint.helper import decode_base64_"
},
{
"path": "iopaint/model/plms_sampler.py",
"chars": 11833,
"preview": "# From: https://github.com/CompVis/latent-diffusion/blob/main/ldm/models/diffusion/plms.py\nimport torch\nimport numpy as "
},
{
"path": "iopaint/model/power_paint/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/power_paint/pipeline_powerpaint.py",
"chars": 56634,
"preview": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "iopaint/model/power_paint/power_paint.py",
"chars": 3658,
"preview": "from PIL import Image\nimport PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\n\nfrom ..base import DiffusionIn"
},
{
"path": "iopaint/model/power_paint/power_paint_v2.py",
"chars": 6964,
"preview": "from itertools import chain\n\nimport PIL.Image\nimport cv2\nimport torch\nfrom iopaint.model.original_sd_configs import get_"
},
{
"path": "iopaint/model/power_paint/powerpaint_tokenizer.py",
"chars": 9439,
"preview": "import copy\nimport random\nfrom typing import Any, List, Union\nfrom transformers import CLIPTokenizer\n\nfrom iopaint.schem"
},
{
"path": "iopaint/model/power_paint/v2/BrushNet_CA.py",
"chars": 49516,
"preview": "from dataclasses import dataclass\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom diffuser"
},
{
"path": "iopaint/model/power_paint/v2/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/model/power_paint/v2/pipeline_PowerPaint_Brushnet_CA.py",
"chars": 77816,
"preview": "import inspect\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nimport numpy as np\nimport PIL.Image\nimport"
},
{
"path": "iopaint/model/power_paint/v2/unet_2d_blocks.py",
"chars": 12496,
"preview": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "iopaint/model/power_paint/v2/unet_2d_condition.py",
"chars": 17567,
"preview": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Lic"
},
{
"path": "iopaint/model/sd.py",
"chars": 4447,
"preview": "import PIL.Image\nimport cv2\nimport torch\nfrom loguru import logger\n\nfrom .base import DiffusionInpaintModel\nfrom .helper"
},
{
"path": "iopaint/model/sdxl.py",
"chars": 3993,
"preview": "import os\n\nimport PIL.Image\nimport cv2\nimport torch\nfrom diffusers import AutoencoderKL\nfrom loguru import logger\n\nfrom "
},
{
"path": "iopaint/model/utils.py",
"chars": 37847,
"preview": "import gc\nimport math\nimport random\nimport traceback\nfrom typing import Any\n\nimport torch\nimport numpy as np\nimport coll"
},
{
"path": "iopaint/model/zits.py",
"chars": 16303,
"preview": "import os\nimport time\n\nimport cv2\nimport torch\nimport torch.nn.functional as F\n\nfrom iopaint.helper import get_cache_pat"
},
{
"path": "iopaint/model_manager.py",
"chars": 10320,
"preview": "from typing import List, Dict\n\nimport torch\nfrom loguru import logger\nimport numpy as np\n\nfrom iopaint.download import s"
},
{
"path": "iopaint/plugins/__init__.py",
"chars": 2471,
"preview": "from typing import Dict\n\nfrom loguru import logger\n\nfrom .anime_seg import AnimeSeg\nfrom .gfpgan_plugin import GFPGANPlu"
},
{
"path": "iopaint/plugins/anime_seg.py",
"chars": 13779,
"preview": "import cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nfrom PIL import Image\n\n"
},
{
"path": "iopaint/plugins/base_plugin.py",
"chars": 730,
"preview": "from loguru import logger\nimport numpy as np\n\nfrom iopaint.schema import RunPluginRequest\n\n\nclass BasePlugin:\n name: "
},
{
"path": "iopaint/plugins/basicsr/LICENSE",
"chars": 11349,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "iopaint/plugins/basicsr/__init__.py",
"chars": 926,
"preview": "\"\"\"\nAdapted from https://github.com/XPixelGroup/BasicSR\nLicense: Apache-2.0\n\nAs of Feb 2024, `basicsr` appears to be unm"
},
{
"path": "iopaint/plugins/basicsr/arch_util.py",
"chars": 2531,
"preview": "from typing import Type, List, Union\n\nimport torch\nfrom torch import nn as nn\nfrom torch.nn import init as init\nfrom tor"
},
{
"path": "iopaint/plugins/basicsr/img_util.py",
"chars": 6195,
"preview": "import cv2\nimport math\nimport numpy as np\nimport os\nimport torch\nfrom torchvision.utils import make_grid\n\n\ndef img2tenso"
},
{
"path": "iopaint/plugins/basicsr/rrdbnet_arch.py",
"chars": 4875,
"preview": "import torch\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\nfrom .arch_util import default_init_weight"
},
{
"path": "iopaint/plugins/briarmbg.py",
"chars": 15021,
"preview": "# copy from: https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4/blob/main/briarmbg.py\nimport cv2\nimport torch\nimport tor"
},
{
"path": "iopaint/plugins/briarmbg2.py",
"chars": 100066,
"preview": "# copy from https://huggingface.co/briaai/RMBG-2.0/tree/main\nimport os\nimport math\nimport numpy as np\n\nimport torch\nimpo"
},
{
"path": "iopaint/plugins/facexlib/.gitignore",
"chars": 1837,
"preview": ".vscode\n*.pth\n*.png\n*.jpg\nversion.py\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C ext"
},
{
"path": "iopaint/plugins/facexlib/__init__.py",
"chars": 61,
"preview": "# flake8: noqa\nfrom .detection import *\nfrom .utils import *\n"
},
{
"path": "iopaint/plugins/facexlib/detection/__init__.py",
"chars": 1313,
"preview": "import torch\nfrom copy import deepcopy\n\nfrom ..utils import load_file_from_url\nfrom .retinaface import RetinaFace\n\n\ndef "
},
{
"path": "iopaint/plugins/facexlib/detection/align_trans.py",
"chars": 7941,
"preview": "import cv2\nimport numpy as np\n\nfrom .matlab_cp2tform import get_similarity_transform_for_cv2\n\n# reference facial points,"
},
{
"path": "iopaint/plugins/facexlib/detection/matlab_cp2tform.py",
"chars": 8109,
"preview": "import numpy as np\nfrom numpy.linalg import inv, lstsq\nfrom numpy.linalg import matrix_rank as rank\nfrom numpy.linalg im"
},
{
"path": "iopaint/plugins/facexlib/detection/retinaface.py",
"chars": 13928,
"preview": "import cv2\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom PIL import Image\nf"
},
{
"path": "iopaint/plugins/facexlib/detection/retinaface_net.py",
"chars": 6281,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef conv_bn(inp, oup, stride=1, leaky=0):\n retur"
},
{
"path": "iopaint/plugins/facexlib/detection/retinaface_utils.py",
"chars": 16362,
"preview": "import numpy as np\nimport torch\nimport torchvision\nfrom itertools import product as product\nfrom math import ceil\n\n\nclas"
},
{
"path": "iopaint/plugins/facexlib/parsing/__init__.py",
"chars": 1009,
"preview": "import torch\n\nfrom ..utils import load_file_from_url\nfrom .bisenet import BiSeNet\nfrom .parsenet import ParseNet\n\n\ndef i"
},
{
"path": "iopaint/plugins/facexlib/parsing/bisenet.py",
"chars": 5190,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .resnet import ResNet18\n\n\nclass ConvBNReLU(nn.M"
},
{
"path": "iopaint/plugins/facexlib/parsing/parsenet.py",
"chars": 6477,
"preview": "\"\"\"Modified from https://github.com/chaofengc/PSFRGAN\n\"\"\"\nimport numpy as np\nimport torch.nn as nn\nfrom torch.nn import "
},
{
"path": "iopaint/plugins/facexlib/parsing/resnet.py",
"chars": 2357,
"preview": "import torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolu"
},
{
"path": "iopaint/plugins/facexlib/utils/__init__.py",
"chars": 330,
"preview": "from .face_utils import align_crop_face_landmarks, compute_increased_bbox, get_valid_bboxes, paste_face_back\nfrom .misc "
},
{
"path": "iopaint/plugins/facexlib/utils/face_restoration_helper.py",
"chars": 18774,
"preview": "import cv2\nimport numpy as np\nimport os\nimport torch\nfrom torchvision.transforms.functional import normalize\n\nfrom ..det"
},
{
"path": "iopaint/plugins/facexlib/utils/face_utils.py",
"chars": 8725,
"preview": "import cv2\nimport numpy as np\nimport torch\n\n\ndef compute_increased_bbox(bbox, increase_area, preserve_aspect=True):\n "
},
{
"path": "iopaint/plugins/facexlib/utils/misc.py",
"chars": 4081,
"preview": "import cv2\nimport os\nimport os.path as osp\nimport torch\nfrom torch.hub import download_url_to_file, get_dir\nfrom urllib."
},
{
"path": "iopaint/plugins/gfpgan/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/plugins/gfpgan/archs/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/plugins/gfpgan/archs/gfpganv1_clean_arch.py",
"chars": 13560,
"preview": "import math\nimport random\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .stylegan2_clean_"
},
{
"path": "iopaint/plugins/gfpgan/archs/restoreformer_arch.py",
"chars": 23875,
"preview": "\"\"\"Modified from https://github.com/wzhouxiff/RestoreFormer\"\"\"\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nim"
},
{
"path": "iopaint/plugins/gfpgan/archs/stylegan2_clean_arch.py",
"chars": 14902,
"preview": "import math\nimport random\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom iopaint.plugins.b"
},
{
"path": "iopaint/plugins/gfpgan_plugin.py",
"chars": 2147,
"preview": "import cv2\nimport numpy as np\nfrom loguru import logger\n\nfrom iopaint.helper import download_model\nfrom iopaint.plugins."
},
{
"path": "iopaint/plugins/gfpganer.py",
"chars": 5550,
"preview": "import os\n\nimport cv2\nimport torch\nfrom torchvision.transforms.functional import normalize\nfrom torch.hub import get_dir"
},
{
"path": "iopaint/plugins/interactive_seg.py",
"chars": 5483,
"preview": "import hashlib\nfrom typing import List\n\nimport numpy as np\nimport torch\nfrom loguru import logger\n\nfrom iopaint.helper i"
},
{
"path": "iopaint/plugins/realesrgan.py",
"chars": 17801,
"preview": "import math\n\nimport cv2\nimport numpy as np\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom loguru"
},
{
"path": "iopaint/plugins/remove_bg.py",
"chars": 3530,
"preview": "import os\nimport cv2\nimport numpy as np\nfrom loguru import logger\nimport torch\nfrom torch.hub import get_dir\n\nfrom iopai"
},
{
"path": "iopaint/plugins/restoreformer.py",
"chars": 1494,
"preview": "import cv2\nimport numpy as np\nfrom loguru import logger\n\nfrom iopaint.helper import download_model\nfrom iopaint.plugins."
},
{
"path": "iopaint/plugins/segment_anything/__init__.py",
"chars": 420,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/build_sam.py",
"chars": 7864,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/__init__.py",
"chars": 385,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/common.py",
"chars": 1479,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/image_encoder.py",
"chars": 14407,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/image_encoder_hq.py",
"chars": 14795,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/mask_decoder.py",
"chars": 15490,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/prompt_encoder.py",
"chars": 8594,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/sam.py",
"chars": 7225,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/sam_hq.py",
"chars": 7474,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/modeling/tiny_vit_sam.py",
"chars": 26182,
"preview": "# --------------------------------------------------------\n# TinyViT Model Architecture\n# Copyright (c) 2022 Microsoft\n#"
},
{
"path": "iopaint/plugins/segment_anything/modeling/transformer.py",
"chars": 8396,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/predictor.py",
"chars": 11845,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/predictor_hq.py",
"chars": 12157,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/utils/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything/utils/transforms.py",
"chars": 4054,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/build_sam.py",
"chars": 9648,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/backbones/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/backbones/hieradet.py",
"chars": 9873,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/backbones/image_encoder.py",
"chars": 4705,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/backbones/utils.py",
"chars": 3053,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/memory_attention.py",
"chars": 5483,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/memory_encoder.py",
"chars": 5644,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/position_encoding.py",
"chars": 8361,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam/mask_decoder.py",
"chars": 12645,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam/prompt_encoder.py",
"chars": 6992,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam/transformer.py",
"chars": 11446,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam2_base.py",
"chars": 46887,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/modeling/sam2_utils.py",
"chars": 13170,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/sam2_image_predictor.py",
"chars": 19208,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/utils/__init__.py",
"chars": 197,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/utils/misc.py",
"chars": 3635,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/plugins/segment_anything2/utils/transforms.py",
"chars": 2775,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the li"
},
{
"path": "iopaint/runtime.py",
"chars": 2496,
"preview": "# https://github.com/huggingface/huggingface_hub/blob/5a12851f54bf614be39614034ed3a9031922d297/src/huggingface_hub/utils"
},
{
"path": "iopaint/schema.py",
"chars": 16019,
"preview": "import random\nfrom enum import Enum\nfrom pathlib import Path\nfrom typing import Optional, Literal, List\n\nfrom loguru imp"
},
{
"path": "iopaint/tests/.gitignore",
"chars": 20,
"preview": "*_result.png\nresult/"
},
{
"path": "iopaint/tests/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "iopaint/tests/test_adjust_mask.py",
"chars": 716,
"preview": "import cv2\nfrom iopaint.helper import adjust_mask\nfrom iopaint.tests.utils import current_dir, save_dir\n\nmask_p = curren"
},
{
"path": "iopaint/tests/test_anytext.py",
"chars": 1461,
"preview": "import os\n\nfrom iopaint.tests.utils import check_device, get_config, assert_equal\n\nos.environ[\"PYTORCH_ENABLE_MPS_FALLBA"
},
{
"path": "iopaint/tests/test_brushnet.py",
"chars": 3632,
"preview": "import os\n\nfrom iopaint.const import SD_BRUSHNET_CHOICES\nfrom iopaint.tests.utils import check_device, get_config, asser"
},
{
"path": "iopaint/tests/test_controlnet.py",
"chars": 3388,
"preview": "import os\n\nfrom iopaint.const import SD_CONTROLNET_CHOICES\nfrom iopaint.tests.utils import current_dir, check_device, ge"
},
{
"path": "iopaint/tests/test_instruct_pix2pix.py",
"chars": 1194,
"preview": "from pathlib import Path\n\nimport pytest\nimport torch\n\nfrom iopaint.model_manager import ModelManager\nfrom iopaint.schema"
},
{
"path": "iopaint/tests/test_load_img.py",
"chars": 552,
"preview": "from iopaint.helper import load_img\nfrom iopaint.tests.utils import current_dir\n\npng_img_p = current_dir / \"image.png\"\nj"
}
]
// ... and 102 more files (download for full content)
About this extraction
This page contains the full source code of the Sanster/IOPaint GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 302 files (2.2 MB), approximately 601.3k tokens, and a symbol index with 2187 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.