nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add mv-adapter ig2mv

Browse Source
This commit is contained in:
Denis Kochetov
2025-05-30 07:38:39 +00:00
parent ea90580926
commit 76179ca7fe
75 changed files with 18743 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Gen_3D_Modules/MV_Adapter/.gitignore
+23
View File
@@ -0,0 +1,23 @@
*.png
*.jpg
*.jpeg
*.gif
*.bmp
*.tiff
*.ico
*.webp
#glb
*.glb
*.gltf
*.obj
*.mtl
*.ply
*.stl
*.dae
#pycache
__pycache__
#checkpoints
checkpoints/
Gen_3D_Modules/MV_Adapter/LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Gen_3D_Modules/MV_Adapter/README.md
+471
View File
@@ -0,0 +1,471 @@
# MV-Adapter: Multi-view Consistent Image Generation Made Easy🚀
## 🏠 <a href="https://huanngzh.github.io/MV-Adapter-Page/" target="_blank">Project Page</a> | <a href="https://arxiv.org/abs/2412.03632" target="_blank">Paper</a> | <a href="https://huggingface.co/huanngzh/mv-adapter" target="_blank">Model Weights</a> | [Demo](https://huggingface.co/collections/huanngzh/mv-adapter-spaces-677e497578747fd734a1b999) | <a href="https://github.com/huanngzh/ComfyUI-MVAdapter" target="_blank">ComfyUI</a>
![teaser](assets/doc/teaser.jpg)
MV-Adapter is a **versatile plug-and-play adapter** that adapt T2I models and their derivatives to multi-view generators.
Highlight Features: Generate multi-view images
- with 768 Resolution using SDXL
- using personalized models (e.g. <a href="https://civitai.com/models/112902/dreamshaper-xl" target="_blank">DreamShaper</a>), distilled models (e.g. <a href="https://huggingface.co/docs/diffusers/api/pipelines/latent_consistency_models" target="_blank">LCM</a>), or extensions (e.g. <a href="https://github.com/lllyasviel/ControlNet" target="_blank">ControlNet</a>)
- from text or image condition
- can be guided by geometry for texture generation
## 🔥 Updates
* [2025-05-15] Release full pipeline for text-to-texture and image-to-texture generation. [See [guidelines](#usage-texture-generation)]
* [2025-04-23] Release dataset ([Objaverse-Ortho10View](https://huggingface.co/datasets/huanngzh/Objaverse-Ortho10View) and [Objaverse-Rand6View](https://huggingface.co/datasets/huanngzh/Objaverse-Rand6View)) and training code. [See [guidelines](#️-training)]
* [2025-03-31] Release text/image-conditioned 3D texture generation demos on [Text2Texture](https://huggingface.co/spaces/VAST-AI/MV-Adapter-Text2Texture) and [Image2Texture](https://huggingface.co/spaces/VAST-AI/MV-Adapter-Img2Texture). Feel free to try them!
* [2025-03-17] Release model weights for partial-image conditioned geometry-to-multiview generation, which can be used to generate textured 3D scenes combined with [MIDI](https://github.com/VAST-AI-Research/MIDI-3D). [See [guidelines](#partial-image--geometry-to-multiview)]
* [2025-03-07] Release model weights for geometry-guided multi-view generation. [See [guidelines](#text-geometry-to-multiview-generation)]
* [2024-12-27] Release model weights, gradio demo, inference scripts and comfyui of text-/image- to multi-view generation models.
## TOC
- [MV-Adapter: Multi-view Consistent Image Generation Made Easy🚀](#mv-adapter-multi-view-consistent-image-generation-made-easy)
- [🏠 Project Page | Paper | Model Weights | Demo | ComfyUI](#-project-page--paper--model-weights--demo--comfyui)
- [🔥 Updates](#-updates)
- [TOC](#toc)
- [Model Zoo \& Demos](#model-zoo--demos)
- [Installation](#installation)
- [Notes](#notes)
- [System Requirements](#system-requirements)
- [Usage: Multiview Generation](#usage-multiview-generation)
- [Launch Demo](#launch-demo)
- [Text to Multiview Generation](#text-to-multiview-generation)
- [Image to Multiview Generation](#image-to-multiview-generation)
- [Inference Scripts](#inference-scripts)
- [Text to Multiview Generation](#text-to-multiview-generation-1)
- [Image to Multiview Generation](#image-to-multiview-generation-1)
- [Text-Geometry to Multiview Generation](#text-geometry-to-multiview-generation)
- [Image-Geometry to Multiview Generation](#image-geometry-to-multiview-generation)
- [Partial Image + Geometry to Multiview](#partial-image--geometry-to-multiview)
- [ComfyUI](#comfyui)
- [Usage: Texture Generation](#usage-texture-generation)
- [Citation](#citation)
## Model Zoo & Demos
No need to download manually. Running the scripts will download model weights automatically.
Notes: Running MV-Adapter for SDXL may need higher GPU memory and more time, but produce higher-quality and higher-resolution results. On the other hand, running its SD2.1 variant needs lower computing cost, but shows a bit lower performance.
| Model | Base Model | HF Weights | Demo Link |
| :-------------------------: | :--------: | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | :-------------------------------------------------------------------------------------------------------------------------------------------: |
| Text-to-Multiview | SD2.1 | [mvadapter_t2mv_sd21.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_t2mv_sd21.safetensors) | |
| Text-to-Multiview | SDXL | [mvadapter_t2mv_sdxl.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_t2mv_sdxl.safetensors) | [General](https://huggingface.co/spaces/VAST-AI/MV-Adapter-T2MV-SDXL) / [Anime](https://huggingface.co/spaces/huanngzh/MV-Adapter-T2MV-Anime) |
| Image-to-Multiview | SD2.1 | [mvadapter_i2mv_sd21.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_i2mv_sd21.safetensors) | |
| Image-to-Multiview | SDXL | [mvadapter_i2mv_sdxl.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_i2mv_sdxl.safetensors) | [Demo](https://huggingface.co/spaces/VAST-AI/MV-Adapter-I2MV-SDXL) |
| Text-Geometry-to-Multiview | SDXL | [mvadapter_tg2mv_sdxl.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_tg2mv_sdxl.safetensors) | [Demo](https://huggingface.co/spaces/VAST-AI/MV-Adapter-Text2Texture) |
| Image-Geometry-to-Multiview | SDXL | [mvadapter_ig2mv_sdxl.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_ig2mv_sdxl.safetensors) / [mvadapter_ig2mv_partial_sdxl.safetensors](https://huggingface.co/huanngzh/mv-adapter/resolve/main/mvadapter_ig2mv_partial_sdxl.safetensors) | [Demo](https://huggingface.co/spaces/VAST-AI/MV-Adapter-Img2Texture) |
| Image-to-Arbitrary-Views | SDXL | | |
## Installation
Clone the repo first:
```Bash
git clone https://github.com/huanngzh/MV-Adapter.git
cd MV-Adapter
```
(Optional) Create a fresh conda env:
```Bash
conda create -n mvadapter python=3.10
conda activate mvadapter
```
Install necessary packages (torch > 2):
```Bash
# pytorch (select correct CUDA version)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
# other dependencies
pip install -r requirements.txt
```
For texture generation, you need to install `CV-CUDA` according to [CVCUDA/CV-CUDA](https://github.com/CVCUDA/CV-CUDA?tab=readme-ov-file#installation).
## Notes
### System Requirements
In the model zoo of MV-Adapter, running image-to-multiview generation has the highest system requirements, which requires about 14G GPU memory.
## Usage: Multiview Generation
### Launch Demo
#### Text to Multiview Generation
**With SDXL:**
```Bash
python -m scripts.gradio_demo_t2mv --base_model "stabilityai/stable-diffusion-xl-base-1.0"
```
![demo_t2mv_1](assets/doc/demo_t2mv_1.png)
![demo_t2mv_2](assets/doc/demo_t2mv_2.png)
> Reminder: When switching the demo to another base model, delete the `gradio_cached_examples` directory, otherwise it will affect the examples results of the next demo.
**With anime-themed <a href="https://huggingface.co/cagliostrolab/animagine-xl-3.1" target="_blank">Animagine XL 3.1</a>:**
```Bash
python -m scripts.gradio_demo_t2mv --base_model "cagliostrolab/animagine-xl-3.1"
```
![demo_t2mv_anime_1](assets/doc/demo_t2mv_anime_1.png)
![demo_t2mv_anime_2](assets/doc/demo_t2mv_anime_2.png)
**With general <a href="https://huggingface.co/Lykon/dreamshaper-xl-1-0" target="_blank">Dreamshaper</a>:**
```Bash
python -m scripts.gradio_demo_t2mv --base_model "Lykon/dreamshaper-xl-1-0" --scheduler ddpm
```
![demo_t2mv_dreamshaper_1](assets/doc/demo_t2mv_dreamshaper_1.png)
You can also specify a new diffusers-format text-to-image diffusion model using `--base_model`. Note that it should be the model name in huggingface, such as `stabilityai/stable-diffusion-xl-base-1.0`, or a local path refer to a text-to-image pipeline directory. Note that if you specify `latent-consistency/lcm-sdxl` to use latent consistency models, please add `--scheduler lcm` to the command.
#### Image to Multiview Generation
**With SDXL:**
```Bash
python -m scripts.gradio_demo_i2mv
```
![demo_i2mv_1](assets/doc/demo_i2mv_1.png)
![demo_i2mv_2](assets/doc/demo_i2mv_2.png)
### Inference Scripts
We recommend that experienced users check the files in the scripts directory to adjust the parameters appropriately to try the best "card drawing" results.
#### Text to Multiview Generation
Note that you can specify a diffusers-format text-to-image diffusion model as the base model using `--base_model xxx`. It should be the model name in huggingface, such as `stabilityai/stable-diffusion-xl-base-1.0`, or a local path refer to a text-to-image pipeline directory.
**With SDXL:**
```Bash
python -m scripts.inference_t2mv_sdxl --text "an astronaut riding a horse" \
--seed 42 \
--output output.png
```
**With personalized models:**
anime-themed <a href="https://huggingface.co/cagliostrolab/animagine-xl-3.1" target="_blank">Animagine XL 3.1</a>
```Bash
python -m scripts.inference_t2mv_sdxl --base_model "cagliostrolab/animagine-xl-3.1" \
--text "1girl, izayoi sakuya, touhou, solo, maid headdress, maid, apron, short sleeves, dress, closed mouth, white apron, serious face, upper body, masterpiece, best quality, very aesthetic, absurdres" \
--seed 0 \
--output output.png
```
general <a href="https://huggingface.co/Lykon/dreamshaper-xl-1-0" target="_blank">Dreamshaper</a>
```Bash
python -m scripts.inference_t2mv_sdxl --base_model "Lykon/dreamshaper-xl-1-0" \
--scheduler ddpm \
--text "the warrior Aragorn from Lord of the Rings, film grain, 8k hd" \
--seed 0 \
--output output.png
```
realistic <a href="https://huggingface.co/stablediffusionapi/real-dream-sdxl" target="_blank">real-dream-sdxl</a>
```Bash
python -m scripts.inference_t2mv_sdxl --base_model "stablediffusionapi/real-dream-sdxl" \
--scheduler ddpm \
--text "macro shot, parrot, colorful, dark shot, film grain, extremely detailed" \
--seed 42 \
--output output.png
```
**With <a href="https://huggingface.co/latent-consistency/lcm-sdxl" target="_blank">LCM</a>:**
```Bash
python -m scripts.inference_t2mv_sdxl --unet_model "latent-consistency/lcm-sdxl" \
--scheduler lcm \
--text "Samurai koala bear" \
--num_inference_steps 8 \
--seed 42 \
--output output.png
```
**With LoRA:**
stylized lora <a href="https://huggingface.co/goofyai/3d_render_style_xl" target="_blank">3d_render_style_xl</a>
```Bash
python -m scripts.inference_t2mv_sdxl --lora_model "goofyai/3d_render_style_xl/3d_render_style_xl.safetensors" \
--text "3d style, a fox with flowers around it" \
--seed 20 \
--lora_scale 1.0 \
--output output.png
```
**With ControlNet:**
Scribble to Multiview with <a href="https://huggingface.co/xinsir/controlnet-scribble-sdxl-1.0" target="_blank">controlnet-scribble-sdxl-1.0</a>
```Bash
python -m scripts.inference_scribble2mv_sdxl --text "A 3D model of Finn the Human from the animated television series Adventure Time. He is wearing his iconic blue shirt and green backpack and has a neutral expression on his face. He is standing in a relaxed pose with his left foot slightly forward and his right foot back. His arms are at his sides and his head is turned slightly to the right. The model is made up of simple shapes and has a stylized, cartoon-like appearance. It is textured to resemble the character's appearance in the show." \
--seed 0 \
--output output.png \
--guidance_scale 5.0 \
--controlnet_images "assets/demo/scribble2mv/color_0000.webp" "assets/demo/scribble2mv/color_0001.webp" "assets/demo/scribble2mv/color_0002.webp" "assets/demo/scribble2mv/color_0003.webp" "assets/demo/scribble2mv/color_0004.webp" "assets/demo/scribble2mv/color_0005.webp" \
--controlnet_conditioning_scale 0.7
```
**With SD2.1:**
> SD2.1 has lower demand for computing resources and higher inference speed, but a bit lower performance than SDXL.
> In our tests, ddpm scheduler works better than other schedulers here.
```Bash
python -m scripts.inference_t2mv_sd --text "a corgi puppy" \
--seed 42 --scheduler ddpm \
--output output.png
```
#### Image to Multiview Generation
**With SDXL:**
```Bash
python -m scripts.inference_i2mv_sdxl \
--image assets/demo/i2mv/A_decorative_figurine_of_a_young_anime-style_girl.png \
--text "A decorative figurine of a young anime-style girl" \
--seed 21 --output output.png --remove_bg
```
**With LCM:**
```Bash
python -m scripts.inference_i2mv_sdxl \
--unet_model "latent-consistency/lcm-sdxl" \
--scheduler lcm \
--image assets/demo/i2mv/A_juvenile_emperor_penguin_chick.png \
--text "A juvenile emperor penguin chick" \
--num_inference_steps 8 \
--seed 0 --output output.png --remove_bg
```
**With SD2.1:** (lower demand for computing resources and higher inference speed)
> In our tests, ddpm scheduler works better than other schedulers here.
```Bash
python -m scripts.inference_i2mv_sd \
--image assets/demo/i2mv/A_decorative_figurine_of_a_young_anime-style_girl.png \
--text "A decorative figurine of a young anime-style girl" \
--output output.png --remove_bg --scheduler ddpm
```
#### Text-Geometry to Multiview Generation
**Importantly**, when using geometry-condition generation, please make sure that the orientation of the mesh you provide is consistent with the following example. Otherwise, you need to adjust the angles in the scripts when rendering the view.
**With SDXL:**
```Bash
python -m scripts.inference_tg2mv_sdxl \
--mesh assets/demo/tg2mv/ac9d4e4f44f34775ad46878ba8fbfd86.glb \
--text "Mater, a rusty and beat-up tow truck from the 2006 Disney/Pixar animated film 'Cars', with a rusty brown exterior, big blue eyes."
```
![tg2mv_example_out](assets/demo/tg2mv/ac9d4e4f44f34775ad46878ba8fbfd86_mv.png)
```Bash
python -m scripts.inference_tg2mv_sdxl \
--mesh assets/demo/tg2mv/b5f0f0f33e3644d1ba73576ceb486d42.glb \
--text "Optimus Prime, a character from Transformers, with blue, red and gray colors, and has a flame-like pattern on the body"
```
#### Image-Geometry to Multiview Generation
**With SDXL:**
```Bash
python -m scripts.inference_ig2mv_sdxl \
--image assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.jpeg \
--mesh assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.glb \
--output output.png --remove_bg
```
![example_ig2mv_out](assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c_mv.png)
#### Partial Image + Geometry to Multiview
**With SDXL:**
```Bash
python -m scripts.inference_ig2mv_partial_sdxl \
--image assets/demo/ig2mv/cartoon_style_table.png \
--mesh assets/demo/ig2mv/cartoon_style_table.glb \
--output output.png
```
Example input:
<img src="assets/demo/ig2mv/cartoon_style_table.png" alt="partial input" style="width: 20%">
Example output:
![example_partial_ig2mv](assets/demo/ig2mv/cartoon_style_table_mv.png)
The above command will save a `*_transform.json` file in the output dir, which contains transformation information like this:
```json
{
"offset": [
0.7446051140100826,
-0.3421213991056582,
0.1360104325533671
],
"scale": 1.0086087120792668
}
```
You can use it to transform your mesh into the canonical space, map the generated multi-view images onto the mesh, and then re-transform the mesh back to the original spatial position.
### ComfyUI
Please check <a href="https://github.com/huanngzh/ComfyUI-MVAdapter" target="_blank">ComfyUI-MVAdapter Repo</a> for details.
**Text to Multiview Generation**
![comfyui_t2mv](assets/doc/comfyui_t2mv.png)
**Text to Multiview Generation with LoRA**
![comfyui_t2mv_lora](assets/doc/comfyui_t2mv_lora.png)
**Image to Multiview Generation**
![comfyui_i2mv](assets/doc/comfyui_i2mv.png)
## Usage: Texture Generation
**Prepare Models**
Download pre-trained [RealESRGAN](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth) for upscaling images and [LaMa](https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt) for view in-painting.
```Bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth -O ./checkpoints/RealESRGAN_x2plus.pth
wget https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt -O ./checkpoints/big-lama.pt
```
**Usage**
> All in one script
Text-conditioned texture generation:
```Bash
python -m scripts.texture_t2tex \
--mesh assets/demo/tg2mv/ac9d4e4f44f34775ad46878ba8fbfd86.glb \
--text "Mater, a rusty and beat-up tow truck from the 2006 Disney/Pixar animated film 'Cars', with a rusty brown exterior, big blue eyes." \
--save_dir outputs --save_name t2tex_sample
```
Image-conditioned texture generation:
```Bash
python -m scripts.texture_i2tex \
--image assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.jpeg \
--mesh assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.glb \
--save_dir outputs --save_name i2tex_sample \
--remove_bg
```
It will save the textured model into `<save_dir>/<save_name>_shaded.glb`.
## 📊 Dataset
Our training dataset, rendered from [Objaverse](https://huggingface.co/datasets/allenai/objaverse), can be downloaded from [Objaverse-Ortho10View](https://huggingface.co/datasets/huanngzh/Objaverse-Ortho10View) and [Objaverse-Rand6View](https://huggingface.co/datasets/huanngzh/Objaverse-Rand6View).
* [Objaverse-Ortho10View](https://huggingface.co/datasets/huanngzh/Objaverse-Ortho10View) contains 10 orthographic views of 1024x1024 resolution, and is used as ground truth.
* [Objaverse-Rand6View](https://huggingface.co/datasets/huanngzh/Objaverse-Rand6View) contains 6 randomly distributed views, and is used as reference image conditions.
Please refer to their dataset cards to extract the data files, and organize them into the following structures:
```Bash
data
├── texture_ortho10view_easylight_objaverse # Objaverse-Ortho10View
│ ├── 00
│ │ ├── 00a4d2b0c4c240289ed456e87d8b9e02
│ ...
├── texture_rand_easylight_objaverse # Objaverse-Rand6View
│ ├── 00
│ │ ├── 00a4d2b0c4c240289ed456e87d8b9e02
│ ...
├── objaverse_list_6w.json # objaverse ids
└── objaverse_short_captions.json # id to captions
```
## 🏋️ Training
The key training code can be found in `mvadapter/systems`:
* `MVAdapterTextSDXLSystem` in `mvadapter_text_sdxl.py` is used for text or text+geometry conditioned multi-view generation.
* `MVAdapterImageSDXLSystem` in `mvadapter_image_sdxl.py` is used for image or image+geometry conditioned multi-view generation.
The specific training commands are as follows:
For text to 6 view generation:
```Bash
python launch.py --config configs/view-guidance/mvadapter_t2mv_sdxl.yaml --train --gpu 0,1,2,3,4,5,6,7
```
For single image to 6 view generation:
```Bash
python launch.py --config configs/view-guidance/mvadapter_i2mv_sdxl.yaml --train --gpu 0,1,2,3,4,5,6,7
```
For single image to 2/3/4/6 view generation:
```Bash
python launch.py --config configs/view-guidance/mvadapter_i2mv_sdxl_aug_quantity.yaml --train --gpu 0,1,2,3,4,5,6,7
```
For text + geometry to 6 view generation:
```Bash
python launch.py --config configs/geometry-guidance/mvadapter_tg2mv_sdxl.yaml --train --gpu 0,1,2,3,4,5,6,7
```
For single image + geometry to 6 view generation:
```Bash
python launch.py --config configs/geometry-guidance/mvadapter_ig2mv_sdxl.yaml --train --gpu 0,1,2,3,4,5,6,7
```
For single partial image + geometry to 6 view generation (used for texture generation conditioned on occluded image, for example, used in [MIDI-3D](https://github.com/VAST-AI-Research/MIDI-3D)):
```Bash
python launch.py --config configs/geometry-guidance/mvadapter_ig2mv_partialimg_sdxl.yaml --train --gpu 0,1,2,3,4,5,6,7
```
## Citation
```
@article{huang2024mvadapter,
title={MV-Adapter: Multi-view Consistent Image Generation Made Easy},
author={Huang, Zehuan and Guo, Yuanchen and Wang, Haoran and Yi, Ran and Ma, Lizhuang and Cao, Yan-Pei and Sheng, Lu},
journal={arXiv preprint arXiv:2412.03632},
year={2024}
}
```
Gen_3D_Modules/MV_Adapter/configs/geometry-guidance/mvadapter_ig2mv_partialimg_sdxl.yaml
+123
View File
@@ -0,0 +1,123 @@
name: ig2mv
tag: "r768-ortho-nv6-ig2mv-partialimg-dcrowcol-sdxl"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: data/texture_ortho6view_easylight
scene_list: data/objaverse_list_6w.json
background_color: gray
image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["position", "normal"]
position_offset: 0.5
position_scale: 1.0
reference_root_dir: ["data/texture_rand_easylight_objaverse"]
reference_scene_list: ["data/objaverse_list_6w.json"]
reference_image_modality: color
reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
reference_mask_aug: true
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv", "to_out_ref"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv", "_ref"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/geometry-guidance/mvadapter_ig2mv_sdxl.yaml
+122
View File
@@ -0,0 +1,122 @@
name: ig2mv
tag: "r768-ortho-nv6-ig2mv-dcrowcol-sdxl"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: data/texture_ortho6view_easylight
scene_list: data/objaverse_list_6w.json
background_color: gray
image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["position", "normal"]
position_offset: 0.5
position_scale: 1.0
reference_root_dir: ["data/texture_rand_easylight_objaverse"]
reference_scene_list: ["data/objaverse_list_6w.json"]
reference_image_modality: color
reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv", "to_out_ref"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv", "_ref"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/geometry-guidance/mvadapter_tg2mv_sdxl.yaml
+117
View File
@@ -0,0 +1,117 @@
name: tg2mv
tag: "r768-ortho-nv6-tg2mv-dcrowcol-sdxl"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: data/texture_ortho10view_easylight_objaverse
scene_list: data/objaverse_list_6w.json
background_color: gray
image_names: ["0000", "0001", "0002", "0003", "0008", "0009"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["position", "normal"]
position_offset: 0.5
position_scale: 1.0
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_text_sdxl.MVAdapterTextSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowColSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/view-guidance/mvadapter_i2mv_sd21.yaml
+122
View File
@@ -0,0 +1,122 @@
name: i2mv
tag: "r512-ortho-nv6-ele0-sd21"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: "data/texture_ortho10view_easylight_objaverse"
scene_list: "data/objaverse_list_6w.json"
background_color: gray
image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["plucker"]
plucker_offset: 1.0
plucker_scale: 2.0
reference_root_dir: ["data/texture_rand_easylight_objaverse"]
reference_scene_list: ["data/objaverse_list_6w.json"]
reference_image_modality: color
reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 512
width: 512
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_image_sd.MVAdapterImageSDSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-2-1-base"
pretrained_vae_name_or_path: null
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv", "to_out_ref"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv", "_ref"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 5
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/view-guidance/mvadapter_i2mv_sdxl.yaml
+122
View File
@@ -0,0 +1,122 @@
name: i2mv
tag: "r768-ortho-nv6-ele0-sdxl"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: "data/texture_ortho10view_easylight_objaverse"
scene_list: "data/objaverse_list_6w.json"
background_color: gray
image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["plucker"]
plucker_offset: 1.0
plucker_scale: 2.0
reference_root_dir: ["data/texture_rand_easylight_objaverse"]
reference_scene_list: ["data/objaverse_list_6w.json"]
reference_image_modality: color
reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv", "to_out_ref"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv", "_ref"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/view-guidance/mvadapter_i2mv_sdxl_aug_quantity.yaml
+123
View File
@@ -0,0 +1,123 @@
name: i2mv
tag: "r768-ortho-nv6-ele0-sdxl-aug-quantity"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: "data/texture_ortho10view_easylight_objaverse"
scene_list: "data/objaverse_list_6w.json"
background_color: gray
image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
image_modality: color
num_views: 6
random_view_list: [[0, 1, 2, 3, 4, 5], [0, 2, 3, 4], [0, 2, 3], [0, 3]]
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["plucker"]
plucker_offset: 1.0
plucker_scale: 2.0
reference_root_dir: ["data/texture_rand_easylight_objaverse"]
reference_scene_list: ["data/objaverse_list_6w.json"]
reference_image_modality: color
reference_image_names: ["0000", "0001", "0002", "0003", "0004"]
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_image_sdxl.MVAdapterImageSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv", "to_out_ref"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv", "_ref"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/view-guidance/mvadapter_t2mv_sd21.yaml
+117
View File
@@ -0,0 +1,117 @@
name: t2mv
tag: "r512-ortho-nv6-ele0-sd"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: "data/texture_ortho10view_easylight_objaverse"
scene_list: "data/objaverse_list_6w.json"
background_color: gray
image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["plucker"]
plucker_offset: 1.0
plucker_scale: 2.0
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 512
width: 512
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_text_sd.MVAdapterTextSDSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-2-1-base"
pretrained_vae_name_or_path: null
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/configs/view-guidance/mvadapter_t2mv_sdxl.yaml
+117
View File
@@ -0,0 +1,117 @@
name: t2mv
tag: "r768-ortho-nv6-ele0-sdxl"
exp_root_dir: "outputs"
seed: 42
data_cls: mvadapter.data.multiview.MultiviewDataModule
data:
root_dir: "data/texture_ortho10view_easylight_objaverse"
scene_list: "data/objaverse_list_6w.json"
background_color: gray
image_names: ["0000", "0004", "0001", "0002", "0003", "0005"]
image_modality: color
num_views: 6
prompt_db_path: data/objaverse_short_captions.json
return_prompt: true
projection_type: ORTHO
source_image_modality: ["plucker"]
plucker_offset: 1.0
plucker_scale: 2.0
train_indices: [0, -8]
val_indices: [-8, null]
test_indices: [-8, null]
height: 768
width: 768
batch_size: 1
num_workers: 16
system_cls: mvadapter.systems.mvadapter_text_sdxl.MVAdapterTextSDXLSystem
system:
check_train_every_n_steps: 1000
cleanup_after_validation_step: true
cleanup_after_test_step: true
# Model / Adapter
pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: null
init_adapter_kwargs:
# Multi-view adapter
self_attn_processor: "mvadapter.models.attention_processor.DecoupledMVRowSelfAttnProcessor2_0"
# Condition encoder
cond_in_channels: 6
# For training
copy_attn_weights: true
zero_init_module_keys: ["to_out_mv"]
# Training
train_cond_encoder: true
trainable_modules: ["_mv"]
prompt_drop_prob: 0.1
image_drop_prob: 0.1
cond_drop_prob: 0.1
# Noise sampler
shift_noise: true
shift_noise_mode: interpolated
shift_noise_scale: 8
# Evaluation
eval_seed: 42
eval_num_inference_steps: 30
eval_guidance_scale: 3.0
eval_height: ${data.height}
eval_width: ${data.width}
# optimizer definition
# you can set different learning rates separately for each group of parameters, but note that if you do this you should specify EVERY trainable parameters
optimizer:
name: AdamW
args:
lr: 5e-5
betas: [0.9, 0.999]
weight_decay: 0.01
params:
cond_encoder:
lr: 5e-5
unet:
lr: 5e-5
scheduler:
name: SequentialLR
interval: step
schedulers:
- name: LinearLR
interval: step
args:
start_factor: 1e-6
end_factor: 1.0
total_iters: 2000
- name: ConstantLR
interval: step
args:
factor: 1.0
total_iters: 9999999
milestones: [2000]
trainer:
max_epochs: 10
log_every_n_steps: 10
num_sanity_val_steps: 1
val_check_interval: 2000
enable_progress_bar: true
precision: bf16-mixed
gradient_clip_val: 1.0
strategy: ddp
accumulate_grad_batches: 1
checkpoint:
save_last: true # whether to save at each validation time
save_top_k: -1
every_n_epochs: 9999 # do not save at all for debug purpose
Gen_3D_Modules/MV_Adapter/launch.py
+221
View File
@@ -0,0 +1,221 @@
import argparse
import contextlib
import logging
import os
import sys
class ColoredFilter(logging.Filter):
"""
A logging filter to add color to certain log levels.
"""
RESET = "\033[0m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
COLORS = {
"WARNING": YELLOW,
"INFO": GREEN,
"DEBUG": BLUE,
"CRITICAL": MAGENTA,
"ERROR": RED,
}
RESET = "\x1b[0m"
def __init__(self):
super().__init__()
def filter(self, record):
if record.levelname in self.COLORS:
color_start = self.COLORS[record.levelname]
record.levelname = f"{color_start}[{record.levelname}]"
record.msg = f"{record.msg}{self.RESET}"
return True
def main(args, extras) -> None:
# set CUDA_VISIBLE_DEVICES if needed, then import pytorch-lightning
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
env_gpus_str = os.environ.get("CUDA_VISIBLE_DEVICES", None)
env_gpus = list(env_gpus_str.split(",")) if env_gpus_str else []
selected_gpus = [0]
# Always rely on CUDA_VISIBLE_DEVICES if specific GPU ID(s) are specified.
# As far as Pytorch Lightning is concerned, we always use all available GPUs
# (possibly filtered by CUDA_VISIBLE_DEVICES).
devices = -1
if len(env_gpus) > 0:
# CUDA_VISIBLE_DEVICES was set already, e.g. within SLURM srun or higher-level script.
n_gpus = len(env_gpus)
else:
selected_gpus = list(args.gpu.split(","))
n_gpus = len(selected_gpus)
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
import pytorch_lightning as pl
import torch
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger, WandbLogger
from pytorch_lightning.utilities.rank_zero import rank_zero_only
if args.typecheck:
from jaxtyping import install_import_hook
install_import_hook("mvadapter", "typeguard.typechecked")
from mvadapter.systems.base import BaseSystem
from mvadapter.utils.callbacks import (
CodeSnapshotCallback,
ConfigSnapshotCallback,
CustomProgressBar,
ProgressCallback,
)
from mvadapter.utils.config import ExperimentConfig, load_config
from mvadapter.utils.core import find
from mvadapter.utils.misc import get_rank, time_recorder
from mvadapter.utils.typing import Optional
logger = logging.getLogger("pytorch_lightning")
if args.verbose:
logger.setLevel(logging.DEBUG)
if args.benchmark:
time_recorder.enable(True)
for handler in logger.handlers:
if handler.stream == sys.stderr: # type: ignore
handler.setFormatter(logging.Formatter("%(levelname)s %(message)s"))
handler.addFilter(ColoredFilter())
# parse YAML config to OmegaConf
cfg: ExperimentConfig
cfg = load_config(args.config, cli_args=extras, n_gpus=n_gpus)
dm = find(cfg.data_cls)(cfg.data)
system: BaseSystem = find(cfg.system_cls)(
cfg.system, resumed=cfg.resume is not None
)
system.set_save_dir(os.path.join(cfg.trial_dir, "save"))
callbacks = []
if args.train:
callbacks += [
ModelCheckpoint(
dirpath=os.path.join(cfg.trial_dir, "ckpts"), **cfg.checkpoint
),
LearningRateMonitor(logging_interval="step"),
# CodeSnapshotCallback(
# os.path.join(cfg.trial_dir, "code"), use_version=False
# ),
ConfigSnapshotCallback(
args.config,
cfg,
os.path.join(cfg.trial_dir, "configs"),
use_version=False,
),
CustomProgressBar(refresh_rate=1),
]
def write_to_text(file, lines):
with open(file, "w") as f:
for line in lines:
f.write(line + "\n")
loggers = []
if args.train:
# make tensorboard logging dir to suppress warning
rank_zero_only(
lambda: os.makedirs(os.path.join(cfg.trial_dir, "tb_logs"), exist_ok=True)
)()
loggers += [
TensorBoardLogger(cfg.trial_dir, name="tb_logs"),
]
if args.wandb:
wandb_logger = WandbLogger(
project="MV-Adapter", name=f"{cfg.name}-{cfg.tag}"
)
system._wandb_logger = wandb_logger
loggers += [wandb_logger]
rank_zero_only(
lambda: write_to_text(
os.path.join(cfg.trial_dir, "cmd.txt"),
["python " + " ".join(sys.argv), str(args)],
)
)()
trainer = Trainer(
callbacks=callbacks,
logger=loggers,
inference_mode=False,
accelerator="gpu",
devices=devices,
**cfg.trainer,
)
# set a different seed for each device
# NOTE: use trainer.global_rank instead of get_rank() to avoid getting the local rank
pl.seed_everything(cfg.seed + trainer.global_rank, workers=True)
def set_system_status(system: BaseSystem, ckpt_path: Optional[str]):
if ckpt_path is None:
return
ckpt = torch.load(ckpt_path, map_location="cpu")
system.set_resume_status(ckpt["epoch"], ckpt["global_step"])
if args.train:
trainer.fit(system, datamodule=dm, ckpt_path=cfg.resume)
trainer.test(system, datamodule=dm)
elif args.validate:
# manually set epoch and global_step as they cannot be automatically resumed
set_system_status(system, cfg.resume)
trainer.validate(system, datamodule=dm, ckpt_path=cfg.resume)
elif args.test:
# manually set epoch and global_step as they cannot be automatically resumed
set_system_status(system, cfg.resume)
trainer.test(system, datamodule=dm, ckpt_path=cfg.resume)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True, help="path to config file")
parser.add_argument(
"--gpu",
default="0",
help="GPU(s) to be used. 0 means use the 1st available GPU. "
"1,2 means use the 2nd and 3rd available GPU. "
"If CUDA_VISIBLE_DEVICES is set before calling `launch.py`, "
"this argument is ignored and all available GPUs are always used.",
)
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument("--train", action="store_true")
group.add_argument("--validate", action="store_true")
group.add_argument("--test", action="store_true")
parser.add_argument("--wandb", action="store_true", help="if true, log to wandb")
parser.add_argument(
"--verbose", action="store_true", help="if true, set logging level to DEBUG"
)
parser.add_argument(
"--benchmark",
action="store_true",
help="if true, set to benchmark mode to record running times",
)
parser.add_argument(
"--typecheck",
action="store_true",
help="whether to enable dynamic type checking",
)
args, extras = parser.parse_known_args()
main(args, extras)
Gen_3D_Modules/MV_Adapter/mvadapter/__init__.py
View File
Gen_3D_Modules/MV_Adapter/mvadapter/data/multiview.py
+561
View File
@@ -0,0 +1,561 @@
import json
import os
import random
from dataclasses import dataclass, field
import cv2
import numpy as np
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from ..utils.config import parse_structured
from ..utils.geometry import (
get_plucker_embeds_from_cameras,
get_plucker_embeds_from_cameras_ortho,
get_position_map_from_depth,
get_position_map_from_depth_ortho,
)
from ..utils.typing import *
def _parse_scene_list_single(scene_list_path: str, root_data_dir: str):
all_scenes = []
if scene_list_path.endswith(".json"):
with open(scene_list_path) as f:
for p in json.loads(f.read()):
if "/" in p:
all_scenes.append(os.path.join(root_data_dir, p))
else:
all_scenes.append(os.path.join(root_data_dir, p[:2], p))
elif scene_list_path.endswith(".txt"):
with open(scene_list_path) as f:
for p in f.readlines():
p = p.strip()
if "/" in p:
all_scenes.append(os.path.join(root_data_dir, p))
else:
all_scenes.append(os.path.join(root_data_dir, p[:2], p))
else:
raise NotImplementedError
return all_scenes
def _parse_scene_list(
scene_list_path: Union[str, List[str]], root_data_dir: Union[str, List[str]]
):
all_scenes = []
if isinstance(scene_list_path, str):
scene_list_path = [scene_list_path]
if isinstance(root_data_dir, str):
root_data_dir = [root_data_dir]
for scene_list_path_, root_data_dir_ in zip(scene_list_path, root_data_dir):
all_scenes += _parse_scene_list_single(scene_list_path_, root_data_dir_)
return all_scenes
def _parse_reference_scene_list(reference_scenes: List[str], all_scenes: List[str]):
all_ids = set(scene.split("/")[-1] for scene in all_scenes)
ref_ids = set(scene.split("/")[-1] for scene in reference_scenes)
common_ids = ref_ids.intersection(all_ids)
all_scenes = [scene for scene in all_scenes if scene.split("/")[-1] in common_ids]
all_ids = {scene.split("/")[-1]: idx for idx, scene in enumerate(all_scenes)}
ref_scenes = [
scene for scene in reference_scenes if scene.split("/")[-1] in all_ids
]
sorted_ref_scenes = sorted(ref_scenes, key=lambda x: all_ids[x.split("/")[-1]])
scene2ref = {
scene: ref_scene for scene, ref_scene in zip(all_scenes, sorted_ref_scenes)
}
return all_scenes, scene2ref
@dataclass
class MultiviewDataModuleConfig:
root_dir: Any = ""
scene_list: Any = ""
image_suffix: str = "webp"
background_color: Union[str, float] = "gray"
image_names: List[str] = field(default_factory=lambda: [])
image_modality: str = "render"
num_views: int = 1
random_view_list: Optional[List[List[int]]] = None
prompt_db_path: Optional[str] = None
return_prompt: bool = False
use_empty_prompt: bool = False
prompt_prefix: Optional[Any] = None
return_one_prompt: bool = True
projection_type: str = "ORTHO"
# source conditions
source_image_modality: Any = "position"
use_camera_space_normal: bool = False
position_offset: float = 0.5
position_scale: float = 1.0
plucker_offset: float = 1.0
plucker_scale: float = 2.0
# reference image
reference_root_dir: Optional[Any] = None
reference_scene_list: Optional[Any] = None
reference_image_modality: str = "render"
reference_image_names: List[str] = field(default_factory=lambda: [])
reference_augment_resolutions: Optional[List[int]] = None
reference_mask_aug: bool = False
repeat: int = 1 # for debugging purpose
train_indices: Optional[Tuple[Any, Any]] = None
val_indices: Optional[Tuple[Any, Any]] = None
test_indices: Optional[Tuple[Any, Any]] = None
height: int = 768
width: int = 768
batch_size: int = 1
eval_batch_size: int = 1
num_workers: int = 16
class MultiviewDataset(Dataset):
def __init__(self, cfg: Any, split: str = "train") -> None:
super().__init__()
assert split in ["train", "val", "test"]
self.cfg: MultiviewDataModuleConfig = cfg
self.all_scenes = _parse_scene_list(self.cfg.scene_list, self.cfg.root_dir)
if (
self.cfg.reference_root_dir is not None
and self.cfg.reference_scene_list is not None
):
reference_scenes = _parse_scene_list(
self.cfg.reference_scene_list, self.cfg.reference_root_dir
)
self.all_scenes, self.reference_scenes = _parse_reference_scene_list(
reference_scenes, self.all_scenes
)
else:
self.reference_scenes = None
self.split = split
if self.split == "train" and self.cfg.train_indices is not None:
self.all_scenes = self.all_scenes[
self.cfg.train_indices[0] : self.cfg.train_indices[1]
]
self.all_scenes = self.all_scenes * self.cfg.repeat
elif self.split == "val" and self.cfg.val_indices is not None:
self.all_scenes = self.all_scenes[
self.cfg.val_indices[0] : self.cfg.val_indices[1]
]
elif self.split == "test" and self.cfg.test_indices is not None:
self.all_scenes = self.all_scenes[
self.cfg.test_indices[0] : self.cfg.test_indices[1]
]
if self.cfg.prompt_db_path is not None:
self.prompt_db = json.load(open(self.cfg.prompt_db_path))
else:
self.prompt_db = None
def __len__(self):
return len(self.all_scenes)
def get_bg_color(self, bg_color):
if bg_color == "white":
bg_color = np.array([1.0, 1.0, 1.0], dtype=np.float32)
elif bg_color == "black":
bg_color = np.array([0.0, 0.0, 0.0], dtype=np.float32)
elif bg_color == "gray":
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
elif bg_color == "random":
bg_color = np.random.rand(3)
elif bg_color == "random_gray":
bg_color = random.uniform(0.3, 0.7)
bg_color = np.array([bg_color] * 3, dtype=np.float32)
elif isinstance(bg_color, float):
bg_color = np.array([bg_color] * 3, dtype=np.float32)
elif isinstance(bg_color, list) or isinstance(bg_color, tuple):
bg_color = np.array(bg_color, dtype=np.float32)
else:
raise NotImplementedError
return bg_color
def load_image(
self,
image: Union[str, Image.Image],
height: int,
width: int,
background_color: torch.Tensor,
rescale: bool = False,
mask_aug: bool = False,
):
if isinstance(image, str):
image = Image.open(image)
image = image.resize((width, height))
image = torch.from_numpy(np.array(image)).float() / 255.0
if mask_aug:
alpha = image[:, :, 3] # Extract alpha channel
h, w = alpha.shape
y_indices, x_indices = torch.where(alpha > 0.5)
if len(y_indices) > 0 and len(x_indices) > 0:
idx = torch.randint(len(y_indices), (1,)).item()
y_center = y_indices[idx].item()
x_center = x_indices[idx].item()
mask_h = random.randint(h // 8, h // 4)
mask_w = random.randint(w // 8, w // 4)
y1 = max(0, y_center - mask_h // 2)
y2 = min(h, y_center + mask_h // 2)
x1 = max(0, x_center - mask_w // 2)
x2 = min(w, x_center + mask_w // 2)
alpha[y1:y2, x1:x2] = 0.0
image[:, :, 3] = alpha
image = image[:, :, :3] * image[:, :, 3:4] + background_color * (
1 - image[:, :, 3:4]
)
if rescale:
image = image * 2.0 - 1.0
return image
def load_normal_image(
self,
path,
height,
width,
background_color,
camera_space: bool = False,
c2w: Optional[torch.FloatTensor] = None,
):
image = Image.open(path).resize((width, height), resample=Image.NEAREST)
image = torch.from_numpy(np.array(image)).float() / 255.0
alpha = image[:, :, 3:4]
image = image[:, :, :3]
if camera_space:
w2c = torch.linalg.inv(c2w)[:3, :3]
image = (
F.normalize(((image * 2 - 1)[:, :, None, :] * w2c).sum(-1), dim=-1)
* 0.5
+ 0.5
)
image = image * alpha + background_color * (1 - alpha)
return image
def load_depth(self, path, height, width):
depth = cv2.imread(path, cv2.IMREAD_UNCHANGED)
depth = cv2.resize(depth, (width, height), interpolation=cv2.INTER_NEAREST)
depth = torch.from_numpy(depth[..., 0:1]).float()
mask = torch.ones_like(depth)
mask[depth > 1000.0] = 0.0 # depth = 65535 is the invalid value
depth[~(mask > 0.5)] = 0.0
return depth, mask
def retrieve_prompt(self, scene_dir):
assert self.prompt_db is not None
source_id = os.path.basename(scene_dir)
return self.prompt_db.get(source_id, "")
def __getitem__(self, index):
background_color = torch.as_tensor(self.get_bg_color(self.cfg.background_color))
scene_dir = self.all_scenes[index]
with open(os.path.join(scene_dir, "meta.json")) as f:
meta = json.load(f)
name2loc = {loc["index"]: loc for loc in meta["locations"]}
# target multi-view images
image_paths = [
os.path.join(
scene_dir, f"{self.cfg.image_modality}_{f}.{self.cfg.image_suffix}"
)
for f in self.cfg.image_names
]
images = [
self.load_image(
p,
height=self.cfg.height,
width=self.cfg.width,
background_color=background_color,
)
for p in image_paths
]
images = torch.stack(images, dim=0).permute(0, 3, 1, 2)
# camera
c2w = [
torch.as_tensor(name2loc[name]["transform_matrix"])
for name in self.cfg.image_names
]
c2w = torch.stack(c2w, dim=0)
if self.cfg.projection_type == "PERSP":
camera_angle_x = (
meta.get("camera_angle_x", None)
or meta["locations"][0]["camera_angle_x"]
)
focal_length = 0.5 * self.cfg.width / np.tan(0.5 * camera_angle_x)
intrinsics = (
torch.as_tensor(
[
[focal_length, 0.0, 0.5 * self.cfg.width],
[0.0, focal_length, 0.5 * self.cfg.height],
[0.0, 0.0, 1.0],
]
)
.unsqueeze(0)
.float()
.repeat(len(self.cfg.image_names), 1, 1)
)
elif self.cfg.projection_type == "ORTHO":
ortho_scale = (
meta.get("ortho_scale", None) or meta["locations"][0]["ortho_scale"]
)
# source conditions
source_image_modality = self.cfg.source_image_modality
if isinstance(source_image_modality, str):
source_image_modality = [source_image_modality]
source_images = []
for modality in source_image_modality:
if modality == "position":
depth_masks = [
self.load_depth(
os.path.join(scene_dir, f"depth_{f}.exr"),
self.cfg.height,
self.cfg.width,
)
for f in self.cfg.image_names
]
depths = torch.stack([d for d, _ in depth_masks])
masks = torch.stack([m for _, m in depth_masks])
c2w_ = c2w.clone()
c2w_[:, :, 1:3] *= -1
if self.cfg.projection_type == "PERSP":
position_maps = get_position_map_from_depth(
depths,
masks,
intrinsics,
c2w_,
image_wh=(self.cfg.width, self.cfg.height),
)
elif self.cfg.projection_type == "ORTHO":
position_maps = get_position_map_from_depth_ortho(
depths,
masks,
c2w_,
ortho_scale,
image_wh=(self.cfg.width, self.cfg.height),
)
position_maps = (
(position_maps + self.cfg.position_offset) / self.cfg.position_scale
).clamp(0.0, 1.0)
source_images.append(position_maps)
elif modality == "normal":
normal_maps = [
self.load_normal_image(
os.path.join(
scene_dir, f"{modality}_{f}.{self.cfg.image_suffix}"
),
height=self.cfg.height,
width=self.cfg.width,
background_color=background_color,
camera_space=self.cfg.use_camera_space_normal,
c2w=c,
)
for c, f in zip(c2w, self.cfg.image_names)
]
source_images.append(torch.stack(normal_maps, dim=0))
elif modality == "plucker":
if self.cfg.projection_type == "ORTHO":
plucker_embed = get_plucker_embeds_from_cameras_ortho(
c2w, [ortho_scale] * len(c2w), self.cfg.width
)
elif self.cfg.projection_type == "PERSP":
plucker_embed = get_plucker_embeds_from_cameras(
c2w, [camera_angle_x] * len(c2w), self.cfg.width
)
else:
raise NotImplementedError
plucker_embed = plucker_embed.permute(0, 2, 3, 1)
plucker_embed = (
(plucker_embed + self.cfg.plucker_offset) / self.cfg.plucker_scale
).clamp(0.0, 1.0)
source_images.append(plucker_embed)
else:
raise NotImplementedError
source_images = torch.cat(source_images, dim=-1).permute(0, 3, 1, 2)
rv = {"rgb": images, "c2w": c2w, "source_rgb": source_images}
num_images = len(self.cfg.image_names)
# prompt
if self.cfg.return_prompt:
if self.cfg.use_empty_prompt:
prompt = ""
else:
prompt = self.retrieve_prompt(scene_dir)
prompts = [prompt] * num_images
if self.cfg.prompt_prefix is not None:
prompt_prefix = self.cfg.prompt_prefix
if isinstance(prompt_prefix, str):
prompt_prefix = [prompt_prefix] * num_images
for i, prompt in enumerate(prompts):
prompts[i] = f"{prompt_prefix[i]} {prompt}"
if self.cfg.return_one_prompt:
rv.update({"prompts": prompts[0]})
else:
rv.update({"prompts": prompts})
# reference image
if self.reference_scenes is not None:
reference_scene_dir = self.reference_scenes[scene_dir]
reference_image_paths = [
os.path.join(
reference_scene_dir,
f"{self.cfg.reference_image_modality}_{f}.{self.cfg.image_suffix}",
)
for f in self.cfg.reference_image_names
]
reference_image_path = random.choice(reference_image_paths)
if self.cfg.reference_augment_resolutions is None:
reference_image = self.load_image(
reference_image_path,
height=self.cfg.height,
width=self.cfg.width,
background_color=background_color,
mask_aug=self.cfg.reference_mask_aug,
).permute(2, 0, 1)
rv.update({"reference_rgb": reference_image})
else:
random_resolution = random.choice(
self.cfg.reference_augment_resolutions
)
reference_image_ = Image.open(reference_image_path).resize(
(random_resolution, random_resolution)
)
reference_image = self.load_image(
reference_image_,
height=self.cfg.height,
width=self.cfg.width,
background_color=background_color,
mask_aug=self.cfg.reference_mask_aug,
).permute(2, 0, 1)
rv.update({"reference_rgb": reference_image})
return rv
def collate(self, batch):
batch = torch.utils.data.default_collate(batch)
pack = lambda t: t.view(-1, *t.shape[2:])
if self.cfg.random_view_list is not None:
indices = random.choice(self.cfg.random_view_list)
else:
indices = list(range(self.cfg.num_views))
num_views = len(indices)
for k in batch.keys():
if k in ["rgb", "source_rgb", "c2w"]:
batch[k] = batch[k][:, indices]
batch[k] = pack(batch[k])
for k in ["prompts"]:
if not self.cfg.return_one_prompt:
batch[k] = [item for pair in zip(*batch[k]) for item in pair]
batch.update(
{
"num_views": num_views,
# For SDXL
"original_size": (self.cfg.height, self.cfg.width),
"target_size": (self.cfg.height, self.cfg.width),
"crops_coords_top_left": (0, 0),
}
)
return batch
class MultiviewDataModule(pl.LightningDataModule):
cfg: MultiviewDataModuleConfig
def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None:
super().__init__()
self.cfg = parse_structured(MultiviewDataModuleConfig, cfg)
def setup(self, stage=None) -> None:
if stage in [None, "fit"]:
self.train_dataset = MultiviewDataset(self.cfg, "train")
if stage in [None, "fit", "validate"]:
self.val_dataset = MultiviewDataset(self.cfg, "val")
if stage in [None, "test", "predict"]:
self.test_dataset = MultiviewDataset(self.cfg, "test")
def prepare_data(self):
pass
def train_dataloader(self) -> DataLoader:
return DataLoader(
self.train_dataset,
batch_size=self.cfg.batch_size,
num_workers=self.cfg.num_workers,
shuffle=True,
collate_fn=self.train_dataset.collate,
)
def val_dataloader(self) -> DataLoader:
return DataLoader(
self.val_dataset,
batch_size=self.cfg.eval_batch_size,
num_workers=self.cfg.num_workers,
shuffle=False,
collate_fn=self.val_dataset.collate,
)
def test_dataloader(self) -> DataLoader:
return DataLoader(
self.test_dataset,
batch_size=self.cfg.eval_batch_size,
num_workers=self.cfg.num_workers,
shuffle=False,
collate_fn=self.test_dataset.collate,
)
def predict_dataloader(self) -> DataLoader:
return self.test_dataloader()
if __name__ == "__main__":
import torchvision
from omegaconf import OmegaConf
config_file = "configs/view-guidance/mvadapter_t2mv_flux.yaml"
data_cfg = OmegaConf.load(config_file)["data"]
cfg: MultiviewDataModuleConfig = MultiviewDataModuleConfig(**data_cfg)
data_module = MultiviewDataModule(cfg)
data_module.setup()
for batch in data_module.test_dataloader():
# ref_rgb = batch["reference_rgb"] # bchw
rgb = batch["rgb"]
source_rgb = batch["source_rgb"]
print(batch["prompts"])
print(rgb.shape, source_rgb.shape)
torchvision.utils.save_image(rgb[:, :3], "debug_rgb.png")
torchvision.utils.save_image(source_rgb[:, :3], "debug_source_rgb.png")
exit(0)
Gen_3D_Modules/MV_Adapter/mvadapter/loaders/__init__.py
+1
View File
@@ -0,0 +1 @@
from .custom_adapter import CustomAdapterMixin
Gen_3D_Modules/MV_Adapter/mvadapter/loaders/custom_adapter.py
+132
View File
@@ -0,0 +1,132 @@
import os
from typing import Dict, Optional, Union
import safetensors
import torch
from diffusers.utils import _get_model_file, logging
from safetensors import safe_open
from huggingface_hub import snapshot_download
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class CustomAdapterMixin:
def init_custom_adapter(self, *args, **kwargs):
self._init_custom_adapter(*args, **kwargs)
def _init_custom_adapter(self, *args, **kwargs):
raise NotImplementedError
def load_custom_adapter(
self,
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
weight_name: str,
subfolder: Optional[str] = None,
local_cache_dir: Optional[str] = None,
**kwargs,
):
# Load the main state dict first.
cache_dir = kwargs.pop("cache_dir", None)
force_download = kwargs.pop("force_download", False)
proxies = kwargs.pop("proxies", None)
local_files_only = kwargs.pop("local_files_only", None)
token = kwargs.pop("token", None)
revision = kwargs.pop("revision", None)
user_agent = {
"file_type": "attn_procs_weights",
"framework": "pytorch",
}
if not isinstance(pretrained_model_name_or_path_or_dict, dict):
# First try to load from local cache if provided
if local_cache_dir and os.path.exists(os.path.join(local_cache_dir, weight_name)):
model_file = os.path.join(local_cache_dir, weight_name)
logger.info(f"Loading adapter from local cache: {model_file}")
else:
try:
model_file = _get_model_file(
pretrained_model_name_or_path_or_dict,
weights_name=weight_name,
subfolder=subfolder,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
local_files_only=local_files_only,
token=token,
revision=revision,
user_agent=user_agent,
)
except (OSError, EnvironmentError) as e:
# If file not found and we have a custom cache dir, try to download
if local_cache_dir:
logger.info(f"Adapter not found. Downloading {pretrained_model_name_or_path_or_dict} to {local_cache_dir}")
try:
# Create directory if it doesn't exist
os.makedirs(local_cache_dir, exist_ok=True)
# Download the repository to local cache
snapshot_download(
repo_id=pretrained_model_name_or_path_or_dict,
local_dir=local_cache_dir,
ignore_patterns=["*.yaml", "*.json", "*.py", ".png", ".jpg", ".gif"]
)
# Now try to load from the local cache
model_file = os.path.join(local_cache_dir, weight_name)
if not os.path.exists(model_file):
raise EnvironmentError(f"File {weight_name} not found even after download to {local_cache_dir}")
except Exception as download_error:
logger.error(f"Failed to download adapter: {download_error}")
raise e # Re-raise original error
else:
raise e # Re-raise original error
if weight_name.endswith(".safetensors"):
state_dict = {}
with safe_open(model_file, framework="pt", device="cpu") as f:
for key in f.keys():
state_dict[key] = f.get_tensor(key)
else:
state_dict = torch.load(model_file, map_location="cpu")
else:
state_dict = pretrained_model_name_or_path_or_dict
self._load_custom_adapter(state_dict)
def _load_custom_adapter(self, state_dict):
raise NotImplementedError
def save_custom_adapter(
self,
save_directory: Union[str, os.PathLike],
weight_name: str,
safe_serialization: bool = False,
**kwargs,
):
if os.path.isfile(save_directory):
logger.error(
f"Provided path ({save_directory}) should be a directory, not a file"
)
return
if safe_serialization:
def save_function(weights, filename):
return safetensors.torch.save_file(
weights, filename, metadata={"format": "pt"}
)
else:
save_function = torch.save
# Save the model
state_dict = self._save_custom_adapter(**kwargs)
save_function(state_dict, os.path.join(save_directory, weight_name))
logger.info(
f"Custom adapter weights saved in {os.path.join(save_directory, weight_name)}"
)
def _save_custom_adapter(self):
raise NotImplementedError
Gen_3D_Modules/MV_Adapter/mvadapter/models/__init__.py
View File
Gen_3D_Modules/MV_Adapter/mvadapter/models/attention_processor.py
+743
View File
@@ -0,0 +1,743 @@
import math
from typing import Callable, List, Optional, Union
import torch
import torch.nn.functional as F
from diffusers.models.attention_processor import Attention
from diffusers.models.unets import UNet2DConditionModel
from diffusers.utils import deprecate, logging
from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
from einops import rearrange, repeat
from torch import nn
def default_set_attn_proc_func(
name: str,
hidden_size: int,
cross_attention_dim: Optional[int],
ori_attn_proc: object,
) -> object:
return ori_attn_proc
def set_unet_2d_condition_attn_processor(
unet: UNet2DConditionModel,
set_self_attn_proc_func: Callable = default_set_attn_proc_func,
set_cross_attn_proc_func: Callable = default_set_attn_proc_func,
set_custom_attn_proc_func: Callable = default_set_attn_proc_func,
set_self_attn_module_names: Optional[List[str]] = None,
set_cross_attn_module_names: Optional[List[str]] = None,
set_custom_attn_module_names: Optional[List[str]] = None,
) -> None:
do_set_processor = lambda name, module_names: (
any([name.startswith(module_name) for module_name in module_names])
if module_names is not None
else True
) # prefix match
attn_procs = {}
for name, attn_processor in unet.attn_processors.items():
# set attn_processor by default, if module_names is None
set_self_attn_processor = do_set_processor(name, set_self_attn_module_names)
set_cross_attn_processor = do_set_processor(name, set_cross_attn_module_names)
set_custom_attn_processor = do_set_processor(name, set_custom_attn_module_names)
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
is_custom = "attn_mid_blocks" in name or "attn_post_blocks" in name
if is_custom:
attn_procs[name] = (
set_custom_attn_proc_func(name, hidden_size, None, attn_processor)
if set_custom_attn_processor
else attn_processor
)
else:
cross_attention_dim = (
None
if name.endswith("attn1.processor")
else unet.config.cross_attention_dim
)
if cross_attention_dim is None or "motion_modules" in name:
# self attention
attn_procs[name] = (
set_self_attn_proc_func(
name, hidden_size, cross_attention_dim, attn_processor
)
if set_self_attn_processor
else attn_processor
)
else:
# cross attention
attn_procs[name] = (
set_cross_attn_proc_func(
name, hidden_size, cross_attention_dim, attn_processor
)
if set_cross_attn_processor
else attn_processor
)
unet.set_attn_processor(attn_procs)
class DecoupledMVRowSelfAttnProcessor2_0(torch.nn.Module):
r"""
Attention processor for Decoupled Row-wise Self-Attention and Image Cross-Attention for PyTorch 2.0.
"""
def __init__(
self,
query_dim: int,
inner_dim: int,
num_views: int = 1,
name: Optional[str] = None,
use_mv: bool = True,
use_ref: bool = False,
):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
"DecoupledMVRowSelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
)
super().__init__()
self.num_views = num_views
self.name = name # NOTE: need for image cross-attention
self.use_mv = use_mv
self.use_ref = use_ref
if self.use_mv:
self.to_q_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_k_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_v_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_out_mv = nn.ModuleList(
[
nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
nn.Dropout(0.0),
]
)
if self.use_ref:
self.to_q_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_k_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_v_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_out_ref = nn.ModuleList(
[
nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
nn.Dropout(0.0),
]
)
def __call__(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.FloatTensor] = None,
temb: Optional[torch.FloatTensor] = None,
mv_scale: float = 1.0,
ref_hidden_states: Optional[torch.FloatTensor] = None,
ref_scale: float = 1.0,
cache_hidden_states: Optional[List[torch.FloatTensor]] = None,
use_mv: bool = True,
use_ref: bool = True,
num_views: Optional[int] = None,
*args,
**kwargs,
) -> torch.FloatTensor:
"""
New args:
mv_scale (float): scale for multi-view self-attention.
ref_hidden_states (torch.FloatTensor): reference encoder hidden states for image cross-attention.
ref_scale (float): scale for image cross-attention.
cache_hidden_states (List[torch.FloatTensor]): cache hidden states from reference unet.
"""
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
if num_views is not None:
self.num_views = num_views
# NEW: cache hidden states for reference unet
if cache_hidden_states is not None:
cache_hidden_states[self.name] = hidden_states.clone()
# NEW: whether to use multi-view attention and image cross-attention
use_mv = self.use_mv and use_mv
use_ref = self.use_ref and use_ref
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape
if encoder_hidden_states is None
else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(
attention_mask, sequence_length, batch_size
)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(
batch_size, attn.heads, -1, attention_mask.shape[-1]
)
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
1, 2
)
query = attn.to_q(hidden_states)
# NEW: for decoupled multi-view attention
if use_mv:
query_mv = self.to_q_mv(hidden_states)
# NEW: for decoupled reference cross attention
if use_ref:
query_ref = self.to_q_ref(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(
encoder_hidden_states
)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
hidden_states = F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states = hidden_states.to(query.dtype)
####### Decoupled multi-view self-attention ########
if use_mv:
key_mv = self.to_k_mv(encoder_hidden_states)
value_mv = self.to_v_mv(encoder_hidden_states)
query_mv = query_mv.view(batch_size, -1, attn.heads, head_dim)
key_mv = key_mv.view(batch_size, -1, attn.heads, head_dim)
value_mv = value_mv.view(batch_size, -1, attn.heads, head_dim)
height = width = math.isqrt(sequence_length)
# row self-attention
query_mv = rearrange(
query_mv,
"(b nv) (ih iw) h c -> (b nv ih) iw h c",
nv=self.num_views,
ih=height,
iw=width,
).transpose(1, 2)
key_mv = rearrange(
key_mv,
"(b nv) (ih iw) h c -> b ih (nv iw) h c",
nv=self.num_views,
ih=height,
iw=width,
)
key_mv = (
key_mv.repeat_interleave(self.num_views, dim=0)
.view(batch_size * height, -1, attn.heads, head_dim)
.transpose(1, 2)
)
value_mv = rearrange(
value_mv,
"(b nv) (ih iw) h c -> b ih (nv iw) h c",
nv=self.num_views,
ih=height,
iw=width,
)
value_mv = (
value_mv.repeat_interleave(self.num_views, dim=0)
.view(batch_size * height, -1, attn.heads, head_dim)
.transpose(1, 2)
)
hidden_states_mv = F.scaled_dot_product_attention(
query_mv,
key_mv,
value_mv,
dropout_p=0.0,
is_causal=False,
)
hidden_states_mv = rearrange(
hidden_states_mv,
"(b nv ih) h iw c -> (b nv) (ih iw) (h c)",
nv=self.num_views,
ih=height,
)
hidden_states_mv = hidden_states_mv.to(query.dtype)
# linear proj
hidden_states_mv = self.to_out_mv[0](hidden_states_mv)
# dropout
hidden_states_mv = self.to_out_mv[1](hidden_states_mv)
if use_ref:
reference_hidden_states = ref_hidden_states[self.name]
key_ref = self.to_k_ref(reference_hidden_states)
value_ref = self.to_v_ref(reference_hidden_states)
query_ref = query_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
key_ref = key_ref.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value_ref = value_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
hidden_states_ref = F.scaled_dot_product_attention(
query_ref, key_ref, value_ref, dropout_p=0.0, is_causal=False
)
hidden_states_ref = hidden_states_ref.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states_ref = hidden_states_ref.to(query.dtype)
# linear proj
hidden_states_ref = self.to_out_ref[0](hidden_states_ref)
# dropout
hidden_states_ref = self.to_out_ref[1](hidden_states_ref)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if use_mv:
hidden_states = hidden_states + hidden_states_mv * mv_scale
if use_ref:
hidden_states = hidden_states + hidden_states_ref * ref_scale
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
def set_num_views(self, num_views: int) -> None:
self.num_views = num_views
class DecoupledMVRowColSelfAttnProcessor2_0(torch.nn.Module):
r"""
Attention processor for Decoupled Row-wise Self-Attention and Image Cross-Attention for PyTorch 2.0.
"""
def __init__(
self,
query_dim: int,
inner_dim: int,
num_views: int = 1,
name: Optional[str] = None,
use_mv: bool = True,
use_ref: bool = False,
):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
"DecoupledMVRowSelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
)
super().__init__()
self.num_views = num_views
self.name = name # NOTE: need for image cross-attention
self.use_mv = use_mv
self.use_ref = use_ref
if self.use_mv:
self.to_q_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_k_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_v_mv = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_out_mv = nn.ModuleList(
[
nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
nn.Dropout(0.0),
]
)
if self.use_ref:
self.to_q_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_k_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_v_ref = nn.Linear(
in_features=query_dim, out_features=inner_dim, bias=False
)
self.to_out_ref = nn.ModuleList(
[
nn.Linear(in_features=inner_dim, out_features=query_dim, bias=True),
nn.Dropout(0.0),
]
)
def __call__(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.FloatTensor] = None,
temb: Optional[torch.FloatTensor] = None,
mv_scale: float = 1.0,
ref_hidden_states: Optional[torch.FloatTensor] = None,
ref_scale: float = 1.0,
cache_hidden_states: Optional[List[torch.FloatTensor]] = None,
use_mv: bool = True,
use_ref: bool = True,
num_views: Optional[int] = None,
*args,
**kwargs,
) -> torch.FloatTensor:
"""
New args:
mv_scale (float): scale for multi-view self-attention.
ref_hidden_states (torch.FloatTensor): reference encoder hidden states for image cross-attention.
ref_scale (float): scale for image cross-attention.
cache_hidden_states (List[torch.FloatTensor]): cache hidden states from reference unet.
"""
if len(args) > 0 or kwargs.get("scale", None) is not None:
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
deprecate("scale", "1.0.0", deprecation_message)
if num_views is not None:
self.num_views = num_views
# NEW: cache hidden states for reference unet
if cache_hidden_states is not None:
cache_hidden_states[self.name] = hidden_states.clone()
# NEW: whether to use multi-view attention and image cross-attention
use_mv = self.use_mv and use_mv
use_ref = self.use_ref and use_ref
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(
batch_size, channel, height * width
).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape
if encoder_hidden_states is None
else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(
attention_mask, sequence_length, batch_size
)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(
batch_size, attn.heads, -1, attention_mask.shape[-1]
)
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
1, 2
)
query = attn.to_q(hidden_states)
# NEW: for decoupled multi-view attention
if use_mv:
query_mv = self.to_q_mv(hidden_states)
# NEW: for decoupled reference cross attention
if use_ref:
query_ref = self.to_q_ref(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(
encoder_hidden_states
)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
hidden_states = F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states = hidden_states.to(query.dtype)
####### Decoupled multi-view self-attention ########
if use_mv:
key_mv = self.to_k_mv(encoder_hidden_states)
value_mv = self.to_v_mv(encoder_hidden_states)
query_mv = query_mv.view(batch_size, -1, attn.heads, head_dim)
key_mv = key_mv.view(batch_size, -1, attn.heads, head_dim)
value_mv = value_mv.view(batch_size, -1, attn.heads, head_dim)
height = width = math.isqrt(sequence_length)
query_mv = rearrange(
query_mv,
"(b nv) (ih iw) h c -> b nv ih iw h c",
nv=self.num_views,
ih=height,
iw=width,
)
key_mv = rearrange(
key_mv,
"(b nv) (ih iw) h c -> b nv ih iw h c",
nv=self.num_views,
ih=height,
iw=width,
)
value_mv = rearrange(
value_mv,
"(b nv) (ih iw) h c -> b nv ih iw h c",
nv=self.num_views,
ih=height,
iw=width,
)
# row-wise attention for view 0123 (front, right, back, left)
query_mv_0123 = rearrange(
query_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
)
key_mv_0123 = rearrange(
key_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
)
value_mv_0123 = rearrange(
value_mv[:, 0:4], "b nv ih iw h c -> (b ih) h (nv iw) c"
)
hidden_states_mv_0123 = F.scaled_dot_product_attention(
query_mv_0123,
key_mv_0123,
value_mv_0123,
dropout_p=0.0,
is_causal=False,
)
hidden_states_mv_0123 = rearrange(
hidden_states_mv_0123,
"(b ih) h (nv iw) c -> b nv (ih iw) (h c)",
ih=height,
iw=height,
)
# col-wise attention for view 0245 (front, back, top, bottom)
# flip first
query_mv_0245 = torch.cat(
[
torch.flip(query_mv[:, [0]], [3]), # horizontal flip
query_mv[:, [2, 4, 5]],
],
dim=1,
)
key_mv_0245 = torch.cat(
[
torch.flip(key_mv[:, [0]], [3]), # horizontal flip
key_mv[:, [2, 4, 5]],
],
dim=1,
)
value_mv_0245 = torch.cat(
[
torch.flip(value_mv[:, [0]], [3]), # horizontal flip
value_mv[:, [2, 4, 5]],
],
dim=1,
)
# attention
query_mv_0245 = rearrange(
query_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c"
)
key_mv_0245 = rearrange(key_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c")
value_mv_0245 = rearrange(
value_mv_0245, "b nv ih iw h c -> (b iw) h (nv ih) c"
)
hidden_states_mv_0245 = F.scaled_dot_product_attention(
query_mv_0245,
key_mv_0245,
value_mv_0245,
dropout_p=0.0,
is_causal=False,
)
# flip back
hidden_states_mv_0245 = rearrange(
hidden_states_mv_0245,
"(b iw) h (nv ih) c -> b nv ih iw (h c)",
ih=height,
iw=height,
)
hidden_states_mv_0245 = torch.cat(
[
torch.flip(hidden_states_mv_0245[:, [0]], [3]), # horizontal flip
hidden_states_mv_0245[:, [1, 2, 3]],
],
dim=1,
)
hidden_states_mv_0245 = hidden_states_mv_0245.view(
hidden_states_mv_0245.shape[0],
hidden_states_mv_0245.shape[1],
-1,
hidden_states_mv_0245.shape[-1],
)
# combine row and col
hidden_states_mv = torch.stack(
[
(hidden_states_mv_0123[:, 0] + hidden_states_mv_0245[:, 0]) / 2,
hidden_states_mv_0123[:, 1],
(hidden_states_mv_0123[:, 2] + hidden_states_mv_0245[:, 1]) / 2,
hidden_states_mv_0123[:, 3],
hidden_states_mv_0245[:, 2],
hidden_states_mv_0245[:, 3],
],
dim=1,
)
hidden_states_mv = hidden_states_mv.view(
-1, hidden_states_mv.shape[-2], hidden_states_mv.shape[-1]
)
hidden_states_mv = hidden_states_mv.to(query.dtype)
# linear proj
hidden_states_mv = self.to_out_mv[0](hidden_states_mv)
# dropout
hidden_states_mv = self.to_out_mv[1](hidden_states_mv)
if use_ref:
reference_hidden_states = ref_hidden_states[self.name]
key_ref = self.to_k_ref(reference_hidden_states)
value_ref = self.to_v_ref(reference_hidden_states)
query_ref = query_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
key_ref = key_ref.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value_ref = value_ref.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
hidden_states_ref = F.scaled_dot_product_attention(
query_ref, key_ref, value_ref, dropout_p=0.0, is_causal=False
)
hidden_states_ref = hidden_states_ref.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states_ref = hidden_states_ref.to(query.dtype)
# linear proj
hidden_states_ref = self.to_out_ref[0](hidden_states_ref)
# dropout
hidden_states_ref = self.to_out_ref[1](hidden_states_ref)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if use_mv:
hidden_states = hidden_states + hidden_states_mv * mv_scale
if use_ref:
hidden_states = hidden_states + hidden_states_ref * ref_scale
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(
batch_size, channel, height, width
)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
def set_num_views(self, num_views: int) -> None:
self.num_views = num_views
Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_i2mv_sd.py
+777
View File
@@ -0,0 +1,777 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from typing import Any, Callable, Dict, List, Optional, Union
import PIL
import torch
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.models import AutoencoderKL, T2IAdapter, UNet2DConditionModel
from diffusers.pipelines.stable_diffusion.pipeline_output import (
StableDiffusionPipelineOutput,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
StableDiffusionPipeline,
rescale_noise_cfg,
retrieve_timesteps,
)
from diffusers.pipelines.stable_diffusion.safety_checker import (
StableDiffusionSafetyChecker,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import deprecate, is_torch_xla_available, logging
from diffusers.utils.torch_utils import randn_tensor
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTokenizer,
CLIPVisionModelWithProjection,
)
from ..loaders import CustomAdapterMixin
from ..models.attention_processor import (
DecoupledMVRowSelfAttnProcessor2_0,
set_unet_2d_condition_attn_processor,
)
if is_torch_xla_available():
import torch_xla.core.xla_model as xm
XLA_AVAILABLE = True
else:
XLA_AVAILABLE = False
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
def retrieve_latents(
encoder_output: torch.Tensor,
generator: Optional[torch.Generator] = None,
sample_mode: str = "sample",
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
return encoder_output.latent_dist.mode()
elif hasattr(encoder_output, "latents"):
return encoder_output.latents
else:
raise AttributeError("Could not access latents of provided encoder_output")
class MVAdapterI2MVSDPipeline(StableDiffusionPipeline, CustomAdapterMixin):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = False,
):
super().__init__(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
image_encoder=image_encoder,
requires_safety_checker=requires_safety_checker,
)
self.control_image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor,
do_convert_rgb=True,
do_normalize=False,
)
# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.prepare_latents
def prepare_image_latents(
self,
image,
timestep,
batch_size,
num_images_per_prompt,
dtype,
device,
generator=None,
add_noise=True,
):
if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
raise ValueError(
f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
)
image = image.to(device=device, dtype=dtype)
batch_size = batch_size * num_images_per_prompt
if image.shape[1] == 4:
init_latents = image
else:
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
elif isinstance(generator, list):
init_latents = [
retrieve_latents(
self.vae.encode(image[i : i + 1]), generator=generator[i]
)
for i in range(batch_size)
]
init_latents = torch.cat(init_latents, dim=0)
else:
init_latents = retrieve_latents(
self.vae.encode(image), generator=generator
)
init_latents = self.vae.config.scaling_factor * init_latents
if (
batch_size > init_latents.shape[0]
and batch_size % init_latents.shape[0] == 0
):
# expand init_latents for batch_size
additional_image_per_prompt = batch_size // init_latents.shape[0]
init_latents = torch.cat(
[init_latents] * additional_image_per_prompt, dim=0
)
elif (
batch_size > init_latents.shape[0]
and batch_size % init_latents.shape[0] != 0
):
raise ValueError(
f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
)
else:
init_latents = torch.cat([init_latents], dim=0)
if add_noise:
shape = init_latents.shape
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# get latents
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
latents = init_latents
return latents
def prepare_control_image(
self,
image,
width,
height,
batch_size,
num_images_per_prompt,
device,
dtype,
do_classifier_free_guidance=False,
num_empty_images=0, # for concat in batch like ImageDream
):
assert hasattr(
self, "control_image_processor"
), "control_image_processor is not initialized"
image = self.control_image_processor.preprocess(
image, height=height, width=width
).to(dtype=torch.float32)
if num_empty_images > 0:
image = torch.cat(
[image, torch.zeros_like(image[:num_empty_images])], dim=0
)
image_batch_size = image.shape[0]
if image_batch_size == 1:
repeat_by = batch_size
else:
# image batch size is the same as prompt batch size
repeat_by = num_images_per_prompt # always 1 for control image
image = image.repeat_interleave(repeat_by, dim=0)
image = image.to(device=device, dtype=dtype)
if do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
timesteps: List[int] = None,
sigmas: List[float] = None,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[
Union[
Callable[[int, int, Dict], None],
PipelineCallback,
MultiPipelineCallbacks,
]
] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# NEW
mv_scale: float = 1.0,
# Camera or geometry condition
control_image: Optional[PipelineImageInput] = None,
control_conditioning_scale: Optional[float] = 1.0,
control_conditioning_factor: float = 1.0,
# Image condition
reference_image: Optional[PipelineImageInput] = None,
reference_conditioning_scale: Optional[float] = 1.0,
# Optional. controlnet
controlnet_image: Optional[PipelineImageInput] = None,
controlnet_conditioning_scale: Optional[float] = 1.0,
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The width in pixels of the generated image.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
sigmas (`List[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor is generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
provided, text embeddings are generated from the `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
guidance_rescale (`float`, *optional*, defaults to 0.0):
Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
the output of the pre-final layer will be used for computing the prompt embeddings.
callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
callback_on_step_end_tensor_inputs (`List`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
Examples:
Returns:
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
otherwise a `tuple` is returned where the first element is a list with the generated images and the
second element is a list of `bool`s indicating whether the corresponding generated image contains
"not-safe-for-work" (nsfw) content.
"""
callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", None)
if callback is not None:
deprecate(
"callback",
"1.0.0",
"Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
)
if callback_steps is not None:
deprecate(
"callback_steps",
"1.0.0",
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
)
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
# to deal with lora scaling and other possible forward hooks
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
height,
width,
callback_steps,
negative_prompt,
prompt_embeds,
negative_prompt_embeds,
ip_adapter_image,
ip_adapter_image_embeds,
callback_on_step_end_tensor_inputs,
)
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._interrupt = False
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# 3. Encode input prompt
lora_scale = (
self.cross_attention_kwargs.get("scale", None)
if self.cross_attention_kwargs is not None
else None
)
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt,
device,
num_images_per_prompt,
self.do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
)
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
image_embeds = self.prepare_ip_adapter_image_embeds(
ip_adapter_image,
ip_adapter_image_embeds,
device,
batch_size * num_images_per_prompt,
self.do_classifier_free_guidance,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device, timesteps, sigmas
)
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 6.1 Add image embeds for IP-Adapter
added_cond_kwargs = (
{"image_embeds": image_embeds}
if (ip_adapter_image is not None or ip_adapter_image_embeds is not None)
else None
)
# 6.2 Optionally get Guidance Scale Embedding
timestep_cond = None
if self.unet.config.time_cond_proj_dim is not None:
guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
batch_size * num_images_per_prompt
)
timestep_cond = self.get_guidance_scale_embedding(
guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
).to(device=device, dtype=latents.dtype)
# Preprocess reference image
reference_image = self.image_processor.preprocess(reference_image)
reference_latents = self.prepare_image_latents(
reference_image,
timesteps[:1].repeat(batch_size * num_images_per_prompt), # no use
batch_size,
1,
prompt_embeds.dtype,
device,
generator,
add_noise=False,
)
ref_timesteps = torch.zeros_like(timesteps[0])
ref_hidden_states = {}
with torch.no_grad():
self.unet(
reference_latents,
ref_timesteps,
encoder_hidden_states=prompt_embeds[-1:],
cross_attention_kwargs={
"cache_hidden_states": ref_hidden_states,
"use_mv": False,
"use_ref": False,
},
return_dict=False,
)
ref_hidden_states = {
k: v.repeat_interleave(num_images_per_prompt, dim=0)
for k, v in ref_hidden_states.items()
}
if self.do_classifier_free_guidance:
ref_hidden_states = {
k: torch.cat([torch.zeros_like(v), v], dim=0)
for k, v in ref_hidden_states.items()
}
cross_attention_kwargs = {
"num_views": num_images_per_prompt,
"mv_scale": mv_scale,
"ref_hidden_states": {k: v.clone() for k, v in ref_hidden_states.items()},
"ref_scale": reference_conditioning_scale,
**(self.cross_attention_kwargs or {}),
}
# Preprocess control image
control_image_feature = self.prepare_control_image(
image=control_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
control_image_feature = control_image_feature.to(
device=device, dtype=latents.dtype
)
adapter_state = self.cond_encoder(control_image_feature)
for i, state in enumerate(adapter_state):
adapter_state[i] = state * control_conditioning_scale
# Preprocess controlnet image if provided
do_controlnet = controlnet_image is not None and hasattr(self, "controlnet")
if do_controlnet:
controlnet_image = self.prepare_control_image(
image=controlnet_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
controlnet_image = controlnet_image.to(device=device, dtype=latents.dtype)
# 7. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
# expand the latents if we are doing classifier free guidance
latent_model_input = (
torch.cat([latents] * 2)
if self.do_classifier_free_guidance
else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
if i < int(num_inference_steps * control_conditioning_factor):
down_intrablock_additional_residuals = [
state.clone() for state in adapter_state
]
else:
down_intrablock_additional_residuals = None
unet_add_kwargs = {}
# Do controlnet if provided
if do_controlnet:
down_block_res_samples, mid_block_res_sample = self.controlnet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
controlnet_cond=controlnet_image,
conditioning_scale=controlnet_conditioning_scale,
guess_mode=False,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)
unet_add_kwargs.update(
{
"down_block_additional_residuals": down_block_res_samples,
"mid_block_additional_residual": mid_block_res_sample,
}
)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs=cross_attention_kwargs,
down_intrablock_additional_residuals=down_intrablock_additional_residuals,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
**unet_add_kwargs,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (
noise_pred_text - noise_pred_uncond
)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(
noise_pred,
noise_pred_text,
guidance_rescale=self.guidance_rescale,
)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop(
"negative_prompt_embeds", negative_prompt_embeds
)
# call the callback, if provided
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if XLA_AVAILABLE:
xm.mark_step()
if not output_type == "latent":
image = self.vae.decode(
latents / self.vae.config.scaling_factor,
return_dict=False,
generator=generator,
)[0]
image, has_nsfw_concept = self.run_safety_checker(
image, device, prompt_embeds.dtype
)
else:
image = latents
has_nsfw_concept = None
if has_nsfw_concept is None:
do_denormalize = [True] * image.shape[0]
else:
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
image = self.image_processor.postprocess(
image, output_type=output_type, do_denormalize=do_denormalize
)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(
images=image, nsfw_content_detected=has_nsfw_concept
)
### NEW: adapters ###
def _init_custom_adapter(
self,
# Multi-view adapter
num_views: int = 1,
self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
# Condition encoder
cond_in_channels: int = 6,
# For training
copy_attn_weights: bool = True,
zero_init_module_keys: List[str] = [],
):
# Condition encoder
self.cond_encoder = T2IAdapter(
in_channels=cond_in_channels,
channels=self.unet.config.block_out_channels,
num_res_blocks=self.unet.config.layers_per_block,
downscale_factor=8,
)
# set custom attn processor for multi-view attention
self.unet: UNet2DConditionModel
set_unet_2d_condition_attn_processor(
self.unet,
set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=True,
use_ref=True,
),
set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=False,
use_ref=False,
),
)
# copy decoupled attention weights from original unet
if copy_attn_weights:
state_dict = self.unet.state_dict()
for key in state_dict.keys():
if "_mv" in key:
compatible_key = key.replace("_mv", "").replace("processor.", "")
elif "_ref" in key:
compatible_key = key.replace("_ref", "").replace("processor.", "")
else:
compatible_key = key
is_zero_init_key = any([k in key for k in zero_init_module_keys])
if is_zero_init_key:
state_dict[key] = torch.zeros_like(state_dict[compatible_key])
else:
state_dict[key] = state_dict[compatible_key].clone()
self.unet.load_state_dict(state_dict)
def _load_custom_adapter(self, state_dict):
self.unet.load_state_dict(state_dict, strict=False)
self.cond_encoder.load_state_dict(state_dict, strict=False)
def _save_custom_adapter(
self,
include_keys: Optional[List[str]] = None,
exclude_keys: Optional[List[str]] = None,
):
def include_fn(k):
is_included = False
if include_keys is not None:
is_included = is_included or any([key in k for key in include_keys])
if exclude_keys is not None:
is_included = is_included and not any(
[key in k for key in exclude_keys]
)
return is_included
state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
state_dict.update(self.cond_encoder.state_dict())
return state_dict
Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_i2mv_sdxl.py
+962
View File
@@ -0,0 +1,962 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import PIL
import torch
import torch.nn as nn
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.models import (
AutoencoderKL,
ImageProjection,
T2IAdapter,
UNet2DConditionModel,
)
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import (
StableDiffusionXLPipelineOutput,
)
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
StableDiffusionXLPipeline,
rescale_noise_cfg,
retrieve_timesteps,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import deprecate, logging
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTextModelWithProjection,
CLIPTokenizer,
CLIPVisionModelWithProjection,
)
from ..loaders import CustomAdapterMixin
from ..models.attention_processor import (
DecoupledMVRowSelfAttnProcessor2_0,
set_unet_2d_condition_attn_processor,
)
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
def retrieve_latents(
encoder_output: torch.Tensor,
generator: Optional[torch.Generator] = None,
sample_mode: str = "sample",
):
if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
return encoder_output.latent_dist.sample(generator)
elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
return encoder_output.latent_dist.mode()
elif hasattr(encoder_output, "latents"):
return encoder_output.latents
else:
raise AttributeError("Could not access latents of provided encoder_output")
class MVAdapterI2MVSDXLPipeline(StableDiffusionXLPipeline, CustomAdapterMixin):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
text_encoder_2: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
add_watermarker: Optional[bool] = None,
):
super().__init__(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
unet=unet,
scheduler=scheduler,
image_encoder=image_encoder,
feature_extractor=feature_extractor,
force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
add_watermarker=add_watermarker,
)
self.control_image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor,
do_convert_rgb=True,
do_normalize=False,
)
# Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.prepare_latents
def prepare_image_latents(
self,
image,
timestep,
batch_size,
num_images_per_prompt,
dtype,
device,
generator=None,
add_noise=True,
):
if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
raise ValueError(
f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
)
latents_mean = latents_std = None
if (
hasattr(self.vae.config, "latents_mean")
and self.vae.config.latents_mean is not None
):
latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1)
if (
hasattr(self.vae.config, "latents_std")
and self.vae.config.latents_std is not None
):
latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1)
# Offload text encoder if `enable_model_cpu_offload` was enabled
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.text_encoder_2.to("cpu")
torch.cuda.empty_cache()
image = image.to(device=device, dtype=dtype)
batch_size = batch_size * num_images_per_prompt
if image.shape[1] == 4:
init_latents = image
else:
# make sure the VAE is in float32 mode, as it overflows in float16
if self.vae.config.force_upcast:
image = image.float()
self.vae.to(dtype=torch.float32)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
elif isinstance(generator, list):
if image.shape[0] < batch_size and batch_size % image.shape[0] == 0:
image = torch.cat([image] * (batch_size // image.shape[0]), dim=0)
elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0:
raise ValueError(
f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} "
)
init_latents = [
retrieve_latents(
self.vae.encode(image[i : i + 1]), generator=generator[i]
)
for i in range(batch_size)
]
init_latents = torch.cat(init_latents, dim=0)
else:
init_latents = retrieve_latents(
self.vae.encode(image), generator=generator
)
if self.vae.config.force_upcast:
self.vae.to(dtype)
init_latents = init_latents.to(dtype)
if latents_mean is not None and latents_std is not None:
latents_mean = latents_mean.to(device=device, dtype=dtype)
latents_std = latents_std.to(device=device, dtype=dtype)
init_latents = (
(init_latents - latents_mean)
* self.vae.config.scaling_factor
/ latents_std
)
else:
init_latents = self.vae.config.scaling_factor * init_latents
if (
batch_size > init_latents.shape[0]
and batch_size % init_latents.shape[0] == 0
):
# expand init_latents for batch_size
additional_image_per_prompt = batch_size // init_latents.shape[0]
init_latents = torch.cat(
[init_latents] * additional_image_per_prompt, dim=0
)
elif (
batch_size > init_latents.shape[0]
and batch_size % init_latents.shape[0] != 0
):
raise ValueError(
f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
)
else:
init_latents = torch.cat([init_latents], dim=0)
if add_noise:
shape = init_latents.shape
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# get latents
init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
latents = init_latents
return latents
def prepare_control_image(
self,
image,
width,
height,
batch_size,
num_images_per_prompt,
device,
dtype,
do_classifier_free_guidance=False,
num_empty_images=0, # for concat in batch like ImageDream
):
assert hasattr(
self, "control_image_processor"
), "control_image_processor is not initialized"
image = self.control_image_processor.preprocess(
image, height=height, width=width
).to(dtype=torch.float32)
if num_empty_images > 0:
image = torch.cat(
[image, torch.zeros_like(image[:num_empty_images])], dim=0
)
image_batch_size = image.shape[0]
if image_batch_size == 1:
repeat_by = batch_size
else:
# image batch size is the same as prompt batch size
repeat_by = num_images_per_prompt # always 1 for control image
image = image.repeat_interleave(repeat_by, dim=0)
image = image.to(device=device, dtype=dtype)
if do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
timesteps: List[int] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
original_size: Optional[Tuple[int, int]] = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
target_size: Optional[Tuple[int, int]] = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# NEW
mv_scale: float = 1.0,
# Camera or geometry condition
control_image: Optional[PipelineImageInput] = None,
control_conditioning_scale: Optional[float] = 1.0,
control_conditioning_factor: float = 1.0,
# Image condition
reference_image: Optional[PipelineImageInput] = None,
reference_conditioning_scale: Optional[float] = 1.0,
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The width in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
denoising_end (`float`, *optional*):
When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
completed before it is intentionally prematurely terminated. As a result, the returned sample will
still retain a substantial amount of noise as determined by the discrete timesteps selected by the
scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
"Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
guidance_scale (`float`, *optional*, defaults to 5.0):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
negative_prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will ge generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generate image. Choose between
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
of a plain tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
guidance_rescale (`float`, *optional*, defaults to 0.0):
Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
[Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
Guidance rescale factor should fix overexposure when using zero terminal SNR.
original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
callback_on_step_end_tensor_inputs (`List`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
Examples:
Returns:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
`tuple`. When returning a tuple, the first element is a list with the generated images.
"""
callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", None)
if callback is not None:
deprecate(
"callback",
"1.0.0",
"Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
if callback_steps is not None:
deprecate(
"callback_steps",
"1.0.0",
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
# 0. Default height and width to unet
height = height or self.default_sample_size * self.vae_scale_factor
width = width or self.default_sample_size * self.vae_scale_factor
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
prompt_2,
height,
width,
callback_steps,
negative_prompt,
negative_prompt_2,
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
ip_adapter_image,
ip_adapter_image_embeds,
callback_on_step_end_tensor_inputs,
)
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._denoising_end = denoising_end
self._interrupt = False
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# 3. Encode input prompt
lora_scale = (
self.cross_attention_kwargs.get("scale", None)
if self.cross_attention_kwargs is not None
else None
)
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt=prompt,
prompt_2=prompt_2,
device=device,
num_images_per_prompt=num_images_per_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
negative_prompt=negative_prompt,
negative_prompt_2=negative_prompt_2,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device, timesteps
)
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Prepare added time ids & embeddings
add_text_embeds = pooled_prompt_embeds
if self.text_encoder_2 is None:
text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
else:
text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
add_time_ids = self._get_add_time_ids(
original_size,
crops_coords_top_left,
target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
if negative_original_size is not None and negative_target_size is not None:
negative_add_time_ids = self._get_add_time_ids(
negative_original_size,
negative_crops_coords_top_left,
negative_target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
else:
negative_add_time_ids = add_time_ids
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
add_text_embeds = torch.cat(
[negative_pooled_prompt_embeds, add_text_embeds], dim=0
)
add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
prompt_embeds = prompt_embeds.to(device)
add_text_embeds = add_text_embeds.to(device)
add_time_ids = add_time_ids.to(device).repeat(
batch_size * num_images_per_prompt, 1
)
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
image_embeds = self.prepare_ip_adapter_image_embeds(
ip_adapter_image,
ip_adapter_image_embeds,
device,
batch_size * num_images_per_prompt,
self.do_classifier_free_guidance,
)
# Preprocess reference image
reference_image = self.image_processor.preprocess(reference_image)
reference_latents = self.prepare_image_latents(
reference_image,
timesteps[:1].repeat(batch_size * num_images_per_prompt), # no use
batch_size,
1,
prompt_embeds.dtype,
device,
generator,
add_noise=False,
)
with torch.no_grad():
ref_timesteps = torch.zeros_like(timesteps[0])
ref_hidden_states = {}
self.unet(
reference_latents,
ref_timesteps,
encoder_hidden_states=prompt_embeds[-1:],
added_cond_kwargs={
"text_embeds": add_text_embeds[-1:],
"time_ids": add_time_ids[-1:],
},
cross_attention_kwargs={
"cache_hidden_states": ref_hidden_states,
"use_mv": False,
"use_ref": False,
},
return_dict=False,
)
ref_hidden_states = {
k: v.repeat_interleave(num_images_per_prompt, dim=0)
for k, v in ref_hidden_states.items()
}
if self.do_classifier_free_guidance:
ref_hidden_states = {
k: torch.cat([torch.zeros_like(v), v], dim=0)
for k, v in ref_hidden_states.items()
}
cross_attention_kwargs = {
"mv_scale": mv_scale,
"ref_hidden_states": {k: v.clone() for k, v in ref_hidden_states.items()},
"ref_scale": reference_conditioning_scale,
"num_views": num_images_per_prompt,
**(self.cross_attention_kwargs or {}),
}
# Preprocess control image
control_image_feature = self.prepare_control_image(
image=control_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
control_image_feature = control_image_feature.to(
device=device, dtype=latents.dtype
)
adapter_state = self.cond_encoder(control_image_feature)
for i, state in enumerate(adapter_state):
adapter_state[i] = state * control_conditioning_scale
# 8. Denoising loop
num_warmup_steps = max(
len(timesteps) - num_inference_steps * self.scheduler.order, 0
)
# 8.1 Apply denoising_end
if (
self.denoising_end is not None
and isinstance(self.denoising_end, float)
and self.denoising_end > 0
and self.denoising_end < 1
):
discrete_timestep_cutoff = int(
round(
self.scheduler.config.num_train_timesteps
- (self.denoising_end * self.scheduler.config.num_train_timesteps)
)
)
num_inference_steps = len(
list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
)
timesteps = timesteps[:num_inference_steps]
# 9. Optionally get Guidance Scale Embedding
timestep_cond = None
if self.unet.config.time_cond_proj_dim is not None:
guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
batch_size * num_images_per_prompt
)
timestep_cond = self.get_guidance_scale_embedding(
guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
).to(device=device, dtype=latents.dtype)
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
# expand the latents if we are doing classifier free guidance
latent_model_input = (
torch.cat([latents] * 2)
if self.do_classifier_free_guidance
else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
added_cond_kwargs = {
"text_embeds": add_text_embeds,
"time_ids": add_time_ids,
}
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
added_cond_kwargs["image_embeds"] = image_embeds
if i < int(num_inference_steps * control_conditioning_factor):
down_intrablock_additional_residuals = [
state.clone() for state in adapter_state
]
else:
down_intrablock_additional_residuals = None
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs=cross_attention_kwargs,
down_intrablock_additional_residuals=down_intrablock_additional_residuals,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (
noise_pred_text - noise_pred_uncond
)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(
noise_pred,
noise_pred_text,
guidance_rescale=self.guidance_rescale,
)
# compute the previous noisy sample x_t -> x_t-1
latents_dtype = latents.dtype
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
)[0]
if latents.dtype != latents_dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
latents = latents.to(latents_dtype)
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop(
"negative_prompt_embeds", negative_prompt_embeds
)
add_text_embeds = callback_outputs.pop(
"add_text_embeds", add_text_embeds
)
negative_pooled_prompt_embeds = callback_outputs.pop(
"negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
)
add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
negative_add_time_ids = callback_outputs.pop(
"negative_add_time_ids", negative_add_time_ids
)
# call the callback, if provided
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent":
# make sure the VAE is in float32 mode, as it overflows in float16
needs_upcasting = (
self.vae.dtype == torch.float16 and self.vae.config.force_upcast
)
if needs_upcasting:
self.upcast_vae()
latents = latents.to(
next(iter(self.vae.post_quant_conv.parameters())).dtype
)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
self.vae = self.vae.to(latents.dtype)
# unscale/denormalize the latents
# denormalize with the mean and std if available and not None
has_latents_mean = (
hasattr(self.vae.config, "latents_mean")
and self.vae.config.latents_mean is not None
)
has_latents_std = (
hasattr(self.vae.config, "latents_std")
and self.vae.config.latents_std is not None
)
if has_latents_mean and has_latents_std:
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, 4, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(self.vae.config.latents_std)
.view(1, 4, 1, 1)
.to(latents.device, latents.dtype)
)
latents = (
latents * latents_std / self.vae.config.scaling_factor
+ latents_mean
)
else:
latents = latents / self.vae.config.scaling_factor
image = self.vae.decode(latents, return_dict=False)[0]
# cast back to fp16 if needed
if needs_upcasting:
self.vae.to(dtype=torch.float16)
else:
image = latents
if not output_type == "latent":
# apply watermark if available
if self.watermark is not None:
image = self.watermark.apply_watermark(image)
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)
return StableDiffusionXLPipelineOutput(images=image)
### NEW: adapters ###
def _init_custom_adapter(
self,
# Multi-view adapter
num_views: int = 1,
self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
# Condition encoder
cond_in_channels: int = 6,
# For training
copy_attn_weights: bool = True,
zero_init_module_keys: List[str] = [],
):
# Condition encoder
self.cond_encoder = T2IAdapter(
in_channels=cond_in_channels,
channels=(320, 640, 1280, 1280),
num_res_blocks=2,
downscale_factor=16,
adapter_type="full_adapter_xl",
)
# set custom attn processor for multi-view attention and image cross-attention
self.unet: UNet2DConditionModel
set_unet_2d_condition_attn_processor(
self.unet,
set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=True,
use_ref=True,
),
set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=False,
use_ref=False,
),
)
# copy decoupled attention weights from original unet
if copy_attn_weights:
state_dict = self.unet.state_dict()
for key in state_dict.keys():
if "_mv" in key:
compatible_key = key.replace("_mv", "").replace("processor.", "")
elif "_ref" in key:
compatible_key = key.replace("_ref", "").replace("processor.", "")
else:
compatible_key = key
is_zero_init_key = any([k in key for k in zero_init_module_keys])
if is_zero_init_key:
state_dict[key] = torch.zeros_like(state_dict[compatible_key])
else:
state_dict[key] = state_dict[compatible_key].clone()
self.unet.load_state_dict(state_dict)
def _load_custom_adapter(self, state_dict):
self.unet.load_state_dict(state_dict, strict=False)
self.cond_encoder.load_state_dict(state_dict, strict=False)
def _save_custom_adapter(
self,
include_keys: Optional[List[str]] = None,
exclude_keys: Optional[List[str]] = None,
):
def include_fn(k):
is_included = False
if include_keys is not None:
is_included = is_included or any([key in k for key in include_keys])
if exclude_keys is not None:
is_included = is_included and not any(
[key in k for key in exclude_keys]
)
return is_included
state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
state_dict.update(self.cond_encoder.state_dict())
return state_dict
Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_t2mv_sd.py
+634
View File
@@ -0,0 +1,634 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from typing import Any, Callable, Dict, List, Optional, Union
import torch
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.models import AutoencoderKL, T2IAdapter, UNet2DConditionModel
from diffusers.pipelines.stable_diffusion.pipeline_output import (
StableDiffusionPipelineOutput,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
StableDiffusionPipeline,
rescale_noise_cfg,
retrieve_timesteps,
)
from diffusers.pipelines.stable_diffusion.safety_checker import (
StableDiffusionSafetyChecker,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import deprecate, is_torch_xla_available, logging
from diffusers.utils.torch_utils import randn_tensor
from packaging import version
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTokenizer,
CLIPVisionModelWithProjection,
)
from ..loaders import CustomAdapterMixin
from ..models.attention_processor import (
DecoupledMVRowSelfAttnProcessor2_0,
set_unet_2d_condition_attn_processor,
)
if is_torch_xla_available():
import torch_xla.core.xla_model as xm
XLA_AVAILABLE = True
else:
XLA_AVAILABLE = False
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class MVAdapterT2MVSDPipeline(StableDiffusionPipeline, CustomAdapterMixin):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
image_encoder: CLIPVisionModelWithProjection = None,
requires_safety_checker: bool = True,
):
super().__init__(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
image_encoder=image_encoder,
requires_safety_checker=requires_safety_checker,
)
self.control_image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor,
do_convert_rgb=True,
do_normalize=False,
)
def prepare_control_image(
self,
image,
width,
height,
batch_size,
num_images_per_prompt,
device,
dtype,
do_classifier_free_guidance=False,
):
assert hasattr(
self, "control_image_processor"
), "control_image_processor is not initialized"
image = self.control_image_processor.preprocess(
image, height=height, width=width
).to(dtype=torch.float32)
image_batch_size = image.shape[0]
if image_batch_size == 1:
repeat_by = batch_size
else:
# image batch size is the same as prompt batch size
repeat_by = num_images_per_prompt # always 1 for control image
image = image.repeat_interleave(repeat_by, dim=0)
image = image.to(device=device, dtype=dtype)
if do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
timesteps: List[int] = None,
sigmas: List[float] = None,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[
Union[
Callable[[int, int, Dict], None],
PipelineCallback,
MultiPipelineCallbacks,
]
] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# NEW
mv_scale: float = 1.0,
# Camera or geometry condition
control_image: Optional[PipelineImageInput] = None,
control_conditioning_scale: Optional[float] = 1.0,
control_conditioning_factor: float = 1.0,
# Optional. controlnet
controlnet_image: Optional[PipelineImageInput] = None,
controlnet_conditioning_scale: Optional[float] = 1.0,
**kwargs,
):
r"""
The call function to the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
The width in pixels of the generated image.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
sigmas (`List[float]`, *optional*):
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
will be used.
guidance_scale (`float`, *optional*, defaults to 7.5):
A higher guidance scale value encourages the model to generate images closely linked to the text
`prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide what to not include in image generation. If not defined, you need to
pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
latents (`torch.Tensor`, *optional*):
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor is generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
provided, text embeddings are generated from the `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
[`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
guidance_rescale (`float`, *optional*, defaults to 0.0):
Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
using zero terminal SNR.
clip_skip (`int`, *optional*):
Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
the output of the pre-final layer will be used for computing the prompt embeddings.
callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
callback_on_step_end_tensor_inputs (`List`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
Examples:
Returns:
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
otherwise a `tuple` is returned where the first element is a list with the generated images and the
second element is a list of `bool`s indicating whether the corresponding generated image contains
"not-safe-for-work" (nsfw) content.
"""
callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", None)
if callback is not None:
deprecate(
"callback",
"1.0.0",
"Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
)
if callback_steps is not None:
deprecate(
"callback_steps",
"1.0.0",
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
)
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
# to deal with lora scaling and other possible forward hooks
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
height,
width,
callback_steps,
negative_prompt,
prompt_embeds,
negative_prompt_embeds,
ip_adapter_image,
ip_adapter_image_embeds,
callback_on_step_end_tensor_inputs,
)
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._interrupt = False
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# 3. Encode input prompt
lora_scale = (
self.cross_attention_kwargs.get("scale", None)
if self.cross_attention_kwargs is not None
else None
)
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt,
device,
num_images_per_prompt,
self.do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
)
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
image_embeds = self.prepare_ip_adapter_image_embeds(
ip_adapter_image,
ip_adapter_image_embeds,
device,
batch_size * num_images_per_prompt,
self.do_classifier_free_guidance,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device, timesteps, sigmas
)
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 6.1 Add image embeds for IP-Adapter
added_cond_kwargs = (
{"image_embeds": image_embeds}
if (ip_adapter_image is not None or ip_adapter_image_embeds is not None)
else None
)
# 6.2 Optionally get Guidance Scale Embedding
timestep_cond = None
if self.unet.config.time_cond_proj_dim is not None:
guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
batch_size * num_images_per_prompt
)
timestep_cond = self.get_guidance_scale_embedding(
guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
).to(device=device, dtype=latents.dtype)
# Preprocess control image
control_image_feature = self.prepare_control_image(
image=control_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
control_image_feature = control_image_feature.to(
device=device, dtype=latents.dtype
)
adapter_state = self.cond_encoder(control_image_feature)
for i, state in enumerate(adapter_state):
adapter_state[i] = state * control_conditioning_scale
# Preprocess controlnet image if provided
do_controlnet = controlnet_image is not None and hasattr(self, "controlnet")
if do_controlnet:
controlnet_image = self.prepare_control_image(
image=controlnet_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
controlnet_image = controlnet_image.to(device=device, dtype=latents.dtype)
# 7. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
# expand the latents if we are doing classifier free guidance
latent_model_input = (
torch.cat([latents] * 2)
if self.do_classifier_free_guidance
else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
if i < int(num_inference_steps * control_conditioning_factor):
down_intrablock_additional_residuals = [
state.clone() for state in adapter_state
]
else:
down_intrablock_additional_residuals = None
unet_add_kwargs = {}
# Do controlnet if provided
if do_controlnet:
down_block_res_samples, mid_block_res_sample = self.controlnet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
controlnet_cond=controlnet_image,
conditioning_scale=controlnet_conditioning_scale,
guess_mode=False,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)
unet_add_kwargs.update(
{
"down_block_additional_residuals": down_block_res_samples,
"mid_block_additional_residual": mid_block_res_sample,
}
)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs={
"mv_scale": mv_scale,
"num_views": num_images_per_prompt,
**(self.cross_attention_kwargs or {}),
},
down_intrablock_additional_residuals=down_intrablock_additional_residuals,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
**unet_add_kwargs,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (
noise_pred_text - noise_pred_uncond
)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(
noise_pred,
noise_pred_text,
guidance_rescale=self.guidance_rescale,
)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop(
"negative_prompt_embeds", negative_prompt_embeds
)
# call the callback, if provided
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if XLA_AVAILABLE:
xm.mark_step()
if not output_type == "latent":
image = self.vae.decode(
latents / self.vae.config.scaling_factor,
return_dict=False,
generator=generator,
)[0]
image, has_nsfw_concept = self.run_safety_checker(
image, device, prompt_embeds.dtype
)
else:
image = latents
has_nsfw_concept = None
if has_nsfw_concept is None:
do_denormalize = [True] * image.shape[0]
else:
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
image = self.image_processor.postprocess(
image, output_type=output_type, do_denormalize=do_denormalize
)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(
images=image, nsfw_content_detected=has_nsfw_concept
)
### NEW: adapters ###
def _init_custom_adapter(
self,
# Multi-view adapter
num_views: int = 1,
self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
# Condition encoder
cond_in_channels: int = 6,
# For training
copy_attn_weights: bool = True,
zero_init_module_keys: List[str] = [],
):
# Condition encoder
self.cond_encoder = T2IAdapter(
in_channels=cond_in_channels,
channels=self.unet.config.block_out_channels,
num_res_blocks=self.unet.config.layers_per_block,
downscale_factor=8,
)
# set custom attn processor for multi-view attention
self.unet: UNet2DConditionModel
set_unet_2d_condition_attn_processor(
self.unet,
set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=True,
use_ref=False,
),
set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=False,
use_ref=False,
),
)
# copy decoupled attention weights from original unet
if copy_attn_weights:
state_dict = self.unet.state_dict()
for key in state_dict.keys():
if "_mv" in key:
compatible_key = key.replace("_mv", "").replace("processor.", "")
else:
compatible_key = key
is_zero_init_key = any([k in key for k in zero_init_module_keys])
if is_zero_init_key:
state_dict[key] = torch.zeros_like(state_dict[compatible_key])
else:
state_dict[key] = state_dict[compatible_key].clone()
self.unet.load_state_dict(state_dict)
def _load_custom_adapter(self, state_dict):
self.unet.load_state_dict(state_dict, strict=False)
self.cond_encoder.load_state_dict(state_dict, strict=False)
def _save_custom_adapter(
self,
include_keys: Optional[List[str]] = None,
exclude_keys: Optional[List[str]] = None,
):
def include_fn(k):
is_included = False
if include_keys is not None:
is_included = is_included or any([key in k for key in include_keys])
if exclude_keys is not None:
is_included = is_included and not any(
[key in k for key in exclude_keys]
)
return is_included
state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
state_dict.update(self.cond_encoder.state_dict())
return state_dict
Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_mvadapter_t2mv_sdxl.py
+801
View File
@@ -0,0 +1,801 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import torch
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.models import AutoencoderKL, T2IAdapter, UNet2DConditionModel
from diffusers.pipelines.stable_diffusion_xl.pipeline_output import (
StableDiffusionXLPipelineOutput,
)
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
StableDiffusionXLPipeline,
rescale_noise_cfg,
retrieve_timesteps,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import deprecate, logging
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTextModelWithProjection,
CLIPTokenizer,
CLIPVisionModelWithProjection,
)
from ..loaders import CustomAdapterMixin
from ..models.attention_processor import (
DecoupledMVRowSelfAttnProcessor2_0,
set_unet_2d_condition_attn_processor,
)
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class MVAdapterT2MVSDXLPipeline(StableDiffusionXLPipeline, CustomAdapterMixin):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
text_encoder_2: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
image_encoder: CLIPVisionModelWithProjection = None,
feature_extractor: CLIPImageProcessor = None,
force_zeros_for_empty_prompt: bool = True,
add_watermarker: Optional[bool] = None,
):
super().__init__(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
unet=unet,
scheduler=scheduler,
image_encoder=image_encoder,
feature_extractor=feature_extractor,
force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
add_watermarker=add_watermarker,
)
self.control_image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor,
do_convert_rgb=True,
do_normalize=False,
)
def prepare_control_image(
self,
image,
width,
height,
batch_size,
num_images_per_prompt,
device,
dtype,
do_classifier_free_guidance=False,
):
assert hasattr(
self, "control_image_processor"
), "control_image_processor is not initialized"
image = self.control_image_processor.preprocess(
image, height=height, width=width
).to(dtype=torch.float32)
image_batch_size = image.shape[0]
if image_batch_size == 1:
repeat_by = batch_size
else:
# image batch size is the same as prompt batch size
repeat_by = num_images_per_prompt # always 1 for control image
image = image.repeat_interleave(repeat_by, dim=0)
image = image.to(device=device, dtype=dtype)
if do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt_2: Optional[Union[str, List[str]]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
timesteps: List[int] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 5.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
ip_adapter_image_embeds: Optional[List[torch.FloatTensor]] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
guidance_rescale: float = 0.0,
original_size: Optional[Tuple[int, int]] = None,
crops_coords_top_left: Tuple[int, int] = (0, 0),
target_size: Optional[Tuple[int, int]] = None,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
clip_skip: Optional[int] = None,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
# NEW
mv_scale: float = 1.0,
# Camera or geometry condition
control_image: Optional[PipelineImageInput] = None,
control_conditioning_scale: Optional[float] = 1.0,
control_conditioning_factor: float = 1.0,
# Optional. controlnet
controlnet_image: Optional[PipelineImageInput] = None,
controlnet_conditioning_scale: Optional[float] = 1.0,
**kwargs,
):
r"""
Function invoked when calling the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
used in both text-encoders
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The width in pixels of the generated image. This is set to 1024 by default for the best results.
Anything below 512 pixels won't work well for
[stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
and checkpoints that are not specifically fine-tuned on low resolutions.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
passed will be used. Must be in descending order.
denoising_end (`float`, *optional*):
When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
completed before it is intentionally prematurely terminated. As a result, the returned sample will
still retain a substantial amount of noise as determined by the discrete timesteps selected by the
scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
"Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
guidance_scale (`float`, *optional*, defaults to 5.0):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
negative_prompt_2 (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
`text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will ge generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
If not provided, pooled text embeddings will be generated from `prompt` input argument.
negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
input argument.
ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
ip_adapter_image_embeds (`List[torch.FloatTensor]`, *optional*):
Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
provided, embeddings are computed from the `ip_adapter_image` input argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generate image. Choose between
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
of a plain tuple.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
guidance_rescale (`float`, *optional*, defaults to 0.0):
Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
[Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
Guidance rescale factor should fix overexposure when using zero terminal SNR.
original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
`crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
`crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
For most cases, `target_size` should be set to the desired height and width of the generated image. If
not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a specific image resolution. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's
micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
To negatively condition the generation process based on a target image resolution. It should be as same
as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of
[https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more
information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
callback_on_step_end_tensor_inputs (`List`, *optional*):
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
`._callback_tensor_inputs` attribute of your pipeline class.
Examples:
Returns:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
[`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
`tuple`. When returning a tuple, the first element is a list with the generated images.
"""
callback = kwargs.pop("callback", None)
callback_steps = kwargs.pop("callback_steps", None)
if callback is not None:
deprecate(
"callback",
"1.0.0",
"Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
if callback_steps is not None:
deprecate(
"callback_steps",
"1.0.0",
"Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
)
# 0. Default height and width to unet
height = height or self.default_sample_size * self.vae_scale_factor
width = width or self.default_sample_size * self.vae_scale_factor
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
prompt_2,
height,
width,
callback_steps,
negative_prompt,
negative_prompt_2,
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
ip_adapter_image,
ip_adapter_image_embeds,
callback_on_step_end_tensor_inputs,
)
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._denoising_end = denoising_end
self._interrupt = False
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# 3. Encode input prompt
lora_scale = (
self.cross_attention_kwargs.get("scale", None)
if self.cross_attention_kwargs is not None
else None
)
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt=prompt,
prompt_2=prompt_2,
device=device,
num_images_per_prompt=num_images_per_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
negative_prompt=negative_prompt,
negative_prompt_2=negative_prompt_2,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
lora_scale=lora_scale,
clip_skip=self.clip_skip,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device, timesteps
)
# 5. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Prepare added time ids & embeddings
add_text_embeds = pooled_prompt_embeds
if self.text_encoder_2 is None:
text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
else:
text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
add_time_ids = self._get_add_time_ids(
original_size,
crops_coords_top_left,
target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
if negative_original_size is not None and negative_target_size is not None:
negative_add_time_ids = self._get_add_time_ids(
negative_original_size,
negative_crops_coords_top_left,
negative_target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
else:
negative_add_time_ids = add_time_ids
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
add_text_embeds = torch.cat(
[negative_pooled_prompt_embeds, add_text_embeds], dim=0
)
add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
prompt_embeds = prompt_embeds.to(device)
add_text_embeds = add_text_embeds.to(device)
add_time_ids = add_time_ids.to(device).repeat(
batch_size * num_images_per_prompt, 1
)
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
image_embeds = self.prepare_ip_adapter_image_embeds(
ip_adapter_image,
ip_adapter_image_embeds,
device,
batch_size * num_images_per_prompt,
self.do_classifier_free_guidance,
)
# Preprocess control image
control_image_feature = self.prepare_control_image(
image=control_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
control_image_feature = control_image_feature.to(
device=device, dtype=latents.dtype
)
adapter_state = self.cond_encoder(control_image_feature)
for i, state in enumerate(adapter_state):
adapter_state[i] = state * control_conditioning_scale
# Preprocess controlnet image if provided
do_controlnet = controlnet_image is not None and hasattr(self, "controlnet")
if do_controlnet:
controlnet_image = self.prepare_control_image(
image=controlnet_image,
width=width,
height=height,
batch_size=batch_size * num_images_per_prompt,
num_images_per_prompt=1, # NOTE: always 1 for control images
device=device,
dtype=latents.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
controlnet_image = controlnet_image.to(device=device, dtype=latents.dtype)
# 8. Denoising loop
num_warmup_steps = max(
len(timesteps) - num_inference_steps * self.scheduler.order, 0
)
# 8.1 Apply denoising_end
if (
self.denoising_end is not None
and isinstance(self.denoising_end, float)
and self.denoising_end > 0
and self.denoising_end < 1
):
discrete_timestep_cutoff = int(
round(
self.scheduler.config.num_train_timesteps
- (self.denoising_end * self.scheduler.config.num_train_timesteps)
)
)
num_inference_steps = len(
list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
)
timesteps = timesteps[:num_inference_steps]
# 9. Optionally get Guidance Scale Embedding
timestep_cond = None
if self.unet.config.time_cond_proj_dim is not None:
guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
batch_size * num_images_per_prompt
)
timestep_cond = self.get_guidance_scale_embedding(
guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
).to(device=device, dtype=latents.dtype)
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
# expand the latents if we are doing classifier free guidance
latent_model_input = (
torch.cat([latents] * 2)
if self.do_classifier_free_guidance
else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
added_cond_kwargs = {
"text_embeds": add_text_embeds,
"time_ids": add_time_ids,
}
if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
added_cond_kwargs["image_embeds"] = image_embeds
if i < int(num_inference_steps * control_conditioning_factor):
down_intrablock_additional_residuals = [
state.clone() for state in adapter_state
]
else:
down_intrablock_additional_residuals = None
unet_add_kwargs = {}
# Do controlnet if provided
if do_controlnet:
down_block_res_samples, mid_block_res_sample = self.controlnet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
controlnet_cond=controlnet_image,
conditioning_scale=controlnet_conditioning_scale,
guess_mode=False,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)
unet_add_kwargs.update(
{
"down_block_additional_residuals": down_block_res_samples,
"mid_block_additional_residual": mid_block_res_sample,
}
)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=timestep_cond,
cross_attention_kwargs={
"mv_scale": mv_scale,
"num_views": num_images_per_prompt,
**(self.cross_attention_kwargs or {}),
},
down_intrablock_additional_residuals=down_intrablock_additional_residuals,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
**unet_add_kwargs,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (
noise_pred_text - noise_pred_uncond
)
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
noise_pred = rescale_noise_cfg(
noise_pred,
noise_pred_text,
guidance_rescale=self.guidance_rescale,
)
# compute the previous noisy sample x_t -> x_t-1
latents_dtype = latents.dtype
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
)[0]
if latents.dtype != latents_dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
latents = latents.to(latents_dtype)
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop(
"negative_prompt_embeds", negative_prompt_embeds
)
add_text_embeds = callback_outputs.pop(
"add_text_embeds", add_text_embeds
)
negative_pooled_prompt_embeds = callback_outputs.pop(
"negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
)
add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
negative_add_time_ids = callback_outputs.pop(
"negative_add_time_ids", negative_add_time_ids
)
# call the callback, if provided
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if not output_type == "latent":
# make sure the VAE is in float32 mode, as it overflows in float16
needs_upcasting = (
self.vae.dtype == torch.float16 and self.vae.config.force_upcast
)
if needs_upcasting:
self.upcast_vae()
latents = latents.to(
next(iter(self.vae.post_quant_conv.parameters())).dtype
)
elif latents.dtype != self.vae.dtype:
if torch.backends.mps.is_available():
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
self.vae = self.vae.to(latents.dtype)
# unscale/denormalize the latents
# denormalize with the mean and std if available and not None
has_latents_mean = (
hasattr(self.vae.config, "latents_mean")
and self.vae.config.latents_mean is not None
)
has_latents_std = (
hasattr(self.vae.config, "latents_std")
and self.vae.config.latents_std is not None
)
if has_latents_mean and has_latents_std:
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, 4, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(self.vae.config.latents_std)
.view(1, 4, 1, 1)
.to(latents.device, latents.dtype)
)
latents = (
latents * latents_std / self.vae.config.scaling_factor
+ latents_mean
)
else:
latents = latents / self.vae.config.scaling_factor
image = self.vae.decode(latents, return_dict=False)[0]
# cast back to fp16 if needed
if needs_upcasting:
self.vae.to(dtype=torch.float16)
else:
image = latents
if not output_type == "latent":
# apply watermark if available
if self.watermark is not None:
image = self.watermark.apply_watermark(image)
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)
return StableDiffusionXLPipelineOutput(images=image)
### NEW: adapters ###
def _init_custom_adapter(
self,
# Multi-view adapter
num_views: int = 1,
self_attn_processor: Any = DecoupledMVRowSelfAttnProcessor2_0,
# Condition encoder
cond_in_channels: int = 6,
# For training
copy_attn_weights: bool = True,
zero_init_module_keys: List[str] = [],
):
# Condition encoder
self.cond_encoder = T2IAdapter(
in_channels=cond_in_channels,
channels=(320, 640, 1280, 1280),
num_res_blocks=2,
downscale_factor=16,
adapter_type="full_adapter_xl",
)
# set custom attn processor for multi-view attention
self.unet: UNet2DConditionModel
set_unet_2d_condition_attn_processor(
self.unet,
set_self_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=True,
use_ref=False,
),
set_cross_attn_proc_func=lambda name, hs, cad, ap: self_attn_processor(
query_dim=hs,
inner_dim=hs,
num_views=num_views,
name=name,
use_mv=False,
use_ref=False,
),
)
# copy decoupled attention weights from original unet
if copy_attn_weights:
state_dict = self.unet.state_dict()
for key in state_dict.keys():
if "_mv" in key:
compatible_key = key.replace("_mv", "").replace("processor.", "")
else:
compatible_key = key
is_zero_init_key = any([k in key for k in zero_init_module_keys])
if is_zero_init_key:
state_dict[key] = torch.zeros_like(state_dict[compatible_key])
else:
state_dict[key] = state_dict[compatible_key].clone()
self.unet.load_state_dict(state_dict)
def _load_custom_adapter(self, state_dict):
self.unet.load_state_dict(state_dict, strict=False)
self.cond_encoder.load_state_dict(state_dict, strict=False)
def _save_custom_adapter(
self,
include_keys: Optional[List[str]] = None,
exclude_keys: Optional[List[str]] = None,
):
def include_fn(k):
is_included = False
if include_keys is not None:
is_included = is_included or any([key in k for key in include_keys])
if exclude_keys is not None:
is_included = is_included and not any(
[key in k for key in exclude_keys]
)
return is_included
state_dict = {k: v for k, v in self.unet.state_dict().items() if include_fn(k)}
state_dict.update(self.cond_encoder.state_dict())
return state_dict
Gen_3D_Modules/MV_Adapter/mvadapter/pipelines/pipeline_texture.py
+428
View File
@@ -0,0 +1,428 @@
import os
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from spandrel import ModelLoader
from mvadapter.utils import image_to_tensor, make_image_grid, tensor_to_image
from mvadapter.utils.mesh_utils import (
Camera,
CameraProjection,
NVDiffRastContextWrapper,
SmartPainter,
TexturedMesh,
get_camera,
get_orthogonal_camera,
load_mesh,
render,
replace_mesh_texture_and_save,
)
from mvadapter.utils.mesh_utils.mesh_process import process_raw
def clear():
torch.cuda.empty_cache()
@contextmanager
def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
texture_ = mesh.texture
mesh.texture = texture
try:
yield
finally:
mesh.texture = texture_
@dataclass
class ModProcessConfig:
view_upscale: bool = False
view_upscale_factor: int = 2
inpaint_mode: str = "uv" # in ["none", "uv", "view"]
view_inpaint_max_view_score_thresh: float = 0.02
view_inpaint_min_rounds: int = 4
view_inpaint_max_rounds: int = 8
view_inpaint_uv_padding_end: bool = True
@dataclass
class TexturePipelineOutput:
shaded_model_save_path: Optional[str] = None
pbr_model_save_path: Optional[str] = None
class TexturePipeline:
def __init__(self, upscaler_ckpt_path: str, inpaint_ckpt_path: str, device: str):
self.device = device
self.ctx = NVDiffRastContextWrapper(device=self.device)
self.cam_proj = CameraProjection(
pb_backend="torch-cuda", bg_remover=None, device=self.device
)
if upscaler_ckpt_path is not None:
self.upscaler = ModelLoader().load_from_file(upscaler_ckpt_path)
self.upscaler.to(self.device).eval().half()
if inpaint_ckpt_path is not None:
self.inpainter = ModelLoader().load_from_file(inpaint_ckpt_path)
self.inpainter.to(self.device).eval()
self.smart_painter = SmartPainter(self.device)
def load_packed_images(self, packed_image_path: Optional[str]) -> List[Image.Image]:
if packed_image_path is None:
return None
packed_image = Image.open(packed_image_path)
images = np.array_split(np.array(packed_image), 6, axis=1)
images = [Image.fromarray(im) for im in images]
return images
def maybe_upscale_image(
self,
tensor: Optional[torch.FloatTensor],
upscale: bool,
upscale_factor: int,
batched: bool = False,
) -> Optional[torch.FloatTensor]:
if upscale:
with torch.no_grad():
tensor = tensor.permute(0, 3, 1, 2)
if batched:
tensor = self.upscaler(tensor.half()).float()
else:
tensor = torch.concat(
[
self.upscaler(im.unsqueeze(0).half()).float()
for im in tensor
],
dim=0,
)
tensor = tensor.clamp(0, 1).permute(0, 2, 3, 1)
clear()
return tensor
def view_inpaint(
self,
mod_name: str,
mesh: TexturedMesh,
uv_proj: torch.FloatTensor,
uv_valid_mask: torch.BoolTensor,
config: ModProcessConfig,
debug_dir: Optional[str] = None,
) -> Tuple[torch.FloatTensor, torch.BoolTensor]:
def inpaint_func(
image: torch.FloatTensor, mask: torch.FloatTensor
) -> torch.FloatTensor:
with torch.no_grad():
return self.inpainter(image.to(torch.float32), mask.to(torch.float32))
return self.smart_painter(
mod_name,
mesh,
inpaint_func,
uv_proj,
~uv_valid_mask,
max_view_score_thresh=config.view_inpaint_max_view_score_thresh,
min_rounds=config.view_inpaint_min_rounds,
max_rounds=config.view_inpaint_max_rounds,
uv_padding_end=config.view_inpaint_uv_padding_end,
debug_dir=debug_dir,
debug_visualize_details=False,
)
def __call__(
self,
mesh_path: str,
save_dir: str,
save_name: str = "default",
# mesh load settings
move_to_center: bool = False,
front_x: bool = True,
# uv unwarp
uv_unwarp: bool = False,
preprocess_mesh: bool = False,
# projection
uv_size: int = 4096,
# modes
rgb_path: Optional[str] = None,
rgb_process_config: ModProcessConfig = ModProcessConfig(),
base_color_path: Optional[str] = None,
base_color_process_config: ModProcessConfig = ModProcessConfig(),
orm_path: Optional[str] = None,
orm_process_config: ModProcessConfig = ModProcessConfig(),
normal_path: Optional[str] = None,
normal_strength: float = 1.0,
normal_process_config: ModProcessConfig = ModProcessConfig(),
# inpaint
uv_inpaint_use_network: bool = False,
view_inpaint_include_occlusion_boundary: bool = False,
poisson_reprojection: bool = False,
# camera
camera_projection_type: str = "ORTHO",
camera_elevation_deg: List[float] = [0, 0, 0, 0, 89.99, -89.99],
camera_azimuth_deg: List[float] = [0, 90, 180, 270, 180, 180],
camera_distance: float = 1.0,
camera_ortho_scale: float = 1.1,
camera_fov_deg: float = 40,
# debug
debug_mode: bool = False,
):
clear()
if debug_mode:
debug_dir = os.path.join(save_dir, "debug")
os.makedirs(debug_dir, exist_ok=True)
if uv_unwarp:
file_suffix = os.path.splitext(mesh_path)[-1]
mesh_path_new = mesh_path.replace(file_suffix, f"_unwarp{file_suffix}")
process_raw(mesh_path, mesh_path_new, preprocess=preprocess_mesh)
mesh_path = mesh_path_new
mesh: TexturedMesh = load_mesh(
mesh_path,
rescale=True,
move_to_center=move_to_center,
front_x_to_y=front_x,
default_uv_size=uv_size,
device=self.device,
)
# projection
if camera_projection_type == "PERSP":
raise NotImplementedError
elif camera_projection_type == "ORTHO":
cameras = get_orthogonal_camera(
elevation_deg=camera_elevation_deg,
distance=[camera_distance] * 6,
left=-camera_ortho_scale / 2,
right=camera_ortho_scale / 2,
bottom=-camera_ortho_scale / 2,
top=camera_ortho_scale / 2,
azimuth_deg=[x - 90 for x in camera_azimuth_deg], # -y as front
device=self.device,
)
mod_kwargs = {
"rgb": (rgb_path, rgb_process_config),
"base_color": (base_color_path, base_color_process_config),
"orm": (orm_path, orm_process_config),
"normal": (normal_path, normal_process_config),
}
mod_uv_image, mod_uv_tensor = {}, {}
for mod_name, (mod_path, mod_process_config) in mod_kwargs.items():
if mod_path is None:
mod_uv_image[mod_name] = None
mod_uv_tensor[mod_name] = None
continue
mod_images = self.load_packed_images(mod_path)
mod_tensor = image_to_tensor(mod_images, device=self.device)
mod_tensor = self.maybe_upscale_image(
mod_tensor,
mod_process_config.view_upscale,
mod_process_config.view_upscale_factor,
)
if mod_process_config.view_upscale and debug_mode:
make_image_grid(tensor_to_image(mod_tensor, batched=True), rows=1).save(
os.path.join(debug_dir, f"{mod_name}_upscaled.jpg")
)
if mod_name == "normal":
_, height, width, _ = mod_tensor.shape
render_out = render(
self.ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
render_depth=False,
render_normal=True,
render_tangent=True,
)
# compute UV tangent space
vN = render_out.normal
vT = render_out.tangent
vB = torch.cross(vN, vT, dim=-1)
tang_space = F.normalize(torch.stack([vT, vB, vN], dim=-2), dim=-1)
# compute geometry tangent space
vGN = vN
vGT = torch.as_tensor(
[
[1, 0, 0],
[0, 1, 0],
[-1, 0, 0],
[0, -1, 0],
[-1, 0, 0],
[-1, 0, 0],
],
dtype=torch.float32,
device=self.device,
)[:, None, None, :]
vGB = torch.cross(vGN, vGT, dim=-1)
vGT = torch.cross(vGB, vGN, dim=-1)
geo_tang_space = F.normalize(
torch.stack([vGT, vGB, vGN], dim=-2), dim=-1
)
# restore world-space normal from geometry tangent space
mod_tensor = mod_tensor * 2 - 1
mod_tensor = F.normalize(
(
mod_tensor[:, :, :, None, :]
* geo_tang_space.permute(0, 1, 2, 4, 3)
).sum(-1),
dim=-1,
)
# bake world-space normal to UV tangent space
mod_tensor = F.normalize(
(mod_tensor[:, :, :, None, :] * tang_space).sum(-1), dim=-1
)
mod_tensor = (mod_tensor * 0.5 + 0.5).clamp(0, 1)
uv_proj, uv_valid_mask = self.cam_proj(
mod_tensor,
mesh,
cameras,
from_scratch=mod_process_config.inpaint_mode != "none",
poisson_blending=False,
depth_grad_dilation=5,
uv_exp_blend_alpha=3,
uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
aoi_cos_valid_threshold=0.2,
uv_size=uv_size,
return_uv_projection_mask=True,
)
uv_proj[~uv_valid_mask] = torch.as_tensor([0.5, 0.5, 1]).to(uv_proj)
else:
# TODO: tweak depth_grad_dilation
cam_proj_out = self.cam_proj(
mod_tensor,
mesh,
cameras,
from_scratch=mod_process_config.inpaint_mode != "none",
poisson_blending=False,
depth_grad_dilation=5,
depth_grad_threshold=0.1,
uv_exp_blend_alpha=3,
uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
aoi_cos_valid_threshold=0.2,
uv_size=uv_size,
uv_padding=not uv_inpaint_use_network,
return_dict=True,
)
uv_proj, uv_valid_mask, uv_depth_grad = (
cam_proj_out.uv_proj,
cam_proj_out.uv_proj_mask,
cam_proj_out.uv_depth_grad,
)
if uv_inpaint_use_network:
uv_inpaint_mask_input = 1 - uv_valid_mask[None, None].float()
uv_inpaint_image_input = uv_proj[None].permute(0, 3, 1, 2)
with torch.no_grad():
uv_inpaint_result = self.inpainter(
uv_inpaint_image_input, uv_inpaint_mask_input
)[0].permute(1, 2, 0)
clear()
if debug_mode:
make_image_grid(
[
tensor_to_image(uv_proj),
tensor_to_image(uv_valid_mask),
tensor_to_image(uv_inpaint_result),
]
).save(os.path.join(debug_dir, f"{mod_name}_uv_inpaint.jpg"))
uv_proj = uv_inpaint_result.contiguous()
if mod_process_config.inpaint_mode == "view":
if view_inpaint_include_occlusion_boundary:
uv_max_depth_grad = uv_depth_grad.max(dim=0)[0]
uv_valid_mask = uv_valid_mask & (uv_max_depth_grad < 0.1)
uv_proj, uv_valid_mask = self.view_inpaint(
mod_name,
mesh,
uv_proj,
uv_valid_mask,
mod_process_config,
debug_dir=debug_dir if debug_mode else None,
)
if poisson_reprojection:
# up and down
mesh.texture = uv_proj
uv_proj = self.cam_proj(
mod_tensor[4:5],
mesh,
cameras[4:5],
from_scratch=False,
poisson_blending=True,
pb_keep_original_border=True,
depth_grad_dilation=5,
uv_exp_blend_alpha=3,
uv_exp_blend_view_weight=torch.as_tensor([1, 1]),
aoi_cos_valid_threshold=0.2,
uv_size=uv_size,
uv_padding=True,
return_dict=False,
)
# front, sides and back
mesh.texture = uv_proj
uv_proj = self.cam_proj(
mod_tensor[0:4],
mesh,
cameras[0:4],
from_scratch=False,
poisson_blending=True,
pb_keep_original_border=True,
depth_grad_dilation=5,
uv_exp_blend_alpha=3,
uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1]),
aoi_cos_valid_threshold=0.2,
uv_size=uv_size,
uv_padding=True,
return_dict=False,
)
if mod_name == "orm":
uv_proj[:, :, 0] = 1.0
mod_uv_image[mod_name] = tensor_to_image(uv_proj)
mod_uv_tensor[mod_name] = uv_proj
clear()
shaded_model_save_path = None
if mod_uv_image["rgb"] is not None:
shaded_model_save_path = os.path.join(save_dir, f"{save_name}_shaded.glb")
replace_mesh_texture_and_save(
mesh_path,
shaded_model_save_path,
texture=mod_uv_image["rgb"],
backend="gltflib",
task_id=save_name,
)
pbr_model_save_path = None
if mod_uv_image["base_color"] is not None:
pbr_model_save_path = os.path.join(save_dir, f"{save_name}_pbr.glb")
replace_mesh_texture_and_save(
mesh_path,
pbr_model_save_path,
texture=mod_uv_image["base_color"],
metallic_roughness_texture=mod_uv_image["orm"],
normal_texture=mod_uv_image["normal"],
normal_strength=normal_strength,
backend="gltflib",
task_id=save_name,
)
clear()
return TexturePipelineOutput(
shaded_model_save_path=shaded_model_save_path,
pbr_model_save_path=pbr_model_save_path,
)
Gen_3D_Modules/MV_Adapter/mvadapter/schedulers/scheduler_utils.py
+70
View File
@@ -0,0 +1,70 @@
import torch
def get_sigmas(noise_scheduler, timesteps, n_dim=4, dtype=torch.float32, device=None):
sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(device)
timesteps = timesteps.to(device)
step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
def SNR_to_betas(snr):
"""
Converts SNR to betas
"""
# alphas_cumprod = pass
# snr = (alpha / ) ** 2
# alpha_t^2 / (1 - alpha_t^2) = snr
alpha_t = (snr / (1 + snr)) ** 0.5
alphas_cumprod = alpha_t**2
alphas = alphas_cumprod / torch.cat(
[torch.ones(1, device=snr.device), alphas_cumprod[:-1]]
)
betas = 1 - alphas
return betas
def compute_snr(timesteps, noise_scheduler):
"""
Computes SNR as per Min-SNR-Diffusion-Training/guided_diffusion/gaussian_diffusion.py at 521b624bd70c67cee4bdf49225915f5
"""
alphas_cumprod = noise_scheduler.alphas_cumprod
sqrt_alphas_cumprod = alphas_cumprod**0.5
sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
# Expand the tensors.
# Adapted from Min-SNR-Diffusion-Training/guided_diffusion/gaussian_diffusion.py at 521b624bd70c67cee4bdf49225915f5
sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
timesteps
].float()
while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(
device=timesteps.device
)[timesteps].float()
while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
# Compute SNR.
snr = (alpha / sigma) ** 2
return snr
def compute_alpha(timesteps, noise_scheduler):
alphas_cumprod = noise_scheduler.alphas_cumprod
sqrt_alphas_cumprod = alphas_cumprod**0.5
sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
timesteps
].float()
while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
return alpha
Gen_3D_Modules/MV_Adapter/mvadapter/schedulers/scheduling_shift_snr.py
+138
View File
@@ -0,0 +1,138 @@
from typing import Any
import torch
from .scheduler_utils import SNR_to_betas, compute_snr
class ShiftSNRScheduler:
def __init__(
self,
noise_scheduler: Any,
timesteps: Any,
shift_scale: float,
scheduler_class: Any,
):
self.noise_scheduler = noise_scheduler
self.timesteps = timesteps
self.shift_scale = shift_scale
self.scheduler_class = scheduler_class
def _get_shift_scheduler(self):
"""
Prepare scheduler for shifted betas.
:return: A scheduler object configured with shifted betas
"""
snr = compute_snr(self.timesteps, self.noise_scheduler)
shifted_betas = SNR_to_betas(snr / self.shift_scale)
return self.scheduler_class.from_config(
self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
)
def _get_interpolated_shift_scheduler(self):
"""
Prepare scheduler for shifted betas and interpolate with the original betas in log space.
:return: A scheduler object configured with interpolated shifted betas
"""
snr = compute_snr(self.timesteps, self.noise_scheduler)
shifted_snr = snr / self.shift_scale
weighting = self.timesteps.float() / (
self.noise_scheduler.config.num_train_timesteps - 1
)
interpolated_snr = torch.exp(
torch.log(snr) * (1 - weighting) + torch.log(shifted_snr) * weighting
)
shifted_betas = SNR_to_betas(interpolated_snr)
return self.scheduler_class.from_config(
self.noise_scheduler.config, trained_betas=shifted_betas.numpy()
)
@classmethod
def from_scheduler(
cls,
noise_scheduler: Any,
shift_mode: str = "default",
timesteps: Any = None,
shift_scale: float = 1.0,
scheduler_class: Any = None,
):
# Check input
if timesteps is None:
timesteps = torch.arange(0, noise_scheduler.config.num_train_timesteps)
if scheduler_class is None:
scheduler_class = noise_scheduler.__class__
# Create scheduler
shift_scheduler = cls(
noise_scheduler=noise_scheduler,
timesteps=timesteps,
shift_scale=shift_scale,
scheduler_class=scheduler_class,
)
if shift_mode == "default":
return shift_scheduler._get_shift_scheduler()
elif shift_mode == "interpolated":
return shift_scheduler._get_interpolated_shift_scheduler()
else:
raise ValueError(f"Unknown shift_mode: {shift_mode}")
if __name__ == "__main__":
"""
Compare the alpha values for different noise schedulers.
"""
import matplotlib.pyplot as plt
from diffusers import DDPMScheduler
from .scheduler_utils import compute_alpha
# Base
timesteps = torch.arange(0, 1000)
noise_scheduler_base = DDPMScheduler.from_pretrained(
"runwayml/stable-diffusion-v1-5", subfolder="scheduler"
)
alpha = compute_alpha(timesteps, noise_scheduler_base)
plt.plot(timesteps.numpy(), alpha.numpy(), label="Base")
# Kolors
num_train_timesteps_ = 1100
timesteps_ = torch.arange(0, num_train_timesteps_)
noise_kwargs = {"beta_end": 0.014, "num_train_timesteps": num_train_timesteps_}
noise_scheduler_kolors = DDPMScheduler.from_config(
noise_scheduler_base.config, **noise_kwargs
)
alpha = compute_alpha(timesteps_, noise_scheduler_kolors)
plt.plot(timesteps_.numpy(), alpha.numpy(), label="Kolors")
# Shift betas
shift_scale = 8.0
noise_scheduler_shift = ShiftSNRScheduler.from_scheduler(
noise_scheduler_base, shift_mode="default", shift_scale=shift_scale
)
alpha = compute_alpha(timesteps, noise_scheduler_shift)
plt.plot(timesteps.numpy(), alpha.numpy(), label="Shift Noise (scale 8.0)")
# Shift betas (interpolated)
noise_scheduler_inter = ShiftSNRScheduler.from_scheduler(
noise_scheduler_base, shift_mode="interpolated", shift_scale=shift_scale
)
alpha = compute_alpha(timesteps, noise_scheduler_inter)
plt.plot(timesteps.numpy(), alpha.numpy(), label="Interpolated (scale 8.0)")
# ZeroSNR
noise_scheduler = DDPMScheduler.from_config(
noise_scheduler_base.config, rescale_betas_zero_snr=True
)
alpha = compute_alpha(timesteps, noise_scheduler)
plt.plot(timesteps.numpy(), alpha.numpy(), label="ZeroSNR")
plt.legend()
plt.grid()
plt.savefig("check_alpha.png")
Gen_3D_Modules/MV_Adapter/mvadapter/systems/__init__.py
View File
Gen_3D_Modules/MV_Adapter/mvadapter/systems/base.py
+262
View File
@@ -0,0 +1,262 @@
import os
from dataclasses import dataclass, field
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from ..utils.base import Updateable, update_end_if_possible, update_if_possible
from ..utils.config import parse_structured
from ..utils.core import debug, find, info, warn
from ..utils.misc import (
C,
cleanup,
get_device,
get_rank,
load_module_weights,
show_vram_usage,
)
from ..utils.saving import SaverMixin
from ..utils.typing import *
from .utils import parse_optimizer, parse_scheduler
class BaseSystem(pl.LightningModule, Updateable, SaverMixin):
@dataclass
class Config:
optimizer: dict = field(default_factory=dict)
scheduler: Optional[dict] = None
weights: Optional[str] = None
weights_ignore_modules: Optional[List[str]] = None
weights_mapping: Optional[List[Dict[str, str]]] = None
check_train_every_n_steps: int = 0
check_val_limit_rank: int = 8
cleanup_after_validation_step: bool = False
cleanup_after_test_step: bool = False
allow_tf32: bool = True
cfg: Config
def __init__(self, cfg, resumed=False) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self._save_dir: Optional[str] = None
self._resumed: bool = resumed
self._resumed_eval: bool = False
self._resumed_eval_status: dict = {"global_step": 0, "current_epoch": 0}
# weird fix for extra VRAM usage on rank 0
# credit: https://discuss.pytorch.org/t/extra-10gb-memory-on-gpu-0-in-ddp-tutorial/118113
torch.cuda.set_device(get_rank())
torch.cuda.empty_cache()
torch.backends.cuda.matmul.allow_tf32 = self.cfg.allow_tf32
self.configure()
if self.cfg.weights is not None:
self.load_weights(
self.cfg.weights,
self.cfg.weights_ignore_modules,
self.cfg.weights_mapping,
)
self.post_configure()
def load_weights(
self,
weights: str,
ignore_modules: Optional[List[str]] = None,
mapping: Optional[List[Dict[str, str]]] = None,
):
state_dict, epoch, global_step = load_module_weights(
weights,
ignore_modules=ignore_modules,
mapping=mapping,
map_location="cpu",
)
self.load_state_dict(state_dict, strict=False)
# restore step-dependent states
self.do_update_step(epoch, global_step, on_load_weights=True)
def set_resume_status(self, current_epoch: int, global_step: int):
# restore correct epoch and global step in eval
self._resumed_eval = True
self._resumed_eval_status["current_epoch"] = current_epoch
self._resumed_eval_status["global_step"] = global_step
@property
def resumed(self):
# whether from resumed checkpoint
return self._resumed
@property
def true_global_step(self):
if self._resumed_eval:
return self._resumed_eval_status["global_step"]
else:
return self.global_step
@property
def true_current_epoch(self):
if self._resumed_eval:
return self._resumed_eval_status["current_epoch"]
else:
return self.current_epoch
def configure(self) -> None:
pass
def post_configure(self) -> None:
"""
executed after weights are loaded
"""
pass
def C(self, value: Any) -> float:
return C(value, self.true_current_epoch, self.true_global_step)
def configure_optimizers(self):
optim = parse_optimizer(self.cfg.optimizer, self)
ret = {
"optimizer": optim,
}
if self.cfg.scheduler is not None:
ret.update(
{
"lr_scheduler": parse_scheduler(self.cfg.scheduler, optim),
}
)
return ret
def on_fit_start(self) -> None:
if self._save_dir is not None:
info(f"Validation results will be saved to {self._save_dir}")
else:
warn(
f"Saving directory not set for the system, visualization results will not be saved"
)
def training_step(self, batch, batch_idx):
raise NotImplementedError
def check_train(self, batch, **kwargs):
if (
self.global_rank == 0
and self.cfg.check_train_every_n_steps > 0
and self.true_global_step % self.cfg.check_train_every_n_steps == 0
):
self.on_check_train(batch, **kwargs)
def on_check_train(self, batch, outputs, **kwargs):
pass
def validation_step(self, batch, batch_idx):
raise NotImplementedError
def on_validation_epoch_end(self):
pass
def test_step(self, batch, batch_idx):
raise NotImplementedError
def on_test_epoch_end(self):
pass
def on_test_end(self) -> None:
if self._save_dir is not None:
info(f"Test results saved to {self._save_dir}")
def on_predict_start(self) -> None:
pass
def predict_step(self, batch, batch_idx):
pass
def on_predict_epoch_end(self) -> None:
pass
def on_predict_end(self) -> None:
pass
def preprocess_data(self, batch, stage):
pass
"""
Implementing on_after_batch_transfer of DataModule does the same.
But on_after_batch_transfer does not support DP.
"""
def on_train_batch_start(self, batch, batch_idx, unused=0):
self.preprocess_data(batch, "train")
self.dataset = self.trainer.train_dataloader.dataset
update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step)
self.do_update_step(self.true_current_epoch, self.true_global_step)
def on_validation_batch_start(self, batch, batch_idx, dataloader_idx=0):
self.preprocess_data(batch, "validation")
self.dataset = self.trainer.val_dataloaders.dataset
update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step)
self.do_update_step(self.true_current_epoch, self.true_global_step)
def on_test_batch_start(self, batch, batch_idx, dataloader_idx=0):
self.preprocess_data(batch, "test")
self.dataset = self.trainer.test_dataloaders.dataset
update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step)
self.do_update_step(self.true_current_epoch, self.true_global_step)
def on_predict_batch_start(self, batch, batch_idx, dataloader_idx=0):
self.preprocess_data(batch, "predict")
self.dataset = self.trainer.predict_dataloaders.dataset
update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step)
self.do_update_step(self.true_current_epoch, self.true_global_step)
def on_train_batch_end(self, outputs, batch, batch_idx):
self.dataset = self.trainer.train_dataloader.dataset
update_end_if_possible(
self.dataset, self.true_current_epoch, self.true_global_step
)
self.do_update_step_end(self.true_current_epoch, self.true_global_step)
def on_validation_batch_end(self, outputs, batch, batch_idx):
self.dataset = self.trainer.val_dataloaders.dataset
update_end_if_possible(
self.dataset, self.true_current_epoch, self.true_global_step
)
self.do_update_step_end(self.true_current_epoch, self.true_global_step)
if self.cfg.cleanup_after_validation_step:
# cleanup to save vram
cleanup()
def on_test_batch_end(self, outputs, batch, batch_idx):
self.dataset = self.trainer.test_dataloaders.dataset
update_end_if_possible(
self.dataset, self.true_current_epoch, self.true_global_step
)
self.do_update_step_end(self.true_current_epoch, self.true_global_step)
if self.cfg.cleanup_after_test_step:
# cleanup to save vram
cleanup()
def on_predict_batch_end(self, outputs, batch, batch_idx):
self.dataset = self.trainer.predict_dataloaders.dataset
update_end_if_possible(
self.dataset, self.true_current_epoch, self.true_global_step
)
self.do_update_step_end(self.true_current_epoch, self.true_global_step)
if self.cfg.cleanup_after_test_step:
# cleanup to save vram
cleanup()
def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
pass
def on_before_optimizer_step(self, optimizer):
"""
# some gradient-related debugging goes here, example:
from lightning.pytorch.utilities import grad_norm
norms = grad_norm(self.geometry, norm_type=2)
print(norms)
for name, p in self.named_parameters():
if p.grad is None:
info(f"{name} does not receive gradients!")
"""
pass
Gen_3D_Modules/MV_Adapter/mvadapter/systems/mvadapter_image_sd.py
+459
View File
@@ -0,0 +1,459 @@
import os
import random
from dataclasses import dataclass, field
import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers.models import AutoencoderKL
from diffusers.training_utils import compute_snr
from einops import rearrange
from omegaconf import OmegaConf
from PIL import Image
from ..pipelines.pipeline_mvadapter_i2mv_sd import MVAdapterI2MVSDPipeline
from ..schedulers.scheduling_shift_snr import ShiftSNRScheduler
from ..utils.core import find
from ..utils.typing import *
from .base import BaseSystem
from .utils import encode_prompt, vae_encode
class MVAdapterImageSDSystem(BaseSystem):
@dataclass
class Config(BaseSystem.Config):
# Model / Adapter
pretrained_model_name_or_path: str = "stabilityai/stable-diffusion-2-1-base"
pretrained_vae_name_or_path: Optional[str] = None
pretrained_adapter_name_or_path: Optional[str] = None
pretrained_unet_name_or_path: Optional[str] = None
init_adapter_kwargs: Dict[str, Any] = field(default_factory=dict)
use_fp16_vae: bool = True
use_fp16_clip: bool = True
# Training
trainable_modules: List[str] = field(default_factory=list)
train_cond_encoder: bool = True
prompt_drop_prob: float = 0.0
image_drop_prob: float = 0.0
cond_drop_prob: float = 0.0
gradient_checkpointing: bool = False
# Noise sampler
noise_scheduler_kwargs: Dict[str, Any] = field(default_factory=dict)
noise_offset: float = 0.0
input_perturbation: float = 0.0
snr_gamma: Optional[float] = 5.0
prediction_type: Optional[str] = None
shift_noise: bool = False
shift_noise_mode: str = "interpolated"
shift_noise_scale: float = 1.0
# Evaluation
eval_seed: int = 0
eval_num_inference_steps: int = 30
eval_guidance_scale: float = 1.0
eval_height: int = 512
eval_width: int = 512
cfg: Config
def configure(self):
super().configure()
# Prepare pipeline
pipeline_kwargs = {}
if self.cfg.pretrained_vae_name_or_path is not None:
pipeline_kwargs["vae"] = AutoencoderKL.from_pretrained(
self.cfg.pretrained_vae_name_or_path
)
if self.cfg.pretrained_unet_name_or_path is not None:
pipeline_kwargs["unet"] = UNet2DConditionModel.from_pretrained(
self.cfg.pretrained_unet_name_or_path
)
pipeline: MVAdapterI2MVSDPipeline
pipeline = MVAdapterI2MVSDPipeline.from_pretrained(
self.cfg.pretrained_model_name_or_path, **pipeline_kwargs
)
init_adapter_kwargs = OmegaConf.to_container(self.cfg.init_adapter_kwargs)
if "self_attn_processor" in init_adapter_kwargs:
self_attn_processor = init_adapter_kwargs["self_attn_processor"]
if self_attn_processor is not None and isinstance(self_attn_processor, str):
self_attn_processor = find(self_attn_processor)
init_adapter_kwargs["self_attn_processor"] = self_attn_processor
pipeline.init_custom_adapter(**init_adapter_kwargs)
if self.cfg.pretrained_adapter_name_or_path:
pretrained_path = os.path.dirname(self.cfg.pretrained_adapter_name_or_path)
adapter_name = os.path.basename(self.cfg.pretrained_adapter_name_or_path)
pipeline.load_custom_adapter(pretrained_path, weight_name=adapter_name)
noise_scheduler = DDPMScheduler.from_config(
pipeline.scheduler.config, **self.cfg.noise_scheduler_kwargs
)
if self.cfg.shift_noise:
noise_scheduler = ShiftSNRScheduler.from_scheduler(
noise_scheduler,
shift_mode=self.cfg.shift_noise_mode,
shift_scale=self.cfg.shift_noise_scale,
scheduler_class=DDPMScheduler,
)
pipeline.scheduler = noise_scheduler
# Prepare models
self.pipeline: MVAdapterI2MVSDPipeline = pipeline
self.vae = self.pipeline.vae.to(
dtype=torch.float16 if self.cfg.use_fp16_vae else torch.float32
)
self.tokenizer = self.pipeline.tokenizer
self.text_encoder = self.pipeline.text_encoder.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.feature_extractor = self.pipeline.feature_extractor
self.cond_encoder = self.pipeline.cond_encoder
self.unet = self.pipeline.unet
self.noise_scheduler = self.pipeline.scheduler
self.inference_scheduler = DDPMScheduler.from_config(
self.noise_scheduler.config
)
self.pipeline.scheduler = self.inference_scheduler
if self.cfg.prediction_type is not None:
self.noise_scheduler.register_to_config(
prediction_type=self.cfg.prediction_type
)
# Prepare trainable / non-trainable modules
trainable_modules = self.cfg.trainable_modules
if trainable_modules and len(trainable_modules) > 0:
self.unet.requires_grad_(False)
for name, module in self.unet.named_modules():
for trainable_module in trainable_modules:
if trainable_module in name:
module.requires_grad_(True)
else:
self.unet.requires_grad_(True)
self.cond_encoder.requires_grad_(self.cfg.train_cond_encoder)
self.vae.requires_grad_(False)
self.text_encoder.requires_grad_(False)
# Others
# Prepare gradient checkpointing
if self.cfg.gradient_checkpointing:
self.unet.enable_gradient_checkpointing()
def forward(
self,
noisy_latents: Tensor,
conditioning_pixel_values: Tensor,
timesteps: Tensor,
ref_latents: Tensor,
prompts: List[str],
num_views: int,
**kwargs,
) -> Dict[str, Any]:
bsz = noisy_latents.shape[0]
b_samples = bsz // num_views
num_batch_images = num_views
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
input_ids = self.tokenizer(
prompts,
max_length=self.tokenizer.model_max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)["input_ids"].to(noisy_latents.device)
encoder_hidden_states = self.text_encoder(input_ids, return_dict=False)[
0
].float()
prompt_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.prompt_drop_prob
)
image_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.image_drop_prob
)
cond_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device) < self.cfg.cond_drop_prob
)
prompt_drop_mask = prompt_drop_mask | cond_drop_mask
image_drop_mask = image_drop_mask | cond_drop_mask
encoder_hidden_states[prompt_drop_mask] = 0.0
# Process reference latents to obtain reference features
with torch.no_grad():
ref_timesteps = torch.zeros_like(timesteps[:b_samples])
ref_hidden_states = {}
self.unet(
ref_latents,
ref_timesteps,
encoder_hidden_states=encoder_hidden_states,
cross_attention_kwargs={
"cache_hidden_states": ref_hidden_states,
"use_mv": False,
"use_ref": False,
},
return_dict=False,
)
for k, v in ref_hidden_states.items():
v_ = v
v_[image_drop_mask] = 0.0
ref_hidden_states[k] = v_.repeat_interleave(num_batch_images, dim=0)
# Repeat additional embeddings for each image in the batch
encoder_hidden_states = encoder_hidden_states.repeat_interleave(
num_batch_images, dim=0
)
conditioning_features = self.cond_encoder(conditioning_pixel_values)
noise_pred = self.unet(
noisy_latents,
timesteps,
encoder_hidden_states=encoder_hidden_states,
down_intrablock_additional_residuals=conditioning_features,
cross_attention_kwargs={
"ref_hidden_states": ref_hidden_states,
"num_views": num_views,
},
).sample
return {"noise_pred": noise_pred}
def training_step(self, batch, batch_idx):
num_views = batch["num_views"]
vae_max_slice = 8
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
latents = []
for i in range(0, batch["rgb"].shape[0], vae_max_slice):
latents.append(
vae_encode(
self.vae,
batch["rgb"][i : i + vae_max_slice].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
)
latents = torch.cat(latents, dim=0)
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
ref_latents = vae_encode(
self.vae,
batch["reference_rgb"].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
bsz = latents.shape[0]
b_samples = bsz // num_views
noise = torch.randn_like(latents)
if self.cfg.noise_offset is not None:
# # https://www.crosslabs.org//blog/diffusion-with-offset-noise
noise += self.cfg.noise_offset * torch.randn(
(latents.shape[0], latents.shape[1], 1, 1), device=latents.device
)
noise_mask = (
batch["noise_mask"]
if "noise_mask" in batch
else torch.ones((bsz,), dtype=torch.bool, device=latents.device)
)
timesteps = torch.randint(
0,
self.noise_scheduler.config.num_train_timesteps,
(b_samples,),
device=latents.device,
dtype=torch.long,
)
timesteps = timesteps.repeat_interleave(num_views)
timesteps[~noise_mask] = 0
if self.cfg.input_perturbation is not None:
new_noise = noise + self.cfg.input_perturbation * torch.randn_like(noise)
noisy_latents = self.noise_scheduler.add_noise(
latents, new_noise, timesteps
)
else:
noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
noisy_latents[~noise_mask] = latents[~noise_mask]
if self.noise_scheduler.config.prediction_type == "epsilon":
target = noise
elif self.noise_scheduler.config.prediction_type == "v_prediction":
target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
else:
raise ValueError(
f"Unsupported prediction type {self.noise_scheduler.config.prediction_type}"
)
model_pred = self(
noisy_latents, batch["source_rgb"], timesteps, ref_latents, **batch
)["noise_pred"]
model_pred = model_pred[noise_mask]
target = target[noise_mask]
if self.cfg.snr_gamma is None:
loss = F.mse_loss(model_pred, target, reduction="mean")
else:
# Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
# Since we predict the noise instead of x_0, the original formulation is slightly changed.
# This is discussed in Section 4.2 of the same paper.
snr = compute_snr(self.noise_scheduler, timesteps)
if self.noise_scheduler.config.prediction_type == "v_prediction":
# Velocity objective requires that we add one to SNR values before we divide by them.
snr = snr + 1
mse_loss_weights = (
torch.stack(
[snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
).min(dim=1)[0]
/ snr
)
loss = F.mse_loss(model_pred, target, reduction="none")
loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
loss = loss.mean()
self.log("train/loss", loss, prog_bar=True)
# will execute self.on_check_train every self.cfg.check_train_every_n_steps steps
self.check_train(batch)
return {"loss": loss}
def on_train_batch_end(self, outputs, batch, batch_idx):
pass
def get_input_visualizations(self, batch):
return [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(batch["reference_rgb"], "B C H W -> (B H) W C"),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
def get_output_visualizations(self, batch, outputs):
images = [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(batch["reference_rgb"], "B C H W -> (B H) W C"),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
outputs, "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
return images
def generate_images(self, batch, **kwargs):
return self.pipeline(
prompt=batch["prompts"],
control_image=batch["source_rgb"],
num_images_per_prompt=batch["num_views"],
generator=torch.Generator(device=self.device).manual_seed(
self.cfg.eval_seed
),
num_inference_steps=self.cfg.eval_num_inference_steps,
guidance_scale=self.cfg.eval_guidance_scale,
height=self.cfg.eval_height,
width=self.cfg.eval_width,
reference_image=batch["reference_rgb"],
output_type="pt",
).images
def on_save_checkpoint(self, checkpoint):
if self.global_rank == 0:
self.pipeline.save_custom_adapter(
os.path.dirname(self.get_save_dir()),
"custom_adapter.safetensors",
safe_serialization=True,
include_keys=self.cfg.trainable_modules,
)
def on_check_train(self, batch):
self.save_image_grid(
f"it{self.true_global_step}-train.jpg",
self.get_input_visualizations(batch),
name="train_step_input",
step=self.true_global_step,
)
def validation_step(self, batch, batch_idx):
out = self.generate_images(batch)
if (
self.cfg.check_val_limit_rank > 0
and self.global_rank < self.cfg.check_val_limit_rank
):
self.save_image_grid(
f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"validation_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_validation_epoch_end(self):
pass
def test_step(self, batch, batch_idx):
out = self.generate_images(batch)
self.save_image_grid(
f"it{self.true_global_step}-test-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"test_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_test_end(self):
pass
Gen_3D_Modules/MV_Adapter/mvadapter/systems/mvadapter_image_sdxl.py
+506
View File
@@ -0,0 +1,506 @@
import os
import random
from dataclasses import dataclass, field
import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers.models import AutoencoderKL
from diffusers.training_utils import compute_snr
from einops import rearrange
from omegaconf import OmegaConf
from PIL import Image
from ..pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from ..schedulers.scheduling_shift_snr import ShiftSNRScheduler
from ..utils.core import find
from ..utils.typing import *
from .base import BaseSystem
from .utils import encode_prompt, vae_encode
def compute_embeddings(
prompt_batch,
empty_prompt_indices,
text_encoders,
tokenizers,
is_train=True,
**kwargs,
):
original_size = kwargs["original_size"]
target_size = kwargs["target_size"]
crops_coords_top_left = kwargs["crops_coords_top_left"]
for i in range(empty_prompt_indices.shape[0]):
if empty_prompt_indices[i]:
prompt_batch[i] = ""
prompt_embeds, pooled_prompt_embeds = encode_prompt(
prompt_batch, text_encoders, tokenizers, 0, is_train
)
add_text_embeds = pooled_prompt_embeds.to(
device=prompt_embeds.device, dtype=prompt_embeds.dtype
)
# Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids
add_time_ids = list(original_size + crops_coords_top_left + target_size)
add_time_ids = torch.tensor([add_time_ids])
add_time_ids = add_time_ids.repeat(len(prompt_batch), 1)
add_time_ids = add_time_ids.to(
device=prompt_embeds.device, dtype=prompt_embeds.dtype
)
unet_added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
return {"prompt_embeds": prompt_embeds, **unet_added_cond_kwargs}
class MVAdapterImageSDXLSystem(BaseSystem):
@dataclass
class Config(BaseSystem.Config):
# Model / Adapter
pretrained_model_name_or_path: str = "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: Optional[str] = "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: Optional[str] = None
pretrained_unet_name_or_path: Optional[str] = None
init_adapter_kwargs: Dict[str, Any] = field(default_factory=dict)
use_fp16_vae: bool = True
use_fp16_clip: bool = True
# Training
trainable_modules: List[str] = field(default_factory=list)
train_cond_encoder: bool = True
prompt_drop_prob: float = 0.0
image_drop_prob: float = 0.0
cond_drop_prob: float = 0.0
gradient_checkpointing: bool = False
# Noise sampler
noise_scheduler_kwargs: Dict[str, Any] = field(default_factory=dict)
noise_offset: float = 0.0
input_perturbation: float = 0.0
snr_gamma: Optional[float] = 5.0
prediction_type: Optional[str] = None
shift_noise: bool = False
shift_noise_mode: str = "interpolated"
shift_noise_scale: float = 1.0
# Evaluation
eval_seed: int = 0
eval_num_inference_steps: int = 30
eval_guidance_scale: float = 1.0
eval_height: int = 512
eval_width: int = 512
cfg: Config
def configure(self):
super().configure()
# Prepare pipeline
pipeline_kwargs = {}
if self.cfg.pretrained_vae_name_or_path is not None:
pipeline_kwargs["vae"] = AutoencoderKL.from_pretrained(
self.cfg.pretrained_vae_name_or_path
)
if self.cfg.pretrained_unet_name_or_path is not None:
pipeline_kwargs["unet"] = UNet2DConditionModel.from_pretrained(
self.cfg.pretrained_unet_name_or_path
)
pipeline: MVAdapterI2MVSDXLPipeline
pipeline = MVAdapterI2MVSDXLPipeline.from_pretrained(
self.cfg.pretrained_model_name_or_path, **pipeline_kwargs
)
init_adapter_kwargs = OmegaConf.to_container(self.cfg.init_adapter_kwargs)
if "self_attn_processor" in init_adapter_kwargs:
self_attn_processor = init_adapter_kwargs["self_attn_processor"]
if self_attn_processor is not None and isinstance(self_attn_processor, str):
self_attn_processor = find(self_attn_processor)
init_adapter_kwargs["self_attn_processor"] = self_attn_processor
pipeline.init_custom_adapter(**init_adapter_kwargs)
if self.cfg.pretrained_adapter_name_or_path:
pretrained_path = os.path.dirname(self.cfg.pretrained_adapter_name_or_path)
adapter_name = os.path.basename(self.cfg.pretrained_adapter_name_or_path)
pipeline.load_custom_adapter(pretrained_path, weight_name=adapter_name)
noise_scheduler = DDPMScheduler.from_config(
pipeline.scheduler.config, **self.cfg.noise_scheduler_kwargs
)
if self.cfg.shift_noise:
noise_scheduler = ShiftSNRScheduler.from_scheduler(
noise_scheduler,
shift_mode=self.cfg.shift_noise_mode,
shift_scale=self.cfg.shift_noise_scale,
scheduler_class=DDPMScheduler,
)
pipeline.scheduler = noise_scheduler
# Prepare models
self.pipeline: MVAdapterI2MVSDXLPipeline = pipeline
self.vae = self.pipeline.vae.to(
dtype=torch.float16 if self.cfg.use_fp16_vae else torch.float32
)
self.tokenizer = self.pipeline.tokenizer
self.tokenizer_2 = self.pipeline.tokenizer_2
self.text_encoder = self.pipeline.text_encoder.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.text_encoder_2 = self.pipeline.text_encoder_2.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.feature_extractor = self.pipeline.feature_extractor
self.cond_encoder = self.pipeline.cond_encoder
self.unet = self.pipeline.unet
self.noise_scheduler = self.pipeline.scheduler
self.inference_scheduler = DDPMScheduler.from_config(
self.noise_scheduler.config
)
self.pipeline.scheduler = self.inference_scheduler
if self.cfg.prediction_type is not None:
self.noise_scheduler.register_to_config(
prediction_type=self.cfg.prediction_type
)
# Prepare trainable / non-trainable modules
trainable_modules = self.cfg.trainable_modules
if trainable_modules and len(trainable_modules) > 0:
self.unet.requires_grad_(False)
for name, module in self.unet.named_modules():
for trainable_module in trainable_modules:
if trainable_module in name:
module.requires_grad_(True)
else:
self.unet.requires_grad_(True)
self.cond_encoder.requires_grad_(self.cfg.train_cond_encoder)
self.vae.requires_grad_(False)
self.text_encoder.requires_grad_(False)
self.text_encoder_2.requires_grad_(False)
# Others
# Prepare gradient checkpointing
if self.cfg.gradient_checkpointing:
self.unet.enable_gradient_checkpointing()
def forward(
self,
noisy_latents: Tensor,
conditioning_pixel_values: Tensor,
timesteps: Tensor,
ref_latents: Tensor,
prompts: List[str],
num_views: int,
**kwargs,
) -> Dict[str, Any]:
bsz = noisy_latents.shape[0]
b_samples = bsz // num_views
num_batch_images = num_views
prompt_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.prompt_drop_prob
)
image_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.image_drop_prob
)
cond_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device) < self.cfg.cond_drop_prob
)
prompt_drop_mask = prompt_drop_mask | cond_drop_mask
image_drop_mask = image_drop_mask | cond_drop_mask
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
# Here, we compute not just the text embeddings but also the additional embeddings
# needed for the SD XL UNet to operate.
additional_embeds = compute_embeddings(
prompts,
prompt_drop_mask,
[self.text_encoder, self.text_encoder_2],
[self.tokenizer, self.tokenizer_2],
**kwargs,
)
# Process reference latents to obtain reference features
with torch.no_grad():
ref_timesteps = torch.zeros_like(timesteps[:b_samples])
ref_hidden_states = {}
self.unet(
ref_latents,
ref_timesteps,
encoder_hidden_states=additional_embeds["prompt_embeds"],
added_cond_kwargs={
"text_embeds": additional_embeds["text_embeds"],
"time_ids": additional_embeds["time_ids"],
},
cross_attention_kwargs={
"cache_hidden_states": ref_hidden_states,
"use_mv": False,
"use_ref": False,
},
return_dict=False,
)
for k, v in ref_hidden_states.items():
v_ = v
v_[image_drop_mask] = 0.0
ref_hidden_states[k] = v_.repeat_interleave(num_batch_images, dim=0)
# Repeat additional embeddings for each image in the batch
for key, value in additional_embeds.items():
kwargs[key] = value.repeat_interleave(num_batch_images, dim=0)
conditioning_features = self.cond_encoder(conditioning_pixel_values)
added_cond_kwargs = {
"text_embeds": kwargs["text_embeds"],
"time_ids": kwargs["time_ids"],
}
noise_pred = self.unet(
noisy_latents,
timesteps,
encoder_hidden_states=kwargs["prompt_embeds"],
added_cond_kwargs=added_cond_kwargs,
down_intrablock_additional_residuals=conditioning_features,
cross_attention_kwargs={
"ref_hidden_states": ref_hidden_states,
"num_views": num_views,
},
).sample
return {"noise_pred": noise_pred}
def training_step(self, batch, batch_idx):
num_views = batch["num_views"]
vae_max_slice = 8
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
latents = []
for i in range(0, batch["rgb"].shape[0], vae_max_slice):
latents.append(
vae_encode(
self.vae,
batch["rgb"][i : i + vae_max_slice].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
)
latents = torch.cat(latents, dim=0)
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
ref_latents = vae_encode(
self.vae,
batch["reference_rgb"].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
bsz = latents.shape[0]
b_samples = bsz // num_views
noise = torch.randn_like(latents)
if self.cfg.noise_offset is not None:
# # https://www.crosslabs.org//blog/diffusion-with-offset-noise
noise += self.cfg.noise_offset * torch.randn(
(latents.shape[0], latents.shape[1], 1, 1), device=latents.device
)
noise_mask = (
batch["noise_mask"]
if "noise_mask" in batch
else torch.ones((bsz,), dtype=torch.bool, device=latents.device)
)
timesteps = torch.randint(
0,
self.noise_scheduler.config.num_train_timesteps,
(b_samples,),
device=latents.device,
dtype=torch.long,
)
timesteps = timesteps.repeat_interleave(num_views)
timesteps[~noise_mask] = 0
if self.cfg.input_perturbation is not None:
new_noise = noise + self.cfg.input_perturbation * torch.randn_like(noise)
noisy_latents = self.noise_scheduler.add_noise(
latents, new_noise, timesteps
)
else:
noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
noisy_latents[~noise_mask] = latents[~noise_mask]
if self.noise_scheduler.config.prediction_type == "epsilon":
target = noise
elif self.noise_scheduler.config.prediction_type == "v_prediction":
target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
else:
raise ValueError(
f"Unsupported prediction type {self.noise_scheduler.config.prediction_type}"
)
model_pred = self(
noisy_latents, batch["source_rgb"], timesteps, ref_latents, **batch
)["noise_pred"]
model_pred = model_pred[noise_mask]
target = target[noise_mask]
if self.cfg.snr_gamma is None:
loss = F.mse_loss(model_pred, target, reduction="mean")
else:
# Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
# Since we predict the noise instead of x_0, the original formulation is slightly changed.
# This is discussed in Section 4.2 of the same paper.
snr = compute_snr(self.noise_scheduler, timesteps)
if self.noise_scheduler.config.prediction_type == "v_prediction":
# Velocity objective requires that we add one to SNR values before we divide by them.
snr = snr + 1
mse_loss_weights = (
torch.stack(
[snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
).min(dim=1)[0]
/ snr
)
loss = F.mse_loss(model_pred, target, reduction="none")
loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
loss = loss.mean()
self.log("train/loss", loss, prog_bar=True)
# will execute self.on_check_train every self.cfg.check_train_every_n_steps steps
self.check_train(batch)
return {"loss": loss}
def on_train_batch_end(self, outputs, batch, batch_idx):
pass
def get_input_visualizations(self, batch):
return [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(batch["reference_rgb"], "B C H W -> (B H) W C"),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
def get_output_visualizations(self, batch, outputs):
images = [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(batch["reference_rgb"], "B C H W -> (B H) W C"),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
outputs, "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
return images
def generate_images(self, batch, **kwargs):
return self.pipeline(
prompt=batch["prompts"],
control_image=batch["source_rgb"],
num_images_per_prompt=batch["num_views"],
generator=torch.Generator(device=self.device).manual_seed(
self.cfg.eval_seed
),
num_inference_steps=self.cfg.eval_num_inference_steps,
guidance_scale=self.cfg.eval_guidance_scale,
height=self.cfg.eval_height,
width=self.cfg.eval_width,
reference_image=batch["reference_rgb"],
output_type="pt",
).images
def on_save_checkpoint(self, checkpoint):
if self.global_rank == 0:
self.pipeline.save_custom_adapter(
os.path.dirname(self.get_save_dir()),
"custom_adapter.safetensors",
safe_serialization=True,
include_keys=self.cfg.trainable_modules,
)
def on_check_train(self, batch):
self.save_image_grid(
f"it{self.true_global_step}-train.jpg",
self.get_input_visualizations(batch),
name="train_step_input",
step=self.true_global_step,
)
def validation_step(self, batch, batch_idx):
out = self.generate_images(batch)
if (
self.cfg.check_val_limit_rank > 0
and self.global_rank < self.cfg.check_val_limit_rank
):
self.save_image_grid(
f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"validation_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_validation_epoch_end(self):
pass
def test_step(self, batch, batch_idx):
out = self.generate_images(batch)
self.save_image_grid(
f"it{self.true_global_step}-test-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"test_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_test_end(self):
pass
Gen_3D_Modules/MV_Adapter/mvadapter/systems/mvadapter_text_sd.py
+415
View File
@@ -0,0 +1,415 @@
import os
import random
from dataclasses import dataclass, field
import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers.models import AutoencoderKL
from diffusers.training_utils import compute_snr
from einops import rearrange
from omegaconf import OmegaConf
from PIL import Image
from ..pipelines.pipeline_mvadapter_t2mv_sd import MVAdapterT2MVSDPipeline
from ..schedulers.scheduling_shift_snr import ShiftSNRScheduler
from ..utils.core import find
from ..utils.typing import *
from .base import BaseSystem
from .utils import encode_prompt, vae_encode
class MVAdapterTextSDSystem(BaseSystem):
@dataclass
class Config(BaseSystem.Config):
# Model / Adapter
pretrained_model_name_or_path: str = "stabilityai/stable-diffusion-2-1-base"
pretrained_vae_name_or_path: Optional[str] = None
pretrained_adapter_name_or_path: Optional[str] = None
pretrained_unet_name_or_path: Optional[str] = None
init_adapter_kwargs: Dict[str, Any] = field(default_factory=dict)
use_fp16_vae: bool = True
use_fp16_clip: bool = True
# Training
trainable_modules: List[str] = field(default_factory=list)
train_cond_encoder: bool = True
prompt_drop_prob: float = 0.0
image_drop_prob: float = 0.0
cond_drop_prob: float = 0.0
gradient_checkpointing: bool = False
# Noise sampler
noise_scheduler_kwargs: Dict[str, Any] = field(default_factory=dict)
noise_offset: float = 0.0
input_perturbation: float = 0.0
snr_gamma: Optional[float] = 5.0
prediction_type: Optional[str] = None
shift_noise: bool = False
shift_noise_mode: str = "interpolated"
shift_noise_scale: float = 1.0
# Evaluation
eval_seed: int = 0
eval_num_inference_steps: int = 30
eval_guidance_scale: float = 1.0
eval_height: int = 512
eval_width: int = 512
cfg: Config
def configure(self):
super().configure()
# Prepare pipeline
pipeline_kwargs = {}
if self.cfg.pretrained_vae_name_or_path is not None:
pipeline_kwargs["vae"] = AutoencoderKL.from_pretrained(
self.cfg.pretrained_vae_name_or_path
)
if self.cfg.pretrained_unet_name_or_path is not None:
pipeline_kwargs["unet"] = UNet2DConditionModel.from_pretrained(
self.cfg.pretrained_unet_name_or_path
)
pipeline: MVAdapterT2MVSDPipeline
pipeline = MVAdapterT2MVSDPipeline.from_pretrained(
self.cfg.pretrained_model_name_or_path, **pipeline_kwargs
)
init_adapter_kwargs = OmegaConf.to_container(self.cfg.init_adapter_kwargs)
if "self_attn_processor" in init_adapter_kwargs:
self_attn_processor = init_adapter_kwargs["self_attn_processor"]
if self_attn_processor is not None and isinstance(self_attn_processor, str):
self_attn_processor = find(self_attn_processor)
init_adapter_kwargs["self_attn_processor"] = self_attn_processor
pipeline.init_custom_adapter(**init_adapter_kwargs)
if self.cfg.pretrained_adapter_name_or_path:
pretrained_path = os.path.dirname(self.cfg.pretrained_adapter_name_or_path)
adapter_name = os.path.basename(self.cfg.pretrained_adapter_name_or_path)
pipeline.load_custom_adapter(pretrained_path, weight_name=adapter_name)
noise_scheduler = DDPMScheduler.from_config(
pipeline.scheduler.config, **self.cfg.noise_scheduler_kwargs
)
if self.cfg.shift_noise:
noise_scheduler = ShiftSNRScheduler.from_scheduler(
noise_scheduler,
shift_mode=self.cfg.shift_noise_mode,
shift_scale=self.cfg.shift_noise_scale,
scheduler_class=DDPMScheduler,
)
pipeline.scheduler = noise_scheduler
# Prepare models
self.pipeline: MVAdapterT2MVSDPipeline = pipeline
self.vae = self.pipeline.vae.to(
dtype=torch.float16 if self.cfg.use_fp16_vae else torch.float32
)
self.tokenizer = self.pipeline.tokenizer
self.text_encoder = self.pipeline.text_encoder.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.feature_extractor = self.pipeline.feature_extractor
self.cond_encoder = self.pipeline.cond_encoder
self.unet = self.pipeline.unet
self.noise_scheduler = self.pipeline.scheduler
self.inference_scheduler = DDPMScheduler.from_config(
self.noise_scheduler.config
)
self.pipeline.scheduler = self.inference_scheduler
if self.cfg.prediction_type is not None:
self.noise_scheduler.register_to_config(
prediction_type=self.cfg.prediction_type
)
# Prepare trainable / non-trainable modules
trainable_modules = self.cfg.trainable_modules
if trainable_modules and len(trainable_modules) > 0:
self.unet.requires_grad_(False)
for name, module in self.unet.named_modules():
for trainable_module in trainable_modules:
if trainable_module in name:
module.requires_grad_(True)
else:
self.unet.requires_grad_(True)
self.cond_encoder.requires_grad_(self.cfg.train_cond_encoder)
self.vae.requires_grad_(False)
self.text_encoder.requires_grad_(False)
# Others
# Prepare gradient checkpointing
if self.cfg.gradient_checkpointing:
self.unet.enable_gradient_checkpointing()
def forward(
self,
noisy_latents: Tensor,
conditioning_pixel_values: Tensor,
timesteps: Tensor,
prompts: List[str],
num_views: int,
**kwargs,
) -> Dict[str, Any]:
bsz = noisy_latents.shape[0]
b_samples = bsz // num_views
num_batch_images = num_views
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
input_ids = self.tokenizer(
prompts,
max_length=self.tokenizer.model_max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)["input_ids"].to(noisy_latents.device)
encoder_hidden_states = self.text_encoder(input_ids, return_dict=False)[
0
].float()
prompt_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.prompt_drop_prob
)
image_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.image_drop_prob
)
cond_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device) < self.cfg.cond_drop_prob
)
prompt_drop_mask = prompt_drop_mask | cond_drop_mask
image_drop_mask = image_drop_mask | cond_drop_mask
encoder_hidden_states[prompt_drop_mask] = 0.0
# Repeat additional embeddings for each image in the batch
encoder_hidden_states = encoder_hidden_states.repeat_interleave(
num_batch_images, dim=0
)
conditioning_features = self.cond_encoder(conditioning_pixel_values)
noise_pred = self.unet(
noisy_latents,
timesteps,
encoder_hidden_states=encoder_hidden_states,
down_intrablock_additional_residuals=conditioning_features,
cross_attention_kwargs={"num_views": num_views},
).sample
return {"noise_pred": noise_pred}
def training_step(self, batch, batch_idx):
num_views = batch["num_views"]
vae_max_slice = 8
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
latents = []
for i in range(0, batch["rgb"].shape[0], vae_max_slice):
latents.append(
vae_encode(
self.vae,
batch["rgb"][i : i + vae_max_slice].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
)
latents = torch.cat(latents, dim=0)
bsz = latents.shape[0]
b_samples = bsz // num_views
noise = torch.randn_like(latents)
if self.cfg.noise_offset is not None:
# # https://www.crosslabs.org//blog/diffusion-with-offset-noise
noise += self.cfg.noise_offset * torch.randn(
(latents.shape[0], latents.shape[1], 1, 1), device=latents.device
)
noise_mask = (
batch["noise_mask"]
if "noise_mask" in batch
else torch.ones((bsz,), dtype=torch.bool, device=latents.device)
)
timesteps = torch.randint(
0,
self.noise_scheduler.config.num_train_timesteps,
(b_samples,),
device=latents.device,
dtype=torch.long,
)
timesteps = timesteps.repeat_interleave(num_views)
timesteps[~noise_mask] = 0
if self.cfg.input_perturbation is not None:
new_noise = noise + self.cfg.input_perturbation * torch.randn_like(noise)
noisy_latents = self.noise_scheduler.add_noise(
latents, new_noise, timesteps
)
else:
noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
noisy_latents[~noise_mask] = latents[~noise_mask]
if self.noise_scheduler.config.prediction_type == "epsilon":
target = noise
elif self.noise_scheduler.config.prediction_type == "v_prediction":
target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
else:
raise ValueError(
f"Unsupported prediction type {self.noise_scheduler.config.prediction_type}"
)
model_pred = self(noisy_latents, batch["source_rgb"], timesteps, **batch)[
"noise_pred"
]
model_pred = model_pred[noise_mask]
target = target[noise_mask]
if self.cfg.snr_gamma is None:
loss = F.mse_loss(model_pred, target, reduction="mean")
else:
# Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
# Since we predict the noise instead of x_0, the original formulation is slightly changed.
# This is discussed in Section 4.2 of the same paper.
snr = compute_snr(self.noise_scheduler, timesteps)
if self.noise_scheduler.config.prediction_type == "v_prediction":
# Velocity objective requires that we add one to SNR values before we divide by them.
snr = snr + 1
mse_loss_weights = (
torch.stack(
[snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
).min(dim=1)[0]
/ snr
)
loss = F.mse_loss(model_pred, target, reduction="none")
loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
loss = loss.mean()
self.log("train/loss", loss, prog_bar=True)
# will execute self.on_check_train every self.cfg.check_train_every_n_steps steps
self.check_train(batch)
return {"loss": loss}
def on_train_batch_end(self, outputs, batch, batch_idx):
pass
def get_input_visualizations(self, batch):
return [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
def get_output_visualizations(self, batch, outputs):
images = [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
outputs, "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
return images
def generate_images(self, batch, **kwargs):
return self.pipeline(
prompt=batch["prompts"],
control_image=batch["source_rgb"],
num_images_per_prompt=batch["num_views"],
generator=torch.Generator(device=self.device).manual_seed(
self.cfg.eval_seed
),
num_inference_steps=self.cfg.eval_num_inference_steps,
guidance_scale=self.cfg.eval_guidance_scale,
height=self.cfg.eval_height,
width=self.cfg.eval_width,
output_type="pt",
).images
def on_save_checkpoint(self, checkpoint):
if self.global_rank == 0:
self.pipeline.save_custom_adapter(
os.path.dirname(self.get_save_dir()),
"custom_adapter.safetensors",
safe_serialization=True,
include_keys=self.cfg.trainable_modules,
)
def on_check_train(self, batch):
self.save_image_grid(
f"it{self.true_global_step}-train.jpg",
self.get_input_visualizations(batch),
name="train_step_input",
step=self.true_global_step,
)
def validation_step(self, batch, batch_idx):
out = self.generate_images(batch)
if (
self.cfg.check_val_limit_rank > 0
and self.global_rank < self.cfg.check_val_limit_rank
):
self.save_image_grid(
f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"validation_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_validation_epoch_end(self):
pass
def test_step(self, batch, batch_idx):
out = self.generate_images(batch)
self.save_image_grid(
f"it{self.true_global_step}-test-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"test_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_test_end(self):
pass
Gen_3D_Modules/MV_Adapter/mvadapter/systems/mvadapter_text_sdxl.py
+457
View File
@@ -0,0 +1,457 @@
import os
import random
from dataclasses import dataclass, field
import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers.models import AutoencoderKL
from diffusers.training_utils import compute_snr
from einops import rearrange
from omegaconf import OmegaConf
from PIL import Image
from ..pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
from ..schedulers.scheduling_shift_snr import ShiftSNRScheduler
from ..utils.core import find
from ..utils.typing import *
from .base import BaseSystem
from .utils import encode_prompt, vae_encode
def compute_embeddings(
prompt_batch,
empty_prompt_indices,
text_encoders,
tokenizers,
is_train=True,
**kwargs,
):
original_size = kwargs["original_size"]
target_size = kwargs["target_size"]
crops_coords_top_left = kwargs["crops_coords_top_left"]
for i in range(empty_prompt_indices.shape[0]):
if empty_prompt_indices[i]:
prompt_batch[i] = ""
prompt_embeds, pooled_prompt_embeds = encode_prompt(
prompt_batch, text_encoders, tokenizers, 0, is_train
)
add_text_embeds = pooled_prompt_embeds.to(
device=prompt_embeds.device, dtype=prompt_embeds.dtype
)
# Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids
add_time_ids = list(original_size + crops_coords_top_left + target_size)
add_time_ids = torch.tensor([add_time_ids])
add_time_ids = add_time_ids.repeat(len(prompt_batch), 1)
add_time_ids = add_time_ids.to(
device=prompt_embeds.device, dtype=prompt_embeds.dtype
)
unet_added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
return {"prompt_embeds": prompt_embeds, **unet_added_cond_kwargs}
class MVAdapterTextSDXLSystem(BaseSystem):
@dataclass
class Config(BaseSystem.Config):
# Model / Adapter
pretrained_model_name_or_path: str = "stabilityai/stable-diffusion-xl-base-1.0"
pretrained_vae_name_or_path: Optional[str] = "madebyollin/sdxl-vae-fp16-fix"
pretrained_adapter_name_or_path: Optional[str] = None
pretrained_unet_name_or_path: Optional[str] = None
init_adapter_kwargs: Dict[str, Any] = field(default_factory=dict)
use_fp16_vae: bool = True
use_fp16_clip: bool = True
# Training
trainable_modules: List[str] = field(default_factory=list)
train_cond_encoder: bool = True
prompt_drop_prob: float = 0.0
image_drop_prob: float = 0.0
cond_drop_prob: float = 0.0
gradient_checkpointing: bool = False
# Noise sampler
noise_scheduler_kwargs: Dict[str, Any] = field(default_factory=dict)
noise_offset: float = 0.0
input_perturbation: float = 0.0
snr_gamma: Optional[float] = 5.0
prediction_type: Optional[str] = None
shift_noise: bool = False
shift_noise_mode: str = "interpolated"
shift_noise_scale: float = 1.0
# Evaluation
eval_seed: int = 0
eval_num_inference_steps: int = 30
eval_guidance_scale: float = 1.0
eval_height: int = 512
eval_width: int = 512
cfg: Config
def configure(self):
super().configure()
# Prepare pipeline
pipeline_kwargs = {}
if self.cfg.pretrained_vae_name_or_path is not None:
pipeline_kwargs["vae"] = AutoencoderKL.from_pretrained(
self.cfg.pretrained_vae_name_or_path
)
if self.cfg.pretrained_unet_name_or_path is not None:
pipeline_kwargs["unet"] = UNet2DConditionModel.from_pretrained(
self.cfg.pretrained_unet_name_or_path
)
pipeline: MVAdapterT2MVSDXLPipeline
pipeline = MVAdapterT2MVSDXLPipeline.from_pretrained(
self.cfg.pretrained_model_name_or_path, **pipeline_kwargs
)
init_adapter_kwargs = OmegaConf.to_container(self.cfg.init_adapter_kwargs)
if "self_attn_processor" in init_adapter_kwargs:
self_attn_processor = init_adapter_kwargs["self_attn_processor"]
if self_attn_processor is not None and isinstance(self_attn_processor, str):
self_attn_processor = find(self_attn_processor)
init_adapter_kwargs["self_attn_processor"] = self_attn_processor
pipeline.init_custom_adapter(**init_adapter_kwargs)
if self.cfg.pretrained_adapter_name_or_path:
pretrained_path = os.path.dirname(self.cfg.pretrained_adapter_name_or_path)
adapter_name = os.path.basename(self.cfg.pretrained_adapter_name_or_path)
pipeline.load_custom_adapter(pretrained_path, weight_name=adapter_name)
noise_scheduler = DDPMScheduler.from_config(
pipeline.scheduler.config, **self.cfg.noise_scheduler_kwargs
)
if self.cfg.shift_noise:
noise_scheduler = ShiftSNRScheduler.from_scheduler(
noise_scheduler,
shift_mode=self.cfg.shift_noise_mode,
shift_scale=self.cfg.shift_noise_scale,
scheduler_class=DDPMScheduler,
)
pipeline.scheduler = noise_scheduler
# Prepare models
self.pipeline: MVAdapterT2MVSDXLPipeline = pipeline
self.vae = self.pipeline.vae.to(
dtype=torch.float16 if self.cfg.use_fp16_vae else torch.float32
)
self.tokenizer = self.pipeline.tokenizer
self.tokenizer_2 = self.pipeline.tokenizer_2
self.text_encoder = self.pipeline.text_encoder.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.text_encoder_2 = self.pipeline.text_encoder_2.to(
dtype=torch.float16 if self.cfg.use_fp16_clip else torch.float32
)
self.cond_encoder = self.pipeline.cond_encoder
self.unet = self.pipeline.unet
self.noise_scheduler = self.pipeline.scheduler
self.inference_scheduler = DDPMScheduler.from_config(
self.noise_scheduler.config
)
self.pipeline.scheduler = self.inference_scheduler
if self.cfg.prediction_type is not None:
self.noise_scheduler.register_to_config(
prediction_type=self.cfg.prediction_type
)
# Prepare trainable / non-trainable modules
trainable_modules = self.cfg.trainable_modules
if trainable_modules and len(trainable_modules) > 0:
self.unet.requires_grad_(False)
for name, module in self.unet.named_modules():
for trainable_module in trainable_modules:
if trainable_module in name:
module.requires_grad_(True)
else:
self.unet.requires_grad_(True)
self.cond_encoder.requires_grad_(self.cfg.train_cond_encoder)
self.vae.requires_grad_(False)
self.text_encoder.requires_grad_(False)
self.text_encoder_2.requires_grad_(False)
# Others
# Prepare gradient checkpointing
if self.cfg.gradient_checkpointing:
self.unet.enable_gradient_checkpointing()
def forward(
self,
noisy_latents: Tensor,
conditioning_pixel_values: Tensor,
timesteps: Tensor,
prompts: List[str],
num_views: int,
**kwargs,
) -> Dict[str, Any]:
bsz = noisy_latents.shape[0]
b_samples = bsz // num_views
num_batch_images = num_views
prompt_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.prompt_drop_prob
)
image_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device)
< self.cfg.image_drop_prob
)
cond_drop_mask = (
torch.rand(b_samples, device=noisy_latents.device) < self.cfg.cond_drop_prob
)
prompt_drop_mask = prompt_drop_mask | cond_drop_mask
image_drop_mask = image_drop_mask | cond_drop_mask
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
# Here, we compute not just the text embeddings but also the additional embeddings
# needed for the SD XL UNet to operate.
additional_embeds = compute_embeddings(
prompts,
prompt_drop_mask,
[self.text_encoder, self.text_encoder_2],
[self.tokenizer, self.tokenizer_2],
**kwargs,
)
for key, value in additional_embeds.items():
kwargs[key] = value.repeat_interleave(num_batch_images, dim=0)
conditioning_features = self.cond_encoder(conditioning_pixel_values)
added_cond_kwargs = {
"text_embeds": kwargs["text_embeds"],
"time_ids": kwargs["time_ids"],
}
noise_pred = self.unet(
noisy_latents,
timesteps,
encoder_hidden_states=kwargs["prompt_embeds"],
added_cond_kwargs=added_cond_kwargs,
down_intrablock_additional_residuals=conditioning_features,
cross_attention_kwargs={"num_views": num_views},
).sample
return {"noise_pred": noise_pred}
def training_step(self, batch, batch_idx):
num_views = batch["num_views"]
vae_max_slice = 8
with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
latents = []
for i in range(0, batch["rgb"].shape[0], vae_max_slice):
latents.append(
vae_encode(
self.vae,
batch["rgb"][i : i + vae_max_slice].to(self.vae.dtype) * 2 - 1,
sample=True,
apply_scale=True,
).float()
)
latents = torch.cat(latents, dim=0)
bsz = latents.shape[0]
b_samples = bsz // num_views
noise = torch.randn_like(latents)
if self.cfg.noise_offset is not None:
# # https://www.crosslabs.org//blog/diffusion-with-offset-noise
noise += self.cfg.noise_offset * torch.randn(
(latents.shape[0], latents.shape[1], 1, 1), device=latents.device
)
noise_mask = (
batch["noise_mask"]
if "noise_mask" in batch
else torch.ones((bsz,), dtype=torch.bool, device=latents.device)
)
timesteps = torch.randint(
0,
self.noise_scheduler.config.num_train_timesteps,
(b_samples,),
device=latents.device,
dtype=torch.long,
)
timesteps = timesteps.repeat_interleave(num_views)
timesteps[~noise_mask] = 0
if self.cfg.input_perturbation is not None:
new_noise = noise + self.cfg.input_perturbation * torch.randn_like(noise)
noisy_latents = self.noise_scheduler.add_noise(
latents, new_noise, timesteps
)
else:
noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
noisy_latents[~noise_mask] = latents[~noise_mask]
if self.noise_scheduler.config.prediction_type == "epsilon":
target = noise
elif self.noise_scheduler.config.prediction_type == "v_prediction":
target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
else:
raise ValueError(
f"Unsupported prediction type {self.noise_scheduler.config.prediction_type}"
)
model_pred = self(noisy_latents, batch["source_rgb"], timesteps, **batch)[
"noise_pred"
]
model_pred = model_pred[noise_mask]
target = target[noise_mask]
if self.cfg.snr_gamma is None:
loss = F.mse_loss(model_pred, target, reduction="mean")
else:
# Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
# Since we predict the noise instead of x_0, the original formulation is slightly changed.
# This is discussed in Section 4.2 of the same paper.
snr = compute_snr(self.noise_scheduler, timesteps)
if self.noise_scheduler.config.prediction_type == "v_prediction":
# Velocity objective requires that we add one to SNR values before we divide by them.
snr = snr + 1
mse_loss_weights = (
torch.stack(
[snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
).min(dim=1)[0]
/ snr
)
loss = F.mse_loss(model_pred, target, reduction="none")
loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
loss = loss.mean()
self.log("train/loss", loss, prog_bar=True)
# will execute self.on_check_train every self.cfg.check_train_every_n_steps steps
self.check_train(batch)
return {"loss": loss}
def on_train_batch_end(self, outputs, batch, batch_idx):
pass
def get_input_visualizations(self, batch):
return [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
def get_output_visualizations(self, batch, outputs):
images = [
{
"type": "rgb",
"img": rearrange(
batch["source_rgb"],
"(B N) C H W -> (B H) (N W) C",
N=batch["num_views"],
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
batch["rgb"], "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
{
"type": "rgb",
"img": rearrange(
outputs, "(B N) C H W -> (B H) (N W) C", N=batch["num_views"]
),
"kwargs": {"data_format": "HWC"},
},
]
return images
def generate_images(self, batch, **kwargs):
return self.pipeline(
prompt=batch["prompts"],
control_image=batch["source_rgb"],
num_images_per_prompt=batch["num_views"],
generator=torch.Generator(device=self.device).manual_seed(
self.cfg.eval_seed
),
num_inference_steps=self.cfg.eval_num_inference_steps,
guidance_scale=self.cfg.eval_guidance_scale,
height=self.cfg.eval_height,
width=self.cfg.eval_width,
output_type="pt",
).images
def on_save_checkpoint(self, checkpoint):
if self.global_rank == 0:
self.pipeline.save_custom_adapter(
os.path.dirname(self.get_save_dir()),
"custom_adapter.safetensors",
safe_serialization=True,
include_keys=self.cfg.trainable_modules,
)
def on_check_train(self, batch):
self.save_image_grid(
f"it{self.true_global_step}-train.jpg",
self.get_input_visualizations(batch),
name="train_step_input",
step=self.true_global_step,
)
def validation_step(self, batch, batch_idx):
out = self.generate_images(batch)
if (
self.cfg.check_val_limit_rank > 0
and self.global_rank < self.cfg.check_val_limit_rank
):
self.save_image_grid(
f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"validation_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_validation_epoch_end(self):
pass
def test_step(self, batch, batch_idx):
out = self.generate_images(batch)
self.save_image_grid(
f"it{self.true_global_step}-test-{self.global_rank}_{batch_idx}.jpg",
self.get_output_visualizations(batch, out),
name=f"test_step_output_{self.global_rank}_{batch_idx}",
step=self.true_global_step,
)
def on_test_end(self):
pass
Gen_3D_Modules/MV_Adapter/mvadapter/systems/utils.py
+178
View File
@@ -0,0 +1,178 @@
import random
import numpy as np
import torch
import torch.nn as nn
from diffusers import AutoencoderKL
from torch.optim import lr_scheduler
from ..utils.core import debug, find, info, warn
from ..utils.typing import *
"""Diffusers Model Utils"""
def vae_encode(
vae: AutoencoderKL,
pixel_values: Float[Tensor, "B 3 H W"],
sample: bool = True,
apply_scale: bool = True,
):
latent_dist = vae.encode(pixel_values).latent_dist
latents = latent_dist.sample() if sample else latent_dist.mode()
if apply_scale:
latents = latents * vae.config.scaling_factor
return latents
# Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt
def encode_prompt(
prompt_batch, text_encoders, tokenizers, proportion_empty_prompts, is_train=True
):
prompt_embeds_list = []
captions = []
for caption in prompt_batch:
if random.random() < proportion_empty_prompts:
captions.append("")
elif isinstance(caption, str):
captions.append(caption)
elif isinstance(caption, (list, np.ndarray)):
# take a random caption if there are multiple
captions.append(random.choice(caption) if is_train else caption[0])
with torch.no_grad():
for tokenizer, text_encoder in zip(tokenizers, text_encoders):
text_inputs = tokenizer(
captions,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
prompt_embeds = text_encoder(
text_input_ids.to(text_encoder.device),
output_hidden_states=True,
)
# We are only ALWAYS interested in the pooled output of the final text encoder
pooled_prompt_embeds = prompt_embeds[0]
prompt_embeds = prompt_embeds.hidden_states[-2]
bs_embed, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
return prompt_embeds, pooled_prompt_embeds
CLIP_INPUT_MEAN = torch.as_tensor(
[0.48145466, 0.4578275, 0.40821073], dtype=torch.float32
)[None, :, None, None]
CLIP_INPUT_STD = torch.as_tensor(
[0.26862954, 0.26130258, 0.27577711], dtype=torch.float32
)[None, :, None, None]
def normalize_image_for_clip(image: Float[Tensor, "B C H W"]):
return (image - CLIP_INPUT_MEAN.to(image)) / CLIP_INPUT_STD.to(image)
"""Training"""
def get_scheduler(name):
if hasattr(lr_scheduler, name):
return getattr(lr_scheduler, name)
else:
raise NotImplementedError
def getattr_recursive(m, attr):
for name in attr.split("."):
m = getattr(m, name)
return m
def get_parameters(model, name):
module = getattr_recursive(model, name)
if isinstance(module, nn.Module):
return module.parameters()
elif isinstance(module, nn.Parameter):
return module
return []
def parse_optimizer(config, model):
if hasattr(config, "params"):
params = [
{"params": get_parameters(model, name), "name": name, **args}
for name, args in config.params.items()
]
debug(f"Specify optimizer params: {config.params}")
else:
params = model.parameters()
if config.name in ["FusedAdam"]:
import apex
optim = getattr(apex.optimizers, config.name)(params, **config.args)
elif config.name in ["Adam8bit", "AdamW8bit"]:
import bitsandbytes as bnb
optim = bnb.optim.Adam8bit(params, **config.args)
else:
optim = getattr(torch.optim, config.name)(params, **config.args)
return optim
def parse_scheduler_to_instance(config, optimizer):
if config.name == "ChainedScheduler":
schedulers = [
parse_scheduler_to_instance(conf, optimizer) for conf in config.schedulers
]
scheduler = lr_scheduler.ChainedScheduler(schedulers)
elif config.name == "Sequential":
schedulers = [
parse_scheduler_to_instance(conf, optimizer) for conf in config.schedulers
]
scheduler = lr_scheduler.SequentialLR(
optimizer, schedulers, milestones=config.milestones
)
else:
scheduler = getattr(lr_scheduler, config.name)(optimizer, **config.args)
return scheduler
def parse_scheduler(config, optimizer):
interval = config.get("interval", "epoch")
assert interval in ["epoch", "step"]
if config.name == "SequentialLR":
scheduler = {
"scheduler": lr_scheduler.SequentialLR(
optimizer,
[
parse_scheduler(conf, optimizer)["scheduler"]
for conf in config.schedulers
],
milestones=config.milestones,
),
"interval": interval,
}
elif config.name == "ChainedScheduler":
scheduler = {
"scheduler": lr_scheduler.ChainedScheduler(
[
parse_scheduler(conf, optimizer)["scheduler"]
for conf in config.schedulers
]
),
"interval": interval,
}
else:
scheduler = {
"scheduler": get_scheduler(config.name)(optimizer, **config.args),
"interval": interval,
}
return scheduler
Gen_3D_Modules/MV_Adapter/mvadapter/utils/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .mesh_utils.camera import get_camera, get_orthogonal_camera
from .geometry import get_plucker_embeds_from_cameras_ortho
from .saving import make_image_grid, tensor_to_image, image_to_tensor
Gen_3D_Modules/MV_Adapter/mvadapter/utils/base.py
+123
View File
@@ -0,0 +1,123 @@
from dataclasses import dataclass
import torch
import torch.nn as nn
from .config import parse_structured
from .misc import get_device, load_module_weights
from .typing import *
class Configurable:
@dataclass
class Config:
pass
def __init__(self, cfg: Optional[dict] = None) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
class Updateable:
def do_update_step(
self, epoch: int, global_step: int, on_load_weights: bool = False
):
for attr in self.__dir__():
if attr.startswith("_"):
continue
try:
module = getattr(self, attr)
except:
continue # ignore attributes like property, which can't be retrived using getattr?
if isinstance(module, Updateable):
module.do_update_step(
epoch, global_step, on_load_weights=on_load_weights
)
self.update_step(epoch, global_step, on_load_weights=on_load_weights)
def do_update_step_end(self, epoch: int, global_step: int):
for attr in self.__dir__():
if attr.startswith("_"):
continue
try:
module = getattr(self, attr)
except:
continue # ignore attributes like property, which can't be retrived using getattr?
if isinstance(module, Updateable):
module.do_update_step_end(epoch, global_step)
self.update_step_end(epoch, global_step)
def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
# override this method to implement custom update logic
# if on_load_weights is True, you should be careful doing things related to model evaluations,
# as the models and tensors are not guarenteed to be on the same device
pass
def update_step_end(self, epoch: int, global_step: int):
pass
def update_if_possible(module: Any, epoch: int, global_step: int) -> None:
if isinstance(module, Updateable):
module.do_update_step(epoch, global_step)
def update_end_if_possible(module: Any, epoch: int, global_step: int) -> None:
if isinstance(module, Updateable):
module.do_update_step_end(epoch, global_step)
class BaseObject(Updateable):
@dataclass
class Config:
pass
cfg: Config # add this to every subclass of BaseObject to enable static type checking
def __init__(
self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.device = get_device()
self.configure(*args, **kwargs)
def configure(self, *args, **kwargs) -> None:
pass
class BaseModule(nn.Module, Updateable):
@dataclass
class Config:
weights: Optional[str] = None
cfg: Config # add this to every subclass of BaseModule to enable static type checking
def __init__(
self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.device = get_device()
self._non_modules = {}
self.configure(*args, **kwargs)
if self.cfg.weights is not None:
# format: path/to/weights:module_name
weights_path, module_name = self.cfg.weights.split(":")
state_dict, epoch, global_step = load_module_weights(
weights_path, module_name=module_name, map_location="cpu"
)
self.load_state_dict(state_dict)
self.do_update_step(
epoch, global_step, on_load_weights=True
) # restore states
def configure(self, *args, **kwargs) -> None:
pass
def register_non_module(self, name: str, module: nn.Module) -> None:
# non-modules won't be treated as model parameters
self._non_modules[name] = module
def non_module(self, name: str):
return self._non_modules.get(name, None)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/callbacks.py
+158
View File
@@ -0,0 +1,158 @@
import os
import shutil
import subprocess
import pytorch_lightning
from .config import dump_config
from .misc import parse_version
if parse_version(pytorch_lightning.__version__) > parse_version("1.8"):
from pytorch_lightning.callbacks import Callback
else:
from pytorch_lightning.callbacks.base import Callback
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
class VersionedCallback(Callback):
def __init__(self, save_root, version=None, use_version=True):
self.save_root = save_root
self._version = version
self.use_version = use_version
@property
def version(self) -> int:
"""Get the experiment version.
Returns:
The experiment version if specified else the next version.
"""
if self._version is None:
self._version = self._get_next_version()
return self._version
def _get_next_version(self):
existing_versions = []
if os.path.isdir(self.save_root):
for f in os.listdir(self.save_root):
bn = os.path.basename(f)
if bn.startswith("version_"):
dir_ver = os.path.splitext(bn)[0].split("_")[1].replace("/", "")
existing_versions.append(int(dir_ver))
if len(existing_versions) == 0:
return 0
return max(existing_versions) + 1
@property
def savedir(self):
if not self.use_version:
return self.save_root
return os.path.join(
self.save_root,
(
self.version
if isinstance(self.version, str)
else f"version_{self.version}"
),
)
class CodeSnapshotCallback(VersionedCallback):
def __init__(self, save_root, version=None, use_version=True):
super().__init__(save_root, version, use_version)
def get_file_list(self):
return [
b.decode()
for b in set(
subprocess.check_output(
'git ls-files -- ":!:load/*"', shell=True
).splitlines()
)
| set( # hard code, TODO: use config to exclude folders or files
subprocess.check_output(
"git ls-files --others --exclude-standard", shell=True
).splitlines()
)
]
@rank_zero_only
def save_code_snapshot(self):
os.makedirs(self.savedir, exist_ok=True)
for f in self.get_file_list():
if not os.path.exists(f) or os.path.isdir(f):
continue
os.makedirs(os.path.join(self.savedir, os.path.dirname(f)), exist_ok=True)
shutil.copyfile(f, os.path.join(self.savedir, f))
def on_fit_start(self, trainer, pl_module):
try:
self.save_code_snapshot()
except:
rank_zero_warn(
"Code snapshot is not saved. Please make sure you have git installed and are in a git repository."
)
class ConfigSnapshotCallback(VersionedCallback):
def __init__(self, config_path, config, save_root, version=None, use_version=True):
super().__init__(save_root, version, use_version)
self.config_path = config_path
self.config = config
@rank_zero_only
def save_config_snapshot(self):
os.makedirs(self.savedir, exist_ok=True)
dump_config(os.path.join(self.savedir, "parsed.yaml"), self.config)
shutil.copyfile(self.config_path, os.path.join(self.savedir, "raw.yaml"))
def on_fit_start(self, trainer, pl_module):
self.save_config_snapshot()
class CustomProgressBar(TQDMProgressBar):
def get_metrics(self, *args, **kwargs):
# don't show the version number
items = super().get_metrics(*args, **kwargs)
items.pop("v_num", None)
return items
class ProgressCallback(Callback):
def __init__(self, save_path):
super().__init__()
self.save_path = save_path
self._file_handle = None
@property
def file_handle(self):
if self._file_handle is None:
self._file_handle = open(self.save_path, "w")
return self._file_handle
@rank_zero_only
def write(self, msg: str) -> None:
self.file_handle.seek(0)
self.file_handle.truncate()
self.file_handle.write(msg)
self.file_handle.flush()
@rank_zero_only
def on_train_batch_end(self, trainer, pl_module, *args, **kwargs):
self.write(
f"Generation progress: {pl_module.true_global_step / trainer.max_steps * 100:.2f}%"
)
@rank_zero_only
def on_validation_start(self, trainer, pl_module):
self.write(f"Rendering validation image ...")
@rank_zero_only
def on_test_start(self, trainer, pl_module):
self.write(f"Rendering video ...")
@rank_zero_only
def on_predict_start(self, trainer, pl_module):
self.write(f"Exporting mesh assets ...")
Gen_3D_Modules/MV_Adapter/mvadapter/utils/config.py
+140
View File
@@ -0,0 +1,140 @@
import os
from dataclasses import dataclass, field
from datetime import datetime
from omegaconf import OmegaConf
from .core import debug, find, info, warn
from .typing import *
# ============ Register OmegaConf Resolvers ============= #
OmegaConf.register_new_resolver(
"calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n)
)
OmegaConf.register_new_resolver("add", lambda a, b: a + b)
OmegaConf.register_new_resolver("sub", lambda a, b: a - b)
OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
OmegaConf.register_new_resolver("div", lambda a, b: a / b)
OmegaConf.register_new_resolver("idiv", lambda a, b: a // b)
OmegaConf.register_new_resolver("basename", lambda p: os.path.basename(p))
OmegaConf.register_new_resolver("rmspace", lambda s, sub: s.replace(" ", sub))
OmegaConf.register_new_resolver("tuple2", lambda s: [float(s), float(s)])
OmegaConf.register_new_resolver("gt0", lambda s: s > 0)
OmegaConf.register_new_resolver("not", lambda s: not s)
def calc_num_train_steps(num_data, batch_size, max_epochs, num_nodes, num_cards=8):
return int(num_data / (num_nodes * num_cards * batch_size)) * max_epochs
OmegaConf.register_new_resolver("calc_num_train_steps", calc_num_train_steps)
# ======================================================= #
# ============== Automatic Name Resolvers =============== #
def get_naming_convention(cfg):
# TODO
name = f"lrm_{cfg.system.backbone.num_layers}"
return name
# ======================================================= #
@dataclass
class ExperimentConfig:
name: str = "default"
description: str = ""
tag: str = ""
seed: int = 0
use_timestamp: bool = True
timestamp: Optional[str] = None
exp_root_dir: str = "outputs"
### these shouldn't be set manually
exp_dir: str = "outputs/default"
trial_name: str = "exp"
trial_dir: str = "outputs/default/exp"
n_gpus: int = 1
###
resume: Optional[str] = None
data_cls: str = ""
data: dict = field(default_factory=dict)
system_cls: str = ""
system: dict = field(default_factory=dict)
# accept pytorch-lightning trainer parameters
# see https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api
trainer: dict = field(default_factory=dict)
# accept pytorch-lightning checkpoint callback parameters
# see https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
checkpoint: dict = field(default_factory=dict)
def load_config(
*yamls: str, cli_args: list = [], from_string=False, makedirs=True, **kwargs
) -> Any:
if from_string:
parse_func = OmegaConf.create
else:
parse_func = OmegaConf.load
yaml_confs = []
for y in yamls:
conf = parse_func(y)
extends = conf.pop("extends", None)
if extends:
assert os.path.exists(extends), f"File {extends} does not exist."
yaml_confs.append(OmegaConf.load(extends))
yaml_confs.append(conf)
cli_conf = OmegaConf.from_cli(cli_args)
cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs)
OmegaConf.resolve(cfg)
assert isinstance(cfg, DictConfig)
scfg: ExperimentConfig = parse_structured(ExperimentConfig, cfg)
# post processing
# auto naming
if scfg.name == "auto":
scfg.name = get_naming_convention(scfg)
# add timestamp
if not scfg.tag and not scfg.use_timestamp:
raise ValueError("Either tag is specified or use_timestamp is True.")
scfg.trial_name = scfg.tag
# if resume from an existing config, scfg.timestamp should not be None
if scfg.timestamp is None:
scfg.timestamp = ""
if scfg.use_timestamp:
if scfg.n_gpus > 1:
warn(
"Timestamp is disabled when using multiple GPUs, please make sure you have a unique tag."
)
else:
scfg.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S")
# make directories
scfg.trial_name += scfg.timestamp
scfg.exp_dir = os.path.join(scfg.exp_root_dir, scfg.name)
scfg.trial_dir = os.path.join(scfg.exp_dir, scfg.trial_name)
if makedirs:
os.makedirs(scfg.trial_dir, exist_ok=True)
return scfg
def config_to_primitive(config, resolve: bool = True) -> Any:
return OmegaConf.to_container(config, resolve=resolve)
def dump_config(path: str, config) -> None:
with open(path, "w") as fp:
OmegaConf.save(config=config, f=fp)
def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg)
return scfg
Gen_3D_Modules/MV_Adapter/mvadapter/utils/core.py
+29
View File
@@ -0,0 +1,29 @@
import importlib
### grammar sugar for logging utilities ###
import logging
logger = logging.getLogger("pytorch_lightning")
from pytorch_lightning.utilities.rank_zero import (
rank_zero_debug,
rank_zero_info,
rank_zero_only,
)
def find(cls_string):
module_string = ".".join(cls_string.split(".")[:-1])
cls_name = cls_string.split(".")[-1]
module = importlib.import_module(module_string, package=None)
cls = getattr(module, cls_name)
return cls
debug = rank_zero_debug
info = rank_zero_info
@rank_zero_only
def warn(*args, **kwargs):
logger.warn(*args, **kwargs)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/geometry.py
+253
View File
@@ -0,0 +1,253 @@
from typing import List, Optional, Tuple
import numpy as np
import torch
from torch.nn import functional as F
def get_position_map_from_depth(depth, mask, intrinsics, extrinsics, image_wh=None):
"""Compute the position map from the depth map and the camera parameters for a batch of views.
Args:
depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
mask (torch.Tensor): The masks with the shape (B, H, W, 1).
intrinsics (torch.Tensor): The camera intrinsics matrices with the shape (B, 3, 3).
extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
image_wh (Tuple[int, int]): The image width and height.
Returns:
torch.Tensor: The position maps with the shape (B, H, W, 3).
"""
if image_wh is None:
image_wh = depth.shape[2], depth.shape[1]
B, H, W, _ = depth.shape
depth = depth.squeeze(-1)
u_coord, v_coord = torch.meshgrid(
torch.arange(image_wh[0]), torch.arange(image_wh[1]), indexing="xy"
)
u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
# Compute the position map by back-projecting depth pixels to 3D space
x = (
(u_coord - intrinsics[:, 0, 2].unsqueeze(-1).unsqueeze(-1))
* depth
/ intrinsics[:, 0, 0].unsqueeze(-1).unsqueeze(-1)
)
y = (
(v_coord - intrinsics[:, 1, 2].unsqueeze(-1).unsqueeze(-1))
* depth
/ intrinsics[:, 1, 1].unsqueeze(-1).unsqueeze(-1)
)
z = depth
# Concatenate to form the 3D coordinates in the camera frame
camera_coords = torch.stack([x, y, z], dim=-1)
# Apply the extrinsic matrix to get coordinates in the world frame
coords_homogeneous = torch.nn.functional.pad(
camera_coords, (0, 1), "constant", 1.0
) # Add a homogeneous coordinate
world_coords = torch.matmul(
coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
).view(B, H, W, 4)
# Apply the mask to the position map
position_map = world_coords[..., :3] * mask
return position_map
def get_position_map_from_depth_ortho(
depth, mask, extrinsics, ortho_scale, image_wh=None
):
"""Compute the position map from the depth map and the camera parameters for a batch of views
using orthographic projection with a given ortho_scale.
Args:
depth (torch.Tensor): The depth maps with the shape (B, H, W, 1).
mask (torch.Tensor): The masks with the shape (B, H, W, 1).
extrinsics (torch.Tensor): The camera extrinsics matrices with the shape (B, 4, 4).
ortho_scale (torch.Tensor): The scaling factor for the orthographic projection with the shape (B, 1, 1, 1).
image_wh (Tuple[int, int]): Optional. The image width and height.
Returns:
torch.Tensor: The position maps with the shape (B, H, W, 3).
"""
if image_wh is None:
image_wh = depth.shape[2], depth.shape[1]
B, H, W, _ = depth.shape
depth = depth.squeeze(-1)
# Generating grid of coordinates in the image space
u_coord, v_coord = torch.meshgrid(
torch.arange(0, image_wh[0]), torch.arange(0, image_wh[1]), indexing="xy"
)
u_coord = u_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
v_coord = v_coord.type_as(depth).unsqueeze(0).expand(B, -1, -1)
# Compute the position map using orthographic projection with ortho_scale
x = (u_coord - image_wh[0] / 2) * ortho_scale / image_wh[0]
y = (v_coord - image_wh[1] / 2) * ortho_scale / image_wh[1]
z = depth
# Concatenate to form the 3D coordinates in the camera frame
camera_coords = torch.stack([x, y, z], dim=-1)
# Apply the extrinsic matrix to get coordinates in the world frame
coords_homogeneous = torch.nn.functional.pad(
camera_coords, (0, 1), "constant", 1.0
) # Add a homogeneous coordinate
world_coords = torch.matmul(
coords_homogeneous.view(B, -1, 4), extrinsics.transpose(1, 2)
).view(B, H, W, 4)
# Apply the mask to the position map
position_map = world_coords[..., :3] * mask
return position_map
def get_opencv_from_blender(matrix_world, fov=None, image_size=None):
# convert matrix_world to opencv format extrinsics
opencv_world_to_cam = matrix_world.inverse()
opencv_world_to_cam[1, :] *= -1
opencv_world_to_cam[2, :] *= -1
R, T = opencv_world_to_cam[:3, :3], opencv_world_to_cam[:3, 3]
if fov is None: # orthographic camera
return R, T
R, T = R.unsqueeze(0), T.unsqueeze(0)
# convert fov to opencv format intrinsics
focal = 1 / np.tan(fov / 2)
intrinsics = np.diag(np.array([focal, focal, 1])).astype(np.float32)
opencv_cam_matrix = (
torch.from_numpy(intrinsics).unsqueeze(0).float().to(matrix_world.device)
)
opencv_cam_matrix[:, :2, -1] += torch.tensor([image_size / 2, image_size / 2]).to(
matrix_world.device
)
opencv_cam_matrix[:, [0, 1], [0, 1]] *= image_size / 2
return R, T, opencv_cam_matrix
def get_ray_directions(
H: int,
W: int,
focal: float,
principal: Optional[Tuple[float, float]] = None,
use_pixel_centers: bool = True,
) -> torch.Tensor:
"""
Get ray directions for all pixels in camera coordinate.
Args:
H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
Outputs:
directions: (H, W, 3), the direction of the rays in camera coordinate
"""
pixel_center = 0.5 if use_pixel_centers else 0
cx, cy = W / 2, H / 2 if principal is None else principal
i, j = torch.meshgrid(
torch.arange(W, dtype=torch.float32) + pixel_center,
torch.arange(H, dtype=torch.float32) + pixel_center,
indexing="xy",
)
directions = torch.stack(
[(i - cx) / focal, -(j - cy) / focal, -torch.ones_like(i)], -1
)
return F.normalize(directions, dim=-1)
def get_rays(
directions: torch.Tensor, c2w: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Get ray origins and directions from camera coordinates to world coordinates
Args:
directions: (H, W, 3) ray directions in camera coordinates
c2w: (4, 4) camera-to-world transformation matrix
Outputs:
rays_o, rays_d: (H, W, 3) ray origins and directions in world coordinates
"""
# Rotate ray directions from camera coordinate to the world coordinate
rays_d = directions @ c2w[:3, :3].T
rays_o = c2w[:3, 3].expand(rays_d.shape)
return rays_o, rays_d
def compute_plucker_embed(
c2w: torch.Tensor, image_width: int, image_height: int, focal: float
) -> torch.Tensor:
"""
Computes Plucker coordinates for a camera.
Args:
c2w: (4, 4) camera-to-world transformation matrix
image_width: Image width
image_height: Image height
focal: Focal length of the camera
Returns:
plucker: (6, H, W) Plucker embedding
"""
directions = get_ray_directions(image_height, image_width, focal)
rays_o, rays_d = get_rays(directions, c2w)
# Cross product to get Plucker coordinates
cross = torch.cross(rays_o, rays_d, dim=-1)
plucker = torch.cat((rays_d, cross), dim=-1)
return plucker.permute(2, 0, 1)
def get_plucker_embeds_from_cameras(
c2w: List[torch.Tensor], fov: List[float], image_size: int
) -> torch.Tensor:
"""
Given lists of camera transformations and fov, returns the batched plucker embeddings.
Args:
c2w: list of camera-to-world transformation matrices
fov: list of field of view values
image_size: size of the image
Returns:
plucker_embeds: (B, 6, H, W) batched plucker embeddings
"""
plucker_embeds = []
for cam_matrix, cam_fov in zip(c2w, fov):
focal = 0.5 * image_size / np.tan(0.5 * cam_fov)
plucker = compute_plucker_embed(cam_matrix, image_size, image_size, focal)
plucker_embeds.append(plucker)
return torch.stack(plucker_embeds)
def get_plucker_embeds_from_cameras_ortho(
c2w: List[torch.Tensor], ortho_scale: List[float], image_size: int
):
"""
Given lists of camera transformations and fov, returns the batched plucker embeddings.
Parameters:
c2w: list of camera-to-world transformation matrices
fov: list of field of view values
image_size: size of the image
Returns:
plucker_embeds: plucker embeddings (B, 6, H, W)
"""
plucker_embeds = []
# compute pairwise mask and plucker embeddings
for cam_matrix, scale in zip(c2w, ortho_scale):
# blender to opencv to pytorch3d
R, T = get_opencv_from_blender(cam_matrix)
cam_pos = -R.T @ T
view_dir = R.T @ torch.tensor([0, 0, 1]).float().to(cam_matrix.device)
# normalize camera position
cam_pos = F.normalize(cam_pos, dim=0)
plucker = torch.concat([view_dir, cam_pos])
plucker = plucker.unsqueeze(-1).unsqueeze(-1).repeat(1, image_size, image_size)
plucker_embeds.append(plucker)
plucker_embeds = torch.stack(plucker_embeds)
return plucker_embeds
Gen_3D_Modules/MV_Adapter/mvadapter/utils/logging.py
+340
View File
@@ -0,0 +1,340 @@
# coding=utf-8
# Copyright 2024 Optuna, Hugging Face
#
# 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.
""" Logging utilities."""
import logging
import os
import sys
import threading
from logging import CRITICAL # NOQA
from logging import DEBUG # NOQA
from logging import ERROR # NOQA
from logging import FATAL # NOQA
from logging import INFO # NOQA
from logging import NOTSET # NOQA
from logging import WARN # NOQA
from logging import WARNING # NOQA
from typing import Dict, Optional
from tqdm import auto as tqdm_lib
_lock = threading.Lock()
_default_handler: Optional[logging.Handler] = None
log_levels = {
"debug": logging.DEBUG,
"info": logging.INFO,
"warning": logging.WARNING,
"error": logging.ERROR,
"critical": logging.CRITICAL,
}
_default_log_level = logging.INFO
_tqdm_active = True
def _get_default_logging_level() -> int:
"""
If LATEXTURE_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
not - fall back to `_default_log_level`
"""
env_level_str = os.getenv("LATEXTURE_VERBOSITY", None)
if env_level_str:
if env_level_str in log_levels:
return log_levels[env_level_str]
else:
logging.getLogger().warning(
f"Unknown option LATEXTURE_VERBOSITY={env_level_str}, "
f"has to be one of: { ', '.join(log_levels.keys()) }"
)
return _default_log_level
def _get_library_name() -> str:
return __name__.split(".")[0]
def _get_library_root_logger() -> logging.Logger:
return logging.getLogger(_get_library_name())
def _configure_library_root_logger() -> None:
global _default_handler
with _lock:
if _default_handler:
# This library has already configured the library root logger.
return
_default_handler = logging.StreamHandler() # Set sys.stderr as stream.
_default_handler.flush = sys.stderr.flush
# Apply our default configuration to the library root logger.
library_root_logger = _get_library_root_logger()
library_root_logger.addHandler(_default_handler)
library_root_logger.setLevel(_get_default_logging_level())
library_root_logger.propagate = False
enable_explicit_format()
def _reset_library_root_logger() -> None:
global _default_handler
with _lock:
if not _default_handler:
return
library_root_logger = _get_library_root_logger()
library_root_logger.removeHandler(_default_handler)
library_root_logger.setLevel(logging.NOTSET)
_default_handler = None
def get_log_levels_dict() -> Dict[str, int]:
return log_levels
def get_logger(name: Optional[str] = None) -> logging.Logger:
"""
Return a logger with the specified name.
This function is not supposed to be directly accessed unless you are writing a custom diffusers module.
"""
if name is None:
name = _get_library_name()
_configure_library_root_logger()
return logging.getLogger(name)
def get_verbosity() -> int:
"""
Return the current level for the 🤗 Diffusers' root logger as an `int`.
Returns:
`int`:
Logging level integers which can be one of:
- `50`: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
- `40`: `diffusers.logging.ERROR`
- `30`: `diffusers.logging.WARNING` or `diffusers.logging.WARN`
- `20`: `diffusers.logging.INFO`
- `10`: `diffusers.logging.DEBUG`
"""
_configure_library_root_logger()
return _get_library_root_logger().getEffectiveLevel()
def set_verbosity(verbosity: int) -> None:
"""
Set the verbosity level for the 🤗 Diffusers' root logger.
Args:
verbosity (`int`):
Logging level which can be one of:
- `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`
- `diffusers.logging.ERROR`
- `diffusers.logging.WARNING` or `diffusers.logging.WARN`
- `diffusers.logging.INFO`
- `diffusers.logging.DEBUG`
"""
_configure_library_root_logger()
_get_library_root_logger().setLevel(verbosity)
def set_verbosity_info() -> None:
"""Set the verbosity to the `INFO` level."""
return set_verbosity(INFO)
def set_verbosity_warning() -> None:
"""Set the verbosity to the `WARNING` level."""
return set_verbosity(WARNING)
def set_verbosity_debug() -> None:
"""Set the verbosity to the `DEBUG` level."""
return set_verbosity(DEBUG)
def set_verbosity_error() -> None:
"""Set the verbosity to the `ERROR` level."""
return set_verbosity(ERROR)
def disable_default_handler() -> None:
"""Disable the default handler of the 🤗 Diffusers' root logger."""
_configure_library_root_logger()
assert _default_handler is not None
_get_library_root_logger().removeHandler(_default_handler)
def enable_default_handler() -> None:
"""Enable the default handler of the 🤗 Diffusers' root logger."""
_configure_library_root_logger()
assert _default_handler is not None
_get_library_root_logger().addHandler(_default_handler)
def add_handler(handler: logging.Handler) -> None:
"""adds a handler to the HuggingFace Diffusers' root logger."""
_configure_library_root_logger()
assert handler is not None
_get_library_root_logger().addHandler(handler)
def remove_handler(handler: logging.Handler) -> None:
"""removes given handler from the HuggingFace Diffusers' root logger."""
_configure_library_root_logger()
assert handler is not None and handler in _get_library_root_logger().handlers
_get_library_root_logger().removeHandler(handler)
def disable_propagation() -> None:
"""
Disable propagation of the library log outputs. Note that log propagation is disabled by default.
"""
_configure_library_root_logger()
_get_library_root_logger().propagate = False
def enable_propagation() -> None:
"""
Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent
double logging if the root logger has been configured.
"""
_configure_library_root_logger()
_get_library_root_logger().propagate = True
def enable_explicit_format() -> None:
"""
Enable explicit formatting for every 🤗 Diffusers' logger. The explicit formatter is as follows:
```
[LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
```
All handlers currently bound to the root logger are affected by this method.
"""
handlers = _get_library_root_logger().handlers
for handler in handlers:
formatter = logging.Formatter(
"[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s"
)
handler.setFormatter(formatter)
def reset_format() -> None:
"""
Resets the formatting for 🤗 Diffusers' loggers.
All handlers currently bound to the root logger are affected by this method.
"""
handlers = _get_library_root_logger().handlers
for handler in handlers:
handler.setFormatter(None)
def warning_advice(self, *args, **kwargs) -> None:
"""
This method is identical to `logger.warning()`, but if env var LATEXTURE_NO_ADVISORY_WARNINGS=1 is set, this
warning will not be printed
"""
no_advisory_warnings = os.getenv("LATEXTURE_NO_ADVISORY_WARNINGS", False)
if no_advisory_warnings:
return
self.warning(*args, **kwargs)
logging.Logger.warning_advice = warning_advice
class EmptyTqdm:
"""Dummy tqdm which doesn't do anything."""
def __init__(self, *args, **kwargs): # pylint: disable=unused-argument
self._iterator = args[0] if args else None
def __iter__(self):
return iter(self._iterator)
def __getattr__(self, _):
"""Return empty function."""
def empty_fn(*args, **kwargs): # pylint: disable=unused-argument
return
return empty_fn
def __enter__(self):
return self
def __exit__(self, type_, value, traceback):
return
class _tqdm_cls:
def __call__(self, *args, **kwargs):
if _tqdm_active:
return tqdm_lib.tqdm(*args, **kwargs)
else:
return EmptyTqdm(*args, **kwargs)
def set_lock(self, *args, **kwargs):
self._lock = None
if _tqdm_active:
return tqdm_lib.tqdm.set_lock(*args, **kwargs)
def get_lock(self):
if _tqdm_active:
return tqdm_lib.tqdm.get_lock()
tqdm = _tqdm_cls()
def is_progress_bar_enabled() -> bool:
"""Return a boolean indicating whether tqdm progress bars are enabled."""
global _tqdm_active
return bool(_tqdm_active)
def enable_progress_bar() -> None:
"""Enable tqdm progress bar."""
global _tqdm_active
_tqdm_active = True
def disable_progress_bar() -> None:
"""Disable tqdm progress bar."""
global _tqdm_active
_tqdm_active = False
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/__init__.py
+20
View File
@@ -0,0 +1,20 @@
from .camera import (
Camera,
get_c2w,
get_camera,
get_orthogonal_camera,
get_orthogonal_projection_matrix,
get_projection_matrix,
)
from .mesh import TexturedMesh, load_mesh, replace_mesh_texture_and_save
from .projection import CameraProjection, CameraProjectionOutput
from .render import (
DepthControlNetNormalization,
DepthNormalizationStrategy,
NVDiffRastContextWrapper,
RenderOutput,
SimpleNormalization,
Zero123PlusPlusNormalization,
render,
)
from .smart_paint import SmartPainter
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/blend.py
Executable
+324
View File
@@ -0,0 +1,324 @@
import time
from abc import ABC, abstractmethod
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from torch.utils.cpp_extension import load_inline
from .utils import SINGLE_IMAGE_TYPE, image_to_tensor
pb_solver_cpp_source = """
#include <torch/extension.h>
#include <vector>
// CUDA forward declarations
void pb_solver_run_cuda(
torch::Tensor A,
torch::Tensor X,
torch::Tensor B,
torch::Tensor Xbuf,
int num_iters
);
// C++ interface
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
void pb_solver_run(
torch::Tensor A,
torch::Tensor X,
torch::Tensor B,
torch::Tensor Xbuf,
int num_iters
) {
CHECK_INPUT(A);
CHECK_INPUT(X);
CHECK_INPUT(B);
CHECK_INPUT(Xbuf);
pb_solver_run_cuda(A, X, B, Xbuf, num_iters);
return;
}
"""
pb_solver_cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <algorithm>
__global__ void pb_solver_run_cuda_kernel(
torch::PackedTensorAccessor32<int64_t,2,torch::RestrictPtrTraits> A,
torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> X,
torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> B,
torch::PackedTensorAccessor32<float,1,torch::RestrictPtrTraits> Xbuf
) {
const int index = blockIdx.x * blockDim.x + threadIdx.x;
if (index < A.size(0)) {
Xbuf[index] = (X[A[index][0]] + X[A[index][1]] + X[A[index][2]] + X[A[index][3]] + B[index]) / 4.;
}
}
void pb_solver_run_cuda(
torch::Tensor A,
torch::Tensor X,
torch::Tensor B,
torch::Tensor Xbuf,
int num_iters
) {
int batch_size = A.size(0);
const int threads = 1024;
const dim3 blocks((batch_size + threads - 1) / threads);
auto A_ptr = A.packed_accessor32<int64_t,2,torch::RestrictPtrTraits>();
auto X_ptr = X.packed_accessor32<float,1,torch::RestrictPtrTraits>();
auto B_ptr = B.packed_accessor32<float,1,torch::RestrictPtrTraits>();
auto Xbuf_ptr = Xbuf.packed_accessor32<float,1,torch::RestrictPtrTraits>();
for (int i = 0; i < num_iters; ++i) {
pb_solver_run_cuda_kernel<<<blocks, threads>>>(
A_ptr,
X_ptr,
B_ptr,
Xbuf_ptr
);
cudaDeviceSynchronize();
std::swap(X_ptr, Xbuf_ptr);
}
// we may waste an iteration here, but it's fine
return;
}
"""
class PBBackend(ABC):
def __init__(self, *args, **kwargs) -> None:
pass
def solve(self, num_iters, A, X, B, Xbuf) -> None:
pass
try:
@triton.jit
def pb_triton_step_kernel(
A_ptr,
X_ptr,
B_ptr,
Xbuf_ptr,
A_row_stride: tl.constexpr,
n_elements: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0.
block_start = pid * BLOCK_SIZE
offsets = tl.arange(0, BLOCK_SIZE) + block_start
A_start_ptr = A_ptr + block_start * A_row_stride
A_ptrs = A_start_ptr + (
tl.arange(0, BLOCK_SIZE)[:, None] * A_row_stride
+ tl.arange(0, A_row_stride)[None, :]
)
B_ptrs = B_ptr + offsets
Xbuf_ptrs = Xbuf_ptr + offsets
mask = offsets < n_elements
A = tl.load(A_ptrs, mask=mask[:, None], other=0)
X = tl.load(X_ptr + A)
B = tl.load(B_ptrs, mask=mask, other=0.0)
Xout = (tl.sum(X, axis=1) + B) / 4
tl.store(Xbuf_ptrs, Xout, mask=mask)
except:
pb_triton_step_kernel = None
class PBTorchCUDAKernelBackend(PBBackend):
def __init__(self) -> None:
self.kernel = load_inline(
name="pb_solver",
cpp_sources=[pb_solver_cpp_source],
cuda_sources=[pb_solver_cuda_source],
functions=["pb_solver_run"],
verbose=True,
)
def solve(self, num_iters, A, X, B, Xbuf) -> None:
self.kernel.pb_solver_run(A, X, B, Xbuf, num_iters)
class PBTritonBackend(PBBackend):
def step(
self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
) -> None:
n_elements = X.shape[0]
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
assert pb_triton_step_kernel is not None
pb_triton_step_kernel[grid](A, X, B, Xbuf, A.stride(0), n_elements, 1024)
def solve(self, num_iters, A, X, B, Xbuf) -> None:
for _ in range(num_iters):
self.step(X, A, B, Xbuf)
X, Xbuf = Xbuf, X
class PBTorchNativeBackend(PBBackend):
def step(
self, X: torch.Tensor, A: torch.Tensor, B: torch.Tensor, Xbuf: torch.Tensor
) -> None:
X.copy_(X[A].sum(-1).add_(B).div_(4))
def solve(self, num_iters, A, X, B, Xbuf) -> None:
for _ in range(num_iters):
self.step(X, A, B, Xbuf)
class PoissonBlendingSolver:
def __init__(self, backend: str, device: str):
self.backend = backend
if backend == "torch-native":
self.pb_solver = PBTorchNativeBackend()
elif backend == "torch-cuda":
self.pb_solver = PBTorchCUDAKernelBackend()
elif backend == "triton":
self.pb_solver = PBTritonBackend()
else:
raise ValueError(f"Unknown backend: {backend}")
self.device = device
self.lap_kernel = torch.tensor(
[[0, -1, 0], [-1, 4, -1], [0, -1, 0]], device=device, dtype=torch.float32
).view(1, 1, 3, 3)
self.lap_kernel4 = torch.tensor(
[
[[0, -1, 0], [0, 1, 0], [0, 0, 0]],
[[0, 0, 0], [0, 1, 0], [0, -1, 0]],
[[0, 0, 0], [-1, 1, 0], [0, 0, 0]],
[[0, 0, 0], [0, 1, -1], [0, 0, 0]],
],
device=device,
dtype=torch.float32,
).view(4, 1, 3, 3)
self.neighbor_kernel = torch.tensor(
[[0, 1, 0], [1, 0, 1], [0, 1, 0]], device=device, dtype=torch.float32
).view(1, 1, 3, 3)
def __call__(
self,
src: SINGLE_IMAGE_TYPE,
mask: SINGLE_IMAGE_TYPE,
tgt: SINGLE_IMAGE_TYPE,
num_iters: int,
inplace: bool = True,
grad_mode: str = "src",
):
src = image_to_tensor(src, device=self.device)
mask = image_to_tensor(mask, device=self.device)
tgt = image_to_tensor(tgt, device=self.device)
assert src.ndim == 3 and tgt.ndim == 3 and mask.ndim in [2, 3]
if mask.ndim == 3:
mask = mask.mean(-1) > 0.5
else:
mask = mask > 0.5
mask[0, :] = 0
mask[-1, :] = 0
mask[:, 0] = 0
mask[:, -1] = 0
tgt_masked = torch.where(mask[..., None], torch.zeros_like(tgt), tgt)
x, y = mask.nonzero(as_tuple=True)
N = x.shape[0]
index_map = torch.cumsum(mask.reshape(-1).long(), dim=-1).reshape(mask.shape)
index_map[~mask] = 0
if grad_mode == "src":
src_lap = F.conv2d(
src.permute(2, 0, 1)[:, None],
weight=self.lap_kernel,
padding=1,
)[:, 0].permute(1, 2, 0)
lap = src_lap
elif grad_mode == "max":
src_lap = F.conv2d(
src.permute(2, 0, 1)[:, None],
weight=self.lap_kernel4,
padding=1,
)
tgt_lap = F.conv2d(
tgt.permute(2, 0, 1)[:, None],
weight=self.lap_kernel4,
padding=1,
)
lap = (
torch.where(src_lap.abs() > tgt_lap.abs(), src_lap, tgt_lap)
.sum(1)
.permute(1, 2, 0)
)
elif grad_mode == "avg":
src_lap = F.conv2d(
src.permute(2, 0, 1)[:, None],
weight=self.lap_kernel4,
padding=1,
)
tgt_lap = F.conv2d(
tgt.permute(2, 0, 1)[:, None],
weight=self.lap_kernel4,
padding=1,
)
lap = (src_lap + tgt_lap).mul(0.5).sum(1).permute(1, 2, 0)
fq_star = F.conv2d(
tgt_masked.permute(2, 0, 1)[:, None],
weight=self.neighbor_kernel,
padding=1,
)[:, 0].permute(1, 2, 0)
A = torch.zeros(N + 1, 4, device=self.device, dtype=torch.long)
X = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
B = torch.zeros(N + 1, 3, device=self.device, dtype=torch.float32)
A[1:] = torch.stack(
[
index_map[x - 1, y],
index_map[x + 1, y],
index_map[x, y - 1],
index_map[x, y + 1],
],
dim=-1,
)
X[1:] = tgt[x, y]
B[1:] = lap[x, y] + fq_star[x, y]
X_flatten = X.flatten()
B_flatten = B.flatten()
A_flatten = torch.stack([3 * A, 3 * A + 1, 3 * A + 2], dim=1).reshape(-1, 4)
buffer = torch.zeros_like(X_flatten)
self.pb_solver.solve(num_iters, A_flatten, X_flatten, B_flatten, buffer)
if inplace:
tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
else:
tgt = tgt.clone()
tgt[x, y] = X_flatten.view(-1, 3)[1:].clamp(0.0, 1.0)
return tgt
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/camera.py
+223
View File
@@ -0,0 +1,223 @@
import math
from dataclasses import dataclass
from typing import List, Optional, Union
import numpy as np
import torch
import torch.nn.functional as F
import trimesh
from PIL import Image
from torch import BoolTensor, FloatTensor
from .utils import LIST_TYPE
def list_to_pt(
x: LIST_TYPE, dtype: Optional[torch.dtype] = None, device: Optional[str] = None
) -> torch.Tensor:
if isinstance(x, list) or isinstance(x, np.ndarray):
return torch.tensor(x, dtype=dtype, device=device)
return x.to(dtype=dtype)
def get_c2w(
elevation_deg: LIST_TYPE,
distance: LIST_TYPE,
azimuth_deg: Optional[LIST_TYPE],
num_views: Optional[int] = 1,
device: Optional[str] = None,
) -> torch.FloatTensor:
if azimuth_deg is None:
assert (
num_views is not None
), "num_views must be provided if azimuth_deg is None."
azimuth_deg = torch.linspace(
0, 360, num_views + 1, dtype=torch.float32, device=device
)[:-1]
else:
num_views = len(azimuth_deg)
azimuth_deg = list_to_pt(azimuth_deg, dtype=torch.float32, device=device)
elevation_deg = list_to_pt(elevation_deg, dtype=torch.float32, device=device)
camera_distances = list_to_pt(distance, dtype=torch.float32, device=device)
elevation = elevation_deg * math.pi / 180
azimuth = azimuth_deg * math.pi / 180
camera_positions = torch.stack(
[
camera_distances * torch.cos(elevation) * torch.cos(azimuth),
camera_distances * torch.cos(elevation) * torch.sin(azimuth),
camera_distances * torch.sin(elevation),
],
dim=-1,
)
center = torch.zeros_like(camera_positions)
up = torch.tensor([0, 0, 1], dtype=torch.float32, device=device)[None, :].repeat(
num_views, 1
)
lookat = F.normalize(center - camera_positions, dim=-1)
right = F.normalize(torch.cross(lookat, up, dim=-1), dim=-1)
up = F.normalize(torch.cross(right, lookat, dim=-1), dim=-1)
c2w3x4 = torch.cat(
[torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
dim=-1,
)
c2w = torch.cat([c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1)
c2w[:, 3, 3] = 1.0
return c2w
def get_projection_matrix(
fovy_deg: LIST_TYPE,
aspect_wh: float = 1.0,
near: float = 0.1,
far: float = 100.0,
device: Optional[str] = None,
) -> torch.FloatTensor:
fovy_deg = list_to_pt(fovy_deg, dtype=torch.float32, device=device)
batch_size = fovy_deg.shape[0]
fovy = fovy_deg * math.pi / 180
tan_half_fovy = torch.tan(fovy / 2)
projection_matrix = torch.zeros(
batch_size, 4, 4, dtype=torch.float32, device=device
)
projection_matrix[:, 0, 0] = 1 / (aspect_wh * tan_half_fovy)
projection_matrix[:, 1, 1] = -1 / tan_half_fovy
projection_matrix[:, 2, 2] = -(far + near) / (far - near)
projection_matrix[:, 2, 3] = -2 * far * near / (far - near)
projection_matrix[:, 3, 2] = -1
return projection_matrix
def get_orthogonal_projection_matrix(
batch_size: int,
left: float,
right: float,
bottom: float,
top: float,
near: float = 0.1,
far: float = 100.0,
device: Optional[str] = None,
) -> torch.FloatTensor:
projection_matrix = torch.zeros(
batch_size, 4, 4, dtype=torch.float32, device=device
)
projection_matrix[:, 0, 0] = 2 / (right - left)
projection_matrix[:, 1, 1] = -2 / (top - bottom)
projection_matrix[:, 2, 2] = -2 / (far - near)
projection_matrix[:, 0, 3] = -(right + left) / (right - left)
projection_matrix[:, 1, 3] = -(top + bottom) / (top - bottom)
projection_matrix[:, 2, 3] = -(far + near) / (far - near)
projection_matrix[:, 3, 3] = 1
return projection_matrix
@dataclass
class Camera:
c2w: Optional[torch.FloatTensor]
w2c: torch.FloatTensor
proj_mtx: torch.FloatTensor
mvp_mtx: torch.FloatTensor
cam_pos: Optional[torch.FloatTensor]
def __getitem__(self, index):
if isinstance(index, int):
sl = slice(index, index + 1)
elif isinstance(index, slice):
sl = index
elif isinstance(index, list):
sl = index
else:
raise NotImplementedError
return Camera(
c2w=self.c2w[sl] if self.c2w is not None else None,
w2c=self.w2c[sl],
proj_mtx=self.proj_mtx[sl],
mvp_mtx=self.mvp_mtx[sl],
cam_pos=self.cam_pos[sl] if self.cam_pos is not None else None,
)
def to(self, device: Optional[str] = None):
if self.c2w is not None:
self.c2w = self.c2w.to(device)
self.w2c = self.w2c.to(device)
self.proj_mtx = self.proj_mtx.to(device)
self.mvp_mtx = self.mvp_mtx.to(device)
if self.cam_pos is not None:
self.cam_pos = self.cam_pos.to(device)
def __len__(self):
return self.c2w.shape[0]
def get_camera(
elevation_deg: Optional[LIST_TYPE] = None,
distance: Optional[LIST_TYPE] = None,
fovy_deg: Optional[LIST_TYPE] = None,
azimuth_deg: Optional[LIST_TYPE] = None,
num_views: Optional[int] = 1,
c2w: Optional[torch.FloatTensor] = None,
w2c: Optional[torch.FloatTensor] = None,
proj_mtx: Optional[torch.FloatTensor] = None,
aspect_wh: float = 1.0,
near: float = 0.1,
far: float = 100.0,
perturb_camera_position: Optional[float] = None,
device: Optional[str] = None,
):
if w2c is None:
if c2w is None:
c2w = get_c2w(elevation_deg, distance, azimuth_deg, num_views, device)
if perturb_camera_position is not None:
perturbed_pos = (
c2w[:, :3, 3]
+ torch.randn_like(c2w[:, :3, 3]) * perturb_camera_position
)
perturbed_pos = (
F.normalize(perturbed_pos, dim=-1)
* ((c2w[:, :3, 3] ** 2).sum(-1) ** 0.5)[:, None]
)
camera_positions = c2w[:, :3, 3]
w2c = torch.linalg.inv(c2w)
else:
camera_positions = None
c2w = None
if proj_mtx is None:
proj_mtx = get_projection_matrix(
fovy_deg, aspect_wh=aspect_wh, near=near, far=far, device=device
)
mvp_mtx = proj_mtx @ w2c
return Camera(
c2w=c2w, w2c=w2c, proj_mtx=proj_mtx, mvp_mtx=mvp_mtx, cam_pos=camera_positions
)
def get_orthogonal_camera(
elevation_deg: LIST_TYPE,
distance: LIST_TYPE,
left: float,
right: float,
bottom: float,
top: float,
azimuth_deg: Optional[LIST_TYPE] = None,
num_views: Optional[int] = 1,
near: float = 0.1,
far: float = 100.0,
device: Optional[str] = None,
):
c2w = get_c2w(elevation_deg, distance, azimuth_deg, num_views, device)
camera_positions = c2w[:, :3, 3]
w2c = torch.linalg.inv(c2w)
proj_mtx = get_orthogonal_projection_matrix(
batch_size=c2w.shape[0],
left=left,
right=right,
bottom=bottom,
top=top,
near=near,
far=far,
device=device,
)
mvp_mtx = proj_mtx @ w2c
return Camera(
c2w=c2w, w2c=w2c, proj_mtx=proj_mtx, mvp_mtx=mvp_mtx, cam_pos=camera_positions
)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/cv_ops.py
+191
View File
@@ -0,0 +1,191 @@
from typing import Optional
import torch
import torch.nn.functional as F
try:
import cvcuda
torch_to_cvc = lambda x, layout: cvcuda.as_tensor(x, layout)
cvc_to_torch = lambda x, device: torch.tensor(x.cuda(), device=device)
except ImportError:
cvcuda = None
torch_to_cvc = None
cvc_to_torch = None
def inpaint_cvc(
image: torch.Tensor,
mask: torch.Tensor,
padding_size: int,
return_dtype: Optional[torch.dtype] = None,
):
input_dtype = image.dtype
input_device = image.device
image = image.detach()
mask = mask.detach()
if image.dtype != torch.uint8:
image = (image * 255).to(torch.uint8)
if mask.dtype != torch.uint8:
mask = (mask * 255).to(torch.uint8)
image_cvc = torch_to_cvc(image, "HWC")
mask_cvc = torch_to_cvc(mask, "HW")
output_cvc = cvcuda.inpaint(image_cvc, mask_cvc, padding_size)
output = cvc_to_torch(output_cvc, device=input_device)
if return_dtype == torch.uint8 or input_dtype == torch.uint8:
return output
return output.to(dtype=input_dtype) / 255.0
def inpaint_torch(
image: torch.Tensor,
mask: torch.Tensor,
padding_size: int,
return_dtype: Optional[torch.dtype] = None,
):
"""
Inpaint masked areas in an image using PyTorch operations.
Args:
image: Input image tensor [H, W, C]
mask: Binary mask tensor [H, W] where non-zero values indicate pixels to inpaint
padding_size: Size of neighborhood to consider for inpainting
return_dtype: Optional dtype for the output tensor
Returns:
Inpainted image tensor
"""
input_dtype = image.dtype
input_device = image.device
image = image.detach()
mask = mask.detach()
if image.dtype != torch.uint8:
image = (image * 255).to(torch.uint8)
if mask.dtype != torch.uint8:
mask = (mask * 255).to(torch.uint8)
# Convert to float for processing
image_float = image.float() / 255.0
mask_float = (mask > 0).float() # 1 for areas to inpaint, 0 for known areas
# Initialize output with original image
output = image_float.clone()
# Create distance-based weights for neighborhood pixels
kernel_size = min(2 * padding_size + 1, 31) # Limit kernel size to avoid memory issues
y, x = torch.meshgrid(
torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, device=input_device),
torch.arange(-(kernel_size // 2), kernel_size // 2 + 1, device=input_device),
indexing="ij"
)
dist = torch.sqrt(x.float() ** 2 + y.float() ** 2)
weights = 1.0 / (dist + 1e-6)
weights = weights / weights.sum()
weights_kernel = weights.unsqueeze(0).unsqueeze(0)
# For each channel, compute weighted average of non-masked pixels
inpainted = []
for c in range(image_float.shape[-1]):
channel = image_float[..., c]
# Compute weighted sum of valid pixels
valid_values = channel * (1 - mask_float)
weighted_sum = F.conv2d(
valid_values.unsqueeze(0).unsqueeze(0),
weights_kernel,
padding=kernel_size // 2
).squeeze(0).squeeze(0)
# Compute sum of weights for valid pixels
valid_weights = F.conv2d(
(1 - mask_float).unsqueeze(0).unsqueeze(0),
weights_kernel,
padding=kernel_size // 2
).squeeze(0).squeeze(0)
# Avoid division by zero
valid_weights = torch.clamp(valid_weights, min=1e-6)
# Compute weighted average
inpainted_channel = weighted_sum / valid_weights
# Only update masked pixels
inpainted_channel = torch.where(mask_float > 0, inpainted_channel, channel)
inpainted.append(inpainted_channel)
# Combine channels
output = torch.stack(inpainted, dim=-1)
# Convert back to appropriate dtype
if return_dtype == torch.uint8 or input_dtype == torch.uint8:
output = (output * 255).to(torch.uint8)
else:
output = output.to(dtype=input_dtype)
return output
def batch_inpaint_cvc(
images: torch.Tensor,
masks: torch.Tensor,
padding_size: int,
return_dtype: Optional[torch.dtype] = None,
):
output = torch.stack(
[
inpaint_cvc(image, mask, padding_size, return_dtype)
for (image, mask) in zip(images, masks)
],
axis=0,
)
return output
def batch_erode(
masks: torch.Tensor, kernel_size: int, return_dtype: Optional[torch.dtype] = None
):
input_dtype = masks.dtype
input_device = masks.device
masks = masks.detach()
if masks.dtype != torch.uint8:
masks = (masks.float() * 255).to(torch.uint8)
masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
masks_erode_cvc = cvcuda.morphology(
masks_cvc,
cvcuda.MorphologyType.ERODE,
maskSize=(kernel_size, kernel_size),
anchor=(-1, -1),
)
masks_erode = cvc_to_torch(masks_erode_cvc, device=input_device)[..., 0]
if return_dtype == torch.uint8 or input_dtype == torch.uint8:
return masks_erode
return (masks_erode > 0).to(dtype=input_dtype)
def batch_dilate(
masks: torch.Tensor, kernel_size: int, return_dtype: Optional[torch.dtype] = None
):
input_dtype = masks.dtype
input_device = masks.device
masks = masks.detach()
if masks.dtype != torch.uint8:
masks = (masks.float() * 255).to(torch.uint8)
masks_cvc = torch_to_cvc(masks[..., None], "NHWC")
masks_dilate_cvc = cvcuda.morphology(
masks_cvc,
cvcuda.MorphologyType.DILATE,
maskSize=(kernel_size, kernel_size),
anchor=(-1, -1),
)
masks_dilate = cvc_to_torch(masks_dilate_cvc, device=input_device)[..., 0]
if return_dtype == torch.uint8 or input_dtype == torch.uint8:
return masks_dilate
return (masks_dilate > 0).to(dtype=input_dtype)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh.py
+522
View File
@@ -0,0 +1,522 @@
import io
import math
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional, Union
import numpy as np
import nvdiffrast.torch as dr
import torch
import torch.nn.functional as F
import trimesh
from PIL import Image
from torch import BoolTensor, FloatTensor
from .utils import tensor_to_image
def dot(x: torch.FloatTensor, y: torch.FloatTensor) -> torch.FloatTensor:
return torch.sum(x * y, -1, keepdim=True)
@dataclass
class TexturedMesh:
v_pos: torch.FloatTensor
t_pos_idx: torch.LongTensor
# texture coordinates
v_tex: Optional[torch.FloatTensor] = None
t_tex_idx: Optional[torch.LongTensor] = None
# texture map
texture: Optional[torch.FloatTensor] = None
# vertices, faces after vertex merging
_stitched_v_pos: Optional[torch.FloatTensor] = None
_stitched_t_pos_idx: Optional[torch.LongTensor] = None
_v_nrm: Optional[torch.FloatTensor] = None
_v_tang: Optional[torch.FloatTensor] = None
@property
def v_nrm(self) -> torch.FloatTensor:
if self._v_nrm is None:
self._v_nrm = self._compute_vertex_normal()
return self._v_nrm
@property
def v_tang(self) -> torch.FloatTensor:
if self._v_tang is None:
self._v_tang = self._compute_tangent()
return self._v_tang
def set_vertex_normal(self, v_nrm: torch.FloatTensor) -> None:
assert v_nrm.shape == self.v_pos.shape
self._v_nrm = v_nrm.to(self.v_pos)
def set_stitched_mesh(
self, v_pos: torch.FloatTensor, t_pos_idx: torch.LongTensor
) -> None:
self._stitched_v_pos = v_pos
self._stitched_t_pos_idx = t_pos_idx
@property
def stitched_v_pos(self) -> torch.FloatTensor:
if self._stitched_v_pos is None:
print("Warning: Stitched vertices not available, using original vertices!")
return self.v_pos
return self._stitched_v_pos
@property
def stitched_t_pos_idx(self) -> torch.LongTensor:
if self._stitched_t_pos_idx is None:
print("Warning: Stitched faces not available, using original faces!")
return self.t_pos_idx
return self._stitched_t_pos_idx
@property
def uv_size(self) -> Optional[int]:
if self.texture is None:
return None
return self.texture.shape[0]
def _compute_vertex_normal(self) -> torch.FloatTensor:
if self._stitched_v_pos is None or self._stitched_t_pos_idx is None:
print(
"Warning: Stitched vertices and faces not available, computing vertex normals on original mesh, which can be erroneous!"
)
v_pos, t_pos_idx = self.v_pos, self.t_pos_idx
else:
v_pos, t_pos_idx = self._stitched_v_pos, self._stitched_t_pos_idx
i0 = t_pos_idx[:, 0]
i1 = t_pos_idx[:, 1]
i2 = t_pos_idx[:, 2]
v0 = v_pos[i0, :]
v1 = v_pos[i1, :]
v2 = v_pos[i2, :]
face_normals = torch.cross(v1 - v0, v2 - v0)
# Splat face normals to vertices
v_nrm = torch.zeros_like(v_pos)
v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
# Normalize, replace zero (degenerated) normals with some default value
v_nrm = torch.where(
dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
)
v_nrm = F.normalize(v_nrm, dim=1)
if torch.is_anomaly_enabled():
assert torch.all(torch.isfinite(v_nrm))
return v_nrm
def _compute_tangent(self) -> torch.FloatTensor:
vn_idx = [None] * 3
pos = [None] * 3
tex = [None] * 3
for i in range(0, 3):
pos[i] = self.v_pos[self.t_pos_idx[:, i]]
tex[i] = self.v_tex[self.t_tex_idx[:, i]]
vn_idx[i] = self.t_pos_idx[:, i]
tangents = torch.zeros_like(self._v_nrm)
tansum = torch.zeros_like(self._v_nrm)
# Compute tangent space for each triangle
uve1 = tex[1] - tex[0]
uve2 = tex[2] - tex[0]
pe1 = pos[1] - pos[0]
pe2 = pos[2] - pos[0]
nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2]
denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1]
# Avoid division by zero for degenerated texture coordinates
tang = nom / torch.where(
denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6)
)
# Update all 3 vertices
for i in range(0, 3):
idx = vn_idx[i][:, None].repeat(1, 3)
tangents.scatter_add_(0, idx, tang) # tangents[n_i] = tangents[n_i] + tang
tansum.scatter_add_(
0, idx, torch.ones_like(tang)
) # tansum[n_i] = tansum[n_i] + 1
tangents = tangents / tansum
# Normalize and make sure tangent is perpendicular to normal
tangents = F.normalize(tangents, dim=1)
tangents = F.normalize(
tangents - dot(tangents, self._v_nrm) * self._v_nrm, dim=1
)
v_tng = tangents
if torch.is_anomaly_enabled():
assert torch.all(torch.isfinite(v_tng))
return v_tng
def to(self, device: Optional[str] = None):
self.v_pos = self.v_pos.to(device)
self.t_pos_idx = self.t_pos_idx.to(device)
if self.v_tex is not None:
self.v_tex = self.v_tex.to(device)
if self.t_tex_idx is not None:
self.t_tex_idx = self.t_tex_idx.to(device)
if self.texture is not None:
self.texture = self.texture.to(device)
if self._stitched_v_pos is not None:
self._stitched_v_pos = self._stitched_v_pos.to(device)
if self._stitched_t_pos_idx is not None:
self._stitched_t_pos_idx = self._stitched_t_pos_idx.to(device)
if self._v_nrm is not None:
self._v_nrm = self._v_nrm.to(device)
if self._v_tang is not None:
self._v_tang = self._v_tang.to(device)
@contextmanager
def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
texture_ = mesh.texture
mesh.texture = texture
try:
yield
finally:
mesh.texture = texture_
def load_mesh(
mesh_path: str,
rescale: bool = False,
move_to_center: bool = False,
scale: float = 0.5,
flip_uv: bool = True,
merge_vertices: bool = True,
default_uv_size: Optional[int] = None,
shape_init_mesh_up: str = "+y",
shape_init_mesh_front: str = "+x",
front_x_to_y: bool = False,
device: Optional[str] = None,
return_transform: bool = False,
) -> TexturedMesh:
if mesh_path.endswith(".npz"):
class Mesh:
vertices = None
faces = None
data = np.load(mesh_path)
mesh = Mesh()
mesh.vertices = data["vertices"]
mesh.faces = data["faces"]
merge_vertices = False
else:
scene = trimesh.load(mesh_path, force="mesh", process=False)
if isinstance(scene, trimesh.Trimesh):
mesh = scene
elif isinstance(scene, trimesh.scene.Scene):
mesh = trimesh.Trimesh()
for obj in scene.geometry.values():
mesh = trimesh.util.concatenate([mesh, obj])
else:
raise ValueError(f"Unknown mesh type at {mesh_path}.")
vertex_normals = getattr(mesh, "vertex_normals", None)
# move to center
transform_offset = None
if move_to_center:
centroid = mesh.vertices.mean(0)
mesh.vertices = mesh.vertices - centroid
transform_offset = centroid # record the offset
# rescale
transform_scale = None
if rescale:
max_scale = np.abs(mesh.vertices).max()
mesh.vertices = mesh.vertices / max_scale * scale
transform_scale = max_scale / scale
dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
dir2vec = {
"+x": np.array([1, 0, 0]),
"+y": np.array([0, 1, 0]),
"+z": np.array([0, 0, 1]),
"-x": np.array([-1, 0, 0]),
"-y": np.array([0, -1, 0]),
"-z": np.array([0, 0, -1]),
}
if shape_init_mesh_up not in dirs or shape_init_mesh_front not in dirs:
raise ValueError(
f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
)
if shape_init_mesh_up[1] == shape_init_mesh_front[1]:
raise ValueError(
"shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
)
z_, x_ = (
dir2vec[shape_init_mesh_up],
dir2vec[shape_init_mesh_front],
)
y_ = np.cross(z_, x_)
std2mesh = np.stack([x_, y_, z_], axis=0).T
mesh2std = np.linalg.inv(std2mesh)
mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
if vertex_normals is not None:
vertex_normals = np.dot(mesh2std, vertex_normals.T).T
if front_x_to_y:
x = mesh.vertices[:, 1].copy()
y = -mesh.vertices[:, 0].copy()
mesh.vertices[:, 0] = x
mesh.vertices[:, 1] = y
if vertex_normals is not None:
vx = vertex_normals[:, 1].copy()
vy = -vertex_normals[:, 0].copy()
vertex_normals[:, 0] = vx
vertex_normals[:, 1] = vy
v_pos = torch.tensor(mesh.vertices, dtype=torch.float32)
t_pos_idx = torch.tensor(mesh.faces, dtype=torch.int64)
if hasattr(mesh, "visual") and hasattr(mesh.visual, "uv"):
v_tex = torch.tensor(mesh.visual.uv, dtype=torch.float32)
if flip_uv:
v_tex[:, 1] = 1.0 - v_tex[:, 1]
t_tex_idx = t_pos_idx.clone()
if (
default_uv_size is not None
or getattr(mesh.visual.material, "baseColorTexture", None) is None
):
assert default_uv_size is not None
texture = torch.zeros(
(default_uv_size, default_uv_size, 3), dtype=torch.float32
)
else:
texture = torch.tensor(
np.array(mesh.visual.material.baseColorTexture) / 255.0,
dtype=torch.float32,
)[..., :3]
else:
v_tex = None
t_tex_idx = None
texture = None
textured_mesh = TexturedMesh(
v_pos=v_pos,
t_pos_idx=t_pos_idx,
v_tex=v_tex,
t_tex_idx=t_tex_idx,
texture=texture,
)
if vertex_normals is not None:
v_nrm = F.normalize(torch.tensor(vertex_normals, dtype=torch.float32), dim=-1)
textured_mesh.set_vertex_normal(v_nrm)
if merge_vertices and vertex_normals is None:
# only merge vertices when vertex normals are not available
mesh.merge_vertices(merge_tex=True)
textured_mesh.set_stitched_mesh(
torch.tensor(mesh.vertices, dtype=torch.float32),
torch.tensor(mesh.faces, dtype=torch.int64),
)
else:
textured_mesh.set_stitched_mesh(textured_mesh.v_pos, textured_mesh.t_pos_idx)
textured_mesh.to(device)
if return_transform:
return textured_mesh, transform_offset, transform_scale
else:
return textured_mesh
def replace_mesh_texture_and_save_trimesh(
input_path: str,
output_path: str,
texture: Union[Image.Image, np.ndarray, torch.Tensor],
metallic_roughness_texture: Optional[
Union[Image.Image, np.ndarray, torch.Tensor]
] = None,
normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
texture_format: str = "JPEG",
task_id: str = "",
**kwargs,
) -> None:
tmesh = trimesh.load(input_path, force="mesh", process=False)
def convert_image(tensor):
texture = tensor_to_image(tensor)
texture.format = texture_format
return texture
texture = convert_image(texture)
if metallic_roughness_texture is not None:
metallic_roughness_texture = convert_image(metallic_roughness_texture)
if normal_texture is not None:
normal_texture = convert_image(normal_texture)
if not isinstance(tmesh.visual.material, trimesh.visual.material.PBRMaterial):
material = trimesh.visual.texture.PBRMaterial(
baseColorFactor=[1.0, 1.0, 1.0, 1.0],
roughnessFactor=0.9 if metallic_roughness_texture is None else 1.0,
metallicFactor=0.0 if metallic_roughness_texture is None else 1.0,
baseColorTexture=texture,
metallicRoughnessTexture=metallic_roughness_texture,
normalTexture=normal_texture,
)
tmesh.visual = trimesh.visual.TextureVisuals(
uv=tmesh.visual.uv, image=texture, material=material
)
else:
tmesh.visual.material.baseColorTexture = texture
tmesh.visual.material.baseColorFactor = [1.0, 1.0, 1.0, 1.0]
if metallic_roughness_texture is None:
tmesh.visual.material.roughnessFactor = 0.9
tmesh.visual.material.metallicFactor = 0.0
else:
tmesh.visual.material.roughnessFactor = 1.0
tmesh.visual.material.metallicFactor = 1.0
tmesh.visual.material.metallicRoughnessTexture = metallic_roughness_texture
tmesh.visual.material.normalTexture = normal_texture
def p_func(tree, task_id_in=task_id):
if "nodes" in tree and tree["nodes"][0]["name"].find("world") >= 0:
tree["nodes"].pop(0)
tree["asset"]["generator"] = "https://github.com/huanngzh/MV-Adapter"
change_dict = {
"nodes": "node",
"meshes": "mesh",
"materials": "mat",
"images": "image",
}
for k, v in change_dict.items():
if k in tree:
for index in range(len(tree[k])):
tree[k][index]["name"] = f"mvadapter_{v}_{task_id_in}"
tmesh.export(output_path, tree_postprocessor=p_func)
def replace_mesh_texture_and_save_gltflib(
input_path: str,
output_path: str,
texture: Union[Image.Image, np.ndarray, torch.Tensor],
metallic_roughness_texture: Optional[
Union[Image.Image, np.ndarray, torch.Tensor]
] = None,
normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
texture_format: str = "JPEG",
normal_strength: float = 1.0,
task_id: str = "",
**kwargs,
) -> None:
from gltflib import GLTF
from gltflib import Image as gltfImage
from gltflib import NormalTextureInfo, Texture, TextureInfo
def convert_image(tensor):
texture = tensor_to_image(tensor)
texture.format = texture_format
return texture
def add_texture(gltf: GLTF, image: Image.Image) -> int:
jpeg_buffer = io.BytesIO()
image.save(jpeg_buffer, format=texture_format)
new_texture_data = bytearray(jpeg_buffer.getvalue())
_, offset, bytelen = gltf._create_or_extend_glb_resource(new_texture_data)
bufferView = gltf._create_embedded_image_buffer_view(offset, bytelen)
image = gltfImage(mimeType="image/jpeg", bufferView=bufferView)
if gltf.model.images:
gltf.model.images.append(image)
else:
gltf.model.images = [image]
texture = Texture(source=len(gltf.model.images) - 1)
if gltf.model.textures:
gltf.model.textures.append(texture)
else:
gltf.model.textures = [texture]
return len(gltf.model.textures) - 1
gltf = GLTF.load(input_path)
texture = convert_image(texture)
if metallic_roughness_texture is not None:
metallic_roughness_texture = convert_image(metallic_roughness_texture)
if normal_texture is not None:
normal_texture = convert_image(normal_texture)
pbr_material = gltf.model.materials[0].pbrMetallicRoughness
pbr_material.baseColorTexture = TextureInfo(index=add_texture(gltf, texture))
pbr_material.baseColorFactor = [1, 1, 1, 1]
if metallic_roughness_texture is not None:
pbr_material.metallicRoughnessTexture = TextureInfo(
index=add_texture(gltf, metallic_roughness_texture)
)
pbr_material.metallicFactor = 1.0
pbr_material.roughnessFactor = 1.0
else:
pbr_material.metallicFactor = 0.0
pbr_material.roughnessFactor = 0.9
if normal_texture is not None:
gltf.model.materials[0].normalTexture = NormalTextureInfo(
index=add_texture(gltf, normal_texture), scale=normal_strength
)
gltf.model.nodes[0].name = f"mvadapter_node_{task_id}"
gltf.model.meshes[0].name = f"mvadapter_mesh_{task_id}"
gltf.model.materials[0].name = f"mvadapter_material_{task_id}"
gltf.model.asset.generator = "https://github.com/huanngzh/MV-Adapter"
gltf.export(output_path)
def replace_mesh_texture_and_save(
input_path: str,
output_path: str,
texture: Union[Image.Image, np.ndarray, torch.Tensor],
metallic_roughness_texture: Optional[
Union[Image.Image, np.ndarray, torch.Tensor]
] = None,
normal_texture: Optional[Union[Image.Image, np.ndarray, torch.Tensor]] = None,
normal_strength: float = 1.0,
texture_format: str = "JPEG",
task_id: str = "",
backend: str = "trimesh",
) -> None:
if backend == "trimesh":
replace_mesh_texture_and_save_trimesh(
input_path,
output_path,
texture,
metallic_roughness_texture=metallic_roughness_texture,
normal_texture=normal_texture,
texture_format=texture_format,
task_id=task_id,
)
elif backend == "gltflib":
replace_mesh_texture_and_save_gltflib(
input_path,
output_path,
texture,
metallic_roughness_texture=metallic_roughness_texture,
normal_texture=normal_texture,
normal_strength=normal_strength,
task_id=task_id,
)
else:
raise NotImplementedError
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/mesh_process.py
+348
View File
@@ -0,0 +1,348 @@
import numpy as np
import open3d as o3d
import pymeshlab
import torch
import trimesh
from pymeshlab import PercentageValue
### Mesh Utils ###
##### read mesh
def read_mesh_from_path(mesh_path):
ms = pymeshlab.MeshSet()
ms.load_new_mesh(mesh_path)
return ms
def mesh_to_meshlab(vertices, faces):
mesh = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
ms = pymeshlab.MeshSet()
ms.add_mesh(mesh)
return ms
def meshlab_to_mesh(ms):
m = ms.current_mesh()
return m.vertex_matrix(), m.face_matrix(), m.vertex_normal_matrix()
##### decimation
def decimate_quadric_edge_collapse_with_texture(
ms, targetfacenum=None, preservenormal=True, verbose=False
):
# targetfacenum: int, Target number of faces.
# preservenormal: bool, Preserve the normals of the original mesh.
if verbose:
print("Starting decimation ... ")
m = ms.current_mesh()
if targetfacenum is None:
targetfacenum = int(m.face_number() * 0.5)
if verbose:
print("... Initial face number is %d ... " % m.face_number())
ms.meshing_decimation_quadric_edge_collapse_with_texture(
targetfacenum=targetfacenum, preservenormal=preservenormal
)
if verbose:
print("... Decimated face number is %d ... " % m.face_number())
print("Decimation done!\n ")
def decimate_quadric_edge_collapse(
ms, targetfacenum=None, preservenormal=True, verbose=False
):
# targetfacenum: int, Target number of faces.
# preservenormal: bool, Preserve the normals of the original mesh.
if verbose:
print("Starting decimation ... ")
m = ms.current_mesh()
if targetfacenum is None:
targetfacenum = int(m.face_number() * 0.5)
if verbose:
print("... Initial face number is %d ... " % m.face_number())
ms.meshing_decimation_quadric_edge_collapse(
targetfacenum=targetfacenum, preservenormal=preservenormal
)
if verbose:
print("... Decimated face number is %d ... " % m.face_number())
print("Decimation done!\n ")
##### vertex merge
def merge_close_vertices(ms, threshold=0.0001, verbose=False):
# threshold: float, Merge together all the vertices that are nearer than the specified threshold.
if verbose:
print("Starting merge vertices ... ")
m = ms.current_mesh()
if verbose:
print("... Initial vertex number is %d ... " % m.vertex_number())
ms.meshing_merge_close_vertices(threshold=PercentageValue(threshold * 100))
if verbose:
print("... Merged vertex number is %d ... " % m.vertex_number())
print("Merge vertices done!\n ")
##### Island Removal
def remove_isolated_pieces(ms, mincomponentsize=25, diameter=None, verbose=False):
# mincomponentsize: Delete isolated connected components composed by a limited number of triangles
# diameter: Delete isolated connected components whose diameter is smaller than the specified constant
if verbose:
print("Starting remove isolated pieces ... ")
m = ms.current_mesh()
if verbose:
print("... Initial face number is %d ... " % m.face_number())
if diameter is None:
ms.meshing_remove_connected_component_by_face_number(
mincomponentsize=mincomponentsize, removeunref=True
)
else:
ms.meshing_remove_connected_component_by_diameter(
mincomponentdiag=PercentageValue(diameter), removeunref=True
)
if verbose:
print("... Isolated removed face number is %d ... " % m.face_number())
print("Remove isolated pieces done!\n ")
##### hole filling
def fix_hole(ms, maxholesize=30, verbose=False):
# maxholesize: int, Maximum size of the hole to be filled.
if verbose:
print("Starting fix holes ... ")
m = ms.current_mesh()
if verbose:
print("... Initial face number is %d ... " % m.face_number())
ms.meshing_close_holes(maxholesize=maxholesize)
if verbose:
print("... Fixed hole face number is %d ... " % m.face_number())
print("Fix holes done!\n ")
##### repair non manifold edges
def repair_non_manifold(ms, verbose=False):
if verbose:
print("Starting repair non manifold edges ... ")
m = ms.current_mesh()
if verbose:
print("... Initial face number is %d ... " % m.face_number())
ms.meshing_repair_non_manifold_edges()
ms.meshing_repair_non_manifold_vertices(vertdispratio=0.1)
ms.meshing_remove_duplicate_faces()
if verbose:
print("... Fixed non manifold edges face number is %d ... " % m.face_number())
print("Repair non manifold edges done!\n ")
##### laplacian_smooth
def laplacian_smooth(ms, stepsmoothnum=3, verbose=False):
# stepsmoothnum: int, Number of smoothing steps to be performed
if verbose:
print("Starting laplacian smooth ... ")
m = ms.current_mesh()
ms.apply_coord_laplacian_smoothing(stepsmoothnum=stepsmoothnum)
if verbose:
print("Laplacian smooth done!\n ")
##### taubin_smooth
def taubin_smooth(ms, stepsmoothnum=3, verbose=False):
if verbose:
print("Starting Taubin smooth ... ")
m = ms.current_mesh()
ms.apply_coord_taubin_smoothing(stepsmoothnum=stepsmoothnum)
if verbose:
print("Taubin smooth done!\n ")
##### compute_normal
def compute_normal(ms, weightmode="Simple Average", verbose=False):
if verbose:
print("Starting compute_normal_per_vertex ... ")
m = ms.current_mesh()
ms.compute_normal_per_vertex(weightmode=weightmode)
if verbose:
print("compute_normal_per_vertex done!\n ")
### Pre-process Mesh ###
def process_mesh(
vertices,
faces,
threshold=0.0001,
mincomponentRatio=0.02,
targetfacenum=50000,
maxholesize=30,
stepsmoothnum=10,
verbose=False,
):
ms = mesh_to_meshlab(vertices, faces)
### Vertex Merge
merge_close_vertices(ms, threshold=threshold, verbose=verbose)
### Island Removal
faces = ms.current_mesh().face_matrix()
remove_isolated_pieces(
ms, mincomponentsize=int(len(faces) * mincomponentRatio), verbose=verbose
)
### Hole Filling
repair_non_manifold(ms) # repair before fix hole
fix_hole(ms, maxholesize=maxholesize, verbose=verbose)
### Taubin Smoothing
taubin_smooth(ms, stepsmoothnum=stepsmoothnum, verbose=verbose)
vertices, faces, _ = meshlab_to_mesh(ms)
if faces.shape[0] > targetfacenum:
device = o3d.core.Device("CPU:0")
dtype_f = o3d.core.float32
dtype_i = o3d.core.int64
mesh = o3d.t.geometry.TriangleMesh(device)
mesh.vertex.positions = o3d.core.Tensor(
vertices.astype(np.float32), dtype_f, device
)
mesh.triangle.indices = o3d.core.Tensor(faces.astype(np.int64), dtype_i, device)
simplified_mesh = mesh.simplify_quadric_decimation(
target_reduction=1.0 - float(targetfacenum) / faces.shape[0]
)
ms.clear()
vertices = simplified_mesh.vertex.positions.numpy()
faces = simplified_mesh.triangle.indices.numpy()
mesh = pymeshlab.Mesh(vertex_matrix=vertices, face_matrix=faces)
ms.add_mesh(mesh)
### Mesh Simplification/Decimation
# decimate_quadric_edge_collapse(ms, targetfacenum=targetfacenum, verbose=verbose)
taubin_smooth(ms, stepsmoothnum=stepsmoothnum, verbose=verbose)
repair_non_manifold(ms, verbose=verbose)
compute_normal(ms, verbose=verbose)
return meshlab_to_mesh(ms)
### UV Un-Warp ###
def uv_parameterize_uvatlas(
vertices,
faces,
size=1024,
gutter=2.5,
max_stretch=0.1666666716337204,
parallel_partitions=16,
nthreads=0,
):
device = o3d.core.Device("CPU:0")
dtype_f = o3d.core.float32
dtype_i = o3d.core.int64
mesh = o3d.t.geometry.TriangleMesh(device)
mesh.vertex.positions = o3d.core.Tensor(
vertices.astype(np.float32), dtype_f, device
)
mesh.triangle.indices = o3d.core.Tensor(faces.astype(np.int64), dtype_i, device)
mesh.compute_uvatlas(
size=size,
gutter=gutter,
max_stretch=max_stretch,
parallel_partitions=parallel_partitions,
nthreads=nthreads,
)
return mesh.triangle.texture_uvs.numpy() # (#F, 3, 2)
### Pack All ###
def process_raw(mesh_path, save_path, preprocess=True, device="cpu"):
scene = trimesh.load(mesh_path, force="mesh", process=False)
if isinstance(scene, trimesh.Trimesh):
mesh = scene
elif isinstance(scene, trimesh.scene.Scene):
mesh = trimesh.Trimesh()
for obj in scene.geometry.values():
mesh = trimesh.util.concatenate([mesh, obj])
else:
raise ValueError(f"Unknown mesh type at {mesh_path}.")
vertices = mesh.vertices
faces = mesh.faces
mesh_post_process_options = {
"mincomponentRatio": 0.02,
"targetfacenum": 50000,
"maxholesize": 100,
"stepsmoothnum": 10,
"verbose": False,
}
if preprocess:
v_pos, t_pos_idx, normals = process_mesh(
vertices=vertices,
faces=faces,
**mesh_post_process_options,
)
else:
v_pos, t_pos_idx, normals = vertices, faces, mesh.vertex_normals
v_tex_np = (
uv_parameterize_uvatlas(v_pos, t_pos_idx).reshape(-1, 2).astype(np.float32)
)
v_pos = torch.from_numpy(v_pos).to(device=device, dtype=torch.float32)
t_pos_idx = torch.from_numpy(t_pos_idx).to(device=device, dtype=torch.long)
v_tex = torch.from_numpy(v_tex_np).to(device=device, dtype=torch.float32)
normals = torch.from_numpy(normals).to(device=device, dtype=torch.float32)
assert v_tex.shape[0] == t_pos_idx.shape[0] * 3
t_tex_idx = torch.arange(
t_pos_idx.shape[0] * 3,
device=device,
dtype=torch.long,
).reshape(-1, 3)
# uv, index = torch.unique(v_tex, dim=0, return_inverse=True) # 这样实现是2毫秒
# super efficient de-duplication
v_tex_u_uint32 = v_tex_np[..., 0].view(np.uint32)
v_tex_v_uint32 = v_tex_np[..., 1].view(np.uint32)
v_hashed = (v_tex_u_uint32.astype(np.uint64) << 32) | v_tex_v_uint32
v_hashed = torch.from_numpy(v_hashed.view(np.int64)).to(v_pos.device)
t_pos_idx_f3 = torch.arange(
t_pos_idx.shape[0] * 3, device=t_pos_idx.device, dtype=torch.long
).reshape(-1, 3)
v_pos_f3 = v_pos[t_pos_idx].reshape(-1, 3)
normals_f3 = normals[t_pos_idx].reshape(-1, 3)
v_hashed_dedup, inverse_indices = torch.unique(v_hashed, return_inverse=True)
dedup_size, full_size = v_hashed_dedup.shape[0], inverse_indices.shape[0]
indices = torch.scatter_reduce(
torch.full(
[dedup_size],
fill_value=full_size,
device=inverse_indices.device,
dtype=torch.long,
),
index=inverse_indices,
src=torch.arange(full_size, device=inverse_indices.device, dtype=torch.int64),
dim=0,
reduce="amin",
)
v_tex = v_tex[indices]
t_tex_idx = inverse_indices.reshape(-1, 3)
v_pos = v_pos_f3[indices]
normals = normals_f3[indices]
normals = normals.to(dtype=torch.float32, device=device)
# either flip uv or flip texture
# here we flip uv
uv_to_save = v_tex.clone()
uv_to_save[:, 1] = 1.0 - uv_to_save[:, 1]
visual = trimesh.visual.TextureVisuals(uv=uv_to_save.cpu().numpy())
tmesh = trimesh.Trimesh(
vertices=v_pos.cpu().numpy(),
faces=t_tex_idx.cpu().numpy(),
vertex_normals=normals.cpu().numpy(),
visual=visual,
process=False,
)
tmesh.export(save_path)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/projection.py
Executable
+204
View File
@@ -0,0 +1,204 @@
import glob
from dataclasses import dataclass
from typing import List, Optional, Union
import numpy as np
import torch
from einops import rearrange
from PIL import Image
from .blend import PoissonBlendingSolver
from .camera import Camera, get_camera, get_orthogonal_camera
from .mesh import TexturedMesh, load_mesh, replace_mesh_texture_and_save
from .render import NVDiffRastContextWrapper, SimpleNormalization, render
from .seg import SegmentationModel
from .utils import (
IMAGE_TYPE,
LIST_TYPE,
image_to_tensor,
make_image_grid,
tensor_to_image,
)
from .uv import (
ExponentialBlend,
SimpleUVValidityStrategy,
uv_blend,
uv_precompute,
uv_render_attr,
uv_render_geometry,
)
from .warp import compute_warp_field
@dataclass
class CameraProjectionOutput:
uv_proj: torch.FloatTensor
uv_proj_mask: torch.BoolTensor
uv_depth_grad: Optional[torch.FloatTensor]
uv_aoi_cos: Optional[torch.FloatTensor]
class CameraProjection:
def __init__(
self,
pb_backend: str,
bg_remover: Optional[SegmentationModel],
device: str,
context_type: str = "gl",
) -> None:
self.pb_solver = PoissonBlendingSolver(pb_backend, device)
self.ctx = NVDiffRastContextWrapper(device, context_type)
self.bg_remover = bg_remover
self.device = device
def __call__(
self,
images: IMAGE_TYPE,
mesh: TexturedMesh,
cam: Optional[Camera] = None,
fovy_deg: Optional[LIST_TYPE] = None,
masks: Optional[IMAGE_TYPE] = None,
remove_bg: bool = False,
c2w: Optional[torch.FloatTensor] = None,
elevation_deg: Optional[LIST_TYPE] = None,
distance: Optional[LIST_TYPE] = None,
azimuth_deg: Optional[LIST_TYPE] = None,
num_views: Optional[int] = None,
uv_size: int = 2048,
warp_images: bool = False,
images_background: Optional[float] = None,
iou_rejection_threshold: Optional[float] = 0.8,
aoi_cos_valid_threshold: float = 0.3,
depth_grad_dilation: int = 5,
depth_grad_threshold: float = 0.1,
uv_exp_blend_alpha: float = 6,
uv_exp_blend_view_weight: Optional[torch.FloatTensor] = None,
poisson_blending: bool = True,
pb_num_iters: int = 1000,
pb_keep_original_border: bool = True,
from_scratch: bool = False,
uv_padding: bool = True,
return_uv_projection_mask: bool = False,
return_dict: bool = False,
) -> Optional[torch.FloatTensor]:
images_pt = image_to_tensor(images, device=self.device)
assert images_pt.ndim == 4
Nv, H, W, _ = images_pt.shape
if masks is not None:
masks_pt = image_to_tensor(masks, device=self.device)
else:
if remove_bg:
assert self.bg_remover is not None
masks_pt = self.bg_remover(images_pt)
else:
masks_pt = None
if masks_pt is not None and masks_pt.ndim == 4:
masks_pt = masks_pt.mean(-1)
if cam is None:
cam = get_camera(
elevation_deg,
distance,
fovy_deg,
azimuth_deg,
num_views,
c2w,
aspect_wh=W / H,
device=self.device,
)
uv_precompute_output = uv_precompute(
self.ctx, mesh, height=uv_size, width=uv_size
)
uv_render_geometry_output = uv_render_geometry(
self.ctx,
mesh,
cam,
view_height=H,
view_width=W,
uv_precompute_output=uv_precompute_output,
compute_depth_grad=True,
depth_grad_dilation=depth_grad_dilation,
)
# IoU rejection
if masks_pt is not None and iou_rejection_threshold is not None:
given_masks = (masks_pt > 0.5).float()
render_masks = uv_render_geometry_output.view_mask.float()
intersection = given_masks * render_masks
union = given_masks + render_masks - intersection
iou = intersection.sum((1, 2)) / union.sum((1, 2))
iou_min = iou.min()
print(f"Debug: Per view IoU: {iou.tolist()}")
if iou_min < iou_rejection_threshold:
print(
f"Warning: Minimum view IoU {iou_min} below threshold {iou_rejection_threshold}, skipping camera projection!"
)
return None
if warp_images:
# TODO: clean code
assert images_background is not None
render_attr = render(
self.ctx,
mesh,
cam,
height=H,
width=W,
render_attr=True,
attr_background=images_background,
).attr
images_pt = compute_warp_field(
self.ctx.ctx,
images_pt,
render_attr,
n_grid=10,
optim_res=[64, 128],
optim_step_per_res=20,
lambda_reg=2.0,
temp_dir="debug_warp",
verbose=False,
device=self.device,
)
uv_render_attr_output = uv_render_attr(
images=images_pt,
masks=masks_pt,
uv_render_geometry_output=uv_render_geometry_output,
)
uv_blend_output = uv_blend(
uv_precompute_output,
uv_render_geometry_output,
uv_render_attr_output,
uv_validity_strategy=SimpleUVValidityStrategy(
aoi_cos_thresh=aoi_cos_valid_threshold,
depth_grad_thresh=depth_grad_threshold,
),
uv_blend_weight_strategy=ExponentialBlend(
alpha=uv_exp_blend_alpha, view_weight=uv_exp_blend_view_weight
),
empty_value=1.0,
do_uv_padding=uv_padding,
pad_unseen_area=from_scratch,
poisson_blending=poisson_blending,
pb_solver=self.pb_solver,
pb_num_iters=pb_num_iters,
pb_keep_original_border=pb_keep_original_border,
)
if return_dict:
# recommonded new way to get return value
return CameraProjectionOutput(
uv_proj=uv_blend_output.uv_attr_blend,
uv_proj_mask=uv_blend_output.uv_valid_mask_blend,
uv_depth_grad=uv_render_geometry_output.uv_depth_grad,
uv_aoi_cos=uv_render_geometry_output.uv_aoi_cos,
)
else:
if return_uv_projection_mask:
return (
uv_blend_output.uv_attr_blend,
uv_blend_output.uv_valid_mask_blend,
)
return uv_blend_output.uv_attr_blend
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/render.py
+310
View File
@@ -0,0 +1,310 @@
import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional, Union
import numpy as np
import nvdiffrast.torch as dr
import torch
import torch.nn.functional as F
import trimesh
from PIL import Image
from torch import BoolTensor, FloatTensor
from .camera import Camera
from .mesh import TexturedMesh
from .utils import get_clip_space_position, transform_points_homo
@dataclass
class RenderOutput:
attr: Optional[torch.FloatTensor] = None
mask: Optional[torch.BoolTensor] = None
depth: Optional[torch.FloatTensor] = None
normal: Optional[torch.FloatTensor] = None
tangent: Optional[torch.FloatTensor] = None
pos: Optional[torch.FloatTensor] = None
class NVDiffRastContextWrapper:
def __init__(self, device: str, context_type: str = "gl"):
if context_type == "gl":
self.ctx = dr.RasterizeGLContext(device=device)
elif context_type == "cuda":
self.ctx = dr.RasterizeCudaContext(device=device)
else:
raise NotImplementedError
def rasterize(self, pos, tri, resolution, ranges=None, grad_db=True):
"""
Rasterize triangles.
All input tensors must be contiguous and reside in GPU memory except for the ranges tensor that, if specified, has to reside in CPU memory. The output tensors will be contiguous and reside in GPU memory.
Arguments:
glctx Rasterizer context of type RasterizeGLContext or RasterizeCudaContext.
pos Vertex position tensor with dtype torch.float32. To enable range mode, this tensor should have a 2D shape [num_vertices, 4]. To enable instanced mode, use a 3D shape [minibatch_size, num_vertices, 4].
tri Triangle tensor with shape [num_triangles, 3] and dtype torch.int32.
resolution Output resolution as integer tuple (height, width).
ranges In range mode, tensor with shape [minibatch_size, 2] and dtype torch.int32, specifying start indices and counts into tri. Ignored in instanced mode.
grad_db Propagate gradients of image-space derivatives of barycentrics into pos in backward pass. Ignored if using an OpenGL context that was not configured to output image-space derivatives.
Returns:
A tuple of two tensors. The first output tensor has shape [minibatch_size, height, width, 4] and contains the main rasterizer output in order (u, v, z/w, triangle_id). If the OpenGL context was configured to output image-space derivatives of barycentrics, the second output tensor will also have shape [minibatch_size, height, width, 4] and contain said derivatives in order (du/dX, du/dY, dv/dX, dv/dY). Otherwise it will be an empty tensor with shape [minibatch_size, height, width, 0].
"""
return dr.rasterize(
self.ctx,
pos.float().contiguous(),
tri.int().contiguous(),
resolution,
ranges,
grad_db,
)
def interpolate(self, attr, rast, tri, rast_db=None, diff_attrs=None):
"""
Interpolate vertex attributes.
All input tensors must be contiguous and reside in GPU memory. The output tensors will be contiguous and reside in GPU memory.
Arguments:
attr Attribute tensor with dtype torch.float32. Shape is [num_vertices, num_attributes] in range mode, or [minibatch_size, num_vertices, num_attributes] in instanced mode. Broadcasting is supported along the minibatch axis.
rast Main output tensor from rasterize().
tri Triangle tensor with shape [num_triangles, 3] and dtype torch.int32.
rast_db (Optional) Tensor containing image-space derivatives of barycentrics, i.e., the second output tensor from rasterize(). Enables computing image-space derivatives of attributes.
diff_attrs (Optional) List of attribute indices for which image-space derivatives are to be computed. Special value 'all' is equivalent to list [0, 1, ..., num_attributes - 1].
Returns:
A tuple of two tensors. The first output tensor contains interpolated attributes and has shape [minibatch_size, height, width, num_attributes]. If rast_db and diff_attrs were specified, the second output tensor contains the image-space derivatives of the selected attributes and has shape [minibatch_size, height, width, 2 * len(diff_attrs)]. The derivatives of the first selected attribute A will be on channels 0 and 1 as (dA/dX, dA/dY), etc. Otherwise, the second output tensor will be an empty tensor with shape [minibatch_size, height, width, 0].
"""
return dr.interpolate(
attr.float().contiguous(), rast, tri.int().contiguous(), rast_db, diff_attrs
)
def texture(
self,
tex,
uv,
uv_da=None,
mip_level_bias=None,
mip=None,
filter_mode="auto",
boundary_mode="wrap",
max_mip_level=None,
):
"""
Perform texture sampling.
All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory.
Arguments:
tex Texture tensor with dtype torch.float32. For 2D textures, must have shape [minibatch_size, tex_height, tex_width, tex_channels]. For cube map textures, must have shape [minibatch_size, 6, tex_height, tex_width, tex_channels] where tex_width and tex_height are equal. Note that boundary_mode must also be set to 'cube' to enable cube map mode. Broadcasting is supported along the minibatch axis.
uv Tensor containing per-pixel texture coordinates. When sampling a 2D texture, must have shape [minibatch_size, height, width, 2]. When sampling a cube map texture, must have shape [minibatch_size, height, width, 3].
uv_da (Optional) Tensor containing image-space derivatives of texture coordinates. Must have same shape as uv except for the last dimension that is to be twice as long.
mip_level_bias (Optional) Per-pixel bias for mip level selection. If uv_da is omitted, determines mip level directly. Must have shape [minibatch_size, height, width].
mip (Optional) Preconstructed mipmap stack from a texture_construct_mip() call, or a list of tensors specifying a custom mipmap stack. When specifying a custom mipmap stack, the tensors in the list must follow the same format as tex except for width and height that must follow the usual rules for mipmap sizes. The base level texture is still supplied in tex and must not be included in the list. Gradients of a custom mipmap stack are not automatically propagated to base texture but the mipmap tensors will receive gradients of their own. If a mipmap stack is not specified but the chosen filter mode requires it, the mipmap stack is constructed internally and discarded afterwards.
filter_mode Texture filtering mode to be used. Valid values are 'auto', 'nearest', 'linear', 'linear-mipmap-nearest', and 'linear-mipmap-linear'. Mode 'auto' selects 'linear' if neither uv_da or mip_level_bias is specified, and 'linear-mipmap-linear' when at least one of them is specified, these being the highest-quality modes possible depending on the availability of the image-space derivatives of the texture coordinates or direct mip level information.
boundary_mode Valid values are 'wrap', 'clamp', 'zero', and 'cube'. If tex defines a cube map, this must be set to 'cube'. The default mode 'wrap' takes fractional part of texture coordinates. Mode 'clamp' clamps texture coordinates to the centers of the boundary texels. Mode 'zero' virtually extends the texture with all-zero values in all directions.
max_mip_level If specified, limits the number of mipmaps constructed and used in mipmap-based filter modes.
Returns:
A tensor containing the results of the texture sampling with shape [minibatch_size, height, width, tex_channels]. Cube map fetches with invalid uv coordinates (e.g., zero vectors) output all zeros and do not propagate gradients.
"""
return dr.texture(
tex.float(),
uv.float(),
uv_da,
mip_level_bias,
mip,
filter_mode,
boundary_mode,
max_mip_level,
)
def antialias(
self, color, rast, pos, tri, topology_hash=None, pos_gradient_boost=1.0
):
"""
Perform antialiasing.
All input tensors must be contiguous and reside in GPU memory. The output tensor will be contiguous and reside in GPU memory.
Note that silhouette edge determination is based on vertex indices in the triangle tensor. For it to work properly, a vertex belonging to multiple triangles must be referred to using the same vertex index in each triangle. Otherwise, nvdiffrast will always classify the adjacent edges as silhouette edges, which leads to bad performance and potentially incorrect gradients. If you are unsure whether your data is good, check which pixels are modified by the antialias operation and compare to the example in the documentation.
Arguments:
color Input image to antialias with shape [minibatch_size, height, width, num_channels].
rast Main output tensor from rasterize().
pos Vertex position tensor used in the rasterization operation.
tri Triangle tensor used in the rasterization operation.
topology_hash (Optional) Preconstructed topology hash for the triangle tensor. If not specified, the topology hash is constructed internally and discarded afterwards.
pos_gradient_boost (Optional) Multiplier for gradients propagated to pos.
Returns:
A tensor containing the antialiased image with the same shape as color input tensor.
"""
return dr.antialias(
color.float(),
rast,
pos.float(),
tri.int(),
topology_hash,
pos_gradient_boost,
)
class DepthNormalizationStrategy(ABC):
@abstractmethod
def __init__(self, *args, **kwargs):
pass
@abstractmethod
def __call__(
self, depth: torch.FloatTensor, mask: torch.BoolTensor
) -> torch.FloatTensor:
pass
class DepthControlNetNormalization(DepthNormalizationStrategy):
def __init__(
self, far_clip: float = 0.25, near_clip: float = 1.0, bg_value: float = 0.0
):
self.far_clip = far_clip
self.near_clip = near_clip
self.bg_value = bg_value
def __call__(
self, depth: torch.FloatTensor, mask: torch.BoolTensor
) -> torch.FloatTensor:
batch_size = depth.shape[0]
min_depth = depth.view(batch_size, -1).min(dim=-1)[0][:, None, None]
max_depth = depth.view(batch_size, -1).max(dim=-1)[0][:, None, None]
depth = 1.0 - ((depth - min_depth) / (max_depth - min_depth + 1e-5)).clamp(
0.0, 1.0
)
depth = depth * (self.near_clip - self.far_clip) + self.far_clip
depth[~mask] = self.bg_value
return depth
class Zero123PlusPlusNormalization(DepthNormalizationStrategy):
def __init__(self, bg_value: float = 0.8):
self.bg_value = bg_value
def __call__(self, depth: FloatTensor, mask: BoolTensor) -> FloatTensor:
batch_size = depth.shape[0]
min_depth = depth.view(batch_size, -1).min(dim=-1)[0][:, None, None]
max_depth = depth.view(batch_size, -1).max(dim=-1)[0][:, None, None]
depth = ((depth - min_depth) / (max_depth - min_depth + 1e-5)).clamp(0.0, 1.0)
depth[~mask] = self.bg_value
return depth
class SimpleNormalization(DepthNormalizationStrategy):
def __init__(
self,
scale: float = 1.0,
offset: float = -1.0,
clamp: bool = True,
bg_value: float = 1.0,
):
self.scale = scale
self.offset = offset
self.clamp = clamp
self.bg_value = bg_value
def __call__(self, depth: FloatTensor, mask: BoolTensor) -> FloatTensor:
depth = depth * self.scale + self.offset
if self.clamp:
depth = depth.clamp(0.0, 1.0)
depth[~mask] = self.bg_value
return depth
def render(
ctx: NVDiffRastContextWrapper,
mesh: TexturedMesh,
cam: Camera,
height: int,
width: int,
render_attr: bool = True,
render_depth: bool = True,
render_normal: bool = True,
render_tangent: bool = False,
depth_normalization_strategy: DepthNormalizationStrategy = DepthControlNetNormalization(),
attr_background: Union[float, torch.FloatTensor] = 0.5,
antialias_attr=False,
normal_background: Union[float, torch.FloatTensor] = 0.0,
tangent_background: Union[float, torch.FloatTensor] = 0.0,
texture_override=None,
texture_filter_mode: str = "linear",
) -> RenderOutput:
output_dict = {}
v_pos_clip = get_clip_space_position(mesh.v_pos, cam.mvp_mtx)
rast, _ = ctx.rasterize(v_pos_clip, mesh.t_pos_idx, (height, width), grad_db=True)
mask = rast[..., 3] > 0
gb_pos, _ = ctx.interpolate(mesh.v_pos[None], rast, mesh.t_pos_idx)
output_dict.update({"mask": mask, "pos": gb_pos})
if render_depth:
gb_pos_vs = transform_points_homo(gb_pos, cam.w2c)
gb_depth = -gb_pos_vs[..., 2]
# set background pixels to min depth value for correct min/max calculation
gb_depth = torch.where(
mask,
gb_depth,
gb_depth.view(gb_depth.shape[0], -1).min(dim=-1)[0][:, None, None],
)
gb_depth = depth_normalization_strategy(gb_depth, mask)
output_dict["depth"] = gb_depth
if render_attr:
tex_c, _ = ctx.interpolate(mesh.v_tex[None], rast, mesh.t_tex_idx)
texture = (
texture_override[None]
if texture_override is not None
else mesh.texture[None]
)
gb_rgb_fg = ctx.texture(texture, tex_c, filter_mode=texture_filter_mode)
gb_rgb_bg = torch.ones_like(gb_rgb_fg) * attr_background
gb_rgb = torch.where(mask[..., None], gb_rgb_fg, gb_rgb_bg)
if antialias_attr:
gb_rgb = ctx.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx)
output_dict["attr"] = gb_rgb
if render_normal:
gb_nrm, _ = ctx.interpolate(mesh.v_nrm[None], rast, mesh.stitched_t_pos_idx)
gb_nrm = F.normalize(gb_nrm, dim=-1, p=2)
gb_nrm[~mask] = normal_background
output_dict["normal"] = gb_nrm
if render_tangent:
gb_tang, _ = ctx.interpolate(mesh.v_tang[None], rast, mesh.stitched_t_pos_idx)
gb_tang = F.normalize(gb_tang, dim=-1, p=2)
gb_tang[~mask] = tangent_background
output_dict["tangent"] = gb_tang
return RenderOutput(**output_dict)
def tensor_to_image(
data: Union[Image.Image, torch.Tensor, np.ndarray],
batched: bool = False,
format: str = "HWC",
) -> Union[Image.Image, List[Image.Image]]:
if isinstance(data, Image.Image):
return data
if isinstance(data, torch.Tensor):
data = data.detach().cpu().numpy()
if data.dtype == np.float32 or data.dtype == np.float16:
data = (data * 255).astype(np.uint8)
elif data.dtype == np.bool_:
data = data.astype(np.uint8) * 255
assert data.dtype == np.uint8
if format == "CHW":
if batched and data.ndim == 4:
data = data.transpose((0, 2, 3, 1))
elif not batched and data.ndim == 3:
data = data.transpose((1, 2, 0))
if batched:
return [Image.fromarray(d) for d in data]
return Image.fromarray(data)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/seg.py
Executable
+38
View File
@@ -0,0 +1,38 @@
from abc import ABC, abstractmethod
import torch
import transformers
from .utils import IMAGE_TYPE, image_to_tensor
class SegmentationModel(ABC):
def __init__(self, *args, **kwargs):
pass
def __call__(self, images: IMAGE_TYPE) -> torch.FloatTensor:
pass
class RMBGModel(SegmentationModel):
def __init__(self, pretrained_model_name_or_path: str, device: str):
self.model = transformers.AutoModelForImageSegmentation.from_pretrained(
pretrained_model_name_or_path, trust_remote_code=True
).to(device)
self.device = device
def __call__(self, images: IMAGE_TYPE) -> torch.FloatTensor:
images = image_to_tensor(images, device=self.device)
batched = True
if images.ndim == 3:
images = images.unsqueeze(0)
batched = False
out = (
self.model(images.permute(0, 3, 1, 2) - 0.5)[0][0]
.clamp(0.0, 1.0)
.permute(0, 2, 3, 1)
)
if not batched:
out = out.squeeze(0)
return out
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/smart_paint.py
+334
View File
@@ -0,0 +1,334 @@
import glob
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass
from itertools import product
from typing import Any, Dict, List, Optional, Union
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from .camera import Camera, get_camera
from .mesh import (
TexturedMesh,
load_mesh,
mesh_use_texture,
replace_mesh_texture_and_save,
)
from .projection import CameraProjection
from .render import NVDiffRastContextWrapper, render
from .utils import image_to_tensor, make_image_grid, tensor_to_image
from .uv import (
ExponentialBlend,
NVDiffRastContextWrapper,
PoissonBlendingSolver,
SimpleUVValidityStrategy,
uv_blend,
uv_padding,
uv_precompute,
uv_render_attr,
uv_render_geometry,
)
class SmartPainter:
def __init__(self, device: str):
self.device = device
self.cam_proj = CameraProjection(
pb_backend="torch-cuda", bg_remover=None, device=device
)
self.ctx = NVDiffRastContextWrapper(device=self.device)
def __call__(
self,
mod_name: str,
mesh: TexturedMesh,
inpaint_func: callable,
uv_texture: torch.FloatTensor,
uv_inpaint_mask: torch.BoolTensor,
max_view_score_thresh: float = 0.02,
min_rounds: int = 3,
max_rounds: int = 8,
uv_padding_end: bool = True,
debug_dir: Optional[str] = None,
debug_visualize_details: bool = False,
):
def make_view_selection_cam_params():
elevation_anchors = range(-60, 61, 15)
azimuth_anchors = range(0, 360, 30)
distance_anchors = [1.2]
fovy_deg_anchors = [40]
params = list(
product(
elevation_anchors,
azimuth_anchors,
distance_anchors,
fovy_deg_anchors,
)
)
params = list(zip(*params))
return params
elevation_deg, azimuth_deg, distance, fovy_deg = (
make_view_selection_cam_params()
)
view_selection_cams = get_camera(
elevation_deg=list(elevation_deg),
azimuth_deg=list(azimuth_deg),
distance=list(distance),
fovy_deg=list(fovy_deg),
perturb_camera_position=0.1,
device=self.device,
)
texture_update = uv_texture.clone()
uv_valid_mask_update = ~uv_inpaint_mask.clone()
score_map_update = torch.zeros_like(uv_valid_mask_update, dtype=torch.float32)
score_map_update[uv_valid_mask_update] = 1.0
max_view_score = 1
i = 0
while i < min_rounds or (
max_view_score > max_view_score_thresh and i < max_rounds
):
score_map_image = score_map_update.float()[:, :, None].repeat_interleave(
3, dim=-1
)
render_size = 256
with mesh_use_texture(mesh, score_map_image):
render_out = render(
self.ctx,
mesh,
view_selection_cams,
height=render_size,
width=render_size,
attr_background=1.0,
texture_filter_mode="nearest",
)
def get_view_aoi_cos(render_out, cameras):
view_normal = render_out.normal
view_normal_cs = (
view_normal[:, :, :, None, :] * cameras.w2c[:, None, None, :3, :3]
).sum(-1)
view_normal_cs = F.normalize(view_normal_cs, dim=-1, p=2)
view_normal_cs[~render_out.mask] = 0
view_aoi_cos = (
(
view_normal_cs
* torch.as_tensor(
[0.0, 0.0, 1.0],
dtype=view_normal_cs.dtype,
device=view_normal_cs.device,
)[None, None, None]
)
.sum(-1)
.clamp(0, 1)
)
return view_aoi_cos
view_aoi_cos = get_view_aoi_cos(render_out, view_selection_cams)
if debug_dir is not None and debug_visualize_details:
make_image_grid(tensor_to_image(render_out.attr, batched=True)).save(
os.path.join(debug_dir, f"{mod_name}_render_score_{i:02d}.jpg")
)
make_image_grid(tensor_to_image(view_aoi_cos, batched=True)).save(
os.path.join(
debug_dir, f"{mod_name}_render_view_aoi_cos_{i:02d}.jpg"
)
)
view_score = [
(
((attr < 1e-3) & (aoi_cos > 0.1)).sum().item()
+ (
((attr > 1e-3) & (aoi_cos > 0.1)).float()
* (aoi_cos - attr - 0.3).clamp_min(0)
)
.sum()
.item()
)
/ render_size**2
for attr, aoi_cos in zip(render_out.attr[..., 0], view_aoi_cos)
]
max_view_score = np.max(view_score)
best_view = np.argmax(view_score)
inpaint_render_size = 1024
best_cam = view_selection_cams[best_view : best_view + 1]
def shrink_mask(mask: torch.BoolTensor, radius: int) -> torch.BoolTensor:
return (
(
-torch.nn.functional.max_pool2d(
-(mask[None, None].float()),
kernel_size=radius * 2 + 1,
stride=1,
padding=radius,
)
)
.squeeze()
.bool()
)
def enlarge_mask(mask: torch.BoolTensor, radius: int) -> torch.BoolTensor:
return (
torch.nn.functional.max_pool2d(
mask[None, None].float(),
kernel_size=radius * 2 + 1,
stride=1,
padding=radius,
)
.squeeze()
.bool()
)
def blur_mask(mask: torch.BoolTensor, radius: int) -> torch.FloatTensor:
# Create a Gaussian kernel
x = torch.arange(
-radius, radius + 1, dtype=torch.float32, device=mask.device
)
kernel = torch.exp(-(x**2) / (2 * (radius / 2) ** 2))
kernel = kernel[None, None, None]
kernel = kernel * kernel.transpose(2, 3)
kernel = kernel / kernel.sum()
# Apply convolution to blur the mask
blurred_mask = F.conv2d(
mask[None, None].float(), kernel, padding=radius
)
return blurred_mask.squeeze()
def get_occulusion_boundary(
view_depth: torch.FloatTensor, dilation: int, thresh: float
) -> torch.BoolTensor:
grad_x_kernel = torch.tensor(
[[1, 0, -1], [2, 0, -2], [1, 0, -1]],
device=view_depth.device,
dtype=torch.float32,
).view(1, 1, 3, 3)
grad_y_kernel = torch.tensor(
[[1, 2, 1], [0, 0, 0], [-1, -2, -1]],
device=view_depth.device,
dtype=torch.float32,
).view(1, 1, 3, 3)
view_depth_grad_x = F.conv2d(
view_depth[None, None], grad_x_kernel, padding=1
)
view_depth_grad_y = F.conv2d(
view_depth[None, None], grad_y_kernel, padding=1
)
view_depth_grad = (view_depth_grad_x**2 + view_depth_grad_y**2).sqrt()
occ_boundary = (view_depth_grad > thresh)[0][0]
if dilation > 0:
occ_boundary = enlarge_mask(occ_boundary, dilation)
return occ_boundary
with mesh_use_texture(mesh, score_map_image):
render_out = render(
self.ctx,
mesh,
best_cam,
height=inpaint_render_size,
width=inpaint_render_size,
attr_background=1.0,
texture_filter_mode="nearest",
)
view_aoi_cos = get_view_aoi_cos(render_out, best_cam)
inpaint_mask = (render_out.attr[0, :, :, 0] < 1e-3) | (
view_aoi_cos[0] - render_out.attr[0, :, :, 0] > 0.3
)
occ_boundary = get_occulusion_boundary(
render_out.depth[0], dilation=0, thresh=0.1
)
# shrink to remove uv seam
# enlarge the mask gives better results
# remove the occlusion boundary to avoid bleeding
inpaint_mask = enlarge_mask(shrink_mask(inpaint_mask, 3), 5) & ~occ_boundary
if debug_dir is not None and debug_visualize_details:
tensor_to_image(occ_boundary).save(
os.path.join(debug_dir, f"{mod_name}_occ_boundary_{i:02d}.jpg")
)
with mesh_use_texture(mesh, texture_update):
inpaint_image = render(
self.ctx,
mesh,
best_cam,
height=inpaint_render_size,
width=inpaint_render_size,
texture_filter_mode="linear",
).attr[0]
# according to https://github.com/advimman/lama/issues/262
# blur the sharp edge mask gives better results
# but it gives worse results in our case
# inpaint_mask = blur_mask(inpaint_mask, 1)
inpaint_mask_input = inpaint_mask.float()[None, None]
inpaint_image_input = inpaint_image.permute(2, 0, 1)[None]
inpaint_result = inpaint_func(inpaint_image_input, inpaint_mask_input)[
0
].permute(1, 2, 0)
if debug_dir is not None:
make_image_grid(
[
tensor_to_image(inpaint_image),
tensor_to_image(inpaint_mask),
tensor_to_image(inpaint_result),
],
rows=1,
).save(
os.path.join(debug_dir, f"{mod_name}_inpaint_result_{i:02d}.jpg")
)
with mesh_use_texture(mesh, texture_update):
proj_out = self.cam_proj(
inpaint_result[None],
mesh,
best_cam,
masks=inpaint_mask[None].float(),
from_scratch=False,
poisson_blending=False,
depth_grad_dilation=3,
uv_exp_blend_alpha=3,
aoi_cos_valid_threshold=0.1,
uv_size=mesh.uv_size,
uv_padding=True,
iou_rejection_threshold=None,
return_dict=True,
)
uv_proj_inpaint, uv_valid_mask_inpaint = (
proj_out.uv_proj,
proj_out.uv_proj_mask,
)
if debug_dir is not None and debug_visualize_details:
tensor_to_image(uv_valid_mask_inpaint).save(
os.path.join(debug_dir, f"{mod_name}_uv_inpaint_mask_{i:02d}.jpg")
)
texture_update = uv_proj_inpaint
uv_valid_mask_update = uv_valid_mask_inpaint | uv_valid_mask_update
score_map_inpaint = torch.where(
uv_valid_mask_inpaint,
proj_out.uv_aoi_cos[0],
torch.zeros_like(proj_out.uv_aoi_cos[0]),
)
score_map = torch.max(score_map_update, score_map_inpaint)
if debug_dir is not None and debug_visualize_details:
make_image_grid(
[
tensor_to_image(score_map_update),
tensor_to_image(score_map_inpaint),
tensor_to_image(score_map),
]
).save(os.path.join(debug_dir, f"{mod_name}_score_map_{i:02d}.jpg"))
score_map_update = score_map
i += 1
if uv_padding_end:
texture_update = uv_padding(texture_update, uv_valid_mask_update, 3)
return texture_update, uv_valid_mask_update
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/utils.py
+139
View File
@@ -0,0 +1,139 @@
import io
import math
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import List, Optional, Union
import numpy as np
import nvdiffrast.torch as dr
import torch
import torch.nn.functional as F
import trimesh
from PIL import Image
from torch import BoolTensor, FloatTensor
LIST_TYPE = Union[list, np.ndarray, torch.Tensor]
IMAGE_TYPE = Union[Image.Image, List[Image.Image], np.ndarray, torch.Tensor]
SINGLE_IMAGE_TYPE = Union[Image.Image, np.ndarray, torch.Tensor]
def tensor_to_image(
data: Union[Image.Image, torch.Tensor, np.ndarray],
batched: bool = False,
format: str = "HWC",
) -> Union[Image.Image, List[Image.Image]]:
if isinstance(data, Image.Image):
return data
if isinstance(data, torch.Tensor):
data = data.detach().cpu().numpy()
if data.dtype == np.float32 or data.dtype == np.float16:
data = (data * 255).astype(np.uint8)
elif data.dtype == np.bool_:
data = data.astype(np.uint8) * 255
assert data.dtype == np.uint8
if format == "CHW":
if batched and data.ndim == 4:
data = data.transpose((0, 2, 3, 1))
elif not batched and data.ndim == 3:
data = data.transpose((1, 2, 0))
if batched:
return [Image.fromarray(d) for d in data]
return Image.fromarray(data)
def image_to_tensor(image: IMAGE_TYPE, return_type="pt", device: Optional[str] = None):
assert return_type in ["np", "pt"]
batched = True
if isinstance(image, Image.Image):
batched = False
image = [image]
if isinstance(image, list):
image = np.stack([np.array(img) for img in image], axis=0)
image = image.astype(np.float32) / 255.0
if isinstance(image, np.ndarray) and return_type == "pt":
image = torch.tensor(image, device=device)
if isinstance(image, torch.Tensor):
image = image.to(dtype=torch.float32, device=device)
if not batched:
image = image[0]
return image
def largest_factor_near_sqrt(n: int) -> int:
"""
Finds the largest factor of n that is closest to the square root of n.
Args:
n (int): The integer for which to find the largest factor near its square root.
Returns:
int: The largest factor of n that is closest to the square root of n.
"""
sqrt_n = int(math.sqrt(n)) # Get the integer part of the square root
# First, check if the square root itself is a factor
if sqrt_n * sqrt_n == n:
return sqrt_n
# Otherwise, find the largest factor by iterating from sqrt_n downwards
for i in range(sqrt_n, 0, -1):
if n % i == 0:
return i
# If n is 1, return 1
return 1
def make_image_grid(
images: List[Image.Image],
rows: Optional[int] = None,
cols: Optional[int] = None,
resize: Optional[int] = None,
) -> Image.Image:
"""
Prepares a single grid of images. Useful for visualization purposes.
"""
if rows is None and cols is not None:
assert len(images) % cols == 0
rows = len(images) // cols
elif cols is None and rows is not None:
assert len(images) % rows == 0
cols = len(images) // rows
elif rows is None and cols is None:
rows = largest_factor_near_sqrt(len(images))
cols = len(images) // rows
assert len(images) == rows * cols
if resize is not None:
images = [img.resize((resize, resize)) for img in images]
w, h = images[0].size
grid = Image.new("RGB", size=(cols * w, rows * h))
for i, img in enumerate(images):
grid.paste(img, box=(i % cols * w, i // cols * h))
return grid
def get_current_timestamp(fmt: str = "%Y%m%d%H%M%S") -> str:
return datetime.now().strftime(fmt)
def get_clip_space_position(pos: torch.FloatTensor, mvp_mtx: torch.FloatTensor):
pos_homo = torch.cat([pos, torch.ones([pos.shape[0], 1]).to(pos)], dim=-1)
return torch.matmul(pos_homo, mvp_mtx.permute(0, 2, 1))
def transform_points_homo(pos: torch.FloatTensor, mtx: torch.FloatTensor):
batch_size = pos.shape[0]
pos_shape = pos.shape[1:-1]
pos = pos.reshape(batch_size, -1, 3)
pos_homo = torch.cat([pos, torch.ones_like(pos[..., 0:1])], dim=-1)
pos = (pos_homo.unsqueeze(2) * mtx.unsqueeze(1)).sum(-1)[..., :3]
pos = pos.reshape(batch_size, *pos_shape, 3)
return pos
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/uv.py
+468
View File
@@ -0,0 +1,468 @@
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional
import torch
import torch.nn.functional as F
from .blend import PoissonBlendingSolver
from .camera import Camera
from .mesh import TexturedMesh
from .render import NVDiffRastContextWrapper, SimpleNormalization, render
from .utils import IMAGE_TYPE, get_clip_space_position, image_to_tensor
@dataclass
class UVPrecomputeOutput:
height: int
width: int
uv_attr: torch.Tensor
uv_mask: torch.Tensor
uv_pos: torch.Tensor
def uv_precompute(
ctx: NVDiffRastContextWrapper,
mesh: TexturedMesh,
height: int,
width: int,
clone_attr: bool = False,
) -> UVPrecomputeOutput:
uv_clip = mesh.v_tex * 2.0 - 1.0
uv_clip4 = torch.cat(
(
uv_clip,
torch.zeros_like(uv_clip[..., 0:1]),
torch.ones_like(uv_clip[..., 0:1]),
),
dim=-1,
)
rast, _ = ctx.rasterize(uv_clip4[None], mesh.t_tex_idx, (height, width))
uv_mask = rast[0, :, :, 3] > 0
uv_pos, _ = ctx.interpolate(mesh.v_pos[None], rast, mesh.t_pos_idx)
uv_pos = uv_pos[0]
uv_precompute_output = UVPrecomputeOutput(
height=height,
width=width,
uv_attr=mesh.texture.clone() if clone_attr else mesh.texture,
uv_mask=uv_mask,
uv_pos=uv_pos,
)
return uv_precompute_output
@dataclass
class UVRenderGeometryOutput:
uv_pos_proj: torch.Tensor
uv_pos_error: torch.Tensor
uv_aoi_cos: torch.Tensor
uv_pos_ndc: torch.Tensor
view_mask: torch.Tensor
view_normal: torch.Tensor
view_aoi_cos: torch.Tensor
view_position: torch.Tensor
view_depth: torch.Tensor
view_depth_grad: Optional[torch.Tensor] = None
uv_depth_grad: Optional[torch.Tensor] = None
view_attr: Optional[torch.Tensor] = None
def uv_render_geometry(
ctx: NVDiffRastContextWrapper,
mesh: TexturedMesh,
cam: Camera,
view_height: int,
view_width: int,
uv_precompute_output: UVPrecomputeOutput,
grid_sample_mode="bilinear",
compute_depth_grad: bool = False,
depth_grad_dilation: int = 1,
render_attr: bool = False,
) -> UVRenderGeometryOutput:
batch_size = cam.c2w.shape[0]
height, width, _ = uv_precompute_output.uv_pos.shape
uv_pos_clip = get_clip_space_position(
uv_precompute_output.uv_pos.view(-1, 3), cam.mvp_mtx
).view(batch_size, height, width, 4)
uv_pos_ndc = uv_pos_clip[..., :2] / uv_pos_clip[..., 3:4]
render_output = render(
ctx,
mesh,
cam,
view_height,
view_width,
render_attr=render_attr,
render_depth=True,
render_normal=True,
depth_normalization_strategy=SimpleNormalization(
scale=1.0, offset=0.0, clamp=False, bg_value=1e2
),
)
view_position = render_output.pos
view_mask = render_output.mask
view_normal = render_output.normal
view_normal_cs = (
view_normal[:, :, :, None, :] * cam.w2c[:, None, None, :3, :3]
).sum(-1)
view_normal_cs = F.normalize(view_normal_cs, dim=-1, p=2)
view_normal_cs[~render_output.mask] = render_output.normal[~render_output.mask]
view_aoi_cos = (
view_normal_cs
* torch.as_tensor(
[0.0, 0.0, 1.0], dtype=view_normal_cs.dtype, device=view_normal_cs.device
)[None, None, None]
).sum(-1)
view_aoi_cos = view_aoi_cos.clamp(0.0, 1.0)
view_depth = render_output.depth
if compute_depth_grad:
grad_x_kernel = torch.tensor(
[[1, 0, -1], [2, 0, -2], [1, 0, -1]],
device=view_depth.device,
dtype=torch.float32,
).view(1, 1, 3, 3)
grad_y_kernel = torch.tensor(
[[1, 2, 1], [0, 0, 0], [-1, -2, -1]],
device=view_depth.device,
dtype=torch.float32,
).view(1, 1, 3, 3)
view_depth_grad_x = F.conv2d(view_depth[:, None], grad_x_kernel, padding=1)
view_depth_grad_y = F.conv2d(view_depth[:, None], grad_y_kernel, padding=1)
view_depth_grad = (view_depth_grad_x**2 + view_depth_grad_y**2).sqrt()
view_depth_grad = F.max_pool2d(
view_depth_grad,
depth_grad_dilation,
stride=1,
padding=depth_grad_dilation // 2,
)
uv_depth_grad = F.grid_sample(
view_depth_grad, uv_pos_ndc, align_corners=False, mode=grid_sample_mode
).permute(0, 2, 3, 1)[..., 0]
else:
view_depth_grad = None
uv_depth_grad = None
if render_attr:
view_attr = render_output.attr
else:
view_attr = None
uv_pos_proj = F.grid_sample(
view_position.permute(0, 3, 1, 2),
uv_pos_ndc,
align_corners=False,
mode=grid_sample_mode,
).permute(0, 2, 3, 1)
uv_pos_error = ((uv_pos_proj - uv_precompute_output.uv_pos) ** 2).sum(-1).sqrt()
uv_aoi_cos = F.grid_sample(
view_aoi_cos[..., None].permute(0, 3, 1, 2),
uv_pos_ndc,
align_corners=False,
mode=grid_sample_mode,
).permute(0, 2, 3, 1)[..., 0]
return UVRenderGeometryOutput(
uv_pos_proj=uv_pos_proj,
uv_pos_error=uv_pos_error,
uv_aoi_cos=uv_aoi_cos,
uv_pos_ndc=uv_pos_ndc,
view_mask=view_mask,
view_position=view_position,
view_normal=view_normal,
view_aoi_cos=view_aoi_cos,
view_depth=view_depth,
view_depth_grad=view_depth_grad,
uv_depth_grad=uv_depth_grad,
view_attr=view_attr,
)
@dataclass
class UVRenderAttrOutput:
uv_attr_proj: torch.Tensor
uv_mask_proj: Optional[torch.Tensor]
def uv_render_attr(
images: IMAGE_TYPE,
uv_render_geometry_output: UVRenderGeometryOutput,
masks: Optional[IMAGE_TYPE] = None,
grid_sample_mode: str = "bilinear",
) -> UVRenderAttrOutput:
images = image_to_tensor(images, device=uv_render_geometry_output.uv_pos_ndc.device)
uv_attr_proj = F.grid_sample(
images.permute(0, 3, 1, 2),
uv_render_geometry_output.uv_pos_ndc,
align_corners=False,
mode=grid_sample_mode,
).permute(0, 2, 3, 1)
if masks is not None:
masks = image_to_tensor(
masks, device=uv_render_geometry_output.uv_pos_ndc.device
)
if masks.ndim == 4:
masks = masks.mean(-1)
uv_mask_proj = F.grid_sample(
masks[..., None].permute(0, 3, 1, 2),
uv_render_geometry_output.uv_pos_ndc,
align_corners=False,
mode=grid_sample_mode,
).permute(0, 2, 3, 1)[..., 0]
else:
uv_mask_proj = None
return UVRenderAttrOutput(uv_attr_proj=uv_attr_proj, uv_mask_proj=uv_mask_proj)
@dataclass
class UVBlendOutput:
uv_attr_blend: Optional[torch.Tensor]
uv_valid_mask: torch.Tensor
uv_valid_mask_blend: torch.Tensor
uv_blend_weight: torch.Tensor
class UVValidityStrategy(ABC):
@abstractmethod
def __init__(self, *args, **kwargs):
pass
@abstractmethod
def __call__(
self,
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
) -> torch.BoolTensor:
pass
class SimpleUVValidityStrategy(UVValidityStrategy):
def __init__(
self,
pos_error_eps: float = 1e-3,
aoi_cos_thresh: float = 0.1,
mask_thresh: float = 0.9,
depth_grad_thresh: Optional[float] = None,
first_view_dominate: bool = False,
):
self.pos_error_eps = pos_error_eps
self.aoi_cos_thresh = aoi_cos_thresh
self.mask_thresh = mask_thresh
self.depth_grad_thresh = depth_grad_thresh
self.first_view_dominate = first_view_dominate
def __call__(
self,
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
) -> torch.BoolTensor:
uv_pos_valid = uv_render_geometry_output.uv_pos_error < self.pos_error_eps
uv_aoi_cos_valid = uv_render_geometry_output.uv_aoi_cos > self.aoi_cos_thresh
uv_valid_mask = uv_pos_valid & uv_aoi_cos_valid
if self.depth_grad_thresh is not None:
if uv_render_geometry_output.uv_depth_grad is None:
print(
"Warning: Depth gradient is not computed, depth gradient threshold is ignored."
)
else:
uv_valid_mask &= (
uv_render_geometry_output.uv_depth_grad < self.depth_grad_thresh
)
uv_valid_mask &= uv_precompute_output.uv_mask
if (
uv_render_attr_output is not None
and uv_render_attr_output.uv_mask_proj is not None
):
uv_valid_mask &= uv_render_attr_output.uv_mask_proj > self.mask_thresh
else:
print("No view mask provided for UV blending, using all valid pixels")
if self.first_view_dominate:
uv_valid_mask[1:][
uv_valid_mask[0:1].expand(uv_valid_mask.shape[0] - 1, -1, -1)
] = False
return uv_valid_mask
class UVBlendWeightStrategy(ABC):
@abstractmethod
def __init__(self, *args, **kwargs):
pass
@abstractmethod
def __call__(
self,
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
uv_valid_mask: torch.BoolTensor,
) -> torch.FloatTensor:
pass
class ExponentialBlend(UVBlendWeightStrategy):
def __init__(
self,
alpha: float = 1.0,
normalization: str = "linear",
view_weight: Optional[torch.FloatTensor] = None,
):
self.alpha = alpha
self.normalization = normalization
self.view_weight = view_weight
def __call__(
self,
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
uv_valid_mask: torch.BoolTensor,
) -> torch.FloatTensor:
weight = uv_render_geometry_output.uv_aoi_cos
weight = weight * uv_valid_mask.float()
if self.view_weight is not None:
weight = weight ** (self.alpha / self.view_weight[:, None, None].to(weight))
else:
weight = weight**self.alpha
if self.normalization == "linear":
weight_norm = (weight / weight.sum(axis=0, keepdim=True).clamp(1e-5)).clamp(
0.0, 1.0
)
elif self.normalization == "softmax":
weight[~uv_valid_mask] = -1e5
weight_norm = F.softmax(weight, dim=0)
return weight_norm
class RandomChoiceBlend(UVBlendWeightStrategy):
def __init__(self, alpha):
self.alpha = alpha
def __call__(
self,
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
uv_valid_mask: torch.BoolTensor,
) -> torch.FloatTensor:
weight = uv_render_geometry_output.uv_aoi_cos
weight = weight * uv_valid_mask.float()
weight[weight > 0] = torch.rand_like(weight[weight > 0])
weight_norm = (
F.one_hot(weight.max(dim=0).indices, num_classes=weight.shape[0])
.float()
.permute(2, 0, 1)
)
return weight_norm
def uv_padding(attr: torch.FloatTensor, inside_mask: torch.BoolTensor, radius: int):
if attr.min() < -1e-5 or attr.max() > 1 + 1e-5:
print(
"Warning: UV blending result is out of range [0, 1], force clamped before UV padding."
)
from .cv_ops import inpaint_torch
attr_padded = inpaint_torch(attr.clamp(0.0, 1.0), ~inside_mask, radius)
return attr_padded
def uv_blend(
uv_precompute_output: UVPrecomputeOutput,
uv_render_geometry_output: UVRenderGeometryOutput,
uv_render_attr_output: Optional[UVRenderAttrOutput],
uv_validity_strategy: UVValidityStrategy = SimpleUVValidityStrategy(),
uv_blend_weight_strategy: UVBlendWeightStrategy = ExponentialBlend(),
empty_value: float = 0.0,
do_uv_padding: bool = True,
uv_padding_radius: int = 3,
pad_unseen_area: bool = False,
poisson_blending: bool = False,
pb_solver: Optional[PoissonBlendingSolver] = None,
pb_num_iters: int = 1000,
pb_keep_original_border: bool = True,
pb_inplace: bool = False,
pb_grad_mode: str = "src",
) -> UVBlendOutput:
uv_valid_mask = uv_validity_strategy(
uv_precompute_output, uv_render_geometry_output, uv_render_attr_output
)
uv_blend_weight = uv_blend_weight_strategy(
uv_precompute_output,
uv_render_geometry_output,
uv_render_attr_output,
uv_valid_mask,
)
uv_valid_mask_blend = uv_valid_mask.any(dim=0)
if uv_render_attr_output is None:
return UVBlendOutput(
uv_attr_blend=None,
uv_valid_mask=uv_valid_mask,
uv_valid_mask_blend=uv_valid_mask_blend,
uv_blend_weight=uv_blend_weight,
)
uv_attr_blend = (
uv_render_attr_output.uv_attr_proj * uv_blend_weight[..., None]
).sum(axis=0)
if poisson_blending:
assert do_uv_padding
assert pb_solver is not None
uv_attr_blend_padded = uv_padding(
uv_attr_blend, uv_valid_mask_blend, uv_padding_radius
)
if pb_keep_original_border:
pb_tgt = uv_precompute_output.uv_attr
else:
uv_attr_hard_stitch = (
uv_attr_blend * uv_valid_mask_blend[..., None].float()
+ uv_precompute_output.uv_attr
* (~uv_valid_mask_blend)[..., None].float()
)
uv_attr_hard_stitch_padded = uv_padding(
uv_attr_hard_stitch, uv_precompute_output.uv_mask, uv_padding_radius
)
pb_tgt = uv_attr_hard_stitch_padded
uv_attr_blend = pb_solver(
uv_attr_blend_padded,
uv_valid_mask_blend,
pb_tgt,
pb_num_iters,
inplace=pb_inplace,
grad_mode=pb_grad_mode,
)
else:
uv_attr_blend = (
uv_attr_blend * uv_valid_mask_blend[..., None].float()
+ uv_precompute_output.uv_attr * (~uv_valid_mask_blend)[..., None].float()
)
if do_uv_padding:
content_mask = (
uv_valid_mask_blend if pad_unseen_area else uv_precompute_output.uv_mask
)
uv_attr_blend = uv_padding(uv_attr_blend, content_mask, uv_padding_radius)
return UVBlendOutput(
uv_attr_blend=uv_attr_blend,
uv_valid_mask=uv_valid_mask,
uv_valid_mask_blend=uv_valid_mask_blend,
uv_blend_weight=uv_blend_weight,
)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/mesh_utils/warp.py
Executable
+285
View File
@@ -0,0 +1,285 @@
import os
import numpy as np
import nvdiffrast.torch as dr
import torch
import torch.nn.functional as F
import trimesh
from einops import rearrange
from PIL import Image, ImageDraw
from tqdm import tqdm
def draw(img, lines):
"""
lines: n x 4, x1,y2,x2,y2
"""
if isinstance(img, torch.Tensor):
img = img.detach().cpu().numpy()
if isinstance(img, np.ndarray):
if img.dtype == np.float32 or img.dtype == np.float64:
img = (img * 255.0).astype(np.uint8)
img = Image.fromarray(img)
img = img.resize((512, 512), Image.Resampling.BICUBIC)
w, h = img.size
assert w == h
if isinstance(lines, torch.Tensor):
lines = lines.detach().cpu().numpy()
lines = (lines + 1) / 2 * w
img1 = ImageDraw.Draw(img)
for line in lines:
img1.line(line.tolist(), fill="red", width=1)
return img
def construct_grid_mesh(n_grid):
vertices = []
vertices_movable_ids = []
idx = 0
for j in range(n_grid + 1):
for i in range(n_grid + 1):
# 3D fit in trimesh format
if 0 < i < n_grid and 0 < j < n_grid:
vertices_movable_ids.append(idx)
vertices.append([i / n_grid, j / n_grid, 1.0 / 2])
idx += 1
vertices = np.array(vertices)
vertices = 2 * vertices - 1
faces = []
for j in range(n_grid):
for i in range(n_grid):
# clockwise
faces.append(
[
i + j * (n_grid + 1),
i + 1 + j * (n_grid + 1),
i + (j + 1) * (n_grid + 1),
]
)
faces.append(
[
i + 1 + j * (n_grid + 1),
i + 1 + (j + 1) * (n_grid + 1),
i + (j + 1) * (n_grid + 1),
]
)
faces = np.array(faces)
vertices_movable_ids = np.array(vertices_movable_ids)
mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
return mesh, vertices_movable_ids, vertices
def compute_warp_field(
ctx,
src_images_tensor,
tgt_images_tensor,
n_grid,
optim_res,
optim_step_per_res,
lambda_reg,
temp_dir,
verbose,
device,
):
"""
src_images_tensor: 4 H W 4
tgt_images_tensor: 4 H W 4
return: 4 H W 4
"""
lam_reg = lambda_reg
lam_mask = 2
mesh, _, vertices_np = construct_grid_mesh(n_grid)
# vertices = mesh.vertices
faces = mesh.faces
edges = mesh.edges_unique
n_degree = torch.tensor(mesh.vertex_degree)
vertices = torch.tensor(vertices_np, device=device, dtype=torch.float32)
vertices_unopt = vertices.clone()
faces = torch.tensor(faces, device=device)
edges = torch.tensor(edges, device=device)
warped_images_tensor = []
bs = src_images_tensor.shape[0]
for img_idx in range(bs):
src_image_tensor = src_images_tensor[img_idx]
tgt_image_tensor = tgt_images_tensor[img_idx]
if verbose and temp_dir is not None:
vis_dir = os.path.join(temp_dir, f"{img_idx}")
os.makedirs(vis_dir, exist_ok=True)
# prepare for nvdiff rendering
v_origin_homo = torch.cat(
[vertices_unopt, torch.ones_like(vertices_unopt[..., :1])], dim=-1
)
# prepare parameters
v_move_indices = torch.where(n_degree == 6)[0]
vertices_movable = vertices[v_move_indices].detach().clone()
vertices_movable.requires_grad = True
opt = torch.optim.Adam([vertices_movable], lr=0.02)
for resl in optim_res:
rast, _ = dr.rasterize(
ctx,
v_origin_homo[None, ...].float(),
faces.int(),
(resl, resl),
grad_db=True,
)
face_ids = rast[..., 3].long() - 1
assert torch.all(face_ids >= 0)
u = rast[..., 0]
v = rast[..., 1]
pixel_vertice_ids = faces[face_ids]
pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
src_img = F.interpolate(
rearrange(src_image_tensor[..., :3], "H W C -> () C H W"),
size=(resl, resl),
mode="bilinear",
align_corners=False,
antialias=True,
)
tgt_img = F.interpolate(
rearrange(tgt_image_tensor[..., :3], "H W C -> () C H W"),
size=(resl, resl),
mode="bilinear",
align_corners=False,
antialias=True,
)
# src_img, tgt_img: 1 C H W
if verbose:
Image.fromarray(
(tgt_img[0].permute(1, 2, 0).detach().cpu().numpy() * 255.0).astype(
np.uint8
)
).save(os.path.join(vis_dir, f"target_{resl:04d}.png"))
with tqdm(range(optim_step_per_res), disable=not verbose) as pbar:
for i in pbar:
opt.zero_grad()
vertices_all = vertices_unopt.detach().clone()
vertices_all[v_move_indices] = vertices_movable
pixel_vertices = vertices_all[
pixel_vertice_ids
] # 1 x 512 x 512 x 3 x 3
pixel_coords_inter = torch.sum(
pixel_vertices * pixel_bary_coords[..., None], dim=-2
)
src_img_warped = F.grid_sample(
src_img,
pixel_coords_inter[..., :2],
mode="bilinear",
align_corners=False,
)
edge_vertices_all = vertices_all[edges]
edge_vertices_unopt = vertices_unopt[edges]
edge_len_all = torch.linalg.norm(
edge_vertices_all[:, 0, :2] - edge_vertices_all[:, 1, :2],
dim=-1,
)
edge_len_unopt = torch.linalg.norm(
edge_vertices_unopt[:, 0, :2] - edge_vertices_all[:, 1, :2],
dim=-1,
)
reg_loss = ((edge_len_all - edge_len_unopt) ** 2).mean()
# mask_loss = ((src_img_warped[:, 3:4, :, :] - tgt_img[:, 3:4, :, :])**2).mean()
# rgb_loss = (((src_img_warped[:, :3, :, :] - tgt_img[:, :3, :, :]) * tgt_img[:, 3:4, :, :])**2).mean()
# loss = rgb_loss + lam_reg * reg_loss + lam_mask * mask_loss
img_loss = ((src_img_warped - tgt_img) ** 2).mean()
loss = img_loss + lam_reg * reg_loss
loss.backward()
opt.step()
if verbose:
print(img_loss.item(), reg_loss.item())
if verbose:
draw(
src_img[0].permute(1, 2, 0),
edge_vertices_all[:, :, :2].reshape(-1, 4),
).save(os.path.join(vis_dir, f"src_{resl:04d}_{i:03d}.png"))
Image.fromarray(
(
torch.cat(
[
tgt_img,
src_img_warped,
(tgt_img - src_img_warped).abs(),
],
dim=-1,
)[0]
.permute(1, 2, 0)
.detach()
.cpu()
.numpy()
* 255.0
).astype(np.uint8)
).save(os.path.join(vis_dir, f"opt_{resl:04d}_{i:03d}.png"))
warped_this_resl = Image.fromarray(
(
src_img_warped[0]
.permute(1, 2, 0)
.detach()
.cpu()
.numpy()
* 255.0
).astype(np.uint8)
)
warped_this_resl.save(
os.path.join(vis_dir, f"warped_{resl:04d}_{i:03d}.png")
)
# warp on 512 resolution
resl = src_image_tensor.shape[1]
with torch.no_grad():
rast, _ = dr.rasterize(
ctx,
v_origin_homo[None, ...].float(),
faces.int(),
(resl, resl),
grad_db=True,
)
face_ids = rast[..., 3].long() - 1
assert torch.all(face_ids >= 0)
u = rast[..., 0]
v = rast[..., 1]
pixel_vertice_ids = faces[face_ids]
pixel_bary_coords = torch.stack([u, v, 1 - u - v], dim=-1)
vertices_all = vertices_unopt.detach().clone()
vertices_all[v_move_indices] = vertices_movable
pixel_vertices = vertices_all[pixel_vertice_ids] # 1 x 512 x 512 x 3 x 3
pixel_coords_inter = torch.sum(
pixel_vertices * pixel_bary_coords[..., None], dim=-2
)
src_img_warped = rearrange(
F.grid_sample(
rearrange(src_image_tensor, "H W C -> () C H W"),
pixel_coords_inter[..., :2],
mode="bicubic",
align_corners=False,
),
"() C H W -> H W C",
).clamp(0, 1)
warped_images_tensor.append(src_img_warped)
warped_images_tensor = torch.stack(warped_images_tensor, dim=0) # 4 H W 3
return warped_images_tensor
Gen_3D_Modules/MV_Adapter/mvadapter/utils/misc.py
+220
View File
@@ -0,0 +1,220 @@
import gc
import os
import re
import time
from collections import defaultdict
from contextlib import contextmanager
import psutil
import torch
from packaging import version
from .config import config_to_primitive
from .core import debug, find, info, warn
from .typing import *
def parse_version(ver: str):
return version.parse(ver)
def get_rank():
# SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
# therefore LOCAL_RANK needs to be checked first
rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
for key in rank_keys:
rank = os.environ.get(key)
if rank is not None:
return int(rank)
return 0
def get_device():
return torch.device(f"cuda:{get_rank()}")
def load_module_weights(
path, module_name=None, ignore_modules=None, mapping=None, map_location=None
) -> Tuple[dict, int, int]:
if module_name is not None and ignore_modules is not None:
raise ValueError("module_name and ignore_modules cannot be both set")
if map_location is None:
map_location = get_device()
ckpt = torch.load(path, map_location=map_location)
state_dict = ckpt["state_dict"]
if mapping is not None:
state_dict_to_load = {}
for k, v in state_dict.items():
if any([k.startswith(m["to"]) for m in mapping]):
pass
else:
state_dict_to_load[k] = v
for k, v in state_dict.items():
for m in mapping:
if k.startswith(m["from"]):
k_dest = k.replace(m["from"], m["to"])
info(f"Mapping {k} => {k_dest}")
state_dict_to_load[k_dest] = v.clone()
state_dict = state_dict_to_load
state_dict_to_load = state_dict
if ignore_modules is not None:
state_dict_to_load = {}
for k, v in state_dict.items():
ignore = any(
[k.startswith(ignore_module + ".") for ignore_module in ignore_modules]
)
if ignore:
continue
state_dict_to_load[k] = v
if module_name is not None:
state_dict_to_load = {}
for k, v in state_dict.items():
m = re.match(rf"^{module_name}\.(.*)$", k)
if m is None:
continue
state_dict_to_load[m.group(1)] = v
return state_dict_to_load, ckpt["epoch"], ckpt["global_step"]
def C(value: Any, epoch: int, global_step: int) -> float:
if isinstance(value, int) or isinstance(value, float):
pass
else:
value = config_to_primitive(value)
if not isinstance(value, list):
raise TypeError("Scalar specification only supports list, got", type(value))
if len(value) == 3:
value = [0] + value
assert len(value) == 4
start_step, start_value, end_value, end_step = value
if isinstance(end_step, int):
current_step = global_step
value = start_value + (end_value - start_value) * max(
min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
)
elif isinstance(end_step, float):
current_step = epoch
value = start_value + (end_value - start_value) * max(
min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
)
return value
def cleanup():
gc.collect()
torch.cuda.empty_cache()
try:
import tinycudann as tcnn
tcnn.free_temporary_memory()
except:
pass
def finish_with_cleanup(func: Callable):
def wrapper(*args, **kwargs):
out = func(*args, **kwargs)
cleanup()
return out
return wrapper
def _distributed_available():
return torch.distributed.is_available() and torch.distributed.is_initialized()
def barrier():
if not _distributed_available():
return
else:
torch.distributed.barrier()
def broadcast(tensor, src=0):
if not _distributed_available():
return tensor
else:
torch.distributed.broadcast(tensor, src=src)
return tensor
def enable_gradient(model, enabled: bool = True) -> None:
for param in model.parameters():
param.requires_grad_(enabled)
class TimeRecorder:
_instance = None
def __init__(self):
self.items = {}
self.accumulations = defaultdict(list)
self.time_scale = 1000.0 # ms
self.time_unit = "ms"
self.enabled = False
def __new__(cls):
# singleton
if cls._instance is None:
cls._instance = super(TimeRecorder, cls).__new__(cls)
return cls._instance
def enable(self, enabled: bool) -> None:
self.enabled = enabled
def start(self, name: str) -> None:
if not self.enabled:
return
torch.cuda.synchronize()
self.items[name] = time.time()
def end(self, name: str, accumulate: bool = False) -> float:
if not self.enabled or name not in self.items:
return
torch.cuda.synchronize()
start_time = self.items.pop(name)
delta = time.time() - start_time
if accumulate:
self.accumulations[name].append(delta)
t = delta * self.time_scale
info(f"{name}: {t:.2f}{self.time_unit}")
def get_accumulation(self, name: str, average: bool = False) -> float:
if not self.enabled or name not in self.accumulations:
return
acc = self.accumulations.pop(name)
total = sum(acc)
if average:
t = total / len(acc) * self.time_scale
else:
t = total * self.time_scale
info(f"{name} for {len(acc)} times: {t:.2f}{self.time_unit}")
### global time recorder
time_recorder = TimeRecorder()
@contextmanager
def time_recorder_enabled():
enabled = time_recorder.enabled
time_recorder.enable(enabled=True)
try:
yield
finally:
time_recorder.enable(enabled=enabled)
def show_vram_usage(name):
available, total = torch.cuda.mem_get_info()
used = total - available
print(
f"{name}: {used / 1024**2:.1f}MB, {psutil.Process(os.getpid()).memory_info().rss / 1024**2:.1f}MB"
)
Gen_3D_Modules/MV_Adapter/mvadapter/utils/ops.py
+462
View File
@@ -0,0 +1,462 @@
import math
from collections import defaultdict
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd
from .core import debug, find, info, warn
from .typing import *
def dot(x, y):
return torch.sum(x * y, -1, keepdim=True)
def reflect(x, n):
return 2 * dot(x, n) * n - x
ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]]
def scale_tensor(
dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale
):
if inp_scale is None:
inp_scale = (0, 1)
if tgt_scale is None:
tgt_scale = (0, 1)
if isinstance(tgt_scale, Tensor):
assert dat.shape[-1] == tgt_scale.shape[-1]
dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
return dat
class _TruncExp(Function): # pylint: disable=abstract-method
# Implementation from torch-ngp:
# https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, x): # pylint: disable=arguments-differ
ctx.save_for_backward(x)
return torch.exp(x)
@staticmethod
@custom_bwd
def backward(ctx, g): # pylint: disable=arguments-differ
x = ctx.saved_tensors[0]
return g * torch.exp(torch.clamp(x, max=15))
trunc_exp = _TruncExp.apply
def get_activation(name) -> Callable:
if name is None:
return lambda x: x
name = name.lower()
if name == "none":
return lambda x: x
elif name == "lin2srgb":
return lambda x: torch.where(
x > 0.0031308,
torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
12.92 * x,
).clamp(0.0, 1.0)
elif name == "exp":
return lambda x: torch.exp(x)
elif name == "shifted_exp":
return lambda x: torch.exp(x - 1.0)
elif name == "trunc_exp":
return trunc_exp
elif name == "shifted_trunc_exp":
return lambda x: trunc_exp(x - 1.0)
elif name == "sigmoid":
return lambda x: torch.sigmoid(x)
elif name == "tanh":
return lambda x: torch.tanh(x)
elif name == "shifted_softplus":
return lambda x: F.softplus(x - 1.0)
elif name == "scale_-11_01":
return lambda x: x * 0.5 + 0.5
elif name == "negative":
return lambda x: -x
else:
try:
return getattr(F, name)
except AttributeError:
raise ValueError(f"Unknown activation function: {name}")
def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any:
if chunk_size <= 0:
return func(*args, **kwargs)
B = None
for arg in list(args) + list(kwargs.values()):
if isinstance(arg, torch.Tensor):
B = arg.shape[0]
break
assert (
B is not None
), "No tensor found in args or kwargs, cannot determine batch size."
out = defaultdict(list)
out_type = None
# max(1, B) to support B == 0
for i in range(0, max(1, B), chunk_size):
out_chunk = func(
*[
arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
for arg in args
],
**{
k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
for k, arg in kwargs.items()
},
)
if out_chunk is None:
continue
out_type = type(out_chunk)
if isinstance(out_chunk, torch.Tensor):
out_chunk = {0: out_chunk}
elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list):
chunk_length = len(out_chunk)
out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)}
elif isinstance(out_chunk, dict):
pass
else:
print(
f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}."
)
exit(1)
for k, v in out_chunk.items():
v = v if torch.is_grad_enabled() else v.detach()
out[k].append(v)
if out_type is None:
return None
out_merged: Dict[Any, Optional[torch.Tensor]] = {}
for k, v in out.items():
if all([vv is None for vv in v]):
# allow None in return value
out_merged[k] = None
elif all([isinstance(vv, torch.Tensor) for vv in v]):
out_merged[k] = torch.cat(v, dim=0)
else:
raise TypeError(
f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}"
)
if out_type is torch.Tensor:
return out_merged[0]
elif out_type in [tuple, list]:
return out_type([out_merged[i] for i in range(chunk_length)])
elif out_type is dict:
return out_merged
def get_ray_directions(
H: int,
W: int,
focal: Union[float, Tuple[float, float]],
principal: Optional[Tuple[float, float]] = None,
use_pixel_centers: bool = True,
normalize: bool = True,
) -> Float[Tensor, "H W 3"]:
"""
Get ray directions for all pixels in camera coordinate.
Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
ray-tracing-generating-camera-rays/standard-coordinate-systems
Inputs:
H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
Outputs:
directions: (H, W, 3), the direction of the rays in camera coordinate
"""
pixel_center = 0.5 if use_pixel_centers else 0
if isinstance(focal, float):
fx, fy = focal, focal
cx, cy = W / 2, H / 2
else:
fx, fy = focal
assert principal is not None
cx, cy = principal
i, j = torch.meshgrid(
torch.arange(W, dtype=torch.float32) + pixel_center,
torch.arange(H, dtype=torch.float32) + pixel_center,
indexing="xy",
)
directions: Float[Tensor, "H W 3"] = torch.stack(
[(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1
)
if normalize:
directions = F.normalize(directions, dim=-1)
return directions
def get_rays(
directions: Float[Tensor, "... 3"],
c2w: Float[Tensor, "... 4 4"],
keepdim=False,
noise_scale=0.0,
normalize=False,
) -> Tuple[Float[Tensor, "... 3"], Float[Tensor, "... 3"]]:
# Rotate ray directions from camera coordinate to the world coordinate
assert directions.shape[-1] == 3
if directions.ndim == 2: # (N_rays, 3)
if c2w.ndim == 2: # (4, 4)
c2w = c2w[None, :, :]
assert c2w.ndim == 3 # (N_rays, 4, 4) or (1, 4, 4)
rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1) # (N_rays, 3)
rays_o = c2w[:, :3, 3].expand(rays_d.shape)
elif directions.ndim == 3: # (H, W, 3)
assert c2w.ndim in [2, 3]
if c2w.ndim == 2: # (4, 4)
rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
-1
) # (H, W, 3)
rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
elif c2w.ndim == 3: # (B, 4, 4)
rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
-1
) # (B, H, W, 3)
rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
elif directions.ndim == 4: # (B, H, W, 3)
assert c2w.ndim == 3 # (B, 4, 4)
rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
-1
) # (B, H, W, 3)
rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
# add camera noise to avoid grid-like artifect
# https://github.com/ashawkey/stable-dreamfusion/blob/49c3d4fa01d68a4f027755acf94e1ff6020458cc/nerf/utils.py#L373
if noise_scale > 0:
rays_o = rays_o + torch.randn(3, device=rays_o.device) * noise_scale
rays_d = rays_d + torch.randn(3, device=rays_d.device) * noise_scale
if normalize:
rays_d = F.normalize(rays_d, dim=-1)
if not keepdim:
rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
return rays_o, rays_d
def get_projection_matrix(
fovy: Union[float, Float[Tensor, "B"]], aspect_wh: float, near: float, far: float
) -> Float[Tensor, "*B 4 4"]:
if isinstance(fovy, float):
proj_mtx = torch.zeros(4, 4, dtype=torch.float32)
proj_mtx[0, 0] = 1.0 / (math.tan(fovy / 2.0) * aspect_wh)
proj_mtx[1, 1] = -1.0 / math.tan(
fovy / 2.0
) # add a negative sign here as the y axis is flipped in nvdiffrast output
proj_mtx[2, 2] = -(far + near) / (far - near)
proj_mtx[2, 3] = -2.0 * far * near / (far - near)
proj_mtx[3, 2] = -1.0
else:
batch_size = fovy.shape[0]
proj_mtx = torch.zeros(batch_size, 4, 4, dtype=torch.float32)
proj_mtx[:, 0, 0] = 1.0 / (torch.tan(fovy / 2.0) * aspect_wh)
proj_mtx[:, 1, 1] = -1.0 / torch.tan(
fovy / 2.0
) # add a negative sign here as the y axis is flipped in nvdiffrast output
proj_mtx[:, 2, 2] = -(far + near) / (far - near)
proj_mtx[:, 2, 3] = -2.0 * far * near / (far - near)
proj_mtx[:, 3, 2] = -1.0
return proj_mtx
def get_mvp_matrix(
c2w: Float[Tensor, "*B 4 4"], proj_mtx: Float[Tensor, "*B 4 4"]
) -> Float[Tensor, "*B 4 4"]:
# calculate w2c from c2w: R' = Rt, t' = -Rt * t
# mathematically equivalent to (c2w)^-1
if c2w.ndim == 2:
assert proj_mtx.ndim == 2
w2c: Float[Tensor, "4 4"] = torch.zeros(4, 4).to(c2w)
w2c[:3, :3] = c2w[:3, :3].permute(1, 0)
w2c[:3, 3:] = -c2w[:3, :3].permute(1, 0) @ c2w[:3, 3:]
w2c[3, 3] = 1.0
else:
w2c: Float[Tensor, "B 4 4"] = torch.zeros(c2w.shape[0], 4, 4).to(c2w)
w2c[:, :3, :3] = c2w[:, :3, :3].permute(0, 2, 1)
w2c[:, :3, 3:] = -c2w[:, :3, :3].permute(0, 2, 1) @ c2w[:, :3, 3:]
w2c[:, 3, 3] = 1.0
# calculate mvp matrix by proj_mtx @ w2c (mv_mtx)
mvp_mtx = proj_mtx @ w2c
return mvp_mtx
def get_intrinsic_from_fov(fov, H, W, bs=-1):
focal_length = 0.5 * H / np.tan(0.5 * fov)
intrinsic = np.identity(3, dtype=np.float32)
intrinsic[0, 0] = focal_length
intrinsic[1, 1] = focal_length
intrinsic[0, 2] = W / 2.0
intrinsic[1, 2] = H / 2.0
if bs > 0:
intrinsic = intrinsic[None].repeat(bs, axis=0)
return torch.from_numpy(intrinsic)
def binary_cross_entropy(input, target):
"""
F.binary_cross_entropy is not numerically stable in mixed-precision training.
"""
return -(target * torch.log(input) + (1 - target) * torch.log(1 - input)).mean()
def tet_sdf_diff(
vert_sdf: Float[Tensor, "Nv 1"], tet_edges: Integer[Tensor, "Ne 2"]
) -> Float[Tensor, ""]:
sdf_f1x6x2 = vert_sdf[:, 0][tet_edges.reshape(-1)].reshape(-1, 2)
mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1])
sdf_f1x6x2 = sdf_f1x6x2[mask]
sdf_diff = F.binary_cross_entropy_with_logits(
sdf_f1x6x2[..., 0], (sdf_f1x6x2[..., 1] > 0).float()
) + F.binary_cross_entropy_with_logits(
sdf_f1x6x2[..., 1], (sdf_f1x6x2[..., 0] > 0).float()
)
return sdf_diff
def validate_empty_rays(ray_indices, t_start, t_end):
if ray_indices.nelement() == 0:
warn("Empty rays_indices!")
ray_indices = torch.LongTensor([0]).to(ray_indices)
t_start = torch.Tensor([0]).to(ray_indices)
t_end = torch.Tensor([0]).to(ray_indices)
return ray_indices, t_start, t_end
def rays_intersect_bbox(
rays_o: Float[Tensor, "N 3"],
rays_d: Float[Tensor, "N 3"],
radius: Float,
near: Float = 0.0,
valid_thresh: Float = 0.01,
background: bool = False,
):
input_shape = rays_o.shape[:-1]
rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3)
rays_d_valid = torch.where(
rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d
)
if type(radius) in [int, float]:
radius = torch.FloatTensor(
[[-radius, radius], [-radius, radius], [-radius, radius]]
).to(rays_o.device)
radius = (
1.0 - 1.0e-3
) * radius # tighten the radius to make sure the intersection point lies in the bounding box
interx0 = (radius[..., 1] - rays_o) / rays_d_valid
interx1 = (radius[..., 0] - rays_o) / rays_d_valid
t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near)
t_far = torch.maximum(interx0, interx1).amin(dim=-1)
# check wheter a ray intersects the bbox or not
rays_valid = t_far - t_near > valid_thresh
t_near_valid, t_far_valid = t_near[rays_valid], t_far[rays_valid]
global_near = t_near_valid.min().item()
global_far = t_far_valid.max().item()
t_near[torch.where(~rays_valid)] = 0.0
t_far[torch.where(~rays_valid)] = 0.0
t_near = t_near.view(*input_shape, 1)
t_far = t_far.view(*input_shape, 1)
rays_valid = rays_valid.view(*input_shape)
return t_near, t_far, rays_valid
def get_plucker_rays(
rays_o: Float[Tensor, "*N 3"], rays_d: Float[Tensor, "*N 3"]
) -> Float[Tensor, "*N 6"]:
rays_o = F.normalize(rays_o, dim=-1)
rays_d = F.normalize(rays_d, dim=-1)
return torch.cat([rays_o.cross(rays_d), rays_d], dim=-1)
def c2w_to_polar(c2w: Float[Tensor, "4 4"]) -> Tuple[float, float, float]:
cam_pos = c2w[:3, 3]
x, y, z = cam_pos.tolist()
distance = cam_pos.norm().item()
elevation = math.asin(z / distance)
if abs(x) < 1.0e-5 and abs(y) < 1.0e-5:
azimuth = 0
else:
azimuth = math.atan2(y, x)
if azimuth < 0:
azimuth += 2 * math.pi
return elevation, azimuth, distance
def polar_to_c2w(
elevation: float, azimuth: float, distance: float
) -> Float[Tensor, "4 4"]:
"""
Compute L = p - C.
Normalize L.
Compute s = L x u. (cross product)
Normalize s.
Compute u' = s x L.
rotation = [s, u, -l]
"""
z = distance * math.sin(elevation)
x = distance * math.cos(elevation) * math.cos(azimuth)
y = distance * math.cos(elevation) * math.sin(azimuth)
l = -torch.as_tensor([x, y, z]).float()
l = F.normalize(l, dim=0)
u = torch.as_tensor([0.0, 0.0, 1.0]).float()
s = l.cross(u)
s = F.normalize(s, dim=0)
u = s.cross(l)
rot = torch.stack([s, u, -l], dim=0).T
c2w = torch.zeros((4, 4), dtype=torch.float32)
c2w[:3, :3] = rot
c2w[:3, 3] = torch.as_tensor([x, y, z])
c2w[3, 3] = 1
return c2w
def fourier_position_encoding(x, n_freq: int, dim: int):
assert n_freq > 0
input_shape = x.shape
input_ndim = x.ndim
if dim < 0:
dim = input_ndim + dim
bands = 2 ** torch.arange(n_freq, dtype=x.dtype, device=x.device)
for i in range(dim + 1):
bands = bands.unsqueeze(0)
for i in range(input_ndim - dim - 1):
bands = bands.unsqueeze(-1)
x = x.view(*input_shape[: dim + 1], 1, *input_shape[dim + 1 :])
x = torch.cat(
[
torch.sin(bands * x).reshape(
*input_shape[:dim], -1, *input_shape[dim + 1 :]
),
torch.cos(bands * x).reshape(
*input_shape[:dim], -1, *input_shape[dim + 1 :]
),
],
dim=dim,
)
return x
Gen_3D_Modules/MV_Adapter/mvadapter/utils/saving.py
+544
View File
@@ -0,0 +1,544 @@
import json
import math
import os
import re
import shutil
from typing import List, Optional, Union
import cv2
import imageio
import matplotlib.pyplot as plt
import numpy as np
import torch
# import wandb
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap
from PIL import Image, ImageDraw
from .typing import *
def image_to_tensor(
images: List[Image.Image], device: str = "cuda"
) -> torch.FloatTensor:
tensors = []
for image in images:
tensor = torch.from_numpy(np.array(image)).float() / 255.0
tensor = tensor.to(device)
tensors.append(tensor)
return torch.stack(tensors)
def tensor_to_image(
data: Union[Image.Image, torch.Tensor, np.ndarray],
batched: bool = False,
format: str = "HWC",
) -> Union[Image.Image, List[Image.Image]]:
if isinstance(data, Image.Image):
return data
if isinstance(data, torch.Tensor):
data = data.detach().cpu().numpy()
if data.dtype == np.float32 or data.dtype == np.float16:
data = (data * 255).astype(np.uint8)
elif data.dtype == np.bool_:
data = data.astype(np.uint8) * 255
assert data.dtype == np.uint8
if format == "CHW":
if batched and data.ndim == 4:
data = data.transpose((0, 2, 3, 1))
elif not batched and data.ndim == 3:
data = data.transpose((1, 2, 0))
if batched:
return [Image.fromarray(d) for d in data]
return Image.fromarray(data)
def largest_factor_near_sqrt(n: int) -> int:
"""
Finds the largest factor of n that is closest to the square root of n.
Args:
n (int): The integer for which to find the largest factor near its square root.
Returns:
int: The largest factor of n that is closest to the square root of n.
"""
sqrt_n = int(math.sqrt(n)) # Get the integer part of the square root
# First, check if the square root itself is a factor
if sqrt_n * sqrt_n == n:
return sqrt_n
# Otherwise, find the largest factor by iterating from sqrt_n downwards
for i in range(sqrt_n, 0, -1):
if n % i == 0:
return i
# If n is 1, return 1
return 1
def make_image_grid(
images: List[Image.Image],
rows: Optional[int] = None,
cols: Optional[int] = None,
resize: Optional[int] = None,
) -> Image.Image:
"""
Prepares a single grid of images. Useful for visualization purposes.
"""
if rows is None and cols is not None:
assert len(images) % cols == 0
rows = len(images) // cols
elif cols is None and rows is not None:
assert len(images) % rows == 0
cols = len(images) // rows
elif rows is None and cols is None:
rows = largest_factor_near_sqrt(len(images))
cols = len(images) // rows
assert len(images) == rows * cols
if resize is not None:
images = [img.resize((resize, resize)) for img in images]
w, h = images[0].size
grid = Image.new("RGB", size=(cols * w, rows * h))
for i, img in enumerate(images):
grid.paste(img, box=(i % cols * w, i // cols * h))
return grid
class SaverMixin:
_save_dir: Optional[str] = None
_wandb_logger: Optional[Any] = None
def set_save_dir(self, save_dir: str):
self._save_dir = save_dir
def get_save_dir(self):
if self._save_dir is None:
raise ValueError("Save dir is not set")
return self._save_dir
def convert_data(self, data):
if data is None:
return None
elif isinstance(data, np.ndarray):
return data
elif isinstance(data, torch.Tensor):
if data.dtype in [torch.float16, torch.bfloat16]:
data = data.float()
return data.detach().cpu().numpy()
elif isinstance(data, list):
return [self.convert_data(d) for d in data]
elif isinstance(data, dict):
return {k: self.convert_data(v) for k, v in data.items()}
else:
raise TypeError(
"Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
type(data),
)
def get_save_path(self, filename):
save_path = os.path.join(self.get_save_dir(), filename)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
return save_path
DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
DEFAULT_UV_KWARGS = {
"data_format": "HWC",
"data_range": (0, 1),
"cmap": "checkerboard",
}
DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"}
DEFAULT_GRID_KWARGS = {"align": "max"}
def get_rgb_image_(self, img, data_format, data_range, rgba=False):
img = self.convert_data(img)
assert data_format in ["CHW", "HWC"]
if data_format == "CHW":
img = img.transpose(1, 2, 0)
if img.dtype != np.uint8:
img = img.clip(min=data_range[0], max=data_range[1])
img = (
(img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0
).astype(np.uint8)
nc = 4 if rgba else 3
imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)]
imgs = [
(
img_
if img_.shape[-1] == nc
else np.concatenate(
[
img_,
np.zeros(
(img_.shape[0], img_.shape[1], nc - img_.shape[2]),
dtype=img_.dtype,
),
],
axis=-1,
)
)
for img_ in imgs
]
img = np.concatenate(imgs, axis=1)
if rgba:
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
else:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
def _save_rgb_image(
self,
filename,
img,
data_format,
data_range,
name: Optional[str] = None,
step: Optional[int] = None,
):
img = self.get_rgb_image_(img, data_format, data_range)
cv2.imwrite(filename, img)
if name and self._wandb_logger:
self._wandb_logger.log_image(
key=name, images=[self.get_save_path(filename)], step=step
)
def save_rgb_image(
self,
filename,
img,
data_format=DEFAULT_RGB_KWARGS["data_format"],
data_range=DEFAULT_RGB_KWARGS["data_range"],
name: Optional[str] = None,
step: Optional[int] = None,
) -> str:
save_path = self.get_save_path(filename)
self._save_rgb_image(save_path, img, data_format, data_range, name, step)
return save_path
def get_uv_image_(self, img, data_format, data_range, cmap):
img = self.convert_data(img)
assert data_format in ["CHW", "HWC"]
if data_format == "CHW":
img = img.transpose(1, 2, 0)
img = img.clip(min=data_range[0], max=data_range[1])
img = (img - data_range[0]) / (data_range[1] - data_range[0])
assert cmap in ["checkerboard", "color"]
if cmap == "checkerboard":
n_grid = 64
mask = (img * n_grid).astype(int)
mask = (mask[..., 0] + mask[..., 1]) % 2 == 0
img = np.ones((img.shape[0], img.shape[1], 3), dtype=np.uint8) * 255
img[mask] = np.array([255, 0, 255], dtype=np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
elif cmap == "color":
img_ = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
img_[..., 0] = (img[..., 0] * 255).astype(np.uint8)
img_[..., 1] = (img[..., 1] * 255).astype(np.uint8)
img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR)
img = img_
return img
def save_uv_image(
self,
filename,
img,
data_format=DEFAULT_UV_KWARGS["data_format"],
data_range=DEFAULT_UV_KWARGS["data_range"],
cmap=DEFAULT_UV_KWARGS["cmap"],
) -> str:
save_path = self.get_save_path(filename)
img = self.get_uv_image_(img, data_format, data_range, cmap)
cv2.imwrite(save_path, img)
return save_path
def get_grayscale_image_(self, img, data_range, cmap):
img = self.convert_data(img)
img = np.nan_to_num(img)
if data_range is None:
img = (img - img.min()) / (img.max() - img.min())
else:
img = img.clip(data_range[0], data_range[1])
img = (img - data_range[0]) / (data_range[1] - data_range[0])
assert cmap in [None, "jet", "magma", "spectral"]
if cmap == None:
img = (img * 255.0).astype(np.uint8)
img = np.repeat(img[..., None], 3, axis=2)
elif cmap == "jet":
img = (img * 255.0).astype(np.uint8)
img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
elif cmap == "magma":
img = 1.0 - img
base = cm.get_cmap("magma")
num_bins = 256
colormap = LinearSegmentedColormap.from_list(
f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins
)(np.linspace(0, 1, num_bins))[:, :3]
a = np.floor(img * 255.0)
b = (a + 1).clip(max=255.0)
f = img * 255.0 - a
a = a.astype(np.uint16).clip(0, 255)
b = b.astype(np.uint16).clip(0, 255)
img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None]
img = (img * 255.0).astype(np.uint8)
elif cmap == "spectral":
colormap = plt.get_cmap("Spectral")
def blend_rgba(image):
image = image[..., :3] * image[..., -1:] + (
1.0 - image[..., -1:]
) # blend A to RGB
return image
img = colormap(img)
img = blend_rgba(img)
img = (img * 255).astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
def _save_grayscale_image(
self,
filename,
img,
data_range,
cmap,
name: Optional[str] = None,
step: Optional[int] = None,
):
img = self.get_grayscale_image_(img, data_range, cmap)
cv2.imwrite(filename, img)
if name and self._wandb_logger:
self._wandb_logger.log_image(
key=name, images=[self.get_save_path(filename)], step=step
)
def save_grayscale_image(
self,
filename,
img,
data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
name: Optional[str] = None,
step: Optional[int] = None,
) -> str:
save_path = self.get_save_path(filename)
self._save_grayscale_image(save_path, img, data_range, cmap, name, step)
return save_path
def get_image_grid_(self, imgs, align):
if isinstance(imgs[0], list):
return np.concatenate(
[self.get_image_grid_(row, align) for row in imgs], axis=0
)
cols = []
for col in imgs:
assert col["type"] in ["rgb", "uv", "grayscale"]
if col["type"] == "rgb":
rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy()
rgb_kwargs.update(col["kwargs"])
cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs))
elif col["type"] == "uv":
uv_kwargs = self.DEFAULT_UV_KWARGS.copy()
uv_kwargs.update(col["kwargs"])
cols.append(self.get_uv_image_(col["img"], **uv_kwargs))
elif col["type"] == "grayscale":
grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy()
grayscale_kwargs.update(col["kwargs"])
cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs))
if align == "max":
h = max([col.shape[0] for col in cols])
elif align == "min":
h = min([col.shape[0] for col in cols])
elif isinstance(align, int):
h = align
else:
raise ValueError(
f"Unsupported image grid align: {align}, should be min, max, or int"
)
for i in range(len(cols)):
if cols[i].shape[0] != h:
w = int(cols[i].shape[1] * h / cols[i].shape[0])
cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_CUBIC)
return np.concatenate(cols, axis=1)
def save_image_grid(
self,
filename,
imgs,
align=DEFAULT_GRID_KWARGS["align"],
name: Optional[str] = None,
step: Optional[int] = None,
texts: Optional[List[float]] = None,
):
save_path = self.get_save_path(filename)
img = self.get_image_grid_(imgs, align=align)
if texts is not None:
img = Image.fromarray(img)
draw = ImageDraw.Draw(img)
black, white = (0, 0, 0), (255, 255, 255)
for i, text in enumerate(texts):
draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
img = np.asarray(img)
cv2.imwrite(save_path, img)
if name and self._wandb_logger:
self._wandb_logger.log_image(key=name, images=[save_path], step=step)
return save_path
def save_image(self, filename, img) -> str:
save_path = self.get_save_path(filename)
img = self.convert_data(img)
assert img.dtype == np.uint8 or img.dtype == np.uint16
if img.ndim == 3 and img.shape[-1] == 3:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
elif img.ndim == 3 and img.shape[-1] == 4:
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
cv2.imwrite(save_path, img)
return save_path
def save_cubemap(self, filename, img, data_range=(0, 1), rgba=False) -> str:
save_path = self.get_save_path(filename)
img = self.convert_data(img)
assert img.ndim == 4 and img.shape[0] == 6 and img.shape[1] == img.shape[2]
imgs_full = []
for start in range(0, img.shape[-1], 3):
img_ = img[..., start : start + 3]
img_ = np.stack(
[
self.get_rgb_image_(img_[i], "HWC", data_range, rgba=rgba)
for i in range(img_.shape[0])
],
axis=0,
)
size = img_.shape[1]
placeholder = np.zeros((size, size, 3), dtype=np.float32)
img_full = np.concatenate(
[
np.concatenate(
[placeholder, img_[2], placeholder, placeholder], axis=1
),
np.concatenate([img_[1], img_[4], img_[0], img_[5]], axis=1),
np.concatenate(
[placeholder, img_[3], placeholder, placeholder], axis=1
),
],
axis=0,
)
imgs_full.append(img_full)
imgs_full = np.concatenate(imgs_full, axis=1)
cv2.imwrite(save_path, imgs_full)
return save_path
def save_data(self, filename, data) -> str:
data = self.convert_data(data)
if isinstance(data, dict):
if not filename.endswith(".npz"):
filename += ".npz"
save_path = self.get_save_path(filename)
np.savez(save_path, **data)
else:
if not filename.endswith(".npy"):
filename += ".npy"
save_path = self.get_save_path(filename)
np.save(save_path, data)
return save_path
def save_state_dict(self, filename, data) -> str:
save_path = self.get_save_path(filename)
torch.save(data, save_path)
return save_path
def save_img_sequence(
self,
filename,
img_dir,
matcher,
save_format="mp4",
fps=30,
name: Optional[str] = None,
step: Optional[int] = None,
) -> str:
assert save_format in ["gif", "mp4"]
if not filename.endswith(save_format):
filename += f".{save_format}"
save_path = self.get_save_path(filename)
matcher = re.compile(matcher)
img_dir = os.path.join(self.get_save_dir(), img_dir)
imgs = []
for f in os.listdir(img_dir):
if matcher.search(f):
imgs.append(f)
imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
if save_format == "gif":
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
imageio.mimsave(save_path, imgs, fps=fps, palettesize=256)
elif save_format == "mp4":
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
imageio.mimsave(save_path, imgs, fps=fps)
if name and self._wandb_logger:
from .core import warn
warn("Wandb logger does not support video logging yet!")
return save_path
def save_img_sequences(
self,
seq_dir,
matcher,
save_format="mp4",
fps=30,
delete=True,
name: Optional[str] = None,
step: Optional[int] = None,
):
seq_dir_ = os.path.join(self.get_save_dir(), seq_dir)
for f in os.listdir(seq_dir_):
img_dir_ = os.path.join(seq_dir_, f)
if not os.path.isdir(img_dir_):
continue
try:
self.save_img_sequence(
os.path.join(seq_dir, f),
os.path.join(seq_dir, f),
matcher,
save_format=save_format,
fps=fps,
name=f"{name}_{f}",
step=step,
)
if delete:
shutil.rmtree(img_dir_)
except:
from .core import warn
warn(f"Video saving for directory {seq_dir_} failed!")
def save_file(self, filename, src_path, delete=False) -> str:
save_path = self.get_save_path(filename)
shutil.copyfile(src_path, save_path)
if delete:
os.remove(src_path)
return save_path
def save_json(self, filename, payload) -> str:
save_path = self.get_save_path(filename)
with open(save_path, "w") as f:
f.write(json.dumps(payload))
return save_path
Gen_3D_Modules/MV_Adapter/mvadapter/utils/typing.py
+40
View File
@@ -0,0 +1,40 @@
"""
This module contains type annotations for the project, using
1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects
2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors
Two types of typing checking can be used:
1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode)
2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking)
"""
# Basic types
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Literal,
NamedTuple,
NewType,
Optional,
Sized,
Tuple,
Type,
TypeVar,
Union,
)
# Tensor dtype
# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
# Config type
from omegaconf import DictConfig, ListConfig
# PyTorch Tensor type
from torch import Tensor
# Runtime type checking decorator
from typeguard import typechecked as typechecker
Gen_3D_Modules/MV_Adapter/mvadapter_node_utils.py
+302
View File
@@ -0,0 +1,302 @@
"""
Utilities for ComfyUI MV-Adapter nodes
"""
import os
import torch
import numpy as np
from PIL import Image
from typing import Optional, Union, List
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from huggingface_hub import snapshot_download
from .mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from .mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
from .mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from .mvadapter.utils.mesh_utils import (
NVDiffRastContextWrapper,
get_orthogonal_camera,
load_mesh,
render,
)
from .mvadapter.utils import make_image_grid, tensor_to_image
try:
from mmgp import offload, profile_type
MMGP_AVAILABLE = True
except ImportError:
MMGP_AVAILABLE = False
print("Warning: mmgp not available, memory optimization disabled")
def prepare_pipeline(
base_model: str,
vae_model: Optional[str] = None,
unet_model: Optional[str] = None,
lora_model: Optional[str] = None,
adapter_path: str = "huanngzh/mv-adapter",
scheduler: str = "default",
num_views: int = 6,
device: str = "cuda",
dtype: torch.dtype = torch.float16,
use_mmgp: bool = False,
adapter_local_path: Optional[str] = None,
):
"""
Prepare MV-Adapter pipeline
Args:
base_model: Base Stable Diffusion XL model
vae_model: Custom VAE model (optional)
unet_model: Custom UNet model (optional)
lora_model: LoRA model (optional)
adapter_path: Path to adapter weights
scheduler: Scheduler type ("default", "ddpm", "lcm")
num_views: Number of views to generate
device: Device for computation
dtype: Data type
use_mmgp: Use mmgp for memory optimization
adapter_local_path: Local path for adapter download
"""
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
if scheduler != "default":
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(
num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
)
# Load adapter weights
weight_name = "mvadapter_ig2mv_sdxl.safetensors"
pipe.load_custom_adapter(adapter_path, weight_name=weight_name, local_cache_dir=adapter_local_path)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
# Apply mmgp optimization if available and requested
if use_mmgp and MMGP_AVAILABLE:
import gc
torch.cuda.empty_cache()
gc.collect()
all_models = {
"vae": pipe.vae,
"text_encoder": pipe.text_encoder,
"text_encoder_2": pipe.text_encoder_2,
"unet": pipe.unet,
"cond_encoder": pipe.cond_encoder,
}
if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
all_models["image_encoder"] = pipe.image_encoder
offload.profile(all_models, profile_type.HighRAM_LowVRAM)
print("mmgp profiling applied successfully.")
# Enable VAE slicing for memory efficiency
pipe.enable_vae_slicing()
return pipe
def create_bg_remover(device: str = "cuda"):
"""Create background removal function"""
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(device)
transform_image = transforms.Compose([
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
def remove_bg(image):
image_size = image.size
input_images = transform_image(image).unsqueeze(0).to(device)
with torch.no_grad():
preds = birefnet(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(image_size)
image.putalpha(mask)
return image
return remove_bg
def preprocess_image(image: Image.Image, height: int, width: int) -> Image.Image:
"""Preprocess image for MV-Adapter"""
image = np.array(image)
alpha = image[..., 3] > 0
H, W = alpha.shape
# get the bounding box of alpha
y, x = np.where(alpha)
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
image_center = image[y0:y1, x0:x1]
# resize the longer side to H * 0.9
H, W, _ = image_center.shape
if H > W:
W = int(W * (height * 0.9) / H)
H = int(height * 0.9)
else:
H = int(H * (width * 0.9) / W)
W = int(width * 0.9)
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
# pad to H, W
start_h = (height - H) // 2
start_w = (width - W) // 2
image = np.zeros((height, width, 4), dtype=np.uint8)
image[start_h : start_h + H, start_w : start_w + W] = image_center
image = image.astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = (image * 255).clip(0, 255).astype(np.uint8)
image = Image.fromarray(image)
return image
def run_pipeline(
pipe,
mesh_path: str,
num_views: int,
text: str,
image: Union[str, Image.Image],
height: int,
width: int,
num_inference_steps: int,
guidance_scale: float,
seed: int,
remove_bg_fn=None,
reference_conditioning_scale: float = 1.0,
negative_prompt: str = "watermark, ugly, deformed, noisy, blurry, low contrast",
lora_scale: float = 1.0,
device: str = "cuda",
):
"""
Execute multi-view image generation
"""
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
device=device,
)
ctx = NVDiffRastContextWrapper(device=device)
mesh = load_mesh(mesh_path, rescale=True, device=device)
render_out = render(
ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
normal_background=0.0,
)
pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
normal_images = tensor_to_image(
(render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
)
control_images = (
torch.cat(
[
(render_out.pos + 0.5).clamp(0, 1),
(render_out.normal / 2 + 0.5).clamp(0, 1),
],
dim=-1,
)
.permute(0, 3, 1, 2)
.to(device)
)
# Prepare image
reference_image = Image.open(image) if isinstance(image, str) else image
if remove_bg_fn is not None:
reference_image = remove_bg_fn(reference_image)
reference_image = preprocess_image(reference_image, height, width)
elif reference_image.mode == "RGBA":
reference_image = preprocess_image(reference_image, height, width)
pipe_kwargs = {}
if seed != -1 and isinstance(seed, int):
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=reference_image,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, pos_images, normal_images, reference_image
def create_image_grid(images: List[Image.Image], rows: int = 1) -> Image.Image:
"""Создание сетки изображений"""
return make_image_grid(images, rows=rows)
def pil_to_torch_tensor(image: Image.Image) -> torch.Tensor:
"""Конвертация PIL изображения в torch tensor"""
import numpy as np
array = np.array(image)
tensor = torch.from_numpy(array).float() / 255.0
# Add batch dimension if needed
if tensor.dim() == 3:
tensor = tensor.unsqueeze(0)
return tensor
def pils_to_torch_tensor(images: List[Image.Image]) -> torch.Tensor:
"""Конвертация списка PIL изображений в torch tensor"""
tensors = [pil_to_torch_tensor(img).squeeze(0) for img in images]
return torch.stack(tensors, dim=0)
Gen_3D_Modules/MV_Adapter/requirements.txt
+28
View File
@@ -0,0 +1,28 @@
controlnet_aux
diffusers
transformers
peft
numpy
huggingface_hub
accelerate
opencv-python
safetensors
pillow
omegaconf
trimesh
einops
gradio
timm
kornia
scikit-image
sentencepiece
git+https://github.com/NVlabs/nvdiffrast.git
# new added for training
pytorch-lightning
# new added for texturing
spandrel==0.4.1
open3d
pymeshlab
cvcuda_cu12
Gen_3D_Modules/MV_Adapter/scripts/__init__.py
View File
Gen_3D_Modules/MV_Adapter/scripts/gradio_demo_i2mv.py
+214
View File
@@ -0,0 +1,214 @@
import argparse
import random
import gradio as gr
import numpy as np
# import spaces
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from .inference_i2mv_sdxl import prepare_pipeline, remove_bg, run_pipeline
# Device and dtype
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
# Hyperparameters
NUM_VIEWS = 6
HEIGHT = 768
WIDTH = 768
MAX_SEED = np.iinfo(np.int32).max
pipe = prepare_pipeline(
base_model="stabilityai/stable-diffusion-xl-base-1.0",
vae_model="madebyollin/sdxl-vae-fp16-fix",
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=None,
num_views=NUM_VIEWS,
device=device,
dtype=dtype,
)
# remove bg
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
# @spaces.GPU()
def infer(
prompt,
image,
do_rembg=True,
seed=42,
randomize_seed=False,
guidance_scale=3.0,
num_inference_steps=50,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
progress=gr.Progress(track_tqdm=True),
):
if do_rembg:
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, device)
else:
remove_bg_fn = None
if randomize_seed:
seed = random.randint(0, MAX_SEED)
images, preprocessed_image = run_pipeline(
pipe,
num_views=NUM_VIEWS,
text=prompt,
image=image,
height=HEIGHT,
width=WIDTH,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
remove_bg_fn=remove_bg_fn,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
device=device,
)
return images, preprocessed_image, seed
examples = [
[
"A decorative figurine of a young anime-style girl",
"assets/demo/i2mv/A_decorative_figurine_of_a_young_anime-style_girl.png",
True,
21,
],
[
"A juvenile emperor penguin chick",
"assets/demo/i2mv/A_juvenile_emperor_penguin_chick.png",
True,
0,
],
[
"A striped tabby cat with white fur sitting upright",
"assets/demo/i2mv/A_striped_tabby_cat_with_white_fur_sitting_upright.png",
True,
0,
],
]
with gr.Blocks() as demo:
with gr.Row():
gr.Markdown(
f"""# MV-Adapter [Image-to-Multi-View]
Generate 768x768 multi-view images from a single image using SDXL <br>
[[page](https://huanngzh.github.io/MV-Adapter-Page/)] [[repo](https://github.com/huanngzh/MV-Adapter)]
"""
)
with gr.Row():
with gr.Column():
with gr.Row():
input_image = gr.Image(
label="Input Image",
sources=["upload", "webcam", "clipboard"],
type="pil",
)
preprocessed_image = gr.Image(label="Preprocessed Image", type="pil")
prompt = gr.Textbox(
label="Prompt", placeholder="Enter your prompt", value="high quality"
)
do_rembg = gr.Checkbox(label="Remove background", value=True)
run_button = gr.Button("Run")
with gr.Accordion("Advanced Settings", open=False):
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=50,
)
with gr.Row():
guidance_scale = gr.Slider(
label="CFG scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=3.0,
)
with gr.Row():
reference_conditioning_scale = gr.Slider(
label="Image conditioning scale",
minimum=0.0,
maximum=2.0,
step=0.1,
value=1.0,
)
with gr.Row():
negative_prompt = gr.Textbox(
label="Negative prompt",
placeholder="Enter your negative prompt",
value="watermark, ugly, deformed, noisy, blurry, low contrast",
)
with gr.Column():
result = gr.Gallery(
label="Result",
show_label=False,
columns=[3],
rows=[2],
object_fit="contain",
height="auto",
)
with gr.Row():
gr.Examples(
examples=examples,
fn=infer,
inputs=[prompt, input_image, do_rembg, seed],
outputs=[result, preprocessed_image, seed],
cache_examples=True,
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[
prompt,
input_image,
do_rembg,
seed,
randomize_seed,
guidance_scale,
num_inference_steps,
reference_conditioning_scale,
negative_prompt,
],
outputs=[result, preprocessed_image, seed],
)
demo.launch()
Gen_3D_Modules/MV_Adapter/scripts/gradio_demo_t2mv.py
+195
View File
@@ -0,0 +1,195 @@
import argparse
import random
import gradio as gr
import numpy as np
# import spaces
import torch
from .inference_t2mv_sdxl import prepare_pipeline, run_pipeline
# Base model
parser = argparse.ArgumentParser()
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument("--scheduler", type=str, default=None)
args = parser.parse_args()
base_model = args.base_model
scheduler = args.scheduler
# Device and dtype
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
# Hyperparameters
NUM_VIEWS = 6
HEIGHT = 768
WIDTH = 768
MAX_SEED = np.iinfo(np.int32).max
pipe = prepare_pipeline(
base_model=base_model,
vae_model="madebyollin/sdxl-vae-fp16-fix",
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=scheduler,
num_views=NUM_VIEWS,
device=device,
dtype=dtype,
)
# @spaces.GPU()
def infer(
prompt,
seed=42,
randomize_seed=False,
guidance_scale=7.0,
num_inference_steps=50,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
progress=gr.Progress(track_tqdm=True),
):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
images = run_pipeline(
pipe,
num_views=NUM_VIEWS,
text=prompt,
height=HEIGHT,
width=WIDTH,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
negative_prompt=negative_prompt,
device=device,
)
return images, seed
examples = {
"stabilityai/stable-diffusion-xl-base-1.0": [
["An astronaut riding a horse", 42],
["A DSLR photo of a frog wearing a sweater", 42],
],
"cagliostrolab/animagine-xl-3.1": [
[
"1girl, izayoi sakuya, touhou, solo, maid headdress, maid, apron, short sleeves, dress, closed mouth, white apron, serious face, upper body, masterpiece, best quality, very aesthetic, absurdres",
0,
],
[
"1boy, male focus, ikari shinji, neon genesis evangelion, solo, serious face,(masterpiece), (best quality), (ultra-detailed), very aesthetic, illustration, disheveled hair, moist skin, intricate details",
0,
],
[
"1girl, pink hair, pink shirts, smile, shy, masterpiece, anime",
0,
],
],
"Lykon/dreamshaper-xl-1-0": [
["the warrior Aragorn from Lord of the Rings, film grain, 8k hd", 0],
[
"Oil painting, masterpiece, regal, fancy. A well-dressed dog named Puproy Doggerson III wearing reading glasses types an important letter on a typewriter and enjoys a cup of coffee with the newspaper.",
42,
],
],
}
css = """
#col-container {
margin: 0 auto;
max-width: 600px;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(
f"""# MV-Adapter [Text-to-Multi-View]
Generate 768x768 multi-view images using {base_model} <br>
[[page](https://huanngzh.github.io/MV-Adapter-Page/)] [[repo](https://github.com/huanngzh/MV-Adapter)]
"""
)
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
run_button = gr.Button("Run", scale=0)
result = gr.Gallery(
label="Result",
show_label=False,
columns=[3],
rows=[2],
object_fit="contain",
height="auto",
)
with gr.Accordion("Advanced Settings", open=False):
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=50,
)
with gr.Row():
guidance_scale = gr.Slider(
label="CFG scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=7.0,
)
with gr.Row():
negative_prompt = gr.Textbox(
label="Negative prompt",
placeholder="Enter your negative prompt",
value="watermark, ugly, deformed, noisy, blurry, low contrast",
)
if base_model in examples:
gr.Examples(
examples=examples[base_model],
fn=infer,
inputs=[prompt, seed],
outputs=[result, seed],
cache_examples=True,
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[
prompt,
seed,
randomize_seed,
guidance_scale,
num_inference_steps,
negative_prompt,
],
outputs=[result, seed],
)
demo.launch()
Gen_3D_Modules/MV_Adapter/scripts/inference_i2mv_sd.py
+255
View File
@@ -0,0 +1,255 @@
import argparse
import numpy as np
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from mvadapter.pipelines.pipeline_mvadapter_i2mv_sd import MVAdapterI2MVSDPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils.mesh_utils import get_orthogonal_camera
from mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
from mvadapter.utils import make_image_grid
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterI2MVSDPipeline
pipe = MVAdapterI2MVSDPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(num_views=num_views)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_i2mv_sd21.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
# vae slicing for lower memory usage
pipe.enable_vae_slicing()
return pipe
def remove_bg(image, net, transform, device):
image_size = image.size
input_images = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
preds = net(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(image_size)
image.putalpha(mask)
return image
def preprocess_image(image: Image.Image, height, width):
image = np.array(image)
alpha = image[..., 3] > 0
H, W = alpha.shape
# get the bounding box of alpha
y, x = np.where(alpha)
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
image_center = image[y0:y1, x0:x1]
# resize the longer side to H * 0.9
H, W, _ = image_center.shape
if H > W:
W = int(W * (height * 0.9) / H)
H = int(height * 0.9)
else:
H = int(H * (width * 0.9) / W)
W = int(width * 0.9)
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
# pad to H, W
start_h = (height - H) // 2
start_w = (width - W) // 2
image = np.zeros((height, width, 4), dtype=np.uint8)
image[start_h : start_h + H, start_w : start_w + W] = image_center
image = image.astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = (image * 255).clip(0, 255).astype(np.uint8)
image = Image.fromarray(image)
return image
def run_pipeline(
pipe,
num_views,
text,
image,
height,
width,
num_inference_steps,
guidance_scale,
seed,
remove_bg_fn=None,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 0, 0],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 45, 90, 180, 270, 315]],
device=device,
)
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
cameras.c2w, [1.1] * num_views, width
)
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
# Prepare image
reference_image = Image.open(image) if isinstance(image, str) else image
if remove_bg_fn is not None:
reference_image = remove_bg_fn(reference_image)
reference_image = preprocess_image(reference_image, height, width)
elif reference_image.mode == "RGBA":
reference_image = preprocess_image(reference_image, height, width)
pipe_kwargs = {}
if seed != -1 and isinstance(seed, int):
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=reference_image,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, reference_image
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-2-1-base"
)
parser.add_argument("--vae_model", type=str, default=None)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--image", type=str, required=True)
parser.add_argument("--text", type=str, default="high quality")
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--output", type=str, default="output.png")
# Extra
parser.add_argument("--remove_bg", action="store_true", help="Remove background")
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
if args.remove_bg:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(args.device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
else:
remove_bg_fn = None
images, reference_image = run_pipeline(
pipe,
num_views=args.num_views,
text=args.text,
image=args.image,
height=512,
width=512,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
lora_scale=args.lora_scale,
reference_conditioning_scale=args.reference_conditioning_scale,
negative_prompt=args.negative_prompt,
device=args.device,
remove_bg_fn=remove_bg_fn,
)
make_image_grid(images, rows=1).save(args.output)
reference_image.save(args.output.rsplit(".", 1)[0] + "_reference.png")
Gen_3D_Modules/MV_Adapter/scripts/inference_i2mv_sdxl.py
+267
View File
@@ -0,0 +1,267 @@
import argparse
import numpy as np
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
from transformers import AutoModelForImageSegmentation
from mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils.mesh_utils import get_orthogonal_camera
from mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
from mvadapter.utils import make_image_grid
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterI2MVSDXLPipeline
pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(num_views=num_views)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_i2mv_sdxl.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
# vae slicing for lower memory usage
pipe.enable_vae_slicing()
return pipe
def remove_bg(image, net, transform, device):
image_size = image.size
input_images = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
preds = net(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(image_size)
image.putalpha(mask)
return image
def preprocess_image(image: Image.Image, height, width):
image = np.array(image)
alpha = image[..., 3] > 0
H, W = alpha.shape
# get the bounding box of alpha
y, x = np.where(alpha)
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
image_center = image[y0:y1, x0:x1]
# resize the longer side to H * 0.9
H, W, _ = image_center.shape
if H > W:
W = int(W * (height * 0.9) / H)
H = int(height * 0.9)
else:
H = int(H * (width * 0.9) / W)
W = int(width * 0.9)
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
# pad to H, W
start_h = (height - H) // 2
start_w = (width - W) // 2
image = np.zeros((height, width, 4), dtype=np.uint8)
image[start_h : start_h + H, start_w : start_w + W] = image_center
image = image.astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = (image * 255).clip(0, 255).astype(np.uint8)
image = Image.fromarray(image)
return image
def run_pipeline(
pipe,
num_views,
text,
image,
height,
width,
num_inference_steps,
guidance_scale,
seed,
remove_bg_fn=None,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
lora_scale=1.0,
device="cuda",
azimuth_deg=None,
):
# Prepare cameras
if azimuth_deg is None:
azimuth_deg = [0, 45, 90, 180, 270, 315]
cameras = get_orthogonal_camera(
elevation_deg=[0] * num_views,
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in azimuth_deg],
device=device,
)
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
cameras.c2w, [1.1] * num_views, width
)
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
# Prepare image
reference_image = Image.open(image) if isinstance(image, str) else image
if remove_bg_fn is not None:
reference_image = remove_bg_fn(reference_image)
reference_image = preprocess_image(reference_image, height, width)
elif reference_image.mode == "RGBA":
reference_image = preprocess_image(reference_image, height, width)
pipe_kwargs = {}
if seed != -1 and isinstance(seed, int):
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=reference_image,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, reference_image
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--num_views", type=int, default=6) # not used
parser.add_argument(
"--azimuth_deg", type=int, nargs="+", default=[0, 45, 90, 180, 270, 315]
)
parser.add_argument("--image", type=str, required=True)
parser.add_argument("--text", type=str, default="high quality")
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--output", type=str, default="output.png")
# Extra
parser.add_argument("--remove_bg", action="store_true", help="Remove background")
args = parser.parse_args()
num_views = len(args.azimuth_deg)
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=num_views,
device=args.device,
dtype=torch.float16,
)
if args.remove_bg:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(args.device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
else:
remove_bg_fn = None
images, reference_image = run_pipeline(
pipe,
num_views=num_views,
text=args.text,
image=args.image,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
lora_scale=args.lora_scale,
reference_conditioning_scale=args.reference_conditioning_scale,
negative_prompt=args.negative_prompt,
device=args.device,
remove_bg_fn=remove_bg_fn,
azimuth_deg=args.azimuth_deg,
)
make_image_grid(images, rows=1).save(args.output)
reference_image.save(args.output.rsplit(".", 1)[0] + "_reference.png")
Gen_3D_Modules/MV_Adapter/scripts/inference_ig2mv_partial_sdxl.py
+303
View File
@@ -0,0 +1,303 @@
import argparse
import json
import numpy as np
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
from transformers import AutoModelForImageSegmentation
from mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
from mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils import make_image_grid, tensor_to_image
from mvadapter.utils.mesh_utils import (
NVDiffRastContextWrapper,
get_orthogonal_camera,
load_mesh,
render,
)
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterI2MVSDXLPipeline
pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(
num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_ig2mv_partial_sdxl.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
return pipe
def remove_bg(image, net, transform, device):
image_size = image.size
input_images = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
preds = net(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(image_size)
image.putalpha(mask)
return image
def preprocess_image(image: Image.Image, height, width):
image = np.array(image)
alpha = image[..., 3] > 0
H, W = alpha.shape
# get the bounding box of alpha
y, x = np.where(alpha)
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
image_center = image[y0:y1, x0:x1]
# resize the longer side to H * 0.9
H, W, _ = image_center.shape
if H > W:
W = int(W * (height * 0.9) / H)
H = int(height * 0.9)
else:
H = int(H * (width * 0.9) / W)
W = int(width * 0.9)
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
# pad to H, W
start_h = (height - H) // 2
start_w = (width - W) // 2
image = np.zeros((height, width, 4), dtype=np.uint8)
image[start_h : start_h + H, start_w : start_w + W] = image_center
image = image.astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = (image * 255).clip(0, 255).astype(np.uint8)
image = Image.fromarray(image)
return image
def run_pipeline(
pipe,
mesh_path,
num_views,
text,
image,
height,
width,
num_inference_steps,
guidance_scale,
seed,
remove_bg_fn=None,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
device=device,
)
ctx = NVDiffRastContextWrapper(device=device)
mesh, offset, scale = load_mesh(
mesh_path,
rescale=True,
move_to_center=True,
device=device,
return_transform=True,
)
transform_dict = {"offset": offset.tolist(), "scale": scale.tolist()}
render_out = render(
ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
normal_background=0.0,
)
pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
normal_images = tensor_to_image(
(render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
)
control_images = (
torch.cat(
[
(render_out.pos + 0.5).clamp(0, 1),
(render_out.normal / 2 + 0.5).clamp(0, 1),
],
dim=-1,
)
.permute(0, 3, 1, 2)
.to(device)
)
# Prepare image
reference_image = Image.open(image) if isinstance(image, str) else image
if remove_bg_fn is not None:
reference_image = remove_bg_fn(reference_image)
reference_image = preprocess_image(reference_image, height, width)
elif reference_image.mode == "RGBA":
reference_image = preprocess_image(reference_image, height, width)
pipe_kwargs = {}
if seed != -1 and isinstance(seed, int):
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=reference_image,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, pos_images, normal_images, reference_image, transform_dict
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--image", type=str, required=True)
parser.add_argument("--text", type=str, required=False, default="high quality")
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--output", type=str, default="output.png")
# Extra
parser.add_argument("--remove_bg", action="store_true", help="Remove background")
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
if args.remove_bg:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(args.device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
else:
remove_bg_fn = None
images, pos_images, normal_images, reference_image, transform_dict = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=args.num_views,
text=args.text,
image=args.image,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
lora_scale=args.lora_scale,
reference_conditioning_scale=args.reference_conditioning_scale,
negative_prompt=args.negative_prompt,
device=args.device,
remove_bg_fn=remove_bg_fn,
)
make_image_grid(images, rows=1).save(args.output)
make_image_grid(pos_images, rows=1).save(args.output.rsplit(".", 1)[0] + "_pos.png")
make_image_grid(normal_images, rows=1).save(
args.output.rsplit(".", 1)[0] + "_nor.png"
)
reference_image.save(args.output.rsplit(".", 1)[0] + "_reference.png")
with open(args.output.rsplit(".", 1)[0] + "_transform.json", "w") as f:
json.dump(transform_dict, f, indent=4)
Gen_3D_Modules/MV_Adapter/scripts/inference_ig2mv_sdxl.py
+313
View File
@@ -0,0 +1,313 @@
import argparse
import numpy as np
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
from transformers import AutoModelForImageSegmentation
from mmgp import offload, profile_type
from mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
from mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils import make_image_grid, tensor_to_image
from mvadapter.utils.mesh_utils import (
NVDiffRastContextWrapper,
get_orthogonal_camera,
load_mesh,
render,
)
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterI2MVSDXLPipeline
pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(
num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_ig2mv_sdxl.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
import gc
torch.cuda.empty_cache()
gc.collect()
all_models = {
"vae": pipe.vae,
"text_encoder": pipe.text_encoder,
"text_encoder_2": pipe.text_encoder_2,
"unet": pipe.unet,
"cond_encoder": pipe.cond_encoder,
}
if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
all_models["image_encoder"] = pipe.image_encoder
offload.profile(all_models, profile_type.HighRAM_LowVRAM)
print("mmgp profiling applied successfully.")
return pipe
def remove_bg(image, net, transform, device):
image_size = image.size
input_images = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
preds = net(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(image_size)
image.putalpha(mask)
return image
def preprocess_image(image: Image.Image, height, width):
image = np.array(image)
alpha = image[..., 3] > 0
H, W = alpha.shape
# get the bounding box of alpha
y, x = np.where(alpha)
y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
image_center = image[y0:y1, x0:x1]
# resize the longer side to H * 0.9
H, W, _ = image_center.shape
if H > W:
W = int(W * (height * 0.9) / H)
H = int(height * 0.9)
else:
H = int(H * (width * 0.9) / W)
W = int(width * 0.9)
image_center = np.array(Image.fromarray(image_center).resize((W, H)))
# pad to H, W
start_h = (height - H) // 2
start_w = (width - W) // 2
image = np.zeros((height, width, 4), dtype=np.uint8)
image[start_h : start_h + H, start_w : start_w + W] = image_center
image = image.astype(np.float32) / 255.0
image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
image = (image * 255).clip(0, 255).astype(np.uint8)
image = Image.fromarray(image)
return image
def run_pipeline(
pipe,
mesh_path,
num_views,
text,
image,
height,
width,
num_inference_steps,
guidance_scale,
seed,
remove_bg_fn=None,
reference_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
device=device,
)
ctx = NVDiffRastContextWrapper(device=device)
mesh = load_mesh(mesh_path, rescale=True, device=device)
render_out = render(
ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
normal_background=0.0,
)
pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
normal_images = tensor_to_image(
(render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
)
control_images = (
torch.cat(
[
(render_out.pos + 0.5).clamp(0, 1),
(render_out.normal / 2 + 0.5).clamp(0, 1),
],
dim=-1,
)
.permute(0, 3, 1, 2)
.to(device)
)
# Prepare image
reference_image = Image.open(image) if isinstance(image, str) else image
if remove_bg_fn is not None:
reference_image = remove_bg_fn(reference_image)
reference_image = preprocess_image(reference_image, height, width)
elif reference_image.mode == "RGBA":
reference_image = preprocess_image(reference_image, height, width)
pipe_kwargs = {}
if seed != -1 and isinstance(seed, int):
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
reference_image=reference_image,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, pos_images, normal_images, reference_image
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--image", type=str, required=True)
parser.add_argument("--text", type=str, required=False, default="high quality")
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--output", type=str, default="output.png")
# Extra
parser.add_argument("--remove_bg", action="store_true", help="Remove background")
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
if args.remove_bg:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(args.device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
else:
remove_bg_fn = None
images, pos_images, normal_images, reference_image = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=args.num_views,
text=args.text,
image=args.image,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
lora_scale=args.lora_scale,
reference_conditioning_scale=args.reference_conditioning_scale,
negative_prompt=args.negative_prompt,
device=args.device,
remove_bg_fn=remove_bg_fn,
)
make_image_grid(images, rows=1).save(args.output)
make_image_grid(pos_images, rows=1).save(args.output.rsplit(".", 1)[0] + "_pos.png")
make_image_grid(normal_images, rows=1).save(
args.output.rsplit(".", 1)[0] + "_nor.png"
)
reference_image.save(args.output.rsplit(".", 1)[0] + "_reference.png")
Gen_3D_Modules/MV_Adapter/scripts/inference_scribble2mv_sdxl.py
+252
View File
@@ -0,0 +1,252 @@
import argparse
import random
import cv2
import numpy as np
import torch
from controlnet_aux import HEDdetector, PidiNetDetector
from diffusers import (
AutoencoderKL,
ControlNetModel,
DDPMScheduler,
LCMScheduler,
UNet2DConditionModel,
)
from PIL import Image
from mvadapter.pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils.mesh_utils import get_orthogonal_camera
from mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
from mvadapter.utils import make_image_grid
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterT2MVSDXLPipeline
pipe = MVAdapterT2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(num_views=num_views)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_t2mv_sdxl.safetensors"
)
# ControlNet
pipe.controlnet = ControlNetModel.from_pretrained(
"xinsir/controlnet-scribble-sdxl-1.0"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
pipe.controlnet.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
# vae slicing for lower memory usage
pipe.enable_vae_slicing()
return pipe
def nms(x, t, s):
x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
y = np.zeros_like(x)
for f in [f1, f2, f3, f4]:
np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
z = np.zeros_like(y, dtype=np.uint8)
z[y > t] = 255
return z
def preprocess_controlnet_image(image_path, height, width):
image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
image = cv2.resize(image, (width, height))
if image.shape[2] == 4:
alpha_channel = image[:, :, 3] / 255.0
rgb_channels = image[:, :, :3] / 255.0
gray_background = np.ones_like(rgb_channels) * 0.5
image = (
alpha_channel[..., None] * rgb_channels
+ (1 - alpha_channel[..., None]) * gray_background
)
image = (image * 255).astype(np.uint8)
processor = HEDdetector.from_pretrained("lllyasviel/Annotators")
image = processor(image, scribble=False)
# following is some processing to simulate human sketch draw, different threshold can generate different width of lines
image = np.array(image)
image = nms(image, 127, 3)
image = cv2.GaussianBlur(image, (0, 0), 3)
# higher threshold, thiner line
random_val = int(round(random.uniform(0.01, 0.10), 2) * 255)
image[image > random_val] = 255
image[image < 255] = 0
return Image.fromarray(image)
def run_pipeline(
pipe,
num_views,
text,
height,
width,
num_inference_steps,
guidance_scale,
seed,
controlnet_images,
controlnet_conditioning_scale,
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 0, 0],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 45, 90, 180, 270, 315]],
device=device,
)
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
cameras.c2w, [1.1] * num_views, width
)
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
pipe_kwargs = {}
if seed != -1:
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
# Prepare controlnet images
controlnet_image = [
preprocess_controlnet_image(path, height, width) for path in controlnet_images
]
pipe_kwargs.update(
{
"controlnet_image": controlnet_image,
"controlnet_conditioning_scale": controlnet_conditioning_scale,
}
)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, controlnet_image
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=7.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--output", type=str, default="output.png")
parser.add_argument("--controlnet_images", type=str, nargs="+", required=True)
parser.add_argument("--controlnet_conditioning_scale", type=float, default=1.0)
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
images, controlnet_images = run_pipeline(
pipe,
num_views=args.num_views,
text=args.text,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
controlnet_images=args.controlnet_images,
controlnet_conditioning_scale=args.controlnet_conditioning_scale,
lora_scale=args.lora_scale,
device=args.device,
)
make_image_grid(images, rows=1).save(args.output)
make_image_grid(controlnet_images, rows=1).save(
args.output.rsplit(".", 1)[0] + "_controlnet.png"
)
Gen_3D_Modules/MV_Adapter/scripts/inference_t2mv_sd.py
+170
View File
@@ -0,0 +1,170 @@
import argparse
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from mvadapter.pipelines.pipeline_mvadapter_t2mv_sd import MVAdapterT2MVSDPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils.mesh_utils import get_orthogonal_camera
from mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
from mvadapter.utils import make_image_grid
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterT2MVSDPipeline
pipe = MVAdapterT2MVSDPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(num_views=num_views)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_t2mv_sd21.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
# vae slicing for lower memory usage
pipe.enable_vae_slicing()
return pipe
def run_pipeline(
pipe,
num_views,
text,
height,
width,
num_inference_steps,
guidance_scale,
seed,
negative_prompt,
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 0, 0],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 45, 90, 180, 270, 315]],
device=device,
)
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
cameras.c2w, [1.1] * num_views, width
)
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
pipe_kwargs = {}
if seed != -1:
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-2-1-base"
)
parser.add_argument("--vae_model", type=str, default=None)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=7.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--output", type=str, default="output.png")
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
images = run_pipeline(
pipe,
num_views=args.num_views,
text=args.text,
height=512,
width=512,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
negative_prompt=args.negative_prompt,
lora_scale=args.lora_scale,
device=args.device,
)
make_image_grid(images, rows=1).save(args.output)
Gen_3D_Modules/MV_Adapter/scripts/inference_t2mv_sdxl.py
+201
View File
@@ -0,0 +1,201 @@
import argparse
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from mmgp import offload, profile_type
from mvadapter.pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils.mesh_utils import get_orthogonal_camera
from mvadapter.utils.geometry import get_plucker_embeds_from_cameras_ortho
from mvadapter.utils import make_image_grid
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterT2MVSDXLPipeline
pipe = MVAdapterT2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(num_views=num_views)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_t2mv_sdxl.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
import gc
torch.cuda.empty_cache()
gc.collect()
all_models = {
"vae": pipe.vae,
"text_encoder": pipe.text_encoder,
"text_encoder_2": pipe.text_encoder_2,
"unet": pipe.unet,
"cond_encoder": pipe.cond_encoder,
}
if hasattr(pipe, 'image_encoder') and pipe.image_encoder is not None:
all_models["image_encoder"] = pipe.image_encoder
offload.profile(all_models, profile_type.HighRAM_LowVRAM)
print("mmgp profiling applied successfully.")
# vae slicing for lower memory usage
pipe.enable_vae_slicing()
return pipe
def run_pipeline(
pipe,
num_views,
text,
height,
width,
num_inference_steps,
guidance_scale,
seed,
negative_prompt,
lora_scale=1.0,
device="cuda",
azimuth_deg=None,
):
# Prepare cameras
if azimuth_deg is None:
azimuth_deg = [0, 45, 90, 180, 270, 315]
cameras = get_orthogonal_camera(
elevation_deg=[0] * num_views,
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in azimuth_deg],
device=device,
)
plucker_embeds = get_plucker_embeds_from_cameras_ortho(
cameras.c2w, [1.1] * num_views, width
)
control_images = ((plucker_embeds + 1.0) / 2.0).clamp(0, 1)
pipe_kwargs = {}
if seed != -1:
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--num_views", type=int, default=6)
parser.add_argument(
"--azimuth_deg", type=int, nargs="+", default=[0, 45, 90, 180, 270, 315]
)
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=7.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--output", type=str, default="output.png")
args = parser.parse_args()
num_views = len(args.azimuth_deg)
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=num_views,
device=args.device,
dtype=torch.float16,
)
images = run_pipeline(
pipe,
num_views=num_views,
text=args.text,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
negative_prompt=args.negative_prompt,
lora_scale=args.lora_scale,
device=args.device,
azimuth_deg=args.azimuth_deg,
)
make_image_grid(images, rows=1).save(args.output)
Gen_3D_Modules/MV_Adapter/scripts/inference_tg2mv_sdxl.py
+200
View File
@@ -0,0 +1,200 @@
import argparse
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
from mvadapter.pipelines.pipeline_mvadapter_t2mv_sdxl import MVAdapterT2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler
from mvadapter.utils import get_orthogonal_camera, make_image_grid, tensor_to_image
from mvadapter.utils.mesh_utils import NVDiffRastContextWrapper, load_mesh, render
def prepare_pipeline(
base_model,
vae_model,
unet_model,
lora_model,
adapter_path,
scheduler,
num_views,
device,
dtype,
):
# Load vae and unet if provided
pipe_kwargs = {}
if vae_model is not None:
pipe_kwargs["vae"] = AutoencoderKL.from_pretrained(vae_model)
if unet_model is not None:
pipe_kwargs["unet"] = UNet2DConditionModel.from_pretrained(unet_model)
# Prepare pipeline
pipe: MVAdapterT2MVSDXLPipeline
pipe = MVAdapterT2MVSDXLPipeline.from_pretrained(base_model, **pipe_kwargs)
# Load scheduler if provided
scheduler_class = None
if scheduler == "ddpm":
scheduler_class = DDPMScheduler
elif scheduler == "lcm":
scheduler_class = LCMScheduler
pipe.scheduler = ShiftSNRScheduler.from_scheduler(
pipe.scheduler,
shift_mode="interpolated",
shift_scale=8.0,
scheduler_class=scheduler_class,
)
pipe.init_custom_adapter(
num_views=num_views, self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0
)
pipe.load_custom_adapter(
adapter_path, weight_name="mvadapter_tg2mv_sdxl.safetensors"
)
pipe.to(device=device, dtype=dtype)
pipe.cond_encoder.to(device=device, dtype=dtype)
# load lora if provided
if lora_model is not None:
model_, name_ = lora_model.rsplit("/", 1)
pipe.load_lora_weights(model_, weight_name=name_)
return pipe
def run_pipeline(
pipe,
mesh_path,
num_views,
text,
height,
width,
num_inference_steps,
guidance_scale,
seed,
negative_prompt,
lora_scale=1.0,
device="cuda",
):
# Prepare cameras
cameras = get_orthogonal_camera(
elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
distance=[1.8] * num_views,
left=-0.55,
right=0.55,
bottom=-0.55,
top=0.55,
azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
device=device,
)
ctx = NVDiffRastContextWrapper(device=device)
mesh = load_mesh(mesh_path, rescale=True, device=device)
render_out = render(
ctx,
mesh,
cameras,
height=height,
width=width,
render_attr=False,
normal_background=0.0,
)
pos_images = tensor_to_image((render_out.pos + 0.5).clamp(0, 1), batched=True)
normal_images = tensor_to_image(
(render_out.normal / 2 + 0.5).clamp(0, 1), batched=True
)
control_images = (
torch.cat(
[
(render_out.pos + 0.5).clamp(0, 1),
(render_out.normal / 2 + 0.5).clamp(0, 1),
],
dim=-1,
)
.permute(0, 3, 1, 2)
.to(device)
)
pipe_kwargs = {}
if seed != -1:
pipe_kwargs["generator"] = torch.Generator(device=device).manual_seed(seed)
images = pipe(
text,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
num_images_per_prompt=num_views,
control_image=control_images,
control_conditioning_scale=1.0,
negative_prompt=negative_prompt,
cross_attention_kwargs={"scale": lora_scale},
**pipe_kwargs,
).images
return images, pos_images, normal_images
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# Models
parser.add_argument(
"--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0"
)
parser.add_argument(
"--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix"
)
parser.add_argument("--unet_model", type=str, default=None)
parser.add_argument("--scheduler", type=str, default=None)
parser.add_argument("--lora_model", type=str, default=None)
parser.add_argument("--adapter_path", type=str, default="huanngzh/mv-adapter")
parser.add_argument("--num_views", type=int, default=6)
# Device
parser.add_argument("--device", type=str, default="cuda")
# Inference
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=7.0)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument(
"--negative_prompt",
type=str,
default="watermark, ugly, deformed, noisy, blurry, low contrast",
)
parser.add_argument("--lora_scale", type=float, default=1.0)
parser.add_argument("--output", type=str, default="output.png")
args = parser.parse_args()
pipe = prepare_pipeline(
base_model=args.base_model,
vae_model=args.vae_model,
unet_model=args.unet_model,
lora_model=args.lora_model,
adapter_path=args.adapter_path,
scheduler=args.scheduler,
num_views=args.num_views,
device=args.device,
dtype=torch.float16,
)
images, pos_images, normal_images = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=args.num_views,
text=args.text,
height=768,
width=768,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
negative_prompt=args.negative_prompt,
lora_scale=args.lora_scale,
device=args.device,
)
make_image_grid(images, rows=1).save(args.output)
make_image_grid(pos_images, rows=1).save(args.output.rsplit(".", 1)[0] + "_pos.png")
make_image_grid(normal_images, rows=1).save(
args.output.rsplit(".", 1)[0] + "_nor.png"
)
Gen_3D_Modules/MV_Adapter/scripts/texture_i2tex.py
Executable
+104
View File
@@ -0,0 +1,104 @@
import argparse
import os
import sys
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from mvadapter.pipelines.pipeline_texture import ModProcessConfig, TexturePipeline
from mvadapter.utils import make_image_grid
from .inference_ig2mv_sdxl import prepare_pipeline, remove_bg, run_pipeline
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--device", type=str, default="cuda")
# I/O
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--image", type=str, required=True)
parser.add_argument("--text", type=str, default="high quality")
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--save_dir", type=str, default="./output")
parser.add_argument("--save_name", type=str, default="i2tex_sample")
# Extra
parser.add_argument("--reference_conditioning_scale", type=float, default=1.0)
parser.add_argument("--preprocess_mesh", action="store_true")
parser.add_argument("--remove_bg", action="store_true")
args = parser.parse_args()
device = args.device
num_views = 6
# Prepare pipelines
pipe = prepare_pipeline(
base_model="stabilityai/stable-diffusion-xl-base-1.0",
vae_model="madebyollin/sdxl-vae-fp16-fix",
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=None,
num_views=num_views,
device=device,
dtype=torch.float16,
)
if args.remove_bg:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to(args.device)
transform_image = transforms.Compose(
[
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, args.device)
else:
remove_bg_fn = None
texture_pipe = TexturePipeline(
upscaler_ckpt_path="./checkpoints/RealESRGAN_x2plus.pth",
inpaint_ckpt_path="./checkpoints/big-lama.pt",
device=device,
)
print("Pipeline ready.")
os.makedirs(args.save_dir, exist_ok=True)
# 1. run MV-Adapter to generate multi-view images
images, _, _, _ = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=num_views,
text=args.text,
image=args.image,
height=768,
width=768,
num_inference_steps=50,
guidance_scale=3.0,
seed=args.seed,
reference_conditioning_scale=args.reference_conditioning_scale,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
device=device,
remove_bg_fn=remove_bg_fn,
)
mv_path = os.path.join(args.save_dir, f"{args.save_name}.png")
make_image_grid(images, rows=1).save(mv_path)
torch.cuda.empty_cache()
# 2. un-project and complete texture
out = texture_pipe(
mesh_path=args.mesh,
save_dir=args.save_dir,
save_name=args.save_name,
uv_unwarp=True,
preprocess_mesh=args.preprocess_mesh,
uv_size=4096,
rgb_path=mv_path,
rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
)
print(f"Output saved to {out.shaded_model_save_path}")
Gen_3D_Modules/MV_Adapter/scripts/texture_t2tex.py
Executable
+80
View File
@@ -0,0 +1,80 @@
import argparse
import os
import sys
import torch
from mvadapter.pipelines.pipeline_texture import ModProcessConfig, TexturePipeline
from mvadapter.utils import make_image_grid
from .inference_tg2mv_sdxl import prepare_pipeline, run_pipeline
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--device", type=str, default="cuda")
# I/O
parser.add_argument("--mesh", type=str, required=True)
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--seed", type=int, default=-1)
parser.add_argument("--save_dir", type=str, default="./output")
parser.add_argument("--save_name", type=str, default="t2tex_sample")
# Extra
parser.add_argument("--preprocess_mesh", action="store_true")
args = parser.parse_args()
device = args.device
num_views = 6
# Prepare pipelines
pipe = prepare_pipeline(
base_model="stabilityai/stable-diffusion-xl-base-1.0",
vae_model="madebyollin/sdxl-vae-fp16-fix",
unet_model=None,
lora_model=None,
adapter_path="huanngzh/mv-adapter",
scheduler=None,
num_views=num_views,
device=device,
dtype=torch.float16,
)
texture_pipe = TexturePipeline(
upscaler_ckpt_path="./checkpoints/RealESRGAN_x2plus.pth",
inpaint_ckpt_path="./checkpoints/big-lama.pt",
device=device,
)
print("Pipeline ready.")
os.makedirs(args.save_dir, exist_ok=True)
# 1. run MV-Adapter to generate multi-view images
images, pos_images, normal_images = run_pipeline(
pipe,
mesh_path=args.mesh,
num_views=num_views,
text=args.text,
height=768,
width=768,
num_inference_steps=50,
guidance_scale=7.0,
seed=args.seed,
negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
device=device,
)
mv_path = os.path.join(args.save_dir, f"{args.save_name}.png")
make_image_grid(images, rows=1).save(mv_path)
torch.cuda.empty_cache()
# 2. un-project and complete texture
out = texture_pipe(
mesh_path=args.mesh,
save_dir=args.save_dir,
save_name=args.save_name,
uv_unwarp=True,
preprocess_mesh=args.preprocess_mesh,
uv_size=4096,
rgb_path=mv_path,
rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
)
print(f"Output saved to {out.shaded_model_save_path}")
Gen_3D_Modules/MV_Adapter/setup.py
+59
View File
@@ -0,0 +1,59 @@
from setuptools import find_packages, setup
setup(
name="mvadapter",
version="0.1.0",
description="MV-Adapter: A Multi-View Adapter for 3D Generation",
author="Zehuan Huang",
packages=[
"mvadapter",
"mvadapter.loaders",
"mvadapter.models",
"mvadapter.pipelines",
"mvadapter.schedulers",
"mvadapter.utils",
"mvadapter.utils.mesh_utils",
],
install_requires=[
"torch>=2.0.0",
"torchvision",
"controlnet_aux",
"diffusers",
"transformers",
"peft",
"numpy",
"huggingface_hub",
"accelerate",
"opencv-python",
"safetensors",
"pillow",
"omegaconf",
"trimesh",
"einops",
"gradio",
"timm",
"kornia",
"scikit-image",
"sentencepiece",
"spandrel",
"open3d",
"pymeshlab",
"cvcuda_cu12",
],
python_requires=">=3.8",
classifiers=[
"Development Status :: 3 - Alpha",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
],
long_description=open("README.md").read(),
long_description_content_type="text/markdown",
url="https://github.com/huanngzh/MV-Adapter",
project_urls={
"Project Page": "https://huanngzh.github.io/MV-Adapter-Page/",
"Paper": "https://arxiv.org/abs/2412.03632",
"Model Weights": "https://huggingface.co/huanngzh/mv-adapter",
},
)
Gen_3D_Modules/MV_Adapter/test.sh
+138
View File
@@ -0,0 +1,138 @@
#!/bin/bash
# Performance testing script for MV-Adapter - dual run comparison
echo "MV-Adapter Performance Testing - Dual Run"
# Function to get memory usage by process
get_memory_usage() {
local pid=$1
if [ -n "$pid" ] && ps -p $pid > /dev/null 2>&1; then
ps -p $pid -o rss= | awk '{print $1/1024}' 2>/dev/null || echo "0"
else
echo "0"
fi
}
# Function to get VRAM usage
get_vram_usage() {
if command -v nvidia-smi &> /dev/null; then
nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | awk '{print $1}'
else
echo "0"
fi
}
# Function to run single test
run_test() {
local test_name=$1
local use_mmgp=$2
local output_suffix=$3
echo "========================================="
echo "Running $test_name"
echo "========================================="
# Memory monitoring setup
log_file="memory_${output_suffix}.log"
echo "timestamp,ram_mb,vram_mb" > $log_file
# Build command
cmd="python -m scripts.inference_ig2mv_sdxl \
--image assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.jpeg \
--mesh assets/demo/ig2mv/1ccd5c1563ea4f5fb8152eac59dabd5c.glb \
--output output_${output_suffix}.png \
--remove_bg"
if [ "$use_mmgp" = "true" ]; then
cmd="$cmd --mmgp"
fi
# Start test
start_time=$(date +%s.%N)
$cmd &
python_pid=$!
# Monitor memory
max_ram=0
max_vram=0
while ps -p $python_pid > /dev/null 2>&1; do
current_time=$(date +%s.%N)
current_ram=$(get_memory_usage $python_pid)
current_vram=$(get_vram_usage)
echo "${current_time},${current_ram},${current_vram}" >> $log_file
if (( $(echo "$current_ram > $max_ram" | bc -l) )); then
max_ram=$current_ram
fi
if (( $current_vram > $max_vram )); then
max_vram=$current_vram
fi
sleep 1
done
# Wait for completion
wait $python_pid
exit_code=$?
# Calculate execution time
end_time=$(date +%s.%N)
execution_time=$(echo "$end_time - $start_time" | bc -l)
# Display results
echo "Results for $test_name:"
echo " Execution time: ${execution_time} seconds"
echo " Max RAM usage: ${max_ram} MB"
echo " Max VRAM usage: ${max_vram} MB"
echo " Exit code: ${exit_code}"
echo ""
# Save results
{
echo "Test: $test_name"
echo "Date: $(date)"
echo "Execution time: ${execution_time} seconds"
echo "Max RAM usage: ${max_ram} MB"
echo "Max VRAM usage: ${max_vram} MB"
echo "Exit code: ${exit_code}"
echo "Log file: ${log_file}"
echo "========================================="
} >> "dual_test_results.txt"
if [ $exit_code -eq 0 ]; then
echo "✓ $test_name completed successfully"
if [ -f "output_${output_suffix}.png" ]; then
ls -lh "output_${output_suffix}.png"
fi
else
echo "✗ $test_name failed with exit code: ${exit_code}"
fi
}
# Initialize results file
echo "Dual Performance Test Results" > "dual_test_results.txt"
echo "Generated: $(date)" >> "dual_test_results.txt"
echo "=========================================" >> "dual_test_results.txt"
# Run tests
echo "Starting dual performance tests..."
echo ""
# Test 1: Without mmgp
# run_test "Test without mmgp" "false" "no_mmgp"
# Small delay between tests
# sleep 5
# Test 2: With mmgp
run_test "Test with mmgp" "true" "with_mmgp"
echo "========================================="
echo "DUAL TEST COMPLETE"
echo "========================================="
echo "Results saved to: dual_test_results.txt"
echo "Memory logs saved to: memory_no_mmgp.log and memory_with_mmgp.log"
echo "Output images: output_no_mmgp.png and output_with_mmgp.png"
nodes.py
+133
View File
@@ -105,6 +105,12 @@ from TripoSG.pipelines.pipeline_triposg import TripoSGPipeline
from TripoSG.pipelines.pipeline_triposg_scribble import TripoSGScribblePipeline
from Stable3DGen.trellis.backend_config import get_available_backends, get_available_sparse_backends
from Stable3DGen.pipeline_builders import StableGenPipelineBuilder
from MV_Adapter.mvadapter_node_utils import (
prepare_pipeline as mvadapter_prepare_pipeline,
run_pipeline as mvadapter_run_pipeline,
create_bg_remover,
)
from MV_Adapter.mvadapter.utils import make_image_grid
os.environ['SPCONV_ALGO'] = 'native'
@@ -4781,3 +4787,130 @@ class StableGen_StableX_Process_Image:
except Exception as e:
raise Exception(f"[StableGen_StableX_Process_Image] Image processing failed: {str(e)}")
class Load_MVAdapter_Pipeline:
"""Loader pipeline for MV-Adapter (Image to Multi-View)"""
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("DIFFUSERS_PIPE",)
RETURN_NAMES = ("mvadapter_pipe",)
FUNCTION = "load"
CKPT_MVADAPTER_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh", "MV-Adapter")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_model": (["stabilityai/stable-diffusion-xl-base-1.0"],
{"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
"vae_model": (["madebyollin/sdxl-vae-fp16-fix", "None"],
{"default": "madebyollin/sdxl-vae-fp16-fix"}),
"adapter_path": (["huanngzh/mv-adapter"], {"default": "huanngzh/mv-adapter"}),
"scheduler": (["default", "ddpm", "lcm"], {"default": "default"}),
"num_views": ("INT", {"default": 6, "min": 1, "max": 16}),
"use_fp16": ("BOOLEAN", {"default": True}),
"use_mmgp": ("BOOLEAN", {"default": True}),
},
"optional": {
"unet_model": ("STRING", {"default": ""}),
"lora_model": ("STRING", {"default": ""}),
}
}
@classmethod
def load(cls, base_model, vae_model, adapter_path, scheduler, num_views,
use_fp16, use_mmgp, unet_model="", lora_model=""):
dtype = torch.float16 if use_fp16 else torch.float32
vae_model = None if vae_model == "None" else vae_model
unet_model = None if not unet_model else unet_model
lora_model = None if not lora_model else lora_model
pipe = mvadapter_prepare_pipeline(
base_model=base_model,
vae_model=vae_model,
unet_model=unet_model,
lora_model=lora_model,
adapter_path=adapter_path,
scheduler=scheduler,
num_views=num_views,
device=DEVICE_STR,
dtype=dtype,
use_mmgp=use_mmgp,
adapter_local_path=cls.CKPT_MVADAPTER_PATH
)
print("MV-Adapter pipeline loaded successfully")
return (pipe,)
class MVAdapter_Image_To_MultiView:
"""Generate multi-view images from single image and 3D mesh"""
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("multiview_images",)
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mvadapter_pipe": ("DIFFUSERS_PIPE",),
"mesh_path": ("STRING", {"default": ""}),
"reference_image": ("IMAGE",),
"prompt": ("STRING", {"default": "high quality", "multiline": True}),
"negative_prompt": ("STRING", {"default": "watermark, ugly, deformed, noisy, blurry, low contrast", "multiline": True}),
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200}),
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.1}),
"reference_conditioning_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
"height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
"width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
"seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
"remove_background": ("BOOLEAN", {"default": False}),
},
"optional": {
"lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
}
}
def run(self, mvadapter_pipe, mesh_path, reference_image, prompt, negative_prompt,
num_inference_steps, guidance_scale, reference_conditioning_scale,
height, width, seed, remove_background, lora_scale=1.0):
if isinstance(reference_image, torch.Tensor):
reference_images = torch_imgs_to_pils(reference_image)
reference_image = reference_images[0]
if not mesh_path or not os.path.exists(mesh_path):
raise ValueError(f"Mesh path does not exist: {mesh_path}")
remove_bg_fn = None
if remove_background:
remove_bg_fn = create_bg_remover(DEVICE_STR)
num_views = 6
images, pos_images, normal_images, processed_ref_image = mvadapter_run_pipeline(
pipe=mvadapter_pipe,
mesh_path=mesh_path,
num_views=num_views,
text=prompt,
image=reference_image,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
remove_bg_fn=remove_bg_fn,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
lora_scale=lora_scale,
device=DEVICE_STR,
)
grid_image = make_image_grid(images, rows=1)
grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
print(f"Generated multiview images: {grid_tensor.shape}")
return (grid_tensor,)
stablegen_nodes.py
+408
View File
@@ -0,0 +1,408 @@
import torch
import os
import trimesh
import numpy as np
from collections import OrderedDict
from huggingface_hub import snapshot_download
# Import StableGen modules
try:
from .Gen_3D_Modules.Stable3DGen.trellis.pipelines.trellis_image_to_3d import TrellisImageTo3DPipeline as Stable3DGenTrellisImageTo3DPipeline
from .Gen_3D_Modules.Stable3DGen.stablex.pipeline_yoso import YosoPipeline
from .Gen_3D_Modules.Stable3DGen.stablex.controlnetvae import ControlNetVAEModel
from .Gen_3D_Modules.Stable3DGen.pipeline_builders import StableGenPipelineBuilder
except ImportError as e:
print(f"Warning: Could not import StableGen modules: {e}")
# Create dummy functions for fallback
def get_available_backends():
return {'flash_attn': True, 'sdpa': True}
def get_available_sparse_backends():
return {'spconv': True}
def set_attention_backend(backend):
return True
def set_sparse_backend(backend, algo):
return True
class StableGenPipelineBuilder:
@staticmethod
def build_trellis_pipeline(*args, **kwargs):
raise ImportError("StableGen modules not available")
@staticmethod
def build_stablex_pipeline(*args, **kwargs):
raise ImportError("StableGen modules not available")
# Import diffusers components for StableX
try:
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from safetensors.torch import load_file
except ImportError as e:
print(f"Warning: Could not import diffusers/accelerate: {e}")
import json
from .shared_utils.mesh_utils import Mesh
from .shared_utils.image_utils import torch_imgs_to_pils, pils_to_torch_imgs, pil_make_image_grid
ROOT_PATH = os.path.dirname(os.path.realpath(__file__))
CKPT_ROOT_PATH = os.path.join(ROOT_PATH, "Checkpoints")
CKPT_DIFFUSERS_PATH = os.path.join(CKPT_ROOT_PATH, "Diffusers")
WEIGHT_DTYPE = torch.float16
DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
DEVICE = torch.device(DEVICE_STR)
HF_DOWNLOAD_IGNORE = ["*.yaml", "*.json", "*.py", ".png", ".jpg", ".gif"]
class Load_StableGen_Trellis_Pipeline:
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("DIFFUSERS_PIPE",)
RETURN_NAMES = ("trellis_pipe",)
FUNCTION = "load"
_REPO_ID_BASE = "Stable-X"
CKPT_STABLEGEN_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "Stable3DGen")
_MODES = {
"trellis-normal-v0-1": ("trellis-normal-v0-1", 12, 12), # (repo, ss_steps, slat_steps)
}
@classmethod
def INPUT_TYPES(cls):
available_attn, available_sparse = StableGenPipelineBuilder.get_available_backends()
return {
"required": {
"model_name": (list(cls._MODES.keys()),),
"dinov2_model": (["dinov2_vitl14_reg"],),
"use_fp16": ("BOOLEAN", {"default": True}),
"attn_backend": (available_attn,),
"sparse_backend": (available_sparse,),
"spconv_algo": (["implicit_gemm", "native", "auto"],),
"smooth_k": ("BOOLEAN", {"default": True}),
}
}
@classmethod
def _build_pipe(cls, repo: str, dinov2_model: str, use_fp16: bool, attn_backend: str,
sparse_backend: str, spconv_algo: str, smooth_k: bool):
return StableGenPipelineBuilder.build_trellis_pipeline(
repo=repo,
dinov2_model=dinov2_model,
use_fp16=use_fp16,
attn_backend=attn_backend,
sparse_backend=sparse_backend,
spconv_algo=spconv_algo,
smooth_k=smooth_k,
ckpt_path=cls.CKPT_STABLEGEN_PATH
)
def load(self, model_name, dinov2_model, use_fp16, attn_backend, sparse_backend, spconv_algo, smooth_k):
repo, ss_steps, slat_steps = self._MODES[model_name]
pipe = self.__class__._build_pipe(repo, dinov2_model, use_fp16, attn_backend, sparse_backend, spconv_algo, smooth_k)
# Store default steps
pipe.default_ss_steps = ss_steps
pipe.default_slat_steps = slat_steps
return (pipe,)
class Load_StableGen_StableX_Pipeline:
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("DIFFUSERS_PIPE",)
RETURN_NAMES = ("stablex_pipe",)
FUNCTION = "load"
_REPO_ID_BASE = "Stable-X"
CKPT_STABLEGEN_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "Stable3DGen")
_MODES = {
"yoso-normal-v1-8-1": "yoso-normal-v1-8-1",
}
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": (list(cls._MODES.keys()),),
"use_fp16": ("BOOLEAN", {"default": True}),
}
}
@classmethod
def _build_pipe(cls, repo: str, use_fp16: bool):
return StableGenPipelineBuilder.build_stablex_pipeline(
repo=repo,
use_fp16=use_fp16,
ckpt_path=cls.CKPT_STABLEGEN_PATH
)
def load(self, model_name, use_fp16):
repo = self._MODES[model_name]
pipe = self.__class__._build_pipe(repo, use_fp16)
return (pipe,)
class StableGen_Trellis_Image_To_3D:
"""
3D generation pipeline using Trellis model.
"""
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("MESH",)
RETURN_NAMES = ("mesh",)
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"trellis_pipe": ("DIFFUSERS_PIPE",),
"images": ("IMAGE", {"list": True}),
"mode": (["single", "multi"], {"default": "single"}),
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
"ss_guidance_strength": ("FLOAT", {"default": 7.5, "min": 0.0, "step": 0.1}),
"ss_sampling_steps": ("INT", {"default": 12, "min": 1}),
"slat_guidance_strength": ("FLOAT", {"default": 3.0, "min": 0.0, "step": 0.1}),
"slat_sampling_steps": ("INT", {"default": 12, "min": 1}),
"mesh_simplify": ("FLOAT", {"default": 0.95, "min": 0.9, "max": 1.0, "step": 0.01}),
}
}
@torch.no_grad()
def run(
self,
trellis_pipe,
images,
mode="single",
seed=1234,
ss_guidance_strength=7.5,
ss_sampling_steps=12,
slat_guidance_strength=3.0,
slat_sampling_steps=12,
mesh_simplify=0.95
):
if isinstance(images, torch.Tensor):
images = torch_imgs_to_pils(images)
if not isinstance(images, list) or len(images) == 0:
raise Exception(f"[StableGen_Trellis_Image_To_3D] 'images' must be a non-empty list of PIL images. Got type: {type(images)}, len: {len(images) if hasattr(images, '__len__') else 'no len'}")
# Use inference context like in working version
with trellis_pipe.inference_context():
# Prepare input based on mode
if mode == "single":
if len(images) > 1:
print(f"Warning: Single mode selected but {len(images)} images provided. Using first image.")
image_input = images[0]
# Set up generation parameters for single image
pipeline_params = StableGenPipelineBuilder.create_trellis_pipeline_params(
seed=seed,
ss_sampling_steps=ss_sampling_steps,
ss_guidance_strength=ss_guidance_strength,
slat_sampling_steps=slat_sampling_steps,
slat_guidance_strength=slat_guidance_strength,
formats=["mesh"]
)
# Generate 3D mesh using single image method
outputs = trellis_pipe.run(
image_input,
**pipeline_params
)
else: # multi mode
# Set up generation parameters for multi image
pipeline_params = StableGenPipelineBuilder.create_trellis_pipeline_params(
seed=seed,
ss_sampling_steps=ss_sampling_steps,
ss_guidance_strength=ss_guidance_strength,
slat_sampling_steps=slat_sampling_steps,
slat_guidance_strength=slat_guidance_strength,
formats=["mesh"]
)
outputs = trellis_pipe.run_multi_image(
images,
**pipeline_params
)
try:
mesh_output = outputs['mesh'][0]
vertices = mesh_output.vertices.cpu().numpy()
faces = mesh_output.faces.cpu().numpy()
transformation_matrix = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
vertices = vertices @ transformation_matrix
tri_mesh = trimesh.Trimesh(vertices, faces)
if mesh_simplify < 1.0:
try:
target_faces = int(len(faces) * mesh_simplify)
tri_mesh = tri_mesh.simplify_quadric_decimation(target_faces)
except Exception as e:
print(f"Warning: Mesh simplification failed: {e}. Using original mesh.")
# Convert to Mesh object
mesh = Mesh.load_trimesh(given_mesh=tri_mesh)
mesh.auto_normal()
return (mesh,)
except Exception as e:
raise Exception(f"[StableGen_Trellis_Image_To_3D] 3D generation failed: {str(e)}")
# Import MV-Adapter modules
try:
from .Gen_3D_Modules.MV_Adapter.mvadapter_node_utils import (
prepare_pipeline as mvadapter_prepare_pipeline,
run_pipeline as mvadapter_run_pipeline,
create_bg_remover,
)
from .Gen_3D_Modules.MV_Adapter.mvadapter.utils import make_image_grid
MVADAPTER_AVAILABLE = True
except ImportError as e:
print(f"Warning: Could not import MV-Adapter modules: {e}")
MVADAPTER_AVAILABLE = False
class Load_MVAdapter_Pipeline:
"""Loader pipeline for MV-Adapter (Image to Multi-View)"""
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("DIFFUSERS_PIPE",)
RETURN_NAMES = ("mvadapter_pipe",)
FUNCTION = "load"
CKPT_MVADAPTER_PATH = os.path.join(CKPT_DIFFUSERS_PATH, "huanngzh", "MV-Adapter")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"base_model": (["stabilityai/stable-diffusion-xl-base-1.0"],
{"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
"vae_model": (["madebyollin/sdxl-vae-fp16-fix", "None"],
{"default": "madebyollin/sdxl-vae-fp16-fix"}),
"adapter_path": (["huanngzh/mv-adapter"], {"default": "huanngzh/mv-adapter"}),
"scheduler": (["default", "ddpm", "lcm"], {"default": "default"}),
"num_views": ("INT", {"default": 6, "min": 1, "max": 16}),
"use_fp16": ("BOOLEAN", {"default": True}),
"use_mmgp": ("BOOLEAN", {"default": False}),
},
"optional": {
"unet_model": ("STRING", {"default": ""}),
"lora_model": ("STRING", {"default": ""}),
}
}
@classmethod
def load(self, base_model, vae_model, adapter_path, scheduler, num_views,
use_fp16, use_mmgp, unet_model="", lora_model=""):
dtype = torch.float16 if use_fp16 else torch.float32
vae_model = None if vae_model == "None" else vae_model
unet_model = None if not unet_model else unet_model
lora_model = None if not lora_model else lora_model
pipe = mvadapter_prepare_pipeline(
base_model=base_model,
vae_model=vae_model,
unet_model=unet_model,
lora_model=lora_model,
adapter_path=adapter_path,
scheduler=scheduler,
num_views=num_views,
device=DEVICE_STR,
dtype=dtype,
use_mmgp=use_mmgp,
adapter_local_path=cls.CKPT_MVADAPTER_PATH
)
print("MV-Adapter pipeline loaded successfully")
return (pipe,)
class MVAdapter_Image_To_MultiView:
"""Generate multi-view images from single image and 3D mesh"""
CATEGORY = "Comfy3D/Algorithm"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("multiview_images",)
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mvadapter_pipe": ("DIFFUSERS_PIPE",),
"mesh_path": ("STRING", {"default": ""}),
"reference_image": ("IMAGE",),
"prompt": ("STRING", {"default": "high quality", "multiline": True}),
"negative_prompt": ("STRING", {"default": "watermark, ugly, deformed, noisy, blurry, low contrast", "multiline": True}),
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200}),
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.1}),
"reference_conditioning_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
"height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
"width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 8}),
"seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
"remove_background": ("BOOLEAN", {"default": False}),
},
"optional": {
"lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
}
}
def run(self, mvadapter_pipe, mesh_path, reference_image, prompt, negative_prompt,
num_inference_steps, guidance_scale, reference_conditioning_scale,
height, width, seed, remove_background, lora_scale=1.0):
# Convert input image
if isinstance(reference_image, torch.Tensor):
reference_images = torch_imgs_to_pils(reference_image)
reference_image = reference_images[0]
# Check if mesh_path exists
if not mesh_path or not os.path.exists(mesh_path):
raise ValueError(f"Mesh path does not exist: {mesh_path}")
# Create background removal function if needed
remove_bg_fn = None
if remove_background:
remove_bg_fn = create_bg_remover(DEVICE_STR)
# Execute generation
num_views = 6 # Standard number of views
images, pos_images, normal_images, processed_ref_image = mvadapter_run_pipeline(
pipe=mvadapter_pipe,
mesh_path=mesh_path,
num_views=num_views,
text=prompt,
image=reference_image,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
seed=seed,
remove_bg_fn=remove_bg_fn,
reference_conditioning_scale=reference_conditioning_scale,
negative_prompt=negative_prompt,
lora_scale=lora_scale,
device=DEVICE_STR,
)
# Create image grid
grid_image = make_image_grid(images, rows=1)
# Convert PIL image to torch tensor using shared utils
grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
print(f"Generated multiview images: {grid_tensor.shape}")
return (grid_tensor,)