remove unnecessary scripts and update output to multiple images
This commit is contained in:
Denis Kochetov
@@ -1,23 +0,0 @@
|
||||
*.png
|
||||
*.jpg
|
||||
*.jpeg
|
||||
*.gif
|
||||
*.bmp
|
||||
*.tiff
|
||||
*.ico
|
||||
*.webp
|
||||
|
||||
#glb
|
||||
*.glb
|
||||
*.gltf
|
||||
*.obj
|
||||
*.mtl
|
||||
*.ply
|
||||
*.stl
|
||||
*.dae
|
||||
|
||||
#pycache
|
||||
__pycache__
|
||||
|
||||
#checkpoints
|
||||
checkpoints/
|
||||
@@ -1,201 +0,0 @@
|
||||
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.
|
||||
@@ -1,471 +0,0 @@
|
||||
# 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>
|
||||
|
||||
|
||||
|
||||
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"
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
> 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"
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
**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
|
||||
```
|
||||
|
||||
|
||||
|
||||
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
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
### 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."
|
||||
```
|
||||
|
||||
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
|
||||
|
||||
#### 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:
|
||||
|
||||
|
||||
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**
|
||||
|
||||
|
||||
|
||||
**Text to Multiview Generation with LoRA**
|
||||
|
||||
|
||||
|
||||
**Image to Multiview Generation**
|
||||
|
||||
|
||||
|
||||
|
||||
## 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}
|
||||
}
|
||||
```
|
||||
@@ -1,221 +0,0 @@
|
||||
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)
|
||||
@@ -1,561 +0,0 @@
|
||||
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)
|
||||
@@ -1,777 +0,0 @@
|
||||
# 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
|
||||
@@ -1,634 +0,0 @@
|
||||
# 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
|
||||
@@ -1,428 +0,0 @@
|
||||
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,
|
||||
)
|
||||
@@ -1,262 +0,0 @@
|
||||
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
|
||||
@@ -1,459 +0,0 @@
|
||||
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
|
||||
@@ -1,506 +0,0 @@
|
||||
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
|
||||
@@ -1,415 +0,0 @@
|
||||
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
|
||||
@@ -1,457 +0,0 @@
|
||||
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
|
||||
@@ -1,178 +0,0 @@
|
||||
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
|
||||
@@ -130,30 +130,30 @@ def prepare_pipeline(
|
||||
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 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
|
||||
# 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
|
||||
# return remove_bg
|
||||
|
||||
|
||||
def preprocess_image(image: Image.Image, height: int, width: int) -> Image.Image:
|
||||
@@ -251,10 +251,10 @@ def run_pipeline(
|
||||
|
||||
# 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":
|
||||
# if remove_bg_fn is not None:
|
||||
# reference_image = remove_bg_fn(reference_image)
|
||||
# reference_image = preprocess_image(reference_image, height, width)
|
||||
if reference_image.mode == "RGBA":
|
||||
reference_image = preprocess_image(reference_image, height, width)
|
||||
|
||||
pipe_kwargs = {}
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
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
|
||||
@@ -1,214 +0,0 @@
|
||||
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()
|
||||
@@ -1,195 +0,0 @@
|
||||
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()
|
||||
@@ -1,255 +0,0 @@
|
||||
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")
|
||||
@@ -1,267 +0,0 @@
|
||||
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")
|
||||
@@ -1,303 +0,0 @@
|
||||
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)
|
||||
@@ -1,313 +0,0 @@
|
||||
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")
|
||||
@@ -1,252 +0,0 @@
|
||||
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"
|
||||
)
|
||||
@@ -1,170 +0,0 @@
|
||||
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)
|
||||
@@ -1,201 +0,0 @@
|
||||
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)
|
||||
@@ -1,200 +0,0 @@
|
||||
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"
|
||||
)
|
||||
@@ -1,104 +0,0 @@
|
||||
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}")
|
||||
@@ -1,80 +0,0 @@
|
||||
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}")
|
||||
@@ -1,59 +0,0 @@
|
||||
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",
|
||||
},
|
||||
)
|
||||
@@ -1,138 +0,0 @@
|
||||
#!/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"
|
||||
@@ -108,7 +108,6 @@ 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,
|
||||
prepare_t2mv_pipeline as mvadapter_prepare_t2mv_pipeline,
|
||||
run_t2mv_pipeline as mvadapter_run_t2mv_pipeline
|
||||
)
|
||||
@@ -4886,10 +4885,6 @@ class MVAdapter_Image_To_MultiView:
|
||||
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,
|
||||
@@ -4902,19 +4897,22 @@ class MVAdapter_Image_To_MultiView:
|
||||
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_image = make_image_grid(images, rows=1)
|
||||
|
||||
grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
|
||||
# grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
|
||||
|
||||
print(f"Generated multiview images: {grid_tensor.shape}")
|
||||
return (grid_tensor,)
|
||||
# print(f"Generated multiview images: {grid_tensor.shape}")
|
||||
# return (grid_tensor,)
|
||||
# return images
|
||||
|
||||
return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
|
||||
return (return_images,)
|
||||
|
||||
|
||||
class Load_MVAdapter_T2MV_Pipeline:
|
||||
@@ -5033,11 +5031,13 @@ class MVAdapter_Text_To_MultiView:
|
||||
)
|
||||
|
||||
# Create image grid
|
||||
grid_image = make_image_grid(images, rows=1)
|
||||
# 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)
|
||||
# # Convert PIL image to torch tensor using shared utils
|
||||
# grid_tensor = pils_to_torch_imgs([grid_image], device=DEVICE_STR)
|
||||
|
||||
print(f"Generated T2MV multiview images: {grid_tensor.shape}")
|
||||
return (grid_tensor,)
|
||||
# print(f"Generated T2MV multiview images: {grid_tensor.shape}")
|
||||
# return (grid_tensor,)
|
||||
return_images = pils_to_torch_imgs(images, device=DEVICE_STR)
|
||||
return (return_images,)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user