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

Add Unique3D & Era3D workflows

Browse Source
This commit is contained in:
Mr. For Example
2024-06-23 11:48:40 -06:00
parent 340a03babe
commit f3cbd66020
56 changed files with 16009 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.vscode/settings.json
Vendored
+4 -1
View File
@@ -1,3 +1,6 @@
{
"cmake.sourceDirectory": "C:/Users/reall/Softwares/ComfyUI_windows_portable/ComfyUI/custom_nodes/ComfyUI-3D-Pack/diff-gaussian-rasterization"
"cmake.sourceDirectory": "C:/Users/reall/Softwares/ComfyUI_windows_portable/ComfyUI/custom_nodes/ComfyUI-3D-Pack/diff-gaussian-rasterization",
"python.analysis.extraPaths": [
"./gen_3d_modules"
]
}
README.md
+14 -1
View File
@@ -13,10 +13,23 @@ This is an extensive node suite that enables ComfyUI to process 3D inputs (Mesh
</span>
## Currently support:
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 02/06/2024**]
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 07/06/2024**]
- **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
- [tripoSR-layered-diffusion workflow](https://github.com/C0nsumption/Consume-ComfyUI-Workflows/tree/main/assets/tripo_sr/00) by [@Consumption](https://twitter.com/c0nsumption_)
- **Unique3D**: [AiuniAI/Unique3D](https://github.com/AiuniAI/Unique3D)
- Four stages pipeline:
1. Single image to 4 multi-view images with resulution: 256X256
2. Consistent Multi-view images Upscale to 512X512, super resolution to 2048X2048
3. Multi-view images to Normal maps with resulution: 512X512, super resolution to 2048X2048
4. Multi-view images & Normal maps to 3D mesh with texture
- To use the [pure Unique3D workflow](./_Example_Workflows/Unique3D/Unique3D_All_Stages.json), Download Models:
- [img2mvimg](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/img2mvimg) and put it into [./checkpoints/Wuvin/Unique3D/image2mvimage](./checkpoints/Wuvin/Unique3D/image2mvimage)
- [image2normal](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/image2normal) and put it into [./checkpoints/Wuvin/Unique3D/image2normal](./checkpoints/Wuvin/Unique3D/image2normal)
- [fine-tuned controlnet-tile](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/controlnet-tile) and put it into `Your ComfyUI root directory/ComfyUI/models/controlnet`
- [ip-adapter_sd15](https://huggingface.co/h94/IP-Adapter/blob/main/models/ip-adapter_sd15.safetensors) and put it into `Your ComfyUI root directory/ComfyUI/models/ipadapter`
- [RealESRGAN_x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth) and put it into `Your ComfyUI root directory/ComfyUI/models/upscale_models`
- **Era3D Diffusion Model**: [pengHTYX/Era3D](https://github.com/pengHTYX/Era3D)
- Single image to 6 multi-view images & normal maps with resulution: 512X512
- *Note: you need at least 16GB vram to run this model*
_Example_Workflows/Unique3D/0_MV_RGB.json
+488
View File
@@ -0,0 +1,488 @@
{
"last_node_id": 27,
"last_link_id": 40,
"nodes": [
{
"id": 17,
"type": "[Comfy3D] Load Unique3D Custom UNet",
"pos": [
534,
161
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"link": 21
}
],
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
22
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Load Unique3D Custom UNet"
},
"widgets_values": [
"image2mvimage"
]
},
{
"id": 3,
"type": "[Comfy3D] Set Diffusers Pipeline Scheduler",
"pos": [
892,
161
],
"size": {
"0": 412.3726501464844,
"1": 58
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"link": 22
}
],
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
28
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Set Diffusers Pipeline Scheduler"
},
"widgets_values": [
"EulerAncestralDiscreteScheduler"
]
},
{
"id": 23,
"type": "[Comfy3D] Unique3D Diffusion Model",
"pos": [
1374,
159
],
"size": {
"0": 380.4000244140625,
"1": 222
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "unique3d_pipe",
"type": "DIFFUSERS_PIPE",
"link": 28
},
{
"name": "reference_image",
"type": "IMAGE",
"link": 38
}
],
"outputs": [
{
"name": "multiviews",
"type": "IMAGE",
"links": [
30
],
"shape": 3,
"slot_index": 0
},
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Unique3D Diffusion Model"
},
"widgets_values": [
1145,
"fixed",
1.5,
30,
256,
4,
true
]
},
{
"id": 1,
"type": "[Comfy3D] Load Diffusers Pipeline",
"pos": [
56,
161
],
"size": {
"0": 430.03009033203125,
"1": 130
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
21
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
},
"widgets_values": [
"Unique3DImage2MVCustomPipeline",
"Wuvin/Unique3D/image2mvimage",
"",
false
]
},
{
"id": 6,
"type": "LoadImage",
"pos": [
375,
343
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
35
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [
32
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"PriestCat.png",
"image"
]
},
{
"id": 25,
"type": "InvertMask",
"pos": [
723,
419
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 32
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
36
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 22,
"type": "SaveImage",
"pos": [
1830,
159
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 30
}
],
"properties": {},
"widgets_values": [
"Unique3D\\RGB\\rgb"
]
},
{
"id": 26,
"type": "[Comfy3D] Image Add Pure Color Background",
"pos": [
975,
338
],
"size": {
"0": 315,
"1": 126
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 35
},
{
"name": "masks",
"type": "MASK",
"link": 36
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
38,
39
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Image Add Pure Color Background"
},
"widgets_values": [
255,
255,
255
]
},
{
"id": 27,
"type": "ImageScale",
"pos": [
1361,
548
],
"size": {
"0": 315,
"1": 130
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 39
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
40
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageScale"
},
"widgets_values": [
"nearest-exact",
2048,
2048,
"disabled"
]
},
{
"id": 24,
"type": "SaveImage",
"pos": [
1730,
550
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 40
}
],
"properties": {},
"widgets_values": [
"Unique3D\\Normal_Inputs\\refined_sr_rgb_0"
]
}
],
"links": [
[
21,
1,
0,
17,
0,
"DIFFUSERS_PIPE"
],
[
22,
17,
0,
3,
0,
"DIFFUSERS_PIPE"
],
[
28,
3,
0,
23,
0,
"DIFFUSERS_PIPE"
],
[
30,
23,
0,
22,
0,
"IMAGE"
],
[
32,
6,
1,
25,
0,
"MASK"
],
[
35,
6,
0,
26,
0,
"IMAGE"
],
[
36,
25,
0,
26,
1,
"MASK"
],
[
38,
26,
0,
23,
1,
"IMAGE"
],
[
39,
26,
0,
27,
0,
"IMAGE"
],
[
40,
27,
0,
24,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.8264462809917354,
"offset": [
-813.4684892955395,
-48.223766450597736
]
}
},
"version": 0.4
}
_Example_Workflows/Unique3D/1_MV_RGB_Upscale.json
+1476
View File
File diff suppressed because it is too large Load Diff
_Example_Workflows/Unique3D/1_MV_RGB_Upscale_New.json
+1403
View File
File diff suppressed because it is too large Load Diff
_Example_Workflows/Unique3D/2_MV_Normal.json
+560
View File
@@ -0,0 +1,560 @@
{
"last_node_id": 30,
"last_link_id": 46,
"nodes": [
{
"id": 17,
"type": "[Comfy3D] Load Unique3D Custom UNet",
"pos": [
534,
161
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"link": 21
}
],
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
22
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Load Unique3D Custom UNet"
},
"widgets_values": [
"image2normal"
]
},
{
"id": 1,
"type": "[Comfy3D] Load Diffusers Pipeline",
"pos": [
56,
161
],
"size": {
"0": 430.03009033203125,
"1": 130
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
21
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
},
"widgets_values": [
"Unique3DImageCustomPipeline",
"Wuvin/Unique3D/image2normal",
"",
false
]
},
{
"id": 3,
"type": "[Comfy3D] Set Diffusers Pipeline Scheduler",
"pos": [
892,
161
],
"size": {
"0": 412.3726501464844,
"1": 58
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"link": 22
}
],
"outputs": [
{
"name": "pipe",
"type": "DIFFUSERS_PIPE",
"links": [
31
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Set Diffusers Pipeline Scheduler"
},
"widgets_values": [
"EulerAncestralDiscreteScheduler"
]
},
{
"id": 20,
"type": "VHS_LoadImagesPath",
"pos": [
736,
322
],
"size": [
226.8000030517578,
194
],
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
32
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [],
"shape": 3,
"slot_index": 1
},
{
"name": "INT",
"type": "INT",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "VHS_LoadImagesPath"
},
"widgets_values": {
"directory": "C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\Unique3D\\Normal_Inputs",
"image_load_cap": 0,
"skip_first_images": 0,
"select_every_nth": 1,
"choose folder to upload": "image",
"videopreview": {
"hidden": false,
"paused": false,
"params": {
"frame_load_cap": 0,
"skip_first_images": 0,
"filename": "C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\Unique3D\\Normal_Inputs",
"type": "path",
"format": "folder",
"select_every_nth": 1
}
}
}
},
{
"id": 23,
"type": "ImageScale",
"pos": [
1022,
322
],
"size": {
"0": 315,
"1": 130
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 32
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
33
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageScale"
},
"widgets_values": [
"lanczos",
512,
512,
"disabled"
]
},
{
"id": 21,
"type": "[Comfy3D] Unique3D Diffusion Model",
"pos": [
1397,
254
],
"size": {
"0": 380.4000244140625,
"1": 222
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "unique3d_pipe",
"type": "DIFFUSERS_PIPE",
"link": 31
},
{
"name": "reference_image",
"type": "IMAGE",
"link": 33
}
],
"outputs": [
{
"name": "multiviews",
"type": "IMAGE",
"links": [
35
],
"shape": 3,
"slot_index": 0
},
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Unique3D Diffusion Model"
},
"widgets_values": [
1145,
"fixed",
1.5,
30,
512,
4,
false
]
},
{
"id": 24,
"type": "UpscaleModelLoader",
"pos": [
1433,
105
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "UPSCALE_MODEL",
"type": "UPSCALE_MODEL",
"links": [
34
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "UpscaleModelLoader"
},
"widgets_values": [
"realesrganX4plus_v1.pt"
]
},
{
"id": 26,
"type": "[Comfy3D] Rotate Normal Maps Horizontally",
"pos": [
2427,
230
],
"size": {
"0": 315,
"1": 78
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "normal_maps",
"type": "IMAGE",
"link": 44
},
{
"name": "normal_masks",
"type": "MASK",
"link": 46
}
],
"outputs": [
{
"name": "normal_maps",
"type": "IMAGE",
"links": [
38
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Rotate Normal Maps Horizontally"
},
"widgets_values": [
true
]
},
{
"id": 19,
"type": "SaveImage",
"pos": [
2770,
230
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 38
}
],
"properties": {},
"widgets_values": [
"Unique3D/Normals/normal"
]
},
{
"id": 25,
"type": "ImageUpscaleWithModel",
"pos": [
1821,
232
],
"size": {
"0": 241.79998779296875,
"1": 46
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "upscale_model",
"type": "UPSCALE_MODEL",
"link": 34
},
{
"name": "image",
"type": "IMAGE",
"link": 35
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
44,
45
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "ImageUpscaleWithModel"
}
},
{
"id": 30,
"type": "[Comfy3D] Get Masks From Normal Maps",
"pos": [
2093,
299
],
"size": {
"0": 304.79998779296875,
"1": 26
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "normal_maps",
"type": "IMAGE",
"link": 45
}
],
"outputs": [
{
"name": "normal_masks",
"type": "MASK",
"links": [
46
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Get Masks From Normal Maps"
}
}
],
"links": [
[
21,
1,
0,
17,
0,
"DIFFUSERS_PIPE"
],
[
22,
17,
0,
3,
0,
"DIFFUSERS_PIPE"
],
[
31,
3,
0,
21,
0,
"DIFFUSERS_PIPE"
],
[
32,
20,
0,
23,
0,
"IMAGE"
],
[
33,
23,
0,
21,
1,
"IMAGE"
],
[
34,
24,
0,
25,
0,
"UPSCALE_MODEL"
],
[
35,
21,
0,
25,
1,
"IMAGE"
],
[
38,
26,
0,
19,
0,
"IMAGE"
],
[
44,
25,
0,
26,
0,
"IMAGE"
],
[
45,
25,
0,
30,
0,
"IMAGE"
],
[
46,
30,
0,
26,
1,
"MASK"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.8264462809917356,
"offset": [
-1532.094207035033,
229.55273653768023
]
}
},
"version": 0.4
}
_Example_Workflows/Unique3D/3_Mesh_Reconstruction.json
+1633
View File
File diff suppressed because it is too large Load Diff
_Example_Workflows/Unique3D/Unique3D_All_Stages.json
+3993
View File
File diff suppressed because it is too large Load Diff
_Example_Workflows/Unique3D/Unique3D_with_Era3D_All.json
+1479
View File
File diff suppressed because it is too large Load Diff
__init__.py
+16
View File
@@ -2,15 +2,31 @@ import os
import sys
import folder_paths as comfy_paths
from pyhocon import ConfigFactory
import logging
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
MODULE_PATH = os.path.join(ROOT_PATH, "gen_3d_modules")
sys.path.append(ROOT_PATH)
sys.path.append(MODULE_PATH)
import shutil
import __main__
import importlib
import inspect
from .webserver.server import server, set_web_conf
from .utils import setup_logger
# Common formatter for simplicity, adjust as needed
common_formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')
# Setup logging for transformers
setup_logger('transformers', logging.WARNING, [logging.WARNING], [logging.ERROR, logging.CRITICAL], common_formatter)
# Setup logging for diffusers
setup_logger('diffusers_logging', logging.INFO, [logging.INFO, logging.WARNING], [logging.ERROR, logging.CRITICAL], common_formatter)
# Redirect warnings to the logging system
logging.captureWarnings(True)
conf_path = os.path.join(ROOT_PATH, "configs/system.conf")
# Configuration
checkpoints/Wuvin/Unique3D/controlnet-tile/config.json
+52
View File
@@ -0,0 +1,52 @@
{
"_class_name": "ControlNetModel",
"_diffusers_version": "0.27.2",
"_name_or_path": "lllyasviel/control_v11f1e_sd15_tile",
"act_fn": "silu",
"addition_embed_type": null,
"addition_embed_type_num_heads": 64,
"addition_time_embed_dim": null,
"attention_head_dim": 8,
"block_out_channels": [
320,
640,
1280,
1280
],
"class_embed_type": null,
"conditioning_channels": 3,
"conditioning_embedding_out_channels": [
16,
32,
96,
256
],
"controlnet_conditioning_channel_order": "rgb",
"cross_attention_dim": 768,
"down_block_types": [
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D"
],
"downsample_padding": 1,
"encoder_hid_dim": null,
"encoder_hid_dim_type": null,
"flip_sin_to_cos": true,
"freq_shift": 0,
"global_pool_conditions": false,
"in_channels": 4,
"layers_per_block": 2,
"mid_block_scale_factor": 1,
"mid_block_type": "UNetMidBlock2DCrossAttn",
"norm_eps": 1e-05,
"norm_num_groups": 32,
"num_attention_heads": null,
"num_class_embeds": null,
"only_cross_attention": false,
"projection_class_embeddings_input_dim": null,
"resnet_time_scale_shift": "default",
"transformer_layers_per_block": 1,
"upcast_attention": false,
"use_linear_projection": false
}
checkpoints/Wuvin/Unique3D/image2mvimage/feature_extractor/preprocessor_config.json
+44
View File
@@ -0,0 +1,44 @@
{
"_valid_processor_keys": [
"images",
"do_resize",
"size",
"resample",
"do_center_crop",
"crop_size",
"do_rescale",
"rescale_factor",
"do_normalize",
"image_mean",
"image_std",
"do_convert_rgb",
"return_tensors",
"data_format",
"input_data_format"
],
"crop_size": {
"height": 224,
"width": 224
},
"do_center_crop": true,
"do_convert_rgb": true,
"do_normalize": true,
"do_rescale": true,
"do_resize": true,
"image_mean": [
0.48145466,
0.4578275,
0.40821073
],
"image_processor_type": "CLIPImageProcessor",
"image_std": [
0.26862954,
0.26130258,
0.27577711
],
"resample": 3,
"rescale_factor": 0.00392156862745098,
"size": {
"shortest_edge": 224
}
}
checkpoints/Wuvin/Unique3D/image2mvimage/image_encoder/config.json
+23
View File
@@ -0,0 +1,23 @@
{
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"architectures": [
"CLIPVisionModelWithProjection"
],
"attention_dropout": 0.0,
"dropout": 0.0,
"hidden_act": "quick_gelu",
"hidden_size": 1024,
"image_size": 224,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 4096,
"layer_norm_eps": 1e-05,
"model_type": "clip_vision_model",
"num_attention_heads": 16,
"num_channels": 3,
"num_hidden_layers": 24,
"patch_size": 14,
"projection_dim": 768,
"torch_dtype": "float32",
"transformers_version": "4.39.3"
}
checkpoints/Wuvin/Unique3D/image2mvimage/model_index.json
+31
View File
@@ -0,0 +1,31 @@
{
"_class_name": "StableDiffusionImage2MVCustomPipeline",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"condition_offset": true,
"feature_extractor": [
"transformers",
"CLIPImageProcessor"
],
"image_encoder": [
"transformers",
"CLIPVisionModelWithProjection"
],
"requires_safety_checker": true,
"safety_checker": [
null,
null
],
"scheduler": [
"diffusers",
"DDIMScheduler"
],
"unet": [
"diffusers",
"UNet2DConditionModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
checkpoints/Wuvin/Unique3D/image2mvimage/scheduler/scheduler_config.json
+20
View File
@@ -0,0 +1,20 @@
{
"_class_name": "DDIMScheduler",
"_diffusers_version": "0.27.2",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": false,
"clip_sample_range": 1.0,
"dynamic_thresholding_ratio": 0.995,
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"rescale_betas_zero_snr": false,
"sample_max_value": 1.0,
"set_alpha_to_one": false,
"skip_prk_steps": true,
"steps_offset": 1,
"thresholding": false,
"timestep_spacing": "leading",
"trained_betas": null
}
checkpoints/Wuvin/Unique3D/image2mvimage/unet/config.json
+68
View File
@@ -0,0 +1,68 @@
{
"_class_name": "UnifieldWrappedUNet",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"act_fn": "silu",
"addition_embed_type": null,
"addition_embed_type_num_heads": 64,
"addition_time_embed_dim": null,
"attention_head_dim": 8,
"attention_type": "default",
"block_out_channels": [
320,
640,
1280,
1280
],
"center_input_sample": false,
"class_embed_type": null,
"class_embeddings_concat": false,
"conv_in_kernel": 3,
"conv_out_kernel": 3,
"cross_attention_dim": 768,
"cross_attention_norm": null,
"down_block_types": [
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D"
],
"downsample_padding": 1,
"dropout": 0.0,
"dual_cross_attention": false,
"encoder_hid_dim": null,
"encoder_hid_dim_type": null,
"flip_sin_to_cos": true,
"freq_shift": 0,
"in_channels": 8,
"layers_per_block": 2,
"mid_block_only_cross_attention": null,
"mid_block_scale_factor": 1,
"mid_block_type": "UNetMidBlock2DCrossAttn",
"norm_eps": 1e-05,
"norm_num_groups": 32,
"num_attention_heads": null,
"num_class_embeds": 8,
"only_cross_attention": false,
"out_channels": 4,
"projection_class_embeddings_input_dim": null,
"resnet_out_scale_factor": 1.0,
"resnet_skip_time_act": false,
"resnet_time_scale_shift": "default",
"reverse_transformer_layers_per_block": null,
"sample_size": 64,
"time_cond_proj_dim": null,
"time_embedding_act_fn": null,
"time_embedding_dim": null,
"time_embedding_type": "positional",
"timestep_post_act": null,
"transformer_layers_per_block": 1,
"up_block_types": [
"UpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D"
],
"upcast_attention": false,
"use_linear_projection": false
}
checkpoints/Wuvin/Unique3D/image2mvimage/vae/config.json
+34
View File
@@ -0,0 +1,34 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"latents_mean": null,
"latents_std": null,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 256,
"scaling_factor": 0.18215,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
checkpoints/Wuvin/Unique3D/image2normal/feature_extractor/preprocessor_config.json
+44
View File
@@ -0,0 +1,44 @@
{
"_valid_processor_keys": [
"images",
"do_resize",
"size",
"resample",
"do_center_crop",
"crop_size",
"do_rescale",
"rescale_factor",
"do_normalize",
"image_mean",
"image_std",
"do_convert_rgb",
"return_tensors",
"data_format",
"input_data_format"
],
"crop_size": {
"height": 224,
"width": 224
},
"do_center_crop": true,
"do_convert_rgb": true,
"do_normalize": true,
"do_rescale": true,
"do_resize": true,
"image_mean": [
0.48145466,
0.4578275,
0.40821073
],
"image_processor_type": "CLIPImageProcessor",
"image_std": [
0.26862954,
0.26130258,
0.27577711
],
"resample": 3,
"rescale_factor": 0.00392156862745098,
"size": {
"shortest_edge": 224
}
}
checkpoints/Wuvin/Unique3D/image2normal/image_encoder/config.json
+23
View File
@@ -0,0 +1,23 @@
{
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"architectures": [
"CLIPVisionModelWithProjection"
],
"attention_dropout": 0.0,
"dropout": 0.0,
"hidden_act": "quick_gelu",
"hidden_size": 1024,
"image_size": 224,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 4096,
"layer_norm_eps": 1e-05,
"model_type": "clip_vision_model",
"num_attention_heads": 16,
"num_channels": 3,
"num_hidden_layers": 24,
"patch_size": 14,
"projection_dim": 768,
"torch_dtype": "bfloat16",
"transformers_version": "4.39.3"
}
checkpoints/Wuvin/Unique3D/image2normal/model_index.json
+31
View File
@@ -0,0 +1,31 @@
{
"_class_name": "StableDiffusionImageCustomPipeline",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"feature_extractor": [
"transformers",
"CLIPImageProcessor"
],
"image_encoder": [
"transformers",
"CLIPVisionModelWithProjection"
],
"noisy_cond_latents": false,
"requires_safety_checker": true,
"safety_checker": [
null,
null
],
"scheduler": [
"diffusers",
"EulerAncestralDiscreteScheduler"
],
"unet": [
"diffusers",
"UNet2DConditionModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
checkpoints/Wuvin/Unique3D/image2normal/scheduler/scheduler_config.json
+16
View File
@@ -0,0 +1,16 @@
{
"_class_name": "EulerAncestralDiscreteScheduler",
"_diffusers_version": "0.27.2",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": false,
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"rescale_betas_zero_snr": false,
"set_alpha_to_one": false,
"skip_prk_steps": true,
"steps_offset": 1,
"timestep_spacing": "linspace",
"trained_betas": null
}
checkpoints/Wuvin/Unique3D/image2normal/unet/config.json
+68
View File
@@ -0,0 +1,68 @@
{
"_class_name": "UnifieldWrappedUNet",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"act_fn": "silu",
"addition_embed_type": null,
"addition_embed_type_num_heads": 64,
"addition_time_embed_dim": null,
"attention_head_dim": 8,
"attention_type": "default",
"block_out_channels": [
320,
640,
1280,
1280
],
"center_input_sample": false,
"class_embed_type": null,
"class_embeddings_concat": false,
"conv_in_kernel": 3,
"conv_out_kernel": 3,
"cross_attention_dim": 768,
"cross_attention_norm": null,
"down_block_types": [
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D"
],
"downsample_padding": 1,
"dropout": 0.0,
"dual_cross_attention": false,
"encoder_hid_dim": null,
"encoder_hid_dim_type": null,
"flip_sin_to_cos": true,
"freq_shift": 0,
"in_channels": 4,
"layers_per_block": 2,
"mid_block_only_cross_attention": null,
"mid_block_scale_factor": 1,
"mid_block_type": "UNetMidBlock2DCrossAttn",
"norm_eps": 1e-05,
"norm_num_groups": 32,
"num_attention_heads": null,
"num_class_embeds": null,
"only_cross_attention": false,
"out_channels": 4,
"projection_class_embeddings_input_dim": null,
"resnet_out_scale_factor": 1.0,
"resnet_skip_time_act": false,
"resnet_time_scale_shift": "default",
"reverse_transformer_layers_per_block": null,
"sample_size": 64,
"time_cond_proj_dim": null,
"time_embedding_act_fn": null,
"time_embedding_dim": null,
"time_embedding_type": "positional",
"timestep_post_act": null,
"transformer_layers_per_block": 1,
"up_block_types": [
"UpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D"
],
"upcast_attention": false,
"use_linear_projection": false
}
checkpoints/Wuvin/Unique3D/image2normal/vae/config.json
+34
View File
@@ -0,0 +1,34 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.27.2",
"_name_or_path": "lambdalabs/sd-image-variations-diffusers",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"latents_mean": null,
"latents_std": null,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 256,
"scaling_factor": 0.18215,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
checkpoints/Wuvin/Unique3D/v1-inference.yaml
+70
View File
@@ -0,0 +1,70 @@
model:
base_learning_rate: 1.0e-04
target: ldm.models.diffusion.ddpm.LatentDiffusion
params:
linear_start: 0.00085
linear_end: 0.0120
num_timesteps_cond: 1
log_every_t: 200
timesteps: 1000
first_stage_key: "jpg"
cond_stage_key: "txt"
image_size: 64
channels: 4
cond_stage_trainable: false # Note: different from the one we trained before
conditioning_key: crossattn
monitor: val/loss_simple_ema
scale_factor: 0.18215
use_ema: False
scheduler_config: # 10000 warmup steps
target: ldm.lr_scheduler.LambdaLinearScheduler
params:
warm_up_steps: [ 10000 ]
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
f_start: [ 1.e-6 ]
f_max: [ 1. ]
f_min: [ 1. ]
unet_config:
target: ldm.modules.diffusionmodules.openaimodel.UNetModel
params:
image_size: 32 # unused
in_channels: 4
out_channels: 4
model_channels: 320
attention_resolutions: [ 4, 2, 1 ]
num_res_blocks: 2
channel_mult: [ 1, 2, 4, 4 ]
num_heads: 8
use_spatial_transformer: True
transformer_depth: 1
context_dim: 768
use_checkpoint: True
legacy: False
first_stage_config:
target: ldm.models.autoencoder.AutoencoderKL
params:
embed_dim: 4
monitor: val/rec_loss
ddconfig:
double_z: true
z_channels: 4
resolution: 256
in_channels: 3
out_ch: 3
ch: 128
ch_mult:
- 1
- 2
- 4
- 4
num_res_blocks: 2
attn_resolutions: []
dropout: 0.0
lossconfig:
target: torch.nn.Identity
cond_stage_config:
target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
configs/Unique3D_configs/image2mvimage.yaml
+43
View File
@@ -0,0 +1,43 @@
init_config:
# enable controls
enable_cross_attn_lora: False
enable_cross_attn_ip: False
enable_self_attn_lora: False
enable_self_attn_ref: False
enable_multiview_attn: True
# for cross attention
init_cross_attn_lora: False
init_cross_attn_ip: False
cross_attn_lora_rank: 256 # 0 for not enabled
cross_attn_lora_only_kv: False
ipadapter_pretrained_name: "h94/IP-Adapter"
ipadapter_subfolder_name: "models"
ipadapter_weight_name: "ip-adapter_sd15.safetensors"
ipadapter_effect_on: "all" # all, first
# for self attention
init_self_attn_lora: False
self_attn_lora_rank: 256
self_attn_lora_only_kv: False
# for self attention ref
init_self_attn_ref: False
self_attn_ref_position: "attn1"
self_attn_ref_other_model_name: "lambdalabs/sd-image-variations-diffusers"
self_attn_ref_pixel_wise_crosspond: False
self_attn_ref_effect_on: "all"
# for multiview attention
init_multiview_attn: True
multiview_attn_position: "attn1"
use_mv_joint_attn: True
num_modalities: 1
# for unet
init_unet_path: ""
cat_condition: True # cat condition to input
# for cls embedding
init_num_cls_label: 8 # for initialize
cls_labels: [0, 1, 2, 3] # for current task
configs/Unique3D_configs/image2normal.yaml
+44
View File
@@ -0,0 +1,44 @@
pretrained_model_name_or_path: "lambdalabs/sd-image-variations-diffusers"
init_config:
# enable controls
enable_cross_attn_lora: False
enable_cross_attn_ip: False
enable_self_attn_lora: False
enable_self_attn_ref: True
enable_multiview_attn: False
# for cross attention
init_cross_attn_lora: False
init_cross_attn_ip: False
cross_attn_lora_rank: 512 # 0 for not enabled
cross_attn_lora_only_kv: False
ipadapter_pretrained_name: "h94/IP-Adapter"
ipadapter_subfolder_name: "models"
ipadapter_weight_name: "ip-adapter_sd15.safetensors"
ipadapter_effect_on: "all" # all, first
# for self attention
init_self_attn_lora: False
self_attn_lora_rank: 512
self_attn_lora_only_kv: False
# for self attention ref
init_self_attn_ref: True
self_attn_ref_position: "attn1"
self_attn_ref_other_model_name: "lambdalabs/sd-image-variations-diffusers"
self_attn_ref_pixel_wise_crosspond: True
self_attn_ref_effect_on: "all"
# for multiview attention
init_multiview_attn: False
multiview_attn_position: "attn1"
num_modalities: 1
# for unet
init_unet_path: "${pretrained_model_name_or_path}"
init_num_cls_label: 0 # for initialize
cls_labels: [] # for current task
gen_3d_modules/Unique3D/__init__.py
View File
gen_3d_modules/Unique3D/custum_3d_diffusion/__init__.py
View File
gen_3d_modules/Unique3D/custum_3d_diffusion/custum_modules/attention_processors.py
+385
View File
@@ -0,0 +1,385 @@
from typing import Any, Dict, Optional
import torch
from diffusers.models.attention_processor import Attention
def construct_pix2pix_attention(hidden_states_dim, norm_type="none"):
if norm_type == "layernorm":
norm = torch.nn.LayerNorm(hidden_states_dim)
else:
norm = torch.nn.Identity()
attention = Attention(
query_dim=hidden_states_dim,
heads=8,
dim_head=hidden_states_dim // 8,
bias=True,
)
# NOTE: xformers 0.22 does not support batchsize >= 4096
attention.xformers_not_supported = True # hacky solution
return norm, attention
class ExtraAttnProc(torch.nn.Module):
def __init__(
self,
chained_proc,
enabled=False,
name=None,
mode='extract',
with_proj_in=False,
proj_in_dim=768,
target_dim=None,
pixel_wise_crosspond=False,
norm_type="none", # none or layernorm
crosspond_effect_on="all", # all or first
crosspond_chain_pos="parralle", # before or parralle or after
simple_3d=False,
views=4,
) -> None:
super().__init__()
self.enabled = enabled
self.chained_proc = chained_proc
self.name = name
self.mode = mode
self.with_proj_in=with_proj_in
self.proj_in_dim = proj_in_dim
self.target_dim = target_dim or proj_in_dim
self.hidden_states_dim = self.target_dim
self.pixel_wise_crosspond = pixel_wise_crosspond
self.crosspond_effect_on = crosspond_effect_on
self.crosspond_chain_pos = crosspond_chain_pos
self.views = views
self.simple_3d = simple_3d
if self.with_proj_in and self.enabled:
self.in_linear = torch.nn.Linear(self.proj_in_dim, self.target_dim, bias=False)
if self.target_dim == self.proj_in_dim:
self.in_linear.weight.data = torch.eye(proj_in_dim)
else:
self.in_linear = None
if self.pixel_wise_crosspond and self.enabled:
self.crosspond_norm, self.crosspond_attention = construct_pix2pix_attention(self.hidden_states_dim, norm_type=norm_type)
def do_crosspond_attention(self, hidden_states: torch.FloatTensor, other_states: torch.FloatTensor):
hidden_states = self.crosspond_norm(hidden_states)
batch, L, D = hidden_states.shape
assert hidden_states.shape == other_states.shape, f"got {hidden_states.shape} and {other_states.shape}"
# to -> batch * L, 1, D
hidden_states = hidden_states.reshape(batch * L, 1, D)
other_states = other_states.reshape(batch * L, 1, D)
hidden_states_catted = other_states
hidden_states = self.crosspond_attention(
hidden_states,
encoder_hidden_states=hidden_states_catted,
)
return hidden_states.reshape(batch, L, D)
def __call__(
self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None,
ref_dict: dict = None, mode=None, **kwargs
) -> Any:
if not self.enabled:
return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
assert ref_dict is not None
if (mode or self.mode) == 'extract':
ref_dict[self.name] = hidden_states
hidden_states1 = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
if self.pixel_wise_crosspond and self.crosspond_chain_pos == "after":
ref_dict[self.name] = hidden_states1
return hidden_states1
elif (mode or self.mode) == 'inject':
ref_state = ref_dict.pop(self.name)
if self.with_proj_in:
ref_state = self.in_linear(ref_state)
B, L, D = ref_state.shape
if hidden_states.shape[0] == B:
modalities = 1
views = 1
else:
modalities = hidden_states.shape[0] // B // self.views
views = self.views
if self.pixel_wise_crosspond:
if self.crosspond_effect_on == "all":
ref_state = ref_state[:, None].expand(-1, modalities * views, -1, -1).reshape(-1, *ref_state.shape[-2:])
if self.crosspond_chain_pos == "before":
hidden_states = hidden_states + self.do_crosspond_attention(hidden_states, ref_state)
hidden_states1 = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
if self.crosspond_chain_pos == "parralle":
hidden_states1 = hidden_states1 + self.do_crosspond_attention(hidden_states, ref_state)
if self.crosspond_chain_pos == "after":
hidden_states1 = hidden_states1 + self.do_crosspond_attention(hidden_states1, ref_state)
return hidden_states1
else:
assert self.crosspond_effect_on == "first"
# hidden_states [B * modalities * views, L, D]
# ref_state [B, L, D]
ref_state = ref_state[:, None].expand(-1, modalities, -1, -1).reshape(-1, ref_state.shape[-2], ref_state.shape[-1]) # [B * modalities, L, D]
def do_paritial_crosspond(hidden_states, ref_state):
first_view_hidden_states = hidden_states.view(-1, views, hidden_states.shape[1], hidden_states.shape[2])[:, 0] # [B * modalities, L, D]
hidden_states2 = self.do_crosspond_attention(first_view_hidden_states, ref_state) # [B * modalities, L, D]
hidden_states2_padded = torch.zeros_like(hidden_states).reshape(-1, views, hidden_states.shape[1], hidden_states.shape[2])
hidden_states2_padded[:, 0] = hidden_states2
hidden_states2_padded = hidden_states2_padded.reshape(-1, hidden_states.shape[1], hidden_states.shape[2])
return hidden_states2_padded
if self.crosspond_chain_pos == "before":
hidden_states = hidden_states + do_paritial_crosspond(hidden_states, ref_state)
hidden_states1 = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs) # [B * modalities * views, L, D]
if self.crosspond_chain_pos == "parralle":
hidden_states1 = hidden_states1 + do_paritial_crosspond(hidden_states, ref_state)
if self.crosspond_chain_pos == "after":
hidden_states1 = hidden_states1 + do_paritial_crosspond(hidden_states1, ref_state)
return hidden_states1
elif self.simple_3d:
B, L, C = encoder_hidden_states.shape
mv = self.views
encoder_hidden_states = encoder_hidden_states.reshape(B // mv, mv, L, C)
ref_state = ref_state[:, None]
encoder_hidden_states = torch.cat([encoder_hidden_states, ref_state], dim=1)
encoder_hidden_states = encoder_hidden_states.reshape(B // mv, 1, (mv+1) * L, C)
encoder_hidden_states = encoder_hidden_states.repeat(1, mv, 1, 1).reshape(-1, (mv+1) * L, C)
return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
else:
ref_state = ref_state[:, None].expand(-1, modalities * views, -1, -1).reshape(-1, ref_state.shape[-2], ref_state.shape[-1])
encoder_hidden_states = torch.cat([encoder_hidden_states, ref_state], dim=1)
return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
else:
raise NotImplementedError("mode or self.mode is required to be 'extract' or 'inject'")
def add_extra_processor(model: torch.nn.Module, enable_filter=lambda x:True, **kwargs):
return_dict = torch.nn.ModuleDict()
proj_in_dim = kwargs.get('proj_in_dim', False)
kwargs.pop('proj_in_dim', None)
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
if "ref_unet" not in (sub_name + name):
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
new_processor = ExtraAttnProc(
chained_proc=module.get_processor(),
enabled=enable_filter(f"{name}.processor"),
name=f"{name}.processor",
proj_in_dim=proj_in_dim if proj_in_dim else module.cross_attention_dim,
target_dim=module.cross_attention_dim,
**kwargs
)
module.set_processor(new_processor)
return_dict[f"{name}.processor".replace(".", "__")] = new_processor
for name, module in model.named_children():
recursive_add_processors(name, module)
return return_dict
def switch_extra_processor(model, enable_filter=lambda x:True):
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, ExtraAttnProc):
module.enabled = enable_filter(name)
for name, module in model.named_children():
recursive_add_processors(name, module)
class multiviewAttnProc(torch.nn.Module):
def __init__(
self,
chained_proc,
enabled=False,
name=None,
hidden_states_dim=None,
chain_pos="parralle", # before or parralle or after
num_modalities=1,
views=4,
base_img_size=64,
) -> None:
super().__init__()
self.enabled = enabled
self.chained_proc = chained_proc
self.name = name
self.hidden_states_dim = hidden_states_dim
self.num_modalities = num_modalities
self.views = views
self.base_img_size = base_img_size
self.chain_pos = chain_pos
self.diff_joint_attn = True
def __call__(
self,
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.FloatTensor] = None,
**kwargs
) -> torch.Tensor:
if not self.enabled:
return self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
B, L, C = hidden_states.shape
mv = self.views
hidden_states = hidden_states.reshape(B // mv, mv, L, C).reshape(-1, mv * L, C)
hidden_states = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
return hidden_states.reshape(B // mv, mv, L, C).reshape(-1, L, C)
def add_multiview_processor(model: torch.nn.Module, enable_filter=lambda x:True, **kwargs):
return_dict = torch.nn.ModuleDict()
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
if "ref_unet" not in (sub_name + name):
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
new_processor = multiviewAttnProc(
chained_proc=module.get_processor(),
enabled=enable_filter(f"{name}.processor"),
name=f"{name}.processor",
hidden_states_dim=module.inner_dim,
**kwargs
)
module.set_processor(new_processor)
return_dict[f"{name}.processor".replace(".", "__")] = new_processor
for name, module in model.named_children():
recursive_add_processors(name, module)
return return_dict
def switch_multiview_processor(model, enable_filter=lambda x:True):
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
processor = module.get_processor()
if isinstance(processor, multiviewAttnProc):
processor.enabled = enable_filter(f"{name}.processor")
for name, module in model.named_children():
recursive_add_processors(name, module)
class NNModuleWrapper(torch.nn.Module):
def __init__(self, module):
super().__init__()
self.module = module
def forward(self, *args, **kwargs):
return self.module(*args, **kwargs)
def __getattr__(self, name: str):
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self.module, name)
class AttnProcessorSwitch(torch.nn.Module):
def __init__(
self,
proc_dict: dict,
enabled_proc="default",
name=None,
switch_name="default_switch",
):
super().__init__()
self.proc_dict = torch.nn.ModuleDict({k: (v if isinstance(v, torch.nn.Module) else NNModuleWrapper(v)) for k, v in proc_dict.items()})
self.enabled_proc = enabled_proc
self.name = name
self.switch_name = switch_name
self.choose_module(enabled_proc)
def choose_module(self, enabled_proc):
self.enabled_proc = enabled_proc
assert enabled_proc in self.proc_dict.keys()
def __call__(
self,
*args,
**kwargs
) -> torch.FloatTensor:
used_proc = self.proc_dict[self.enabled_proc]
return used_proc(*args, **kwargs)
def add_switch(model: torch.nn.Module, module_filter=lambda x:True, switch_dict_fn=lambda x: {"default": x}, switch_name="default_switch", enabled_proc="default"):
return_dict = torch.nn.ModuleDict()
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
if "ref_unet" not in (sub_name + name):
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
processor = module.get_processor()
if module_filter(processor):
proc_dict = switch_dict_fn(processor)
new_processor = AttnProcessorSwitch(
proc_dict=proc_dict,
enabled_proc=enabled_proc,
name=f"{name}.processor",
switch_name=switch_name,
)
module.set_processor(new_processor)
return_dict[f"{name}.processor".replace(".", "__")] = new_processor
for name, module in model.named_children():
recursive_add_processors(name, module)
return return_dict
def change_switch(model: torch.nn.Module, switch_name="default_switch", enabled_proc="default"):
def recursive_change_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
recursive_change_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
processor = module.get_processor()
if isinstance(processor, AttnProcessorSwitch) and processor.switch_name == switch_name:
processor.choose_module(enabled_proc)
for name, module in model.named_children():
recursive_change_processors(name, module)
########## Hack: Attention fix #############
from diffusers.models.attention import Attention
def forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.FloatTensor] = None,
**cross_attention_kwargs,
) -> torch.Tensor:
r"""
The forward method of the `Attention` class.
Args:
hidden_states (`torch.Tensor`):
The hidden states of the query.
encoder_hidden_states (`torch.Tensor`, *optional*):
The hidden states of the encoder.
attention_mask (`torch.Tensor`, *optional*):
The attention mask to use. If `None`, no mask is applied.
**cross_attention_kwargs:
Additional keyword arguments to pass along to the cross attention.
Returns:
`torch.Tensor`: The output of the attention layer.
"""
# The `Attention` class can call different attention processors / attention functions
# here we simply pass along all tensors to the selected processor class
# For standard processors that are defined here, `**cross_attention_kwargs` is empty
return self.processor(
self,
hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
**cross_attention_kwargs,
)
Attention.forward = forward
gen_3d_modules/Unique3D/custum_3d_diffusion/custum_modules/unifield_processor.py
+460
View File
@@ -0,0 +1,460 @@
from types import FunctionType
from typing import Any, Dict, List
from diffusers import UNet2DConditionModel
import torch
import logging
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, ImageProjection
from diffusers.models.attention_processor import Attention, AttnProcessor, AttnProcessor2_0, XFormersAttnProcessor
from dataclasses import dataclass, field
from diffusers.loaders import IPAdapterMixin
from Unique3D.custum_3d_diffusion.custum_modules.attention_processors import add_extra_processor, switch_extra_processor, add_multiview_processor, switch_multiview_processor, add_switch, change_switch
@dataclass
class AttnConfig:
"""
* CrossAttention: Attention module (inherits knowledge), LoRA module (achieves fine-tuning), IPAdapter module (achieves conceptual control).
* SelfAttention: Attention module (inherits knowledge), LoRA module (achieves fine-tuning), Reference Attention module (achieves pixel-level control).
* Multiview Attention module: Multiview Attention module (achieves multi-view consistency).
* Cross Modality Attention module: Cross Modality Attention module (achieves multi-modality consistency).
For setups:
train_xxx_lr is implemented in the U-Net architecture.
enable_xxx_lora is implemented in the U-Net architecture.
enable_xxx_ip is implemented in the processor and U-Net architecture.
enable_xxx_ref_proj_in is implemented in the processor.
"""
latent_size: int = 64
train_lr: float = 0
# for cross attention
# 0 learning rate for not training
train_cross_attn_lr: float = 0
train_cross_attn_lora_lr: float = 0
train_cross_attn_ip_lr: float = 0 # 0 for not trained
init_cross_attn_lora: bool = False
enable_cross_attn_lora: bool = False
init_cross_attn_ip: bool = False
enable_cross_attn_ip: bool = False
cross_attn_lora_rank: int = 64 # 0 for not enabled
cross_attn_lora_only_kv: bool = False
ipadapter_pretrained_name: str = "h94/IP-Adapter"
ipadapter_subfolder_name: str = "models"
ipadapter_weight_name: str = "ip-adapter-plus_sd15.safetensors"
ipadapter_effect_on: str = "all" # all, first
# for self attention
train_self_attn_lr: float = 0
train_self_attn_lora_lr: float = 0
init_self_attn_lora: bool = False
enable_self_attn_lora: bool = False
self_attn_lora_rank: int = 64
self_attn_lora_only_kv: bool = False
train_self_attn_ref_lr: float = 0
train_ref_unet_lr: float = 0
init_self_attn_ref: bool = False
enable_self_attn_ref: bool = False
self_attn_ref_other_model_name: str = ""
self_attn_ref_position: str = "attn1"
self_attn_ref_pixel_wise_crosspond: bool = False # enable pixel_wise_crosspond in refattn
self_attn_ref_chain_pos: str = "parralle" # before or parralle or after
self_attn_ref_effect_on: str = "all" # all or first, for _crosspond attn
self_attn_ref_zero_init: bool = True
use_simple3d_attn: bool = False
# for multiview attention
init_multiview_attn: bool = False
enable_multiview_attn: bool = False
multiview_attn_position: str = "attn1"
multiview_chain_pose: str = "parralle" # before or parralle or after
num_modalities: int = 1
use_mv_joint_attn: bool = False
# for unet
init_unet_path: str = "runwayml/stable-diffusion-v1-5"
init_num_cls_label: int = 0 # for initialize
cls_labels: List[int] = field(default_factory=lambda: [])
cls_label_type: str = "embedding"
cat_condition: bool = False # cat condition to input
class Configurable:
attn_config: AttnConfig
def set_config(self, attn_config: AttnConfig):
raise NotImplementedError()
def update_config(self, attn_config: AttnConfig):
self.attn_config = attn_config
def do_set_config(self, attn_config: AttnConfig):
self.set_config(attn_config)
for name, module in self.named_modules():
if isinstance(module, Configurable):
if hasattr(module, "do_set_config"):
module.do_set_config(attn_config)
else:
print(f"Warning: {name} has no attribute do_set_config, but is an instance of Configurable")
module.attn_config = attn_config
def do_update_config(self, attn_config: AttnConfig):
self.update_config(attn_config)
for name, module in self.named_modules():
if isinstance(module, Configurable):
if hasattr(module, "do_update_config"):
module.do_update_config(attn_config)
else:
print(f"Warning: {name} has no attribute do_update_config, but is an instance of Configurable")
module.attn_config = attn_config
from diffusers import ModelMixin # Must import ModelMixin for CompiledUNet
class UnifieldWrappedUNet(UNet2DConditionModel):
forward_hook: FunctionType
def forward(self, *args, **kwargs):
if hasattr(self, 'forward_hook'):
return self.forward_hook(super().forward, *args, **kwargs)
return super().forward(*args, **kwargs)
class ConfigurableUNet2DConditionModel(Configurable, IPAdapterMixin):
unet: UNet2DConditionModel
cls_embedding_param_dict = {}
cross_attn_lora_param_dict = {}
self_attn_lora_param_dict = {}
cross_attn_param_dict = {}
self_attn_param_dict = {}
ipadapter_param_dict = {}
ref_attn_param_dict = {}
ref_unet_param_dict = {}
multiview_attn_param_dict = {}
other_param_dict = {}
rev_param_name_mapping = {}
class_labels = []
def set_class_labels(self, class_labels: torch.Tensor):
if self.attn_config.init_num_cls_label != 0:
self.class_labels = class_labels.to(self.unet.device).long()
def __init__(self, init_config: AttnConfig, weight_dtype) -> None:
super().__init__()
self.weight_dtype = weight_dtype
self.set_config(init_config)
def enable_xformers_memory_efficient_attention(self):
self.unet.enable_xformers_memory_efficient_attention
def recursive_add_processors(name: str, module: torch.nn.Module):
for sub_name, child in module.named_children():
recursive_add_processors(f"{name}.{sub_name}", child)
if isinstance(module, Attention):
if hasattr(module, 'xformers_not_supported'):
return
old_processor = module.get_processor()
if isinstance(old_processor, (AttnProcessor, AttnProcessor2_0)):
module.set_use_memory_efficient_attention_xformers(True)
for name, module in self.unet.named_children():
recursive_add_processors(name, module)
def __getattr__(self, name: str) -> Any:
try:
return super().__getattr__(name)
except AttributeError:
return getattr(self.unet, name)
# --- for IPAdapterMixin
def register_modules(self, **kwargs):
for name, module in kwargs.items():
# set models
setattr(self, name, module)
def register_to_config(self, **kwargs):
pass
def unload_ip_adapter(self):
raise NotImplementedError()
# --- for Configurable
def get_refunet(self):
if self.attn_config.self_attn_ref_other_model_name == "self":
return self.unet
else:
return self.unet.ref_unet
def set_config(self, attn_config: AttnConfig):
self.attn_config = attn_config
unet_type = UnifieldWrappedUNet
# class_embed_type = "projection" for 'camera'
# class_embed_type = None for 'embedding'
unet_kwargs = {}
if attn_config.init_num_cls_label > 0:
if attn_config.cls_label_type == "embedding":
unet_kwargs = {
"num_class_embeds": attn_config.init_num_cls_label,
"device_map": None,
"low_cpu_mem_usage": False,
"class_embed_type": None,
}
else:
raise ValueError(f"cls_label_type {attn_config.cls_label_type} is not supported")
self.unet: UnifieldWrappedUNet = unet_type.from_pretrained(
attn_config.init_unet_path, subfolder="unet", torch_dtype=self.weight_dtype,
**unet_kwargs
)
assert isinstance(self.unet, UnifieldWrappedUNet)
self.unet.forward_hook = self.unet_forward_hook
if self.attn_config.cat_condition:
# double in_channels
if self.unet.config.in_channels != 8:
self.unet.register_to_config(in_channels=self.unet.config.in_channels * 2)
# repeate unet.conv_in weight twice
doubled_conv_in = torch.nn.Conv2d(self.unet.conv_in.in_channels * 2, self.unet.conv_in.out_channels, self.unet.conv_in.kernel_size, self.unet.conv_in.stride, self.unet.conv_in.padding)
doubled_conv_in.weight.data = torch.cat([self.unet.conv_in.weight.data, torch.zeros_like(self.unet.conv_in.weight.data)], dim=1)
doubled_conv_in.bias.data = self.unet.conv_in.bias.data
self.unet.conv_in = doubled_conv_in
used_param_ids = set()
if attn_config.init_cross_attn_lora:
# setup lora
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
if attn_config.cross_attn_lora_only_kv:
target_modules=["attn2.to_k", "attn2.to_v"]
else:
target_modules=["attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0"]
lora_config: LoraConfig = LoraConfig(
r=attn_config.cross_attn_lora_rank,
lora_alpha=attn_config.cross_attn_lora_rank,
init_lora_weights="gaussian",
target_modules=target_modules,
)
adapter_name="cross_attn_lora"
self.unet.add_adapter(lora_config, adapter_name=adapter_name)
# update cross_attn_lora_param_dict
self.cross_attn_lora_param_dict = {id(param): param for name, param in self.unet.named_parameters() if adapter_name in name and id(param) not in used_param_ids}
used_param_ids.update(self.cross_attn_lora_param_dict.keys())
if attn_config.init_self_attn_lora:
# setup lora
from peft import LoraConfig
if attn_config.self_attn_lora_only_kv:
target_modules=["attn1.to_k", "attn1.to_v"]
else:
target_modules=["attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0"]
lora_config: LoraConfig = LoraConfig(
r=attn_config.self_attn_lora_rank,
lora_alpha=attn_config.self_attn_lora_rank,
init_lora_weights="gaussian",
target_modules=target_modules,
)
adapter_name="self_attn_lora"
self.unet.add_adapter(lora_config, adapter_name=adapter_name)
# update cross_self_lora_param_dict
self.self_attn_lora_param_dict = {id(param): param for name, param in self.unet.named_parameters() if adapter_name in name and id(param) not in used_param_ids}
used_param_ids.update(self.self_attn_lora_param_dict.keys())
if attn_config.init_num_cls_label != 0:
self.cls_embedding_param_dict = {id(param): param for param in self.unet.class_embedding.parameters()}
used_param_ids.update(self.cls_embedding_param_dict.keys())
self.set_class_labels(torch.tensor(attn_config.cls_labels).long())
if attn_config.init_cross_attn_ip:
self.image_encoder = None
# setup ipadapter
self.load_ip_adapter(
attn_config.ipadapter_pretrained_name,
subfolder=attn_config.ipadapter_subfolder_name,
weight_name=attn_config.ipadapter_weight_name
)
# warp ip_adapter_attn_proc with switch
from diffusers.models.attention_processor import IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0
add_switch(self.unet, module_filter=lambda x: isinstance(x, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0)), switch_dict_fn=lambda x: {"ipadapter": x, "default": XFormersAttnProcessor()}, switch_name="ipadapter_switch", enabled_proc="ipadapter")
# update ipadapter_param_dict
# weights are in attention processors and unet.encoder_hid_proj
self.ipadapter_param_dict = {id(param): param for param in self.unet.encoder_hid_proj.parameters() if id(param) not in used_param_ids}
used_param_ids.update(self.ipadapter_param_dict.keys())
print("DEBUG: ipadapter_param_dict len in encoder_hid_proj", len(self.ipadapter_param_dict))
for name, processor in self.unet.attn_processors.items():
if hasattr(processor, "to_k_ip"):
self.ipadapter_param_dict.update({id(param): param for param in processor.parameters()})
print(f"DEBUG: ipadapter_param_dict len in all", len(self.ipadapter_param_dict))
ref_unet = None
if attn_config.init_self_attn_ref:
# setup reference attention processor
if attn_config.self_attn_ref_other_model_name == "self":
raise NotImplementedError("self reference is not fully implemented")
else:
ref_unet: UNet2DConditionModel = UNet2DConditionModel.from_pretrained(
attn_config.self_attn_ref_other_model_name, subfolder="unet", torch_dtype=self.unet.dtype
)
ref_unet.to(self.unet.device)
if self.attn_config.train_ref_unet_lr == 0:
ref_unet.eval()
ref_unet.requires_grad_(False)
else:
ref_unet.train()
add_extra_processor(
model=ref_unet,
enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"),
mode='extract',
with_proj_in=False,
pixel_wise_crosspond=False,
)
# NOTE: Here require cross_attention_dim in two unet's self attention should be the same
processor_dict = add_extra_processor(
model=self.unet,
enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"),
mode='inject',
with_proj_in=False,
pixel_wise_crosspond=attn_config.self_attn_ref_pixel_wise_crosspond,
crosspond_effect_on=attn_config.self_attn_ref_effect_on,
crosspond_chain_pos=attn_config.self_attn_ref_chain_pos,
simple_3d=attn_config.use_simple3d_attn,
)
self.ref_unet_param_dict = {id(param): param for name, param in ref_unet.named_parameters() if id(param) not in used_param_ids and (attn_config.self_attn_ref_position in name)}
if attn_config.self_attn_ref_chain_pos != "after":
# pop untrainable paramters
for name, param in ref_unet.named_parameters():
if id(param) in self.ref_unet_param_dict and ('up_blocks.3.attentions.2.transformer_blocks.0.' in name):
self.ref_unet_param_dict.pop(id(param))
used_param_ids.update(self.ref_unet_param_dict.keys())
# update ref_attn_param_dict
self.ref_attn_param_dict = {id(param): param for name, param in processor_dict.named_parameters() if id(param) not in used_param_ids}
used_param_ids.update(self.ref_attn_param_dict.keys())
if attn_config.init_multiview_attn:
processor_dict = add_multiview_processor(
model = self.unet,
enable_filter = lambda name: name.endswith(f"{attn_config.multiview_attn_position}.processor"),
num_modalities = attn_config.num_modalities,
base_img_size = attn_config.latent_size,
chain_pos = attn_config.multiview_chain_pose,
)
# update multiview_attn_param_dict
self.multiview_attn_param_dict = {id(param): param for name, param in processor_dict.named_parameters() if id(param) not in used_param_ids}
used_param_ids.update(self.multiview_attn_param_dict.keys())
# initialize cross_attn_param_dict parameters
self.cross_attn_param_dict = {id(param): param for name, param in self.unet.named_parameters() if "attn2" in name and id(param) not in used_param_ids}
used_param_ids.update(self.cross_attn_param_dict.keys())
# initialize self_attn_param_dict parameters
self.self_attn_param_dict = {id(param): param for name, param in self.unet.named_parameters() if "attn1" in name and id(param) not in used_param_ids}
used_param_ids.update(self.self_attn_param_dict.keys())
# initialize other_param_dict parameters
self.other_param_dict = {id(param): param for name, param in self.unet.named_parameters() if id(param) not in used_param_ids}
if ref_unet is not None:
self.unet.ref_unet = ref_unet
self.rev_param_name_mapping = {id(param): name for name, param in self.unet.named_parameters()}
self.update_config(attn_config, force_update=True)
return self.unet
_attn_keys_to_update = ["enable_cross_attn_lora", "enable_cross_attn_ip", "enable_self_attn_lora", "enable_self_attn_ref", "enable_multiview_attn", "cls_labels"]
def update_config(self, attn_config: AttnConfig, force_update=False):
assert isinstance(self.unet, UNet2DConditionModel), "unet must be an instance of UNet2DConditionModel"
need_to_update = False
# update cls_labels
for key in self._attn_keys_to_update:
if getattr(self.attn_config, key) != getattr(attn_config, key):
need_to_update = True
break
if not force_update and not need_to_update:
return
self.set_class_labels(torch.tensor(attn_config.cls_labels).long())
# setup loras
if self.attn_config.init_cross_attn_lora or self.attn_config.init_self_attn_lora:
if attn_config.enable_cross_attn_lora or attn_config.enable_self_attn_lora:
cross_attn_lora_weight = 1. if attn_config.enable_cross_attn_lora > 0 else 0
self_attn_lora_weight = 1. if attn_config.enable_self_attn_lora > 0 else 0
self.unet.set_adapters(["cross_attn_lora", "self_attn_lora"], weights=[cross_attn_lora_weight, self_attn_lora_weight])
else:
self.unet.disable_adapters()
# setup ipadapter
if self.attn_config.init_cross_attn_ip:
if attn_config.enable_cross_attn_ip:
change_switch(self.unet, "ipadapter_switch", "ipadapter")
else:
change_switch(self.unet, "ipadapter_switch", "default")
# setup reference attention processor
if self.attn_config.init_self_attn_ref:
if attn_config.enable_self_attn_ref:
switch_extra_processor(self.unet, enable_filter=lambda name: name.endswith(f"{attn_config.self_attn_ref_position}.processor"))
else:
switch_extra_processor(self.unet, enable_filter=lambda name: False)
# setup multiview attention processor
if self.attn_config.init_multiview_attn:
if attn_config.enable_multiview_attn:
switch_multiview_processor(self.unet, enable_filter=lambda name: name.endswith(f"{attn_config.multiview_attn_position}.processor"))
else:
switch_multiview_processor(self.unet, enable_filter=lambda name: False)
# update cls_labels
for key in self._attn_keys_to_update:
setattr(self.attn_config, key, getattr(attn_config, key))
def unet_forward_hook(self, raw_forward, sample: torch.FloatTensor, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, *args, cross_attention_kwargs=None, condition_latents=None, class_labels=None, noisy_condition_input=False, cond_pixels_clip=None, **kwargs):
if class_labels is None and len(self.class_labels) > 0:
class_labels = self.class_labels.repeat(sample.shape[0] // self.class_labels.shape[0]).to(sample.device)
elif self.attn_config.init_num_cls_label != 0:
assert class_labels is not None, "class_labels should be passed if self.class_labels is empty and self.attn_config.init_num_cls_label is not 0"
if class_labels is not None:
if self.attn_config.cls_label_type == "embedding":
pass
else:
raise ValueError(f"cls_label_type {self.attn_config.cls_label_type} is not supported")
if self.attn_config.init_self_attn_ref and self.attn_config.enable_self_attn_ref:
# NOTE: extra step, extract condition
ref_dict = {}
ref_unet = self.get_refunet().to(sample.device)
assert condition_latents is not None
if self.attn_config.self_attn_ref_other_model_name == "self":
raise NotImplementedError()
else:
with torch.no_grad():
cond_encoder_hidden_states = encoder_hidden_states.reshape(condition_latents.shape[0], -1, *encoder_hidden_states.shape[1:])[:, 0]
if timestep.dim() == 0:
cond_timestep = timestep
else:
cond_timestep = timestep.reshape(condition_latents.shape[0], -1)[:, 0]
ref_unet(condition_latents, cond_timestep, cond_encoder_hidden_states, cross_attention_kwargs=dict(ref_dict=ref_dict))
# NOTE: extra step, inject condition
# Predict the noise residual and compute loss
if cross_attention_kwargs is None:
cross_attention_kwargs = {}
cross_attention_kwargs.update(ref_dict=ref_dict, mode='inject')
elif condition_latents is not None:
if not hasattr(self, 'condition_latents_raised'):
logging.warn("Warning! condition_latents is not None, but self_attn_ref is not enabled! This warning will only be raised once.")
self.condition_latents_raised = True
if self.attn_config.init_cross_attn_ip:
raise NotImplementedError()
if self.attn_config.cat_condition:
assert condition_latents is not None
B = condition_latents.shape[0]
cat_latents = condition_latents.reshape(B, 1, *condition_latents.shape[1:]).repeat(1, sample.shape[0] // B, 1, 1, 1).reshape(*sample.shape)
sample = torch.cat([sample, cat_latents], dim=1)
return raw_forward(sample, timestep, encoder_hidden_states, *args, cross_attention_kwargs=cross_attention_kwargs, class_labels=class_labels, **kwargs)
gen_3d_modules/Unique3D/custum_3d_diffusion/custum_pipeline/unifield_pipeline_img2img.py
+302
View File
@@ -0,0 +1,302 @@
# Copyright 2023 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.
# modified by Wuvin
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionImageVariationPipeline
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker, StableDiffusionPipelineOutput
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
import comfy.utils
class StableDiffusionImageCustomPipeline(
StableDiffusionImageVariationPipeline
):
def __init__(
self,
vae: AutoencoderKL,
image_encoder: CLIPVisionModelWithProjection,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
latents_offset=None,
noisy_cond_latents=False,
):
super().__init__(
vae=vae,
image_encoder=image_encoder,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
requires_safety_checker=requires_safety_checker
)
latents_offset = tuple(latents_offset) if latents_offset is not None else None
self.latents_offset = latents_offset
if latents_offset is not None:
self.register_to_config(latents_offset=latents_offset)
self.noisy_cond_latents = noisy_cond_latents
self.register_to_config(noisy_cond_latents=noisy_cond_latents)
def encode_latents(self, image, device, dtype, height, width):
# support batchsize > 1
if isinstance(image, Image.Image):
image = [image]
image = [img.convert("RGB") for img in image]
images = self.image_processor.preprocess(image, height=height, width=width).to(device, dtype=dtype)
latents = self.vae.encode(images).latent_dist.mode() * self.vae.config.scaling_factor
if self.latents_offset is not None:
return latents - torch.tensor(self.latents_offset).to(latents.device)[None, :, None, None]
else:
return latents
def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance):
dtype = next(self.image_encoder.parameters()).dtype
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(images=image, return_tensors="pt").pixel_values
image = image.to(device=device, dtype=dtype)
image_embeddings = self.image_encoder(image).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
# duplicate image embeddings for each generation per prompt, using mps friendly method
bs_embed, seq_len, _ = image_embeddings.shape
image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1)
image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
if do_classifier_free_guidance:
# NOTE: the same as original code
negative_prompt_embeds = torch.zeros_like(image_embeddings)
# 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
image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])
return image_embeddings
@torch.no_grad()
def __call__(
self,
image: Union[Image.Image, List[Image.Image], torch.FloatTensor],
height: Optional[int] = 1024,
width: Optional[int] = 1024,
height_cond: Optional[int] = 512,
width_cond: Optional[int] = 512,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
upper_left_feature: bool = False,
):
r"""
The call function to the pipeline for generation.
Args:
image (`Image.Image` or `List[Image.Image]` or `torch.FloatTensor`):
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
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. This parameter is modulated by `strength`.
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`.
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.FloatTensor`, *optional*):
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor is generated by sampling using the supplied random `generator`.
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.
callback (`Callable`, *optional*):
A function that calls every `callback_steps` steps during inference. The function is called with the
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
callback_steps (`int`, *optional*, defaults to 1):
The frequency at which the `callback` function is called. If not specified, the callback is called at
every step.
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.
Examples:
```py
from diffusers import StableDiffusionImageVariationPipeline
from PIL import Image
from io import BytesIO
import requests
pipe = StableDiffusionImageVariationPipeline.from_pretrained(
"lambdalabs/sd-image-variations-diffusers", revision="v2.0"
)
pipe = pipe.to("cuda")
url = "https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200"
response = requests.get(url)
image = Image.open(BytesIO(response.content)).convert("RGB")
out = pipe(image, num_images_per_prompt=3, guidance_scale=15)
out["images"][0].save("result.jpg")
```
"""
# 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
# 1. Check inputs. Raise error if not correct
self.check_inputs(image, height, width, callback_steps)
# 2. Define call parameters
if isinstance(image, Image.Image):
batch_size = 1
elif isinstance(image, list):
batch_size = len(image)
else:
batch_size = image.shape[0]
device = self._execution_device
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input image
if isinstance(image, Image.Image) and upper_left_feature:
# only use the first one of four images
emb_image = image.crop((0, 0, image.size[0] // 2, image.size[1] // 2))
else:
emb_image = image
image_embeddings = self._encode_image(emb_image, device, num_images_per_prompt, do_classifier_free_guidance)
cond_latents = self.encode_latents(image, image_embeddings.device, image_embeddings.dtype, height_cond, width_cond)
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
# 5. Prepare latent variables
num_channels_latents = self.unet.config.out_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
image_embeddings.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs.
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Denoising loop
comfy_pbar = comfy.utils.ProgressBar(num_inference_steps)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.noisy_cond_latents:
raise ValueError("Noisy condition latents is not recommended.")
else:
noisy_cond_latents = cond_latents
noisy_cond_latents = torch.cat([torch.zeros_like(noisy_cond_latents), noisy_cond_latents]) if do_classifier_free_guidance else noisy_cond_latents
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings, condition_latents=noisy_cond_latents).sample
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
# 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)
comfy_pbar.update_absolute(i + 1)
self.maybe_free_model_hooks()
if self.latents_offset is not None:
latents = latents + torch.tensor(self.latents_offset).to(latents.device)[None, :, None, None]
if not output_type == "latent":
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
image, has_nsfw_concept = self.run_safety_checker(image, device, image_embeddings.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)
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
if __name__ == "__main__":
pass
gen_3d_modules/Unique3D/custum_3d_diffusion/custum_pipeline/unifield_pipeline_img2mvimg.py
+301
View File
@@ -0,0 +1,301 @@
# Copyright 2023 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.
# modified by Wuvin
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionImageVariationPipeline
from diffusers.schedulers import KarrasDiffusionSchedulers, DDPMScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker, StableDiffusionPipelineOutput
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
import comfy.utils
class StableDiffusionImage2MVCustomPipeline(
StableDiffusionImageVariationPipeline
):
def __init__(
self,
vae: AutoencoderKL,
image_encoder: CLIPVisionModelWithProjection,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPImageProcessor,
requires_safety_checker: bool = True,
latents_offset=None,
noisy_cond_latents=False,
condition_offset=True,
):
super().__init__(
vae=vae,
image_encoder=image_encoder,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
requires_safety_checker=requires_safety_checker
)
latents_offset = tuple(latents_offset) if latents_offset is not None else None
self.latents_offset = latents_offset
if latents_offset is not None:
self.register_to_config(latents_offset=latents_offset)
if noisy_cond_latents:
raise NotImplementedError("Noisy condition latents not supported Now.")
self.condition_offset = condition_offset
self.register_to_config(condition_offset=condition_offset)
def encode_latents(self, image: Image.Image, device, dtype, height, width):
images = self.image_processor.preprocess(image.convert("RGB"), height=height, width=width).to(device, dtype=dtype)
# NOTE: .mode() for condition
latents = self.vae.encode(images).latent_dist.mode() * self.vae.config.scaling_factor
if self.latents_offset is not None and self.condition_offset:
return latents - torch.tensor(self.latents_offset).to(latents.device)[None, :, None, None]
else:
return latents
def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance):
dtype = next(self.image_encoder.parameters()).dtype
if not isinstance(image, torch.Tensor):
image = self.feature_extractor(images=image, return_tensors="pt").pixel_values
image = image.to(device=device, dtype=dtype)
image_embeddings = self.image_encoder(image).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
# duplicate image embeddings for each generation per prompt, using mps friendly method
bs_embed, seq_len, _ = image_embeddings.shape
image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1)
image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
if do_classifier_free_guidance:
# NOTE: the same as original code
negative_prompt_embeds = torch.zeros_like(image_embeddings)
# 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
image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])
return image_embeddings
@torch.no_grad()
def __call__(
self,
image: Union[Image.Image, List[Image.Image], torch.FloatTensor],
height: Optional[int] = 1024,
width: Optional[int] = 1024,
height_cond: Optional[int] = 512,
width_cond: Optional[int] = 512,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 4,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
):
r"""
The call function to the pipeline for generation.
Args:
image (`Image.Image` or `List[Image.Image]` or `torch.FloatTensor`):
Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
[`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
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. This parameter is modulated by `strength`.
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`.
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.FloatTensor`, *optional*):
Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor is generated by sampling using the supplied random `generator`.
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.
callback (`Callable`, *optional*):
A function that calls every `callback_steps` steps during inference. The function is called with the
following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
callback_steps (`int`, *optional*, defaults to 1):
The frequency at which the `callback` function is called. If not specified, the callback is called at
every step.
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.
Examples:
```py
from diffusers import StableDiffusionImageVariationPipeline
from PIL import Image
from io import BytesIO
import requests
pipe = StableDiffusionImageVariationPipeline.from_pretrained(
"lambdalabs/sd-image-variations-diffusers", revision="v2.0"
)
pipe = pipe.to("cuda")
url = "https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200"
response = requests.get(url)
image = Image.open(BytesIO(response.content)).convert("RGB")
out = pipe(image, num_images_per_prompt=3, guidance_scale=15)
out["images"][0].save("result.jpg")
```
"""
# 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
# 1. Check inputs. Raise error if not correct
self.check_inputs(image, height, width, callback_steps)
# 2. Define call parameters
if isinstance(image, Image.Image):
batch_size = 1
elif len(image) == 1:
image = image[0]
batch_size = 1
else:
raise NotImplementedError()
# elif isinstance(image, list):
# batch_size = len(image)
# else:
# batch_size = image.shape[0]
device = self._execution_device
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input image
emb_image = image
image_embeddings = self._encode_image(emb_image, device, num_images_per_prompt, do_classifier_free_guidance)
cond_latents = self.encode_latents(image, image_embeddings.device, image_embeddings.dtype, height_cond, width_cond)
cond_latents = torch.cat([torch.zeros_like(cond_latents), cond_latents]) if do_classifier_free_guidance else cond_latents
image_pixels = self.feature_extractor(images=emb_image, return_tensors="pt").pixel_values
if do_classifier_free_guidance:
image_pixels = torch.cat([torch.zeros_like(image_pixels), image_pixels], dim=0)
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
# 5. Prepare latent variables
num_channels_latents = self.unet.config.out_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
image_embeddings.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs.
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Denoising loop
comfy_pbar = comfy.utils.ProgressBar(num_inference_steps)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings, condition_latents=cond_latents, noisy_condition_input=False, cond_pixels_clip=image_pixels).sample
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
# 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)
comfy_pbar.update_absolute(i + 1)
self.maybe_free_model_hooks()
if self.latents_offset is not None:
latents = latents + torch.tensor(self.latents_offset).to(latents.device)[None, :, None, None]
if not output_type == "latent":
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
image, has_nsfw_concept = self.run_safety_checker(image, device, image_embeddings.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)
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
if __name__ == "__main__":
pass
gen_3d_modules/Unique3D/custum_3d_diffusion/modules.py
+14
View File
@@ -0,0 +1,14 @@
__modules__ = {}
def register(name):
def decorator(cls):
__modules__[name] = cls
return cls
return decorator
def find(name):
return __modules__[name]
from custum_3d_diffusion.trainings import base, image2mvimage_trainer, image2image_trainer
gen_3d_modules/Unique3D/custum_3d_diffusion/trainings/__init__.py
View File
gen_3d_modules/Unique3D/custum_3d_diffusion/trainings/config_classes.py
+35
View File
@@ -0,0 +1,35 @@
from dataclasses import dataclass, field
from typing import List, Optional
@dataclass
class TrainerSubConfig:
trainer_type: str = ""
trainer: dict = field(default_factory=dict)
@dataclass
class ExprimentConfig:
trainers: List[dict] = field(default_factory=lambda: [])
init_config: dict = field(default_factory=dict)
pretrained_model_name_or_path: str = ""
pretrained_unet_state_dict_path: str = ""
# expriments related parameters
linear_beta_schedule: bool = False
zero_snr: bool = False
prediction_type: Optional[str] = None
seed: Optional[int] = None
max_train_steps: int = 1000000
gradient_accumulation_steps: int = 1
learning_rate: float = 1e-4
lr_scheduler: str = "constant"
lr_warmup_steps: int = 500
use_8bit_adam: bool = False
adam_beta1: float = 0.9
adam_beta2: float = 0.999
adam_weight_decay: float = 1e-2
adam_epsilon: float = 1e-08
max_grad_norm: float = 1.0
mixed_precision: Optional[str] = None # ["no", "fp16", "bf16", "fp8"]
skip_training: bool = False
debug: bool = False
gen_3d_modules/Unique3D/custum_3d_diffusion/trainings/utils.py
+25
View File
@@ -0,0 +1,25 @@
from omegaconf import DictConfig, OmegaConf
def parse_structured(fields, cfg) -> DictConfig:
scfg = OmegaConf.structured(fields(**cfg))
return scfg
def load_config(fields, config, extras=None):
if extras is not None:
print("Warning! extra parameter in cli is not verified, may cause erros.")
if isinstance(config, str):
cfg = OmegaConf.load(config)
elif isinstance(config, dict):
cfg = OmegaConf.create(config)
elif isinstance(config, DictConfig):
cfg = config
else:
raise NotImplementedError(f"Unsupported config type {type(config)}")
if extras is not None:
cli_conf = OmegaConf.from_cli(extras)
cfg = OmegaConf.merge(cfg, cli_conf)
OmegaConf.resolve(cfg)
assert isinstance(cfg, DictConfig)
return parse_structured(fields, cfg)
gen_3d_modules/Unique3D/mesh_reconstruction/__init__.py
View File
gen_3d_modules/Unique3D/mesh_reconstruction/func.py
+133
View File
@@ -0,0 +1,133 @@
# modified from https://github.com/Profactor/continuous-remeshing
import torch
import numpy as np
import trimesh
from typing import Tuple
def to_numpy(*args):
def convert(a):
if isinstance(a,torch.Tensor):
return a.detach().cpu().numpy()
assert a is None or isinstance(a,np.ndarray)
return a
return convert(args[0]) if len(args)==1 else tuple(convert(a) for a in args)
def laplacian(
num_verts:int,
edges: torch.Tensor #E,2
) -> torch.Tensor: #sparse V,V
"""create sparse Laplacian matrix"""
V = num_verts
E = edges.shape[0]
#adjacency matrix,
idx = torch.cat([edges, edges.fliplr()], dim=0).type(torch.long).T # (2, 2*E)
ones = torch.ones(2*E, dtype=torch.float32, device=edges.device)
A = torch.sparse.FloatTensor(idx, ones, (V, V))
#degree matrix
deg = torch.sparse.sum(A, dim=1).to_dense()
idx = torch.arange(V, device=edges.device)
idx = torch.stack([idx, idx], dim=0)
D = torch.sparse.FloatTensor(idx, deg, (V, V))
return D - A
def _translation(x, y, z, device):
return torch.tensor([[1., 0, 0, x],
[0, 1, 0, y],
[0, 0, 1, z],
[0, 0, 0, 1]],device=device) #4,4
def _projection(r, device, l=None, t=None, b=None, n=1.0, f=50.0, flip_y=True):
"""
see https://blog.csdn.net/wodownload2/article/details/85069240/
"""
if l is None:
l = -r
if t is None:
t = r
if b is None:
b = -t
p = torch.zeros([4,4],device=device)
p[0,0] = 2*n/(r-l)
p[0,2] = (r+l)/(r-l)
p[1,1] = 2*n/(t-b) * (-1 if flip_y else 1)
p[1,2] = (t+b)/(t-b)
p[2,2] = -(f+n)/(f-n)
p[2,3] = -(2*f*n)/(f-n)
p[3,2] = -1
return p #4,4
def _orthographic(r, device, l=None, t=None, b=None, n=1.0, f=50.0, flip_y=True):
if l is None:
l = -r
if t is None:
t = r
if b is None:
b = -t
o = torch.zeros([4,4],device=device)
o[0,0] = 2/(r-l)
o[0,3] = -(r+l)/(r-l)
o[1,1] = 2/(t-b) * (-1 if flip_y else 1)
o[1,3] = -(t+b)/(t-b)
o[2,2] = -2/(f-n)
o[2,3] = -(f+n)/(f-n)
o[3,3] = 1
return o #4,4
def make_star_cameras(az_count,pol_count,distance:float=10.,r=None,image_size=[512,512],device='cuda'):
if r is None:
r = 1/distance
A = az_count
P = pol_count
C = A * P
phi = torch.arange(0,A) * (2*torch.pi/A)
phi_rot = torch.eye(3,device=device)[None,None].expand(A,1,3,3).clone()
phi_rot[:,0,2,2] = phi.cos()
phi_rot[:,0,2,0] = -phi.sin()
phi_rot[:,0,0,2] = phi.sin()
phi_rot[:,0,0,0] = phi.cos()
theta = torch.arange(1,P+1) * (torch.pi/(P+1)) - torch.pi/2
theta_rot = torch.eye(3,device=device)[None,None].expand(1,P,3,3).clone()
theta_rot[0,:,1,1] = theta.cos()
theta_rot[0,:,1,2] = -theta.sin()
theta_rot[0,:,2,1] = theta.sin()
theta_rot[0,:,2,2] = theta.cos()
mv = torch.empty((C,4,4), device=device)
mv[:] = torch.eye(4, device=device)
mv[:,:3,:3] = (theta_rot @ phi_rot).reshape(C,3,3)
mv = _translation(0, 0, -distance, device) @ mv
return mv, _projection(r,device)
def make_star_cameras_orthographic(az_count,pol_count,distance:float=10.,r=None,image_size=[512,512],device='cuda'):
mv, _ = make_star_cameras(az_count,pol_count,distance,r,image_size,device)
if r is None:
r = 1
return mv, _orthographic(r,device)
def make_sphere(level:int=2,radius=1.,device='cuda') -> Tuple[torch.Tensor,torch.Tensor]:
sphere = trimesh.creation.icosphere(subdivisions=level, radius=1.0, color=None)
vertices = torch.tensor(sphere.vertices, device=device, dtype=torch.float32) * radius
faces = torch.tensor(sphere.faces, device=device, dtype=torch.long)
return vertices,faces
from pytorch3d.renderer import (
FoVOrthographicCameras,
look_at_view_transform,
)
def get_camera(R, T, focal_length=1 / (2**0.5)):
focal_length = 1 / focal_length
camera = FoVOrthographicCameras(device=R.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length)
return camera
def make_star_cameras_orthographic_py3d(azim_list, device, focal=2/1.35, dist=1.1):
R, T = look_at_view_transform(dist, 0, azim_list)
focal_length = 1 / focal
return FoVOrthographicCameras(device=R.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length).to(device)
gen_3d_modules/Unique3D/mesh_reconstruction/opt.py
+190
View File
@@ -0,0 +1,190 @@
# modified from https://github.com/Profactor/continuous-remeshing
import time
import torch
import torch_scatter
from typing import Tuple
from .remesh import calc_edge_length, calc_edges, calc_face_collapses, calc_face_normals, calc_vertex_normals, collapse_edges, flip_edges, pack, prepend_dummies, remove_dummies, split_edges
@torch.no_grad()
def remesh(
vertices_etc:torch.Tensor, #V,D
faces:torch.Tensor, #F,3 long
min_edgelen:torch.Tensor, #V
max_edgelen:torch.Tensor, #V
flip:bool,
max_vertices=1e6
):
# dummies
vertices_etc,faces = prepend_dummies(vertices_etc,faces)
vertices = vertices_etc[:,:3] #V,3
nan_tensor = torch.tensor([torch.nan],device=min_edgelen.device)
min_edgelen = torch.concat((nan_tensor,min_edgelen))
max_edgelen = torch.concat((nan_tensor,max_edgelen))
# collapse
edges,face_to_edge = calc_edges(faces) #E,2 F,3
edge_length = calc_edge_length(vertices,edges) #E
face_normals = calc_face_normals(vertices,faces,normalize=False) #F,3
vertex_normals = calc_vertex_normals(vertices,faces,face_normals) #V,3
face_collapse = calc_face_collapses(vertices,faces,edges,face_to_edge,edge_length,face_normals,vertex_normals,min_edgelen,area_ratio=0.5)
shortness = (1 - edge_length / min_edgelen[edges].mean(dim=-1)).clamp_min_(0) #e[0,1] 0...ok, 1...edgelen=0
priority = face_collapse.float() + shortness
vertices_etc,faces = collapse_edges(vertices_etc,faces,edges,priority)
# split
if vertices.shape[0]<max_vertices:
edges,face_to_edge = calc_edges(faces) #E,2 F,3
vertices = vertices_etc[:,:3] #V,3
edge_length = calc_edge_length(vertices,edges) #E
splits = edge_length > max_edgelen[edges].mean(dim=-1)
vertices_etc,faces = split_edges(vertices_etc,faces,edges,face_to_edge,splits,pack_faces=False)
vertices_etc,faces = pack(vertices_etc,faces)
vertices = vertices_etc[:,:3]
if flip:
edges,_,edge_to_face = calc_edges(faces,with_edge_to_face=True) #E,2 F,3
flip_edges(vertices,faces,edges,edge_to_face,with_border=False)
return remove_dummies(vertices_etc,faces)
def lerp_unbiased(a:torch.Tensor,b:torch.Tensor,weight:float,step:int):
"""lerp with adam's bias correction"""
c_prev = 1-weight**(step-1)
c = 1-weight**step
a_weight = weight*c_prev/c
b_weight = (1-weight)/c
a.mul_(a_weight).add_(b, alpha=b_weight)
class MeshOptimizer:
"""Use this like a pytorch Optimizer, but after calling opt.step(), do vertices,faces = opt.remesh()."""
def __init__(self,
vertices:torch.Tensor, #V,3
faces:torch.Tensor, #F,3
lr=0.3, #learning rate
betas=(0.8,0.8,0), #betas[0:2] are the same as in Adam, betas[2] may be used to time-smooth the relative velocity nu
gammas=(0,0,0), #optional spatial smoothing for m1,m2,nu, values between 0 (no smoothing) and 1 (max. smoothing)
nu_ref=0.3, #reference velocity for edge length controller
edge_len_lims=(.01,.15), #smallest and largest allowed reference edge length
edge_len_tol=.5, #edge length tolerance for split and collapse
gain=.2, #gain value for edge length controller
laplacian_weight=.02, #for laplacian smoothing/regularization
ramp=1, #learning rate ramp, actual ramp width is ramp/(1-betas[0])
grad_lim=10., #gradients are clipped to m1.abs()*grad_lim
remesh_interval=1, #larger intervals are faster but with worse mesh quality
local_edgelen=True, #set to False to use a global scalar reference edge length instead
):
self._vertices = vertices
self._faces = faces
self._lr = lr
self._betas = betas
self._gammas = gammas
self._nu_ref = nu_ref
self._edge_len_lims = edge_len_lims
self._edge_len_tol = edge_len_tol
self._gain = gain
self._laplacian_weight = laplacian_weight
self._ramp = ramp
self._grad_lim = grad_lim
self._remesh_interval = remesh_interval
self._local_edgelen = local_edgelen
self._step = 0
V = self._vertices.shape[0]
# prepare continuous tensor for all vertex-based data
self._vertices_etc = torch.zeros([V,9],device=vertices.device)
self._split_vertices_etc()
self.vertices.copy_(vertices) #initialize vertices
self._vertices.requires_grad_()
self._ref_len.fill_(edge_len_lims[1])
@property
def vertices(self):
return self._vertices
@property
def faces(self):
return self._faces
def _split_vertices_etc(self):
self._vertices = self._vertices_etc[:,:3]
self._m2 = self._vertices_etc[:,3]
self._nu = self._vertices_etc[:,4]
self._m1 = self._vertices_etc[:,5:8]
self._ref_len = self._vertices_etc[:,8]
with_gammas = any(g!=0 for g in self._gammas)
self._smooth = self._vertices_etc[:,:8] if with_gammas else self._vertices_etc[:,:3]
def zero_grad(self):
self._vertices.grad = None
@torch.no_grad()
def step(self):
eps = 1e-8
self._step += 1
# spatial smoothing
edges,_ = calc_edges(self._faces) #E,2
E = edges.shape[0]
edge_smooth = self._smooth[edges] #E,2,S
neighbor_smooth = torch.zeros_like(self._smooth) #V,S
torch_scatter.scatter_mean(src=edge_smooth.flip(dims=[1]).reshape(E*2,-1),index=edges.reshape(E*2,1),dim=0,out=neighbor_smooth)
#apply optional smoothing of m1,m2,nu
if self._gammas[0]:
self._m1.lerp_(neighbor_smooth[:,5:8],self._gammas[0])
if self._gammas[1]:
self._m2.lerp_(neighbor_smooth[:,3],self._gammas[1])
if self._gammas[2]:
self._nu.lerp_(neighbor_smooth[:,4],self._gammas[2])
#add laplace smoothing to gradients
laplace = self._vertices - neighbor_smooth[:,:3]
grad = torch.addcmul(self._vertices.grad, laplace, self._nu[:,None], value=self._laplacian_weight)
#gradient clipping
if self._step>1:
grad_lim = self._m1.abs().mul_(self._grad_lim)
grad.clamp_(min=-grad_lim,max=grad_lim)
# moment updates
lerp_unbiased(self._m1, grad, self._betas[0], self._step)
lerp_unbiased(self._m2, (grad**2).sum(dim=-1), self._betas[1], self._step)
velocity = self._m1 / self._m2[:,None].sqrt().add_(eps) #V,3
speed = velocity.norm(dim=-1) #V
if self._betas[2]:
lerp_unbiased(self._nu,speed,self._betas[2],self._step) #V
else:
self._nu.copy_(speed) #V
# update vertices
ramped_lr = self._lr * min(1,self._step * (1-self._betas[0]) / self._ramp)
self._vertices.add_(velocity * self._ref_len[:,None], alpha=-ramped_lr)
# update target edge length
if self._step % self._remesh_interval == 0:
if self._local_edgelen:
len_change = (1 + (self._nu - self._nu_ref) * self._gain)
else:
len_change = (1 + (self._nu.mean() - self._nu_ref) * self._gain)
self._ref_len *= len_change
self._ref_len.clamp_(*self._edge_len_lims)
def remesh(self, flip:bool=True, poisson=False)->Tuple[torch.Tensor,torch.Tensor]:
min_edge_len = self._ref_len * (1 - self._edge_len_tol)
max_edge_len = self._ref_len * (1 + self._edge_len_tol)
self._vertices_etc,self._faces = remesh(self._vertices_etc,self._faces,min_edge_len,max_edge_len,flip, max_vertices=1e6)
self._split_vertices_etc()
self._vertices.requires_grad_()
return self._vertices, self._faces
gen_3d_modules/Unique3D/mesh_reconstruction/recon.py
+65
View File
@@ -0,0 +1,65 @@
from tqdm import tqdm
from PIL import Image
import numpy as np
import torch
from typing import List
import comfy.utils
from .remesh import calc_vertex_normals
from .opt import MeshOptimizer
from .func import make_star_cameras_orthographic, make_star_cameras_orthographic_py3d
from .render import NormalsRenderer, Pytorch3DNormalsRenderer
from Unique3D.scripts.utils import to_py3d_mesh, init_target
def reconstruct_stage1(pils: List[Image.Image], steps=100, vertices=None, faces=None, start_edge_len=0.1, end_edge_len=0.02, decay=0.995, return_mesh=False, loss_expansion_weight=0.1, gain=0.05):
vertices, faces = vertices.to("cuda"), faces.to("cuda")
assert len(pils) == 4
mv,proj = make_star_cameras_orthographic(4, 1)
renderer = NormalsRenderer(mv,proj,list(pils[0].size))
# cameras = make_star_cameras_orthographic_py3d([0, 270, 180, 90], device="cuda", focal=1., dist=4.0)
# renderer = Pytorch3DNormalsRenderer(cameras, list(pils[0].size), device="cuda")
target_images = init_target(pils, new_bkgd=(0., 0., 0.)) # 4s
# 1. no rotate
target_images = target_images[[0, 3, 2, 1]]
# 2. init from coarse mesh
opt = MeshOptimizer(vertices,faces, local_edgelen=False, gain=gain, edge_len_lims=(end_edge_len, start_edge_len))
vertices = opt.vertices
mask = target_images[..., -1] < 0.5
comfy_pbar = comfy.utils.ProgressBar(steps)
for i in tqdm(range(steps)):
opt.zero_grad()
opt._lr *= decay
normals = calc_vertex_normals(vertices,faces)
images = renderer.render(vertices,normals,faces)
loss_expand = 0.5 * ((vertices+normals).detach() - vertices).pow(2).mean()
t_mask = images[..., -1] > 0.5
loss_target_l2 = (images[t_mask] - target_images[t_mask]).abs().pow(2).mean()
loss_alpha_target_mask_l2 = (images[..., -1][mask] - target_images[..., -1][mask]).pow(2).mean()
loss = loss_target_l2 + loss_alpha_target_mask_l2 + loss_expand * loss_expansion_weight
# out of box
loss_oob = (vertices.abs() > 0.99).float().mean() * 10
loss = loss + loss_oob
loss.backward()
opt.step()
vertices,faces = opt.remesh(poisson=False)
comfy_pbar.update_absolute(i + 1)
vertices, faces = vertices.detach(), faces.detach()
if return_mesh:
return to_py3d_mesh(vertices, faces)
else:
return vertices, faces
gen_3d_modules/Unique3D/mesh_reconstruction/refine.py
+85
View File
@@ -0,0 +1,85 @@
from tqdm import tqdm
from PIL import Image
import torch
from typing import List
import comfy.utils
from .remesh import calc_vertex_normals
from .opt import MeshOptimizer
from .func import make_star_cameras_orthographic, make_star_cameras_orthographic_py3d
from .render import NormalsRenderer, Pytorch3DNormalsRenderer
from Unique3D.scripts.project_mesh import multiview_color_projection, get_cameras_list
from Unique3D.scripts.utils import to_py3d_mesh, from_py3d_mesh, init_target
def run_mesh_refine(vertices, faces, pils: List[Image.Image], steps=100, start_edge_len=0.02, end_edge_len=0.005, decay=0.99, update_normal_interval=10, update_warmup=10, return_mesh=False, process_inputs=False, process_outputs=False):
if process_inputs:
vertices = vertices * 2 / 1.35
vertices[..., [0, 2]] = - vertices[..., [0, 2]]
poission_steps = []
assert len(pils) == 4
mv,proj = make_star_cameras_orthographic(4, 1)
renderer = NormalsRenderer(mv,proj,list(pils[0].size))
# cameras = make_star_cameras_orthographic_py3d([0, 270, 180, 90], device="cuda", focal=1., dist=4.0)
# renderer = Pytorch3DNormalsRenderer(cameras, list(pils[0].size), device="cuda")
target_images = init_target(pils, new_bkgd=(0., 0., 0.)) # 4s
# 1. no rotate
target_images = target_images[[0, 3, 2, 1]]
# 2. init from coarse mesh
opt = MeshOptimizer(vertices,faces, ramp=5, edge_len_lims=(end_edge_len, start_edge_len), local_edgelen=False, laplacian_weight=0.02)
vertices = opt.vertices
alpha_init = None
mask = target_images[..., -1] < 0.5
comfy_pbar = comfy.utils.ProgressBar(steps)
for i in tqdm(range(steps)):
opt.zero_grad()
opt._lr *= decay
normals = calc_vertex_normals(vertices,faces)
images = renderer.render(vertices,normals,faces)
if alpha_init is None:
alpha_init = images.detach()
if i < update_warmup or i % update_normal_interval == 0:
with torch.no_grad():
py3d_mesh = to_py3d_mesh(vertices, faces, normals)
cameras = get_cameras_list(azim_list = [0, 90, 180, 270], device=vertices.device, focal=1.)
_, _, target_normal = from_py3d_mesh(multiview_color_projection(py3d_mesh, pils, cameras_list=cameras, weights=[2.0, 0.8, 1.0, 0.8], confidence_threshold=0.1, complete_unseen=False, below_confidence_strategy='original', reweight_with_cosangle='linear'))
target_normal = target_normal * 2 - 1
target_normal = torch.nn.functional.normalize(target_normal, dim=-1)
debug_images = renderer.render(vertices,target_normal,faces)
d_mask = images[..., -1] > 0.5
loss_debug_l2 = (images[..., :3][d_mask] - debug_images[..., :3][d_mask]).pow(2).mean()
loss_alpha_target_mask_l2 = (images[..., -1][mask] - target_images[..., -1][mask]).pow(2).mean()
loss = loss_debug_l2 + loss_alpha_target_mask_l2
# out of box
loss_oob = (vertices.abs() > 0.99).float().mean() * 10
loss = loss + loss_oob
loss.backward()
opt.step()
vertices,faces = opt.remesh(poisson=(i in poission_steps))
comfy_pbar.update_absolute(i + 1)
vertices, faces = vertices.detach(), faces.detach()
if process_outputs:
vertices = vertices / 2 * 1.35
vertices[..., [0, 2]] = - vertices[..., [0, 2]]
if return_mesh:
return to_py3d_mesh(vertices, faces)
else:
return vertices, faces
gen_3d_modules/Unique3D/mesh_reconstruction/remesh.py
+361
View File
@@ -0,0 +1,361 @@
# modified from https://github.com/Profactor/continuous-remeshing
import torch
import torch.nn.functional as tfunc
import torch_scatter
from typing import Tuple
def prepend_dummies(
vertices:torch.Tensor, #V,D
faces:torch.Tensor, #F,3 long
)->Tuple[torch.Tensor,torch.Tensor]:
"""prepend dummy elements to vertices and faces to enable "masked" scatter operations"""
V,D = vertices.shape
vertices = torch.concat((torch.full((1,D),fill_value=torch.nan,device=vertices.device),vertices),dim=0)
faces = torch.concat((torch.zeros((1,3),dtype=torch.long,device=faces.device),faces+1),dim=0)
return vertices,faces
def remove_dummies(
vertices:torch.Tensor, #V,D - first vertex all nan and unreferenced
faces:torch.Tensor, #F,3 long - first face all zeros
)->Tuple[torch.Tensor,torch.Tensor]:
"""remove dummy elements added with prepend_dummies()"""
return vertices[1:],faces[1:]-1
def calc_edges(
faces: torch.Tensor, # F,3 long - first face may be dummy with all zeros
with_edge_to_face: bool = False
) -> Tuple[torch.Tensor, ...]:
"""
returns Tuple of
- edges E,2 long, 0 for unused, lower vertex index first
- face_to_edge F,3 long
- (optional) edge_to_face shape=E,[left,right],[face,side]
o-<-----e1 e0,e1...edge, e0<e1
| /A L,R....left and right face
| L / | both triangles ordered counter clockwise
| / R | normals pointing out of screen
V/ |
e0---->-o
"""
F = faces.shape[0]
# make full edges, lower vertex index first
face_edges = torch.stack((faces,faces.roll(-1,1)),dim=-1) #F*3,3,2
full_edges = face_edges.reshape(F*3,2)
sorted_edges,_ = full_edges.sort(dim=-1) #F*3,2
# make unique edges
edges,full_to_unique = torch.unique(input=sorted_edges,sorted=True,return_inverse=True,dim=0) #(E,2),(F*3)
E = edges.shape[0]
face_to_edge = full_to_unique.reshape(F,3) #F,3
if not with_edge_to_face:
return edges, face_to_edge
is_right = full_edges[:,0]!=sorted_edges[:,0] #F*3
edge_to_face = torch.zeros((E,2,2),dtype=torch.long,device=faces.device) #E,LR=2,S=2
scatter_src = torch.cartesian_prod(torch.arange(0,F,device=faces.device),torch.arange(0,3,device=faces.device)) #F*3,2
edge_to_face.reshape(2*E,2).scatter_(dim=0,index=(2*full_to_unique+is_right)[:,None].expand(F*3,2),src=scatter_src) #E,LR=2,S=2
edge_to_face[0] = 0
return edges, face_to_edge, edge_to_face
def calc_edge_length(
vertices:torch.Tensor, #V,3 first may be dummy
edges:torch.Tensor, #E,2 long, lower vertex index first, (0,0) for unused
)->torch.Tensor: #E
full_vertices = vertices[edges] #E,2,3
a,b = full_vertices.unbind(dim=1) #E,3
return torch.norm(a-b,p=2,dim=-1)
def calc_face_normals(
vertices:torch.Tensor, #V,3 first vertex may be unreferenced
faces:torch.Tensor, #F,3 long, first face may be all zero
normalize:bool=False,
)->torch.Tensor: #F,3
"""
n
|
c0 corners ordered counterclockwise when
/ \ looking onto surface (in neg normal direction)
c1---c2
"""
full_vertices = vertices[faces] #F,C=3,3
v0,v1,v2 = full_vertices.unbind(dim=1) #F,3
face_normals = torch.cross(v1-v0,v2-v0, dim=1) #F,3
if normalize:
face_normals = tfunc.normalize(face_normals, eps=1e-6, dim=1)
return face_normals #F,3
def calc_vertex_normals(
vertices:torch.Tensor, #V,3 first vertex may be unreferenced
faces:torch.Tensor, #F,3 long, first face may be all zero
face_normals:torch.Tensor=None, #F,3, not normalized
)->torch.Tensor: #F,3
F = faces.shape[0]
if face_normals is None:
face_normals = calc_face_normals(vertices,faces)
vertex_normals = torch.zeros((vertices.shape[0],3,3),dtype=vertices.dtype,device=vertices.device) #V,C=3,3
vertex_normals.scatter_add_(dim=0,index=faces[:,:,None].expand(F,3,3),src=face_normals[:,None,:].expand(F,3,3))
vertex_normals = vertex_normals.sum(dim=1) #V,3
return tfunc.normalize(vertex_normals, eps=1e-6, dim=1)
def calc_face_ref_normals(
faces:torch.Tensor, #F,3 long, 0 for unused
vertex_normals:torch.Tensor, #V,3 first unused
normalize:bool=False,
)->torch.Tensor: #F,3
"""calculate reference normals for face flip detection"""
full_normals = vertex_normals[faces] #F,C=3,3
ref_normals = full_normals.sum(dim=1) #F,3
if normalize:
ref_normals = tfunc.normalize(ref_normals, eps=1e-6, dim=1)
return ref_normals
def pack(
vertices:torch.Tensor, #V,3 first unused and nan
faces:torch.Tensor, #F,3 long, 0 for unused
)->Tuple[torch.Tensor,torch.Tensor]: #(vertices,faces), keeps first vertex unused
"""removes unused elements in vertices and faces"""
V = vertices.shape[0]
# remove unused faces
used_faces = faces[:,0]!=0
used_faces[0] = True
faces = faces[used_faces] #sync
# remove unused vertices
used_vertices = torch.zeros(V,3,dtype=torch.bool,device=vertices.device)
used_vertices.scatter_(dim=0,index=faces,value=True,reduce='add')
used_vertices = used_vertices.any(dim=1)
used_vertices[0] = True
vertices = vertices[used_vertices] #sync
# update used faces
ind = torch.zeros(V,dtype=torch.long,device=vertices.device)
V1 = used_vertices.sum()
ind[used_vertices] = torch.arange(0,V1,device=vertices.device) #sync
faces = ind[faces]
return vertices,faces
def split_edges(
vertices:torch.Tensor, #V,3 first unused
faces:torch.Tensor, #F,3 long, 0 for unused
edges:torch.Tensor, #E,2 long 0 for unused, lower vertex index first
face_to_edge:torch.Tensor, #F,3 long 0 for unused
splits, #E bool
pack_faces:bool=True,
)->Tuple[torch.Tensor,torch.Tensor]: #(vertices,faces)
# c2 c2 c...corners = faces
# . . . . s...side_vert, 0 means no split
# . . .N2 . S...shrunk_face
# . . . . Ni...new_faces
# s2 s1 s2|c2...s1|c1
# . . . . .
# . . . S . .
# . . . . N1 .
# c0...(s0=0)....c1 s0|c0...........c1
#
# pseudo-code:
# S = [s0|c0,s1|c1,s2|c2] example:[c0,s1,s2]
# split = side_vert!=0 example:[False,True,True]
# N0 = split[0]*[c0,s0,s2|c2] example:[0,0,0]
# N1 = split[1]*[c1,s1,s0|c0] example:[c1,s1,c0]
# N2 = split[2]*[c2,s2,s1|c1] example:[c2,s2,s1]
V = vertices.shape[0]
F = faces.shape[0]
S = splits.sum().item() #sync
if S==0:
return vertices,faces
edge_vert = torch.zeros_like(splits, dtype=torch.long) #E
edge_vert[splits] = torch.arange(V,V+S,dtype=torch.long,device=vertices.device) #E 0 for no split, sync
side_vert = edge_vert[face_to_edge] #F,3 long, 0 for no split
split_edges = edges[splits] #S sync
#vertices
split_vertices = vertices[split_edges].mean(dim=1) #S,3
vertices = torch.concat((vertices,split_vertices),dim=0)
#faces
side_split = side_vert!=0 #F,3
shrunk_faces = torch.where(side_split,side_vert,faces) #F,3 long, 0 for no split
new_faces = side_split[:,:,None] * torch.stack((faces,side_vert,shrunk_faces.roll(1,dims=-1)),dim=-1) #F,N=3,C=3
faces = torch.concat((shrunk_faces,new_faces.reshape(F*3,3))) #4F,3
if pack_faces:
mask = faces[:,0]!=0
mask[0] = True
faces = faces[mask] #F',3 sync
return vertices,faces
def collapse_edges(
vertices:torch.Tensor, #V,3 first unused
faces:torch.Tensor, #F,3 long 0 for unused
edges:torch.Tensor, #E,2 long 0 for unused, lower vertex index first
priorities:torch.Tensor, #E float
stable:bool=False, #only for unit testing
)->Tuple[torch.Tensor,torch.Tensor]: #(vertices,faces)
V = vertices.shape[0]
# check spacing
_,order = priorities.sort(stable=stable) #E
rank = torch.zeros_like(order)
rank[order] = torch.arange(0,len(rank),device=rank.device)
vert_rank = torch.zeros(V,dtype=torch.long,device=vertices.device) #V
edge_rank = rank #E
for i in range(3):
torch_scatter.scatter_max(src=edge_rank[:,None].expand(-1,2).reshape(-1),index=edges.reshape(-1),dim=0,out=vert_rank)
edge_rank,_ = vert_rank[edges].max(dim=-1) #E
candidates = edges[(edge_rank==rank).logical_and_(priorities>0)] #E',2
# check connectivity
vert_connections = torch.zeros(V,dtype=torch.long,device=vertices.device) #V
vert_connections[candidates[:,0]] = 1 #start
edge_connections = vert_connections[edges].sum(dim=-1) #E, edge connected to start
vert_connections.scatter_add_(dim=0,index=edges.reshape(-1),src=edge_connections[:,None].expand(-1,2).reshape(-1))# one edge from start
vert_connections[candidates] = 0 #clear start and end
edge_connections = vert_connections[edges].sum(dim=-1) #E, one or two edges from start
vert_connections.scatter_add_(dim=0,index=edges.reshape(-1),src=edge_connections[:,None].expand(-1,2).reshape(-1)) #one or two edges from start
collapses = candidates[vert_connections[candidates[:,1]] <= 2] # E" not more than two connections between start and end
# mean vertices
vertices[collapses[:,0]] = vertices[collapses].mean(dim=1)
# update faces
dest = torch.arange(0,V,dtype=torch.long,device=vertices.device) #V
dest[collapses[:,1]] = dest[collapses[:,0]]
faces = dest[faces] #F,3
c0,c1,c2 = faces.unbind(dim=-1)
collapsed = (c0==c1).logical_or_(c1==c2).logical_or_(c0==c2)
faces[collapsed] = 0
return vertices,faces
def calc_face_collapses(
vertices:torch.Tensor, #V,3 first unused
faces:torch.Tensor, #F,3 long, 0 for unused
edges:torch.Tensor, #E,2 long 0 for unused, lower vertex index first
face_to_edge:torch.Tensor, #F,3 long 0 for unused
edge_length:torch.Tensor, #E
face_normals:torch.Tensor, #F,3
vertex_normals:torch.Tensor, #V,3 first unused
min_edge_length:torch.Tensor=None, #V
area_ratio = 0.5, #collapse if area < min_edge_length**2 * area_ratio
shortest_probability = 0.8
)->torch.Tensor: #E edges to collapse
E = edges.shape[0]
F = faces.shape[0]
# face flips
ref_normals = calc_face_ref_normals(faces,vertex_normals,normalize=False) #F,3
face_collapses = (face_normals*ref_normals).sum(dim=-1)<0 #F
# small faces
if min_edge_length is not None:
min_face_length = min_edge_length[faces].mean(dim=-1) #F
min_area = min_face_length**2 * area_ratio #F
face_collapses.logical_or_(face_normals.norm(dim=-1) < min_area*2) #F
face_collapses[0] = False
# faces to edges
face_length = edge_length[face_to_edge] #F,3
if shortest_probability<1:
#select shortest edge with shortest_probability chance
randlim = round(2/(1-shortest_probability))
rand_ind = torch.randint(0,randlim,size=(F,),device=faces.device).clamp_max_(2) #selected edge local index in face
sort_ind = torch.argsort(face_length,dim=-1,descending=True) #F,3
local_ind = sort_ind.gather(dim=-1,index=rand_ind[:,None])
else:
local_ind = torch.argmin(face_length,dim=-1)[:,None] #F,1 0...2 shortest edge local index in face
edge_ind = face_to_edge.gather(dim=1,index=local_ind)[:,0] #F 0...E selected edge global index
edge_collapses = torch.zeros(E,dtype=torch.long,device=vertices.device)
edge_collapses.scatter_add_(dim=0,index=edge_ind,src=face_collapses.long())
return edge_collapses.bool()
def flip_edges(
vertices:torch.Tensor, #V,3 first unused
faces:torch.Tensor, #F,3 long, first must be 0, 0 for unused
edges:torch.Tensor, #E,2 long, first must be 0, 0 for unused, lower vertex index first
edge_to_face:torch.Tensor, #E,[left,right],[face,side]
with_border:bool=True, #handle border edges (D=4 instead of D=6)
with_normal_check:bool=True, #check face normal flips
stable:bool=False, #only for unit testing
):
V = vertices.shape[0]
E = edges.shape[0]
device=vertices.device
vertex_degree = torch.zeros(V,dtype=torch.long,device=device) #V long
vertex_degree.scatter_(dim=0,index=edges.reshape(E*2),value=1,reduce='add')
neighbor_corner = (edge_to_face[:,:,1] + 2) % 3 #go from side to corner
neighbors = faces[edge_to_face[:,:,0],neighbor_corner] #E,LR=2
edge_is_inside = neighbors.all(dim=-1) #E
if with_border:
# inside vertices should have D=6, border edges D=4, so we subtract 2 for all inside vertices
# need to use float for masks in order to use scatter(reduce='multiply')
vertex_is_inside = torch.ones(V,2,dtype=torch.float32,device=vertices.device) #V,2 float
src = edge_is_inside.type(torch.float32)[:,None].expand(E,2) #E,2 float
vertex_is_inside.scatter_(dim=0,index=edges,src=src,reduce='multiply')
vertex_is_inside = vertex_is_inside.prod(dim=-1,dtype=torch.long) #V long
vertex_degree -= 2 * vertex_is_inside #V long
neighbor_degrees = vertex_degree[neighbors] #E,LR=2
edge_degrees = vertex_degree[edges] #E,2
#
# loss = Sum_over_affected_vertices((new_degree-6)**2)
# loss_change = Sum_over_neighbor_vertices((degree+1-6)**2-(degree-6)**2)
# + Sum_over_edge_vertices((degree-1-6)**2-(degree-6)**2)
# = 2 * (2 + Sum_over_neighbor_vertices(degree) - Sum_over_edge_vertices(degree))
#
loss_change = 2 + neighbor_degrees.sum(dim=-1) - edge_degrees.sum(dim=-1) #E
candidates = torch.logical_and(loss_change<0, edge_is_inside) #E
loss_change = loss_change[candidates] #E'
if loss_change.shape[0]==0:
return
edges_neighbors = torch.concat((edges[candidates],neighbors[candidates]),dim=-1) #E',4
_,order = loss_change.sort(descending=True, stable=stable) #E'
rank = torch.zeros_like(order)
rank[order] = torch.arange(0,len(rank),device=rank.device)
vertex_rank = torch.zeros((V,4),dtype=torch.long,device=device) #V,4
torch_scatter.scatter_max(src=rank[:,None].expand(-1,4),index=edges_neighbors,dim=0,out=vertex_rank)
vertex_rank,_ = vertex_rank.max(dim=-1) #V
neighborhood_rank,_ = vertex_rank[edges_neighbors].max(dim=-1) #E'
flip = rank==neighborhood_rank #E'
if with_normal_check:
# cl-<-----e1 e0,e1...edge, e0<e1
# | /A L,R....left and right face
# | L / | both triangles ordered counter clockwise
# | / R | normals pointing out of screen
# V/ |
# e0---->-cr
v = vertices[edges_neighbors] #E",4,3
v = v - v[:,0:1] #make relative to e0
e1 = v[:,1]
cl = v[:,2]
cr = v[:,3]
n = torch.cross(e1,cl) + torch.cross(cr,e1) #sum of old normal vectors
flip.logical_and_(torch.sum(n*torch.cross(cr,cl),dim=-1)>0) #first new face
flip.logical_and_(torch.sum(n*torch.cross(cl-e1,cr-e1),dim=-1)>0) #second new face
flip_edges_neighbors = edges_neighbors[flip] #E",4
flip_edge_to_face = edge_to_face[candidates,:,0][flip] #E",2
flip_faces = flip_edges_neighbors[:,[[0,3,2],[1,2,3]]] #E",2,3
faces.scatter_(dim=0,index=flip_edge_to_face.reshape(-1,1).expand(-1,3),src=flip_faces.reshape(-1,3))
gen_3d_modules/Unique3D/mesh_reconstruction/render.py
+163
View File
@@ -0,0 +1,163 @@
# modified from https://github.com/Profactor/continuous-remeshing
import nvdiffrast.torch as dr
import torch
from typing import Tuple
def _warmup(glctx, device=None):
device = 'cuda' if device is None else device
#windows workaround for https://github.com/NVlabs/nvdiffrast/issues/59
def tensor(*args, **kwargs):
return torch.tensor(*args, device=device, **kwargs)
pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32)
tri = tensor([[0, 1, 2]], dtype=torch.int32)
dr.rasterize(glctx, pos, tri, resolution=[256, 256])
glctx = dr.RasterizeGLContext(output_db=False, device="cuda")
class NormalsRenderer:
_glctx:dr.RasterizeGLContext = None
def __init__(
self,
mv: torch.Tensor, #C,4,4
proj: torch.Tensor, #C,4,4
image_size: Tuple[int,int],
mvp = None,
device=None,
):
if mvp is None:
self._mvp = proj @ mv #C,4,4
else:
self._mvp = mvp
self._image_size = image_size
self._glctx = glctx
_warmup(self._glctx, device)
def render(self,
vertices: torch.Tensor, #V,3 float
normals: torch.Tensor, #V,3 float in [-1, 1]
faces: torch.Tensor, #F,3 long
) ->torch.Tensor: #C,H,W,4
V = vertices.shape[0]
faces = faces.type(torch.int32)
vert_hom = torch.cat((vertices, torch.ones(V,1,device=vertices.device)),axis=-1) #V,3 -> V,4
vertices_clip = vert_hom @ self._mvp.transpose(-2,-1) #C,V,4
rast_out,_ = dr.rasterize(self._glctx, vertices_clip, faces, resolution=self._image_size, grad_db=False) #C,H,W,4
vert_col = (normals+1)/2 #V,3
col,_ = dr.interpolate(vert_col, rast_out, faces) #C,H,W,3
alpha = torch.clamp(rast_out[..., -1:], max=1) #C,H,W,1
col = torch.concat((col,alpha),dim=-1) #C,H,W,4
col = dr.antialias(col, rast_out, vertices_clip, faces) #C,H,W,4
return col #C,H,W,4
from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh.shader import ShaderBase
from pytorch3d.renderer import (
RasterizationSettings,
MeshRendererWithFragments,
TexturesVertex,
MeshRasterizer,
BlendParams,
FoVOrthographicCameras,
look_at_view_transform,
hard_rgb_blend,
)
class VertexColorShader(ShaderBase):
def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
blend_params = kwargs.get("blend_params", self.blend_params)
texels = meshes.sample_textures(fragments)
return hard_rgb_blend(texels, fragments, blend_params)
def render_mesh_vertex_color(mesh, cameras, H, W, blur_radius=0.0, faces_per_pixel=1, bkgd=(0., 0., 0.), dtype=torch.float32, device="cuda"):
if len(mesh) != len(cameras):
if len(cameras) % len(mesh) == 0:
mesh = mesh.extend(len(cameras))
else:
raise NotImplementedError()
# render requires everything in float16 or float32
input_dtype = dtype
blend_params = BlendParams(1e-4, 1e-4, bkgd)
# Define the settings for rasterization and shading
raster_settings = RasterizationSettings(
image_size=(H, W),
blur_radius=blur_radius,
faces_per_pixel=faces_per_pixel,
clip_barycentric_coords=True,
bin_size=None,
max_faces_per_bin=None,
)
# Create a renderer by composing a rasterizer and a shader
# We simply render vertex colors through the custom VertexColorShader (no lighting, materials are used)
renderer = MeshRendererWithFragments(
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=raster_settings
),
shader=VertexColorShader(
device=device,
cameras=cameras,
blend_params=blend_params
)
)
# render RGB and depth, get mask
with torch.autocast(dtype=input_dtype, device_type=torch.device(device).type):
images, _ = renderer(mesh)
return images # BHW4
class Pytorch3DNormalsRenderer: # 100 times slower!!!
def __init__(self, cameras, image_size, device):
self.cameras = cameras.to(device)
self._image_size = image_size
self.device = device
def render(self,
vertices: torch.Tensor, #V,3 float
normals: torch.Tensor, #V,3 float in [-1, 1]
faces: torch.Tensor, #F,3 long
) ->torch.Tensor: #C,H,W,4
mesh = Meshes(verts=[vertices], faces=[faces], textures=TexturesVertex(verts_features=[(normals + 1) / 2])).to(self.device)
return render_mesh_vertex_color(mesh, self.cameras, self._image_size[0], self._image_size[1], device=self.device)
def save_tensor_to_img(tensor, save_dir):
from PIL import Image
import numpy as np
for idx, img in enumerate(tensor):
img = img[..., :3].cpu().numpy()
img = (img * 255).astype(np.uint8)
img = Image.fromarray(img)
img.save(save_dir + f"{idx}.png")
if __name__ == "__main__":
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from mesh_reconstruction.func import make_star_cameras_orthographic, make_star_cameras_orthographic_py3d
cameras = make_star_cameras_orthographic_py3d([0, 270, 180, 90], device="cuda", focal=1., dist=4.0)
mv,proj = make_star_cameras_orthographic(4, 1)
resolution = 1024
renderer1 = NormalsRenderer(mv,proj, [resolution,resolution], device="cuda")
renderer2 = Pytorch3DNormalsRenderer(cameras, [resolution,resolution], device="cuda")
vertices = torch.tensor([[0,0,0],[0,0,1],[0,1,0],[1,0,0]], device="cuda", dtype=torch.float32)
normals = torch.tensor([[-1,-1,-1],[1,-1,-1],[-1,-1,1],[-1,1,-1]], device="cuda", dtype=torch.float32)
faces = torch.tensor([[0,1,2],[0,1,3],[0,2,3],[1,2,3]], device="cuda", dtype=torch.long)
import time
t0 = time.time()
r1 = renderer1.render(vertices, normals, faces)
print("time r1:", time.time() - t0)
t0 = time.time()
r2 = renderer2.render(vertices, normals, faces)
print("time r2:", time.time() - t0)
for i in range(4):
print((r1[i]-r2[i]).abs().mean(), (r1[i]+r2[i]).abs().mean())
gen_3d_modules/Unique3D/scripts/__init__.py
View File
gen_3d_modules/Unique3D/scripts/all_typing.py
+42
View File
@@ -0,0 +1,42 @@
# code from https://github.com/threestudio-project
"""
This module contains type annotations for the project, using
1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects
2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors
Two types of typing checking can be used:
1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode)
2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking)
"""
# Basic types
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Literal,
NamedTuple,
NewType,
Optional,
Sized,
Tuple,
Type,
TypeVar,
Union,
)
# Tensor dtype
# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
# Config type
from omegaconf import DictConfig
# PyTorch Tensor type
from torch import Tensor
# Runtime type checking decorator
from typeguard import typechecked as typechecker
gen_3d_modules/Unique3D/scripts/mesh_init.py
+145
View File
@@ -0,0 +1,145 @@
from PIL import Image
import torch
import numpy as np
from pytorch3d.structures import Meshes, join_meshes_as_scene
from pytorch3d.renderer import TexturesVertex
from .utils import meshlab_mesh_to_py3dmesh, py3dmesh_to_meshlab_mesh
import pymeshlab
_MAX_THREAD = 8
def fast_geo(front_normal: Image.Image, back_normal: Image.Image, side_normal: Image.Image, clamp=0., init_type="std"):
front_normal = front_normal.resize((192, 192))
back_normal = back_normal.resize((192, 192))
side_normal = side_normal.resize((192, 192))
# build mesh with front back projection # ~3s
side_w_over_h = calc_w_over_h(side_normal)
mesh_front = build_mesh(front_normal, front_normal, clamp_min=clamp, scale=side_w_over_h, init_type=init_type)
mesh_back = build_mesh(back_normal, back_normal, is_back=True, clamp_min=clamp, scale=side_w_over_h, init_type=init_type)
meshes = join_meshes_as_scene([mesh_front, mesh_back])
meshes = fix_border_with_pymeshlab_fast(meshes, poissson_depth=6, simplification=2000)
return meshes
# rgb and depth to mesh
def get_ortho_ray_directions_origins(W, H, use_pixel_centers=True, device="cuda"):
pixel_center = 0.5 if use_pixel_centers else 0
i, j = np.meshgrid(
np.arange(W, dtype=np.float32) + pixel_center,
np.arange(H, dtype=np.float32) + pixel_center,
indexing='xy'
)
i, j = torch.from_numpy(i).to(device), torch.from_numpy(j).to(device)
origins = torch.stack([(i/W-0.5)*2, (j/H-0.5)*2 * H / W, torch.zeros_like(i)], dim=-1) # W, H, 3
directions = torch.stack([torch.zeros_like(i), torch.zeros_like(j), torch.ones_like(i)], dim=-1) # W, H, 3
return origins, directions
def depth_and_color_to_mesh(rgb_BCHW, pred_HWC, valid_HWC=None, is_back=False):
if valid_HWC is None:
valid_HWC = torch.ones_like(pred_HWC).bool()
H, W = rgb_BCHW.shape[-2:]
rgb_BCHW = rgb_BCHW.flip(-2)
pred_HWC = pred_HWC.flip(0)
valid_HWC = valid_HWC.flip(0)
rays_o, rays_d = get_ortho_ray_directions_origins(W, H, device=rgb_BCHW.device)
verts = rays_o + rays_d * pred_HWC # [H, W, 3]
verts = verts.reshape(-1, 3) # [V, 3]
indexes = torch.arange(H * W).reshape(H, W).to(rgb_BCHW.device)
faces1 = torch.stack([indexes[:-1, :-1], indexes[:-1, 1:], indexes[1:, :-1]], dim=-1)
# faces1_valid = valid_HWC[:-1, :-1] | valid_HWC[:-1, 1:] | valid_HWC[1:, :-1]
faces1_valid = valid_HWC[:-1, :-1] & valid_HWC[:-1, 1:] & valid_HWC[1:, :-1]
faces2 = torch.stack([indexes[1:, 1:], indexes[1:, :-1], indexes[:-1, 1:]], dim=-1)
# faces2_valid = valid_HWC[1:, 1:] | valid_HWC[1:, :-1] | valid_HWC[:-1, 1:]
faces2_valid = valid_HWC[1:, 1:] & valid_HWC[1:, :-1] & valid_HWC[:-1, 1:]
faces = torch.cat([faces1[faces1_valid.expand_as(faces1)].reshape(-1, 3), faces2[faces2_valid.expand_as(faces2)].reshape(-1, 3)], dim=0) # (F, 3)
colors = (rgb_BCHW[0].permute((1,2,0)) / 2 + 0.5).reshape(-1, 3) # (V, 3)
if is_back:
verts = verts * torch.tensor([-1, 1, -1], dtype=verts.dtype, device=verts.device)
used_verts = faces.unique()
old_to_new_mapping = torch.zeros_like(verts[..., 0]).long()
old_to_new_mapping[used_verts] = torch.arange(used_verts.shape[0], device=verts.device)
new_faces = old_to_new_mapping[faces]
mesh = Meshes(verts=[verts[used_verts]], faces=[new_faces], textures=TexturesVertex(verts_features=[colors[used_verts]]))
return mesh
def normalmap_to_depthmap(normal_np):
from .normal_to_height_map import estimate_height_map
height = estimate_height_map(normal_np, raw_values=True, thread_count=_MAX_THREAD, target_iteration_count=96)
return height
def transform_back_normal_to_front(normal_pil):
arr = np.array(normal_pil) # in [0, 255]
arr[..., 0] = 255-arr[..., 0]
arr[..., 2] = 255-arr[..., 2]
return Image.fromarray(arr.astype(np.uint8))
def calc_w_over_h(normal_pil):
if isinstance(normal_pil, Image.Image):
arr = np.array(normal_pil)
else:
assert isinstance(normal_pil, np.ndarray)
arr = normal_pil
if arr.shape[-1] == 4:
alpha = arr[..., -1] / 255.
alpha[alpha >= 0.5] = 1
alpha[alpha < 0.5] = 0
else:
alpha = ~(arr.min(axis=-1) >= 250)
h_min, w_min = np.min(np.where(alpha), axis=1)
h_max, w_max = np.max(np.where(alpha), axis=1)
return (w_max - w_min) / (h_max - h_min)
def build_mesh(normal_pil, rgb_pil, is_back=False, clamp_min=-1, scale=0.3, init_type="std", offset=0):
if is_back:
normal_pil = transform_back_normal_to_front(normal_pil)
normal_img = np.array(normal_pil)
rgb_img = np.array(rgb_pil)
if normal_img.shape[-1] == 4:
valid_HWC = normal_img[..., [3]] / 255
elif rgb_img.shape[-1] == 4:
valid_HWC = rgb_img[..., [3]] / 255
else:
raise ValueError("invalid input, either normal or rgb should have alpha channel")
real_height_pix = np.max(np.where(valid_HWC>0.5)[0]) - np.min(np.where(valid_HWC>0.5)[0])
heights = normalmap_to_depthmap(normal_img)
rgb_BCHW = torch.from_numpy(rgb_img[..., :3] / 255.).permute((2,0,1))[None]
valid_HWC[valid_HWC < 0.5] = 0
valid_HWC[valid_HWC >= 0.5] = 1
valid_HWC = torch.from_numpy(valid_HWC).bool()
if init_type == "std":
# accurate but not stable
pred_HWC = torch.from_numpy(heights / heights.max() * (real_height_pix / heights.shape[0]) * scale * 2).float()[..., None]
elif init_type == "thin":
heights = heights - heights.min()
heights = (heights / heights.max() * 0.2)
pred_HWC = torch.from_numpy(heights * scale).float()[..., None]
else:
# stable but not accurate
heights = heights - heights.min()
heights = (heights / heights.max() * (1-offset)) + offset # to [0.2, 1]
pred_HWC = torch.from_numpy(heights * scale).float()[..., None]
# set the boarder pixels to 0 height
import cv2
# edge filter
edge = cv2.Canny((valid_HWC[..., 0] * 255).numpy().astype(np.uint8), 0, 255)
edge = torch.from_numpy(edge).bool()[..., None]
pred_HWC[edge] = 0
valid_HWC[pred_HWC < clamp_min] = False
return depth_and_color_to_mesh(rgb_BCHW.cuda(), pred_HWC.cuda(), valid_HWC.cuda(), is_back)
def fix_border_with_pymeshlab_fast(meshes: Meshes, poissson_depth=6, simplification=0):
ms = pymeshlab.MeshSet()
ms.add_mesh(py3dmesh_to_meshlab_mesh(meshes), "cube_vcolor_mesh")
if simplification > 0:
ms.apply_filter('meshing_decimation_quadric_edge_collapse', targetfacenum=simplification, preservetopology=True)
ms.apply_filter('generate_surface_reconstruction_screened_poisson', threads = 6, depth = poissson_depth, preclean = True)
if simplification > 0:
ms.apply_filter('meshing_decimation_quadric_edge_collapse', targetfacenum=simplification, preservetopology=True)
return meshlab_mesh_to_py3dmesh(ms.current_mesh())
gen_3d_modules/Unique3D/scripts/normal_to_height_map.py
+203
View File
@@ -0,0 +1,203 @@
# code modified from https://github.com/YertleTurtleGit/depth-from-normals
import numpy as np
import cv2 as cv
from multiprocessing.pool import ThreadPool as Pool
from multiprocessing import cpu_count
from typing import Tuple, List, Union
import numba
def calculate_gradients(
normals: np.ndarray, mask: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
horizontal_angle_map = np.arccos(np.clip(normals[:, :, 0], -1, 1))
left_gradients = np.zeros(normals.shape[:2])
left_gradients[mask != 0] = (1 - np.sin(horizontal_angle_map[mask != 0])) * np.sign(
horizontal_angle_map[mask != 0] - np.pi / 2
)
vertical_angle_map = np.arccos(np.clip(normals[:, :, 1], -1, 1))
top_gradients = np.zeros(normals.shape[:2])
top_gradients[mask != 0] = -(1 - np.sin(vertical_angle_map[mask != 0])) * np.sign(
vertical_angle_map[mask != 0] - np.pi / 2
)
return left_gradients, top_gradients
@numba.jit(nopython=True)
def integrate_gradient_field(
gradient_field: np.ndarray, axis: int, mask: np.ndarray
) -> np.ndarray:
heights = np.zeros(gradient_field.shape)
for d1 in numba.prange(heights.shape[1 - axis]):
sum_value = 0
for d2 in range(heights.shape[axis]):
coordinates = (d1, d2) if axis == 1 else (d2, d1)
if mask[coordinates] != 0:
sum_value = sum_value + gradient_field[coordinates]
heights[coordinates] = sum_value
else:
sum_value = 0
return heights
def calculate_heights(
left_gradients: np.ndarray, top_gradients, mask: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
left_heights = integrate_gradient_field(left_gradients, 1, mask)
right_heights = np.fliplr(
integrate_gradient_field(np.fliplr(-left_gradients), 1, np.fliplr(mask))
)
top_heights = integrate_gradient_field(top_gradients, 0, mask)
bottom_heights = np.flipud(
integrate_gradient_field(np.flipud(-top_gradients), 0, np.flipud(mask))
)
return left_heights, right_heights, top_heights, bottom_heights
def combine_heights(*heights: np.ndarray) -> np.ndarray:
return np.mean(np.stack(heights, axis=0), axis=0)
def rotate(matrix: np.ndarray, angle: float) -> np.ndarray:
h, w = matrix.shape[:2]
center = (w / 2, h / 2)
rotation_matrix = cv.getRotationMatrix2D(center, angle, 1.0)
corners = cv.transform(
np.array([[[0, 0], [w, 0], [w, h], [0, h]]]), rotation_matrix
)[0]
_, _, w, h = cv.boundingRect(corners)
rotation_matrix[0, 2] += w / 2 - center[0]
rotation_matrix[1, 2] += h / 2 - center[1]
result = cv.warpAffine(matrix, rotation_matrix, (w, h), flags=cv.INTER_LINEAR)
return result
def rotate_vector_field_normals(normals: np.ndarray, angle: float) -> np.ndarray:
angle = np.radians(angle)
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
rotated_normals = np.empty_like(normals)
rotated_normals[:, :, 0] = (
normals[:, :, 0] * cos_angle - normals[:, :, 1] * sin_angle
)
rotated_normals[:, :, 1] = (
normals[:, :, 0] * sin_angle + normals[:, :, 1] * cos_angle
)
return rotated_normals
def centered_crop(image: np.ndarray, target_resolution: Tuple[int, int]) -> np.ndarray:
return image[
(image.shape[0] - target_resolution[0])
// 2 : (image.shape[0] - target_resolution[0])
// 2
+ target_resolution[0],
(image.shape[1] - target_resolution[1])
// 2 : (image.shape[1] - target_resolution[1])
// 2
+ target_resolution[1],
]
def integrate_vector_field(
vector_field: np.ndarray,
mask: np.ndarray,
target_iteration_count: int,
thread_count: int,
) -> np.ndarray:
shape = vector_field.shape[:2]
angles = np.linspace(0, 90, target_iteration_count, endpoint=False)
def integrate_vector_field_angles(angles: List[float]) -> np.ndarray:
all_combined_heights = np.zeros(shape)
for angle in angles:
rotated_vector_field = rotate_vector_field_normals(
rotate(vector_field, angle), angle
)
rotated_mask = rotate(mask, angle)
left_gradients, top_gradients = calculate_gradients(
rotated_vector_field, rotated_mask
)
(
left_heights,
right_heights,
top_heights,
bottom_heights,
) = calculate_heights(left_gradients, top_gradients, rotated_mask)
combined_heights = combine_heights(
left_heights, right_heights, top_heights, bottom_heights
)
combined_heights = centered_crop(rotate(combined_heights, -angle), shape)
all_combined_heights += combined_heights / len(angles)
return all_combined_heights
with Pool(processes=thread_count) as pool:
heights = pool.map(
integrate_vector_field_angles,
np.array(
np.array_split(angles, thread_count),
dtype=object,
),
)
pool.close()
pool.join()
isotropic_height = np.zeros(shape)
for height in heights:
isotropic_height += height / thread_count
return isotropic_height
def estimate_height_map(
normal_map: np.ndarray,
mask: Union[np.ndarray, None] = None,
height_divisor: float = 1,
target_iteration_count: int = 250,
thread_count: int = cpu_count(),
raw_values: bool = False,
) -> np.ndarray:
if mask is None:
if normal_map.shape[-1] == 4:
mask = normal_map[:, :, 3] / 255
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
else:
mask = np.ones(normal_map.shape[:2], dtype=np.uint8)
normals = ((normal_map[:, :, :3].astype(np.float64) / 255) - 0.5) * 2
heights = integrate_vector_field(
normals, mask, target_iteration_count, thread_count
)
if raw_values:
return heights
heights /= height_divisor
heights[mask > 0] += 1 / 2
heights[mask == 0] = 1 / 2
heights *= 2**16 - 1
if np.min(heights) < 0 or np.max(heights) > 2**16 - 1:
raise OverflowError("Height values are clipping.")
heights = np.clip(heights, 0, 2**16 - 1)
heights = heights.astype(np.uint16)
return heights
gen_3d_modules/Unique3D/scripts/project_mesh.py
+378
View File
@@ -0,0 +1,378 @@
from typing import List
import torch
import numpy as np
from PIL import Image
from pytorch3d.renderer.cameras import look_at_view_transform, OrthographicCameras, CamerasBase
from pytorch3d.renderer.mesh.rasterizer import Fragments
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
RasterizationSettings,
TexturesVertex,
FoVPerspectiveCameras,
FoVOrthographicCameras,
)
from pytorch3d.renderer import MeshRasterizer
def get_camera(world_to_cam, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'):
# pytorch3d expects transforms as row-vectors, so flip rotation: https://github.com/facebookresearch/pytorch3d/issues/1183
R = world_to_cam[:3, :3].t()[None, ...]
T = world_to_cam[:3, 3][None, ...]
if cam_type == 'fov':
camera = FoVPerspectiveCameras(device=world_to_cam.device, R=R, T=T, fov=fov_in_degrees, degrees=True)
else:
focal_length = 1 / focal_length
camera = FoVOrthographicCameras(device=world_to_cam.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length)
return camera
def render_pix2faces_py3d(meshes, cameras, H=512, W=512, blur_radius=0.0, faces_per_pixel=1):
"""
Renders pix2face of visible faces.
:param mesh: Pytorch3d.structures.Meshes
:param cameras: pytorch3d.renderer.Cameras
:param H: target image height
:param W: target image width
:param blur_radius: Float distance in the range [0, 2] used to expand the face
bounding boxes for rasterization. Setting blur radius
results in blurred edges around the shape instead of a
hard boundary. Set to 0 for no blur.
:param faces_per_pixel: (int) Number of faces to keep track of per pixel.
We return the nearest faces_per_pixel faces along the z-axis.
"""
# Define the settings for rasterization and shading
raster_settings = RasterizationSettings(
image_size=(H, W),
blur_radius=blur_radius,
faces_per_pixel=faces_per_pixel
)
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=raster_settings
)
fragments: Fragments = rasterizer(meshes, cameras=cameras)
return {
"pix_to_face": fragments.pix_to_face[..., 0],
}
import nvdiffrast.torch as dr
def _warmup(glctx, device=None):
device = 'cuda' if device is None else device
#windows workaround for https://github.com/NVlabs/nvdiffrast/issues/59
def tensor(*args, **kwargs):
return torch.tensor(*args, device=device, **kwargs)
pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32)
tri = tensor([[0, 1, 2]], dtype=torch.int32)
dr.rasterize(glctx, pos, tri, resolution=[256, 256])
class Pix2FacesRenderer:
def __init__(self, device="cuda"):
self._glctx = dr.RasterizeGLContext(output_db=False, device=device)
self.device = device
_warmup(self._glctx, device)
def transform_vertices(self, meshes: Meshes, cameras: CamerasBase):
vertices = cameras.transform_points_ndc(meshes.verts_padded())
perspective_correct = cameras.is_perspective()
znear = cameras.get_znear()
if isinstance(znear, torch.Tensor):
znear = znear.min().item()
z_clip = None if not perspective_correct or znear is None else znear / 2
if z_clip:
vertices = vertices[vertices[..., 2] >= cameras.get_znear()][None] # clip
vertices = vertices * torch.tensor([-1, -1, 1]).to(vertices)
vertices = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1).to(torch.float32)
return vertices
def render_pix2faces_nvdiff(self, meshes: Meshes, cameras: CamerasBase, H=512, W=512):
meshes = meshes.to(self.device)
cameras = cameras.to(self.device)
vertices = self.transform_vertices(meshes, cameras)
faces = meshes.faces_packed().to(torch.int32)
rast_out,_ = dr.rasterize(self._glctx, vertices, faces, resolution=(H, W), grad_db=False) #C,H,W,4
pix_to_face = rast_out[..., -1].to(torch.int32) - 1
return pix_to_face
pix2faces_renderer = Pix2FacesRenderer()
def get_visible_faces(meshes: Meshes, cameras: CamerasBase, resolution=1024):
# pix_to_face = render_pix2faces_py3d(meshes, cameras, H=resolution, W=resolution)['pix_to_face']
pix_to_face = pix2faces_renderer.render_pix2faces_nvdiff(meshes, cameras, H=resolution, W=resolution)
unique_faces = torch.unique(pix_to_face.flatten())
unique_faces = unique_faces[unique_faces != -1]
return unique_faces
def project_color(meshes: Meshes, cameras: CamerasBase, pil_image: Image.Image, use_alpha=True, eps=0.05, resolution=1024, device="cuda") -> dict:
"""
Projects color from a given image onto a 3D mesh.
Args:
meshes (pytorch3d.structures.Meshes): The 3D mesh object.
cameras (pytorch3d.renderer.cameras.CamerasBase): The camera object.
pil_image (PIL.Image.Image): The input image.
use_alpha (bool, optional): Whether to use the alpha channel of the image. Defaults to True.
eps (float, optional): The threshold for selecting visible faces. Defaults to 0.05.
resolution (int, optional): The resolution of the projection. Defaults to 1024.
device (str, optional): The device to use for computation. Defaults to "cuda".
debug (bool, optional): Whether to save debug images. Defaults to False.
Returns:
dict: A dictionary containing the following keys:
- "new_texture" (TexturesVertex): The updated texture with interpolated colors.
- "valid_verts" (Tensor of [M,3]): The indices of the vertices being projected.
- "valid_colors" (Tensor of [M,3]): The interpolated colors for the valid vertices.
"""
meshes = meshes.to(device)
cameras = cameras.to(device)
image = torch.from_numpy(np.array(pil_image.convert("RGBA")) / 255.).permute((2, 0, 1)).float().to(device) # in CHW format of [0, 1.]
unique_faces = get_visible_faces(meshes, cameras, resolution=resolution)
# visible faces
faces_normals = meshes.faces_normals_packed()[unique_faces]
faces_normals = faces_normals / faces_normals.norm(dim=1, keepdim=True)
world_points = cameras.unproject_points(torch.tensor([[[0., 0., 0.1], [0., 0., 0.2]]]).to(device))[0]
view_direction = world_points[1] - world_points[0]
view_direction = view_direction / view_direction.norm(dim=0, keepdim=True)
# find invalid faces
cos_angles = (faces_normals * view_direction).sum(dim=1)
assert cos_angles.mean() < 0, f"The view direction is not correct. cos_angles.mean()={cos_angles.mean()}"
selected_faces = unique_faces[cos_angles < -eps]
# find verts
faces = meshes.faces_packed()[selected_faces] # [N, 3]
verts = torch.unique(faces.flatten()) # [N, 1]
verts_coordinates = meshes.verts_packed()[verts] # [N, 3]
# compute color
pt_tensor = cameras.transform_points(verts_coordinates)[..., :2] # NDC space points
valid = ~((pt_tensor.isnan()|(pt_tensor<-1)|(1<pt_tensor)).any(dim=1)) # checked, correct
valid_pt = pt_tensor[valid, :]
valid_idx = verts[valid]
valid_color = torch.nn.functional.grid_sample(image[None].flip((-1, -2)), valid_pt[None, :, None, :], align_corners=False, padding_mode="reflection", mode="bilinear")[0, :, :, 0].T.clamp(0, 1) # [N, 4], note that bicubic may give invalid value
alpha, valid_color = valid_color[:, 3:], valid_color[:, :3]
if not use_alpha:
alpha = torch.ones_like(alpha)
# modify color
old_colors = meshes.textures.verts_features_packed()
old_colors[valid_idx] = valid_color * alpha + old_colors[valid_idx] * (1 - alpha)
new_texture = TexturesVertex(verts_features=[old_colors])
valid_verts_normals = meshes.verts_normals_packed()[valid_idx]
valid_verts_normals = valid_verts_normals / valid_verts_normals.norm(dim=1, keepdim=True).clamp_min(0.001)
cos_angles = (valid_verts_normals * view_direction).sum(dim=1)
return {
"new_texture": new_texture,
"valid_verts": valid_idx,
"valid_colors": valid_color,
"valid_alpha": alpha,
"cos_angles": cos_angles,
}
def complete_unseen_vertex_color(meshes: Meshes, valid_index: torch.Tensor) -> dict:
"""
meshes: the mesh with vertex color to be completed.
valid_index: the index of the valid vertices, where valid means colors are fixed. [V, 1]
"""
valid_index = valid_index.to(meshes.device)
colors = meshes.textures.verts_features_packed() # [V, 3]
V = colors.shape[0]
invalid_index = torch.ones_like(colors[:, 0]).bool() # [V]
invalid_index[valid_index] = False
invalid_index = torch.arange(V).to(meshes.device)[invalid_index]
L = meshes.laplacian_packed()
E = torch.sparse_coo_tensor(torch.tensor([list(range(V))] * 2), torch.ones((V,)), size=(V, V)).to(meshes.device)
L = L + E
# E = torch.eye(V, layout=torch.sparse_coo, device=meshes.device)
# L = L + E
colored_count = torch.ones_like(colors[:, 0]) # [V]
colored_count[invalid_index] = 0
L_invalid = torch.index_select(L, 0, invalid_index) # sparse [IV, V]
total_colored = colored_count.sum()
coloring_round = 0
stage = "uncolored"
from tqdm import tqdm
pbar = tqdm(miniters=100)
while stage == "uncolored" or coloring_round > 0:
new_color = torch.matmul(L_invalid, colors * colored_count[:, None]) # [IV, 3]
new_count = torch.matmul(L_invalid, colored_count)[:, None] # [IV, 1]
colors[invalid_index] = torch.where(new_count > 0, new_color / new_count, colors[invalid_index])
colored_count[invalid_index] = (new_count[:, 0] > 0).float()
new_total_colored = colored_count.sum()
if new_total_colored > total_colored:
total_colored = new_total_colored
coloring_round += 1
else:
stage = "colored"
coloring_round -= 1
pbar.update(1)
if coloring_round > 10000:
print("coloring_round > 10000, break")
break
assert not torch.isnan(colors).any()
meshes.textures = TexturesVertex(verts_features=[colors])
return meshes
def multiview_color_projection(meshes: Meshes, image_list: List[Image.Image], cameras_list: List[CamerasBase]=None, camera_focal: float = 2 / 1.35, weights=None, eps=0.05, resolution=1024, device="cuda", reweight_with_cosangle="square", use_alpha=True, confidence_threshold=0.1, complete_unseen=False, below_confidence_strategy="smooth") -> Meshes:
"""
Projects color from a given image onto a 3D mesh.
Args:
meshes (pytorch3d.structures.Meshes): The 3D mesh object, only one mesh.
image_list (PIL.Image.Image): List of images.
cameras_list (list): List of cameras.
camera_focal (float, optional): The focal length of the camera, if cameras_list is not passed. Defaults to 2 / 1.35.
weights (list, optional): List of weights for each image, for ['front', 'front_right', 'right', 'back', 'left', 'front_left']. Defaults to None.
eps (float, optional): The threshold for selecting visible faces. Defaults to 0.05.
resolution (int, optional): The resolution of the projection. Defaults to 1024.
device (str, optional): The device to use for computation. Defaults to "cuda".
reweight_with_cosangle (str, optional): Whether to reweight the color with the angle between the view direction and the vertex normal. Defaults to None.
use_alpha (bool, optional): Whether to use the alpha channel of the image. Defaults to True.
confidence_threshold (float, optional): The threshold for the confidence of the projected color, if final projection weight is less than this, we will use the original color. Defaults to 0.1.
complete_unseen (bool, optional): Whether to complete the unseen vertex color using laplacian. Defaults to False.
Returns:
Meshes: the colored mesh
"""
# 1. preprocess inputs
if image_list is None:
raise ValueError("image_list is None")
if cameras_list is None:
if len(image_list) == 8:
cameras_list = get_8view_cameras(device, focal=camera_focal)
elif len(image_list) == 6:
cameras_list = get_6view_cameras(device, focal=camera_focal)
elif len(image_list) == 4:
cameras_list = get_4view_cameras(device, focal=camera_focal)
elif len(image_list) == 2:
cameras_list = get_2view_cameras(device, focal=camera_focal)
else:
raise ValueError("cameras_list is None, and can not be guessed from image_list")
if weights is None:
if len(image_list) == 8:
weights = [2.0, 0.05, 0.2, 0.02, 1.0, 0.02, 0.2, 0.05]
elif len(image_list) == 6:
weights = [2.0, 0.05, 0.2, 1.0, 0.2, 0.05]
elif len(image_list) == 4:
weights = [2.0, 0.2, 1.0, 0.2]
elif len(image_list) == 2:
weights = [1.0, 1.0]
else:
raise ValueError("weights is None, and can not be guessed from image_list")
# 2. run projection
meshes = meshes.clone().to(device)
if weights is None:
weights = [1. for _ in range(len(cameras_list))]
assert len(cameras_list) == len(image_list) == len(weights)
original_color = meshes.textures.verts_features_packed()
assert not torch.isnan(original_color).any()
texture_counts = torch.zeros_like(original_color[..., :1])
texture_values = torch.zeros_like(original_color)
max_texture_counts = torch.zeros_like(original_color[..., :1])
max_texture_values = torch.zeros_like(original_color)
for camera, image, weight in zip(cameras_list, image_list, weights):
ret = project_color(meshes, camera, image, eps=eps, resolution=resolution, device=device, use_alpha=use_alpha)
if reweight_with_cosangle == "linear":
weight = (ret['cos_angles'].abs() * weight)[:, None]
elif reweight_with_cosangle == "square":
weight = (ret['cos_angles'].abs() ** 2 * weight)[:, None]
if use_alpha:
weight = weight * ret['valid_alpha']
assert weight.min() > -0.0001
texture_counts[ret['valid_verts']] += weight
texture_values[ret['valid_verts']] += ret['valid_colors'] * weight
max_texture_values[ret['valid_verts']] = torch.where(weight > max_texture_counts[ret['valid_verts']], ret['valid_colors'], max_texture_values[ret['valid_verts']])
max_texture_counts[ret['valid_verts']] = torch.max(max_texture_counts[ret['valid_verts']], weight)
# Method2
texture_values = torch.where(texture_counts > confidence_threshold, texture_values / texture_counts, texture_values)
if below_confidence_strategy == "smooth":
texture_values = torch.where(texture_counts <= confidence_threshold, (original_color * (confidence_threshold - texture_counts) + texture_values) / confidence_threshold, texture_values)
elif below_confidence_strategy == "original":
texture_values = torch.where(texture_counts <= confidence_threshold, original_color, texture_values)
else:
raise ValueError(f"below_confidence_strategy={below_confidence_strategy} is not supported")
assert not torch.isnan(texture_values).any()
meshes.textures = TexturesVertex(verts_features=[texture_values])
if complete_unseen:
meshes = complete_unseen_vertex_color(meshes, torch.arange(texture_values.shape[0]).to(device)[texture_counts[:, 0] >= confidence_threshold])
ret_mesh = meshes.detach()
del meshes
return ret_mesh
def get_cameras_list(azim_list, device, focal=2/1.35, dist=1.1):
ret = []
for azim in azim_list:
R, T = look_at_view_transform(dist, 0, azim)
w2c = torch.cat([R[0].T, T[0, :, None]], dim=1)
cameras: OrthographicCameras = get_camera(w2c, focal_length=focal, cam_type='orthogonal').to(device)
ret.append(cameras)
return ret
def get_8view_cameras(device, focal=2/1.35):
return get_cameras_list(azim_list = [180, 225, 270, 315, 0, 45, 90, 135], device=device, focal=focal)
def get_6view_cameras(device, focal=2/1.35):
return get_cameras_list(azim_list = [180, 225, 270, 0, 90, 135], device=device, focal=focal)
def get_4view_cameras(device, focal=2/1.35):
return get_cameras_list(azim_list = [180, 270, 0, 90], device=device, focal=focal)
def get_2view_cameras(device, focal=2/1.35):
return get_cameras_list(azim_list = [180, 0], device=device, focal=focal)
def get_multiple_view_cameras(device, focal=2/1.35, offset=180, num_views=8, dist=1.1):
return get_cameras_list(azim_list = (np.linspace(0, 360, num_views+1)[:-1] + offset) % 360, device=device, focal=focal, dist=dist)
def align_with_alpha_bbox(source_img, target_img, final_size=1024):
# align source_img with target_img using alpha channel
# source_img and target_img are PIL.Image.Image
source_img = source_img.convert("RGBA")
target_img = target_img.convert("RGBA").resize((final_size, final_size))
source_np = np.array(source_img)
target_np = np.array(target_img)
source_alpha = source_np[:, :, 3]
target_alpha = target_np[:, :, 3]
bbox_source_min, bbox_source_max = np.argwhere(source_alpha > 0).min(axis=0), np.argwhere(source_alpha > 0).max(axis=0)
bbox_target_min, bbox_target_max = np.argwhere(target_alpha > 0).min(axis=0), np.argwhere(target_alpha > 0).max(axis=0)
source_content = source_np[bbox_source_min[0]:bbox_source_max[0]+1, bbox_source_min[1]:bbox_source_max[1]+1, :]
# resize source_content to fit in the position of target_content
source_content = Image.fromarray(source_content).resize((bbox_target_max[1]-bbox_target_min[1]+1, bbox_target_max[0]-bbox_target_min[0]+1), resample=Image.BICUBIC)
target_np[bbox_target_min[0]:bbox_target_max[0]+1, bbox_target_min[1]:bbox_target_max[1]+1, :] = np.array(source_content)
return Image.fromarray(target_np)
def load_image_list_from_mvdiffusion(mvdiffusion_path, front_from_pil_or_path=None):
import os
image_list = []
for dir in ['front', 'front_right', 'right', 'back', 'left', 'front_left']:
image_path = os.path.join(mvdiffusion_path, f"rgb_000_{dir}.png")
pil = Image.open(image_path)
if dir == 'front':
if front_from_pil_or_path is not None:
if isinstance(front_from_pil_or_path, str):
replace_pil = Image.open(front_from_pil_or_path)
else:
replace_pil = front_from_pil_or_path
# align replace_pil with pil using bounding box in alpha channel
pil = align_with_alpha_bbox(replace_pil, pil, final_size=1024)
image_list.append(pil)
return image_list
def load_image_list_from_img_grid(img_grid_path, resolution = 1024):
img_list = []
grid = Image.open(img_grid_path)
w, h = grid.size
for row in range(0, h, resolution):
for col in range(0, w, resolution):
img_list.append(grid.crop((col, row, col + resolution, row + resolution)))
return img_list
gen_3d_modules/Unique3D/scripts/utils.py
+344
View File
@@ -0,0 +1,344 @@
import torch
import numpy as np
from PIL import Image
import pymeshlab
import pymeshlab as ml
from pymeshlab import PercentageValue
from pytorch3d.renderer import TexturesVertex
from pytorch3d.structures import Meshes
from rembg import new_session, remove
import torch
import torch.nn.functional as F
from typing import List, Tuple
from PIL import Image
import trimesh
providers = [
('CUDAExecutionProvider', {
'device_id': 0,
'arena_extend_strategy': 'kSameAsRequested',
'gpu_mem_limit': 8 * 1024 * 1024 * 1024,
'cudnn_conv_algo_search': 'HEURISTIC',
})
]
session = new_session(providers=providers)
NEG_PROMPT="sketch, sculpture, hand drawing, outline, single color, NSFW, lowres, bad anatomy,bad hands, text, error, missing fingers, yellow sleeves, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(worst quality:1.4),(low quality:1.4)"
def load_mesh_with_trimesh(file_name, file_type=None):
import trimesh
mesh: trimesh.Trimesh = trimesh.load(file_name, file_type=file_type)
if isinstance(mesh, trimesh.Scene):
assert len(mesh.geometry) > 0
# save to obj first and load again to avoid offset issue
from io import BytesIO
with BytesIO() as f:
mesh.export(f, file_type="obj")
f.seek(0)
mesh = trimesh.load(f, file_type="obj")
if isinstance(mesh, trimesh.Scene):
# we lose texture information here
mesh = trimesh.util.concatenate(
tuple(trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
for g in mesh.geometry.values()))
assert isinstance(mesh, trimesh.Trimesh)
vertices = torch.from_numpy(mesh.vertices).T
faces = torch.from_numpy(mesh.faces).T
colors = None
if mesh.visual is not None:
if hasattr(mesh.visual, 'vertex_colors'):
colors = torch.from_numpy(mesh.visual.vertex_colors)[..., :3].T / 255.
if colors is None:
# print("Warning: no vertex color found in mesh! Filling it with gray.")
colors = torch.ones_like(vertices) * 0.5
return vertices, faces, colors
def meshlab_mesh_to_py3dmesh(mesh: pymeshlab.Mesh) -> Meshes:
verts = torch.from_numpy(mesh.vertex_matrix()).float()
faces = torch.from_numpy(mesh.face_matrix()).long()
colors = torch.from_numpy(mesh.vertex_color_matrix()[..., :3]).float()
textures = TexturesVertex(verts_features=[colors])
return Meshes(verts=[verts], faces=[faces], textures=textures)
def py3dmesh_to_meshlab_mesh(meshes: Meshes) -> pymeshlab.Mesh:
colors_in = F.pad(meshes.textures.verts_features_packed().cpu().float(), [0,1], value=1).numpy().astype(np.float64)
m1 = pymeshlab.Mesh(
vertex_matrix=meshes.verts_packed().cpu().float().numpy().astype(np.float64),
face_matrix=meshes.faces_packed().cpu().long().numpy().astype(np.int32),
v_normals_matrix=meshes.verts_normals_packed().cpu().float().numpy().astype(np.float64),
v_color_matrix=colors_in)
return m1
def to_pyml_mesh(vertices,faces):
m1 = pymeshlab.Mesh(
vertex_matrix=vertices.cpu().float().numpy().astype(np.float64),
face_matrix=faces.cpu().long().numpy().astype(np.int32),
)
return m1
def to_py3d_mesh(vertices, faces, normals=None):
from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh.textures import TexturesVertex
mesh = Meshes(verts=[vertices], faces=[faces], textures=None)
if normals is None:
normals = mesh.verts_normals_packed()
# set normals as vertext colors
mesh.textures = TexturesVertex(verts_features=[normals / 2 + 0.5])
return mesh
def from_py3d_mesh(mesh):
return mesh.verts_list()[0], mesh.faces_list()[0], mesh.textures.verts_features_packed()
def rotate_normalmap_by_angle(normal_map: np.ndarray, angle: float):
"""
rotate along y-axis
normal_map: np.array, shape=(H, W, 3) in [-1, 1]
angle: float, in degree
"""
angle = angle / 180 * np.pi
R = np.array([[np.cos(angle), 0, np.sin(angle)], [0, 1, 0], [-np.sin(angle), 0, np.cos(angle)]])
return np.dot(normal_map.reshape(-1, 3), R.T).reshape(normal_map.shape)
# from view coord to front view world coord
def rotate_normals(normal_pils, return_types='np', rotate_direction=1) -> np.ndarray: # [0, 255]
n_views = len(normal_pils)
ret = []
for idx, rgba_normal in enumerate(normal_pils):
# rotate normal
normal_np = np.array(rgba_normal)[:, :, :3] / 255 # in [-1, 1]
alpha_np = np.array(rgba_normal)[:, :, 3] / 255 # in [0, 1]
normal_np = normal_np * 2 - 1
normal_np = rotate_normalmap_by_angle(normal_np, rotate_direction * idx * (360 / n_views))
normal_np = (normal_np + 1) / 2
normal_np = normal_np * alpha_np[..., None] # make bg black
rgba_normal_np = np.concatenate([normal_np * 255, alpha_np[:, :, None] * 255] , axis=-1)
if return_types == 'np':
ret.append(rgba_normal_np)
elif return_types == 'pil':
ret.append(Image.fromarray(rgba_normal_np.astype(np.uint8)))
else:
raise ValueError(f"return_types should be 'np' or 'pil', but got {return_types}")
return ret
def rotate_normalmap_by_angle_torch(normal_map, angle):
"""
rotate along y-axis
normal_map: torch.Tensor, shape=(H, W, 3) in [-1, 1], device='cuda'
angle: float, in degree
"""
angle = torch.tensor(angle / 180 * np.pi).to(normal_map)
R = torch.tensor([[torch.cos(angle), 0, torch.sin(angle)],
[0, 1, 0],
[-torch.sin(angle), 0, torch.cos(angle)]]).to(normal_map)
return torch.matmul(normal_map.view(-1, 3), R.T).view(normal_map.shape)
def do_rotate(rgba_normal, angle):
rgba_normal = torch.from_numpy(rgba_normal).float().cuda() / 255
rotated_normal_tensor = rotate_normalmap_by_angle_torch(rgba_normal[..., :3] * 2 - 1, angle)
rotated_normal_tensor = (rotated_normal_tensor + 1) / 2
rotated_normal_tensor = rotated_normal_tensor * rgba_normal[:, :, [3]] # make bg black
rgba_normal_np = torch.cat([rotated_normal_tensor * 255, rgba_normal[:, :, [3]] * 255], dim=-1).cpu().numpy()
return rgba_normal_np
def rotate_normals_torch(normal_pils, return_types='np', rotate_direction=1):
n_views = len(normal_pils)
ret = []
for idx, rgba_normal in enumerate(normal_pils):
# rotate normal
angle = rotate_direction * idx * (360 / n_views)
rgba_normal_np = do_rotate(np.array(rgba_normal), angle)
if return_types == 'np':
ret.append(rgba_normal_np)
elif return_types == 'pil':
ret.append(Image.fromarray(rgba_normal_np.astype(np.uint8)))
else:
raise ValueError(f"return_types should be 'np' or 'pil', but got {return_types}")
return ret
def get_normal_map_masks(imgs):
"""Only works for normal"""
if not isinstance(imgs, list):
imgs = [imgs]
rets = []
for img in imgs:
arr = np.array(img)
# remove color from normal map
if arr.shape[-1] == 4:
arr = arr[..., :3]
# calc diffs
base = arr[0, 0]
diffs = np.abs(arr.astype(np.int32) - base.astype(np.int32)).sum(axis=-1)
alpha = (diffs <= 80)
arr[alpha] = 255
alpha = ~alpha
#arr = np.concatenate([arr, alpha[..., None].astype(np.int32) * 255], axis=-1)
#rets.append(Image.fromarray(arr.astype(np.uint8)))
rets.append(torch.from_numpy(alpha))
return rets
def change_bkgd(img_pils, new_bkgd=(0., 0., 0.)):
ret = []
new_bkgd = np.array(new_bkgd).reshape(1, 1, 3)
for rgba_img in img_pils:
img_np = np.array(rgba_img)[:, :, :3] / 255
alpha_np = np.array(rgba_img)[:, :, 3] / 255
ori_bkgd = img_np[:1, :1]
# color = ori_color * alpha + bkgd * (1-alpha)
# ori_color = (color - bkgd * (1-alpha)) / alpha
alpha_np_clamp = np.clip(alpha_np, 1e-6, 1) # avoid divide by zero
ori_img_np = (img_np - ori_bkgd * (1 - alpha_np[..., None])) / alpha_np_clamp[..., None]
img_np = np.where(alpha_np[..., None] > 0.05, ori_img_np * alpha_np[..., None] + new_bkgd * (1 - alpha_np[..., None]), new_bkgd)
rgba_img_np = np.concatenate([img_np * 255, alpha_np[..., None] * 255], axis=-1)
ret.append(Image.fromarray(rgba_img_np.astype(np.uint8)))
return ret
def change_bkgd_to_normal(normal_pils) -> List[Image.Image]:
n_views = len(normal_pils)
ret = []
for idx, rgba_normal in enumerate(normal_pils):
# calcuate background normal
target_bkgd = rotate_normalmap_by_angle(np.array([[[0., 0., 1.]]]), idx * (360 / n_views))
normal_np = np.array(rgba_normal)[:, :, :3] / 255 # in [-1, 1]
alpha_np = np.array(rgba_normal)[:, :, 3] / 255 # in [0, 1]
normal_np = normal_np * 2 - 1
old_bkgd = normal_np[:1,:1]
normal_np[alpha_np > 0.05] = (normal_np[alpha_np > 0.05] - old_bkgd * (1 - alpha_np[alpha_np > 0.05][..., None])) / alpha_np[alpha_np > 0.05][..., None]
normal_np = normal_np * alpha_np[..., None] + target_bkgd * (1 - alpha_np[..., None])
normal_np = (normal_np + 1) / 2
rgba_normal_np = np.concatenate([normal_np * 255, alpha_np[..., None] * 255] , axis=-1)
ret.append(Image.fromarray(rgba_normal_np.astype(np.uint8)))
return ret
def fix_vert_color_glb(mesh_path):
from pygltflib import GLTF2, Material, PbrMetallicRoughness
obj1 = GLTF2().load(mesh_path)
obj1.meshes[0].primitives[0].material = 0
obj1.materials.append(Material(
pbrMetallicRoughness = PbrMetallicRoughness(
baseColorFactor = [1.0, 1.0, 1.0, 1.0],
metallicFactor = 0.,
roughnessFactor = 1.0,
),
emissiveFactor = [0.0, 0.0, 0.0],
doubleSided = True,
))
obj1.save(mesh_path)
def srgb_to_linear(c_srgb):
c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4)
return c_linear.clip(0, 1.)
def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
# convert from pytorch3d meshes to trimesh mesh
vertices = meshes.verts_packed().cpu().float().numpy()
triangles = meshes.faces_packed().cpu().long().numpy()
np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
if save_glb_path.endswith(".glb"):
# rotate 180 along +Y
vertices[:, [0, 2]] = -vertices[:, [0, 2]]
if apply_sRGB_to_LinearRGB:
np_color = srgb_to_linear(np_color)
assert vertices.shape[0] == np_color.shape[0]
assert np_color.shape[1] == 3
assert 0 <= np_color.min() and np_color.max() <= 1, f"min={np_color.min()}, max={np_color.max()}"
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
mesh.remove_unreferenced_vertices()
# save mesh
mesh.export(save_glb_path)
if save_glb_path.endswith(".glb"):
fix_vert_color_glb(save_glb_path)
print(f"saving to {save_glb_path}")
def save_glb_and_video(save_mesh_prefix: str, meshes: Meshes, with_timestamp=True, dist=3.5, azim_offset=180, resolution=512, fov_in_degrees=1 / 1.15, cam_type="ortho", view_padding=60, export_video=True) -> Tuple[str, str]:
import time
if '.' in save_mesh_prefix:
save_mesh_prefix = ".".join(save_mesh_prefix.split('.')[:-1])
if with_timestamp:
save_mesh_prefix = save_mesh_prefix + f"_{int(time.time())}"
ret_mesh = save_mesh_prefix + ".glb"
# optimizied version
save_py3dmesh_with_trimesh_fast(meshes, ret_mesh)
return ret_mesh, None
def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, apply_sub_divide=False, sub_divide_threshold=0.25):
ms = ml.MeshSet()
ms.add_mesh(pyml_mesh, "cube_mesh")
if apply_smooth:
ms.apply_filter("apply_coord_laplacian_smoothing", stepsmoothnum=stepsmoothnum, cotangentweight=False)
if apply_sub_divide: # 5s, slow
ms.apply_filter("meshing_repair_non_manifold_vertices")
ms.apply_filter("meshing_repair_non_manifold_edges", method='Remove Faces')
ms.apply_filter("meshing_surface_subdivision_loop", iterations=2, threshold=PercentageValue(sub_divide_threshold))
return meshlab_mesh_to_py3dmesh(ms.current_mesh())
def expand2square(pil_img, background_color):
width, height = pil_img.size
if width == height:
return pil_img
elif width > height:
result = Image.new(pil_img.mode, (width, width), background_color)
result.paste(pil_img, (0, (width - height) // 2))
return result
else:
result = Image.new(pil_img.mode, (height, height), background_color)
result.paste(pil_img, ((height - width) // 2, 0))
return result
def simple_image_preprocess(input_image, rembg_session=session, background_color=255):
RES = 2048
input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
if input_image.mode != 'RGBA':
image_rem = input_image.convert('RGBA')
input_image = remove(image_rem, alpha_matting=False, session=rembg_session)
arr = np.asarray(input_image)
alpha = np.asarray(input_image)[:, :, -1]
x_nonzero = np.nonzero((alpha > 60).sum(axis=1))
y_nonzero = np.nonzero((alpha > 60).sum(axis=0))
x_min = int(x_nonzero[0].min())
y_min = int(y_nonzero[0].min())
x_max = int(x_nonzero[0].max())
y_max = int(y_nonzero[0].max())
arr = arr[x_min: x_max, y_min: y_max]
input_image = Image.fromarray(arr)
input_image = expand2square(input_image, (background_color, background_color, background_color, 0))
return input_image
def init_target(img_pils, new_bkgd=(0., 0., 0.), device="cuda"):
# Convert the background color to a PyTorch tensor
new_bkgd = torch.tensor(new_bkgd, dtype=torch.float32).view(1, 1, 3).to(device)
# Convert all images to PyTorch tensors and process them
imgs = torch.stack([torch.from_numpy(np.array(img, dtype=np.float32)) for img in img_pils]).to(device) / 255
img_nps = imgs[..., :3]
alpha_nps = imgs[..., 3]
ori_bkgds = img_nps[:, :1, :1]
# Avoid divide by zero and calculate the original image
alpha_nps_clamp = torch.clamp(alpha_nps, 1e-6, 1)
ori_img_nps = (img_nps - ori_bkgds * (1 - alpha_nps.unsqueeze(-1))) / alpha_nps_clamp.unsqueeze(-1)
ori_img_nps = torch.clamp(ori_img_nps, 0, 1)
img_nps = torch.where(alpha_nps.unsqueeze(-1) > 0.05, ori_img_nps * alpha_nps.unsqueeze(-1) + new_bkgd * (1 - alpha_nps.unsqueeze(-1)), new_bkgd)
rgba_img_np = torch.cat([img_nps, alpha_nps.unsqueeze(-1)], dim=-1)
return rgba_img_np
lgm/core/attention.py
+7 -5
View File
@@ -8,10 +8,13 @@
# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
import os
import warnings
import logging
from torch import Tensor
from torch import nn
from .utils import setup_logger
setup_logger('', logging.INFO, [logging.INFO, logging.WARNING], [logging.ERROR, logging.CRITICAL], logging.Formatter('%(message)s'))
XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
try:
@@ -19,14 +22,13 @@ try:
from xformers.ops import memory_efficient_attention, unbind
XFORMERS_AVAILABLE = True
warnings.warn("xFormers is available (Attention)")
logging.warn("xFormers is available (Attention)")
else:
warnings.warn("xFormers is disabled (Attention)")
logging.warn("xFormers is disabled (Attention)")
raise ImportError
except ImportError:
XFORMERS_AVAILABLE = False
warnings.warn("xFormers is not available (Attention)")
logging.warn("xFormers is not available (Attention)")
class Attention(nn.Module):
def __init__(
lgm/core/utils.py
+25
View File
@@ -1,4 +1,6 @@
import numpy as np
import logging
import sys
import torch
import torch.nn as nn
@@ -107,3 +109,26 @@ def grid_distortion(images, strength=0.5):
return images
class WarningFilter(logging.Filter):
def filter(self, record):
if record.levelno == logging.WARNING:
record.msg = f"Warn!: {record.msg}"
return True
def create_handler(stream, levels, formatter):
handler = logging.StreamHandler(stream)
handler.setLevel(min(levels))
handler.addFilter(lambda record: record.levelno in levels)
handler.addFilter(WarningFilter()) # Apply the custom filter
handler.setFormatter(formatter)
return handler
def setup_logger(logger_name, level, stdout_levels, stderr_levels, formatter):
logger = logging.getLogger(logger_name)
for handler in logger.handlers[:]:
logger.removeHandler(handler)
logger.setLevel(level)
stdout_handler = create_handler(sys.stdout, stdout_levels, formatter)
stderr_handler = create_handler(sys.stderr, stderr_levels, formatter)
logger.addHandler(stdout_handler)
logger.addHandler(stderr_handler)
mesh_processer/mesh.py
+1 -1
View File
@@ -466,7 +466,7 @@ class Mesh:
i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]
face_normals = torch.cross(v1 - v0, v2 - v0)
face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
# Splat face normals to vertices
vn = torch.zeros_like(self.v)
nodes.py
+491 -94
View File
@@ -74,7 +74,18 @@ from .era3d.mvdiffusion.pipelines.pipeline_mvdiffusion_unclip import StableUnCLI
from .era3d.mvdiffusion.data.single_image_dataset import SingleImageDataset as Era3DSingleImageDataset
from .era3d.utils.misc import load_config as load_config_era3d
from .shared_utils.image_utils import prepare_torch_img, torch_img_to_pil_rgba, troch_image_dilate
from Unique3D.custum_3d_diffusion.custum_pipeline.unifield_pipeline_img2mvimg import StableDiffusionImage2MVCustomPipeline
from Unique3D.custum_3d_diffusion.custum_pipeline.unifield_pipeline_img2img import StableDiffusionImageCustomPipeline
from Unique3D.scripts.mesh_init import fast_geo
from Unique3D.scripts.utils import from_py3d_mesh, to_py3d_mesh, to_pyml_mesh, simple_clean_mesh
from Unique3D.scripts.project_mesh import multiview_color_projection, get_cameras_list
from Unique3D.mesh_reconstruction.recon import reconstruct_stage1
from Unique3D.mesh_reconstruction.refine import run_mesh_refine
from .shared_utils.image_utils import (
prepare_torch_img, torch_imgs_to_pils, troch_image_dilate,
pils_rgba_to_rgb, pil_make_image_grid, pil_split_image, pils_to_torch_imgs,
)
from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames, resume_or_download_model_from_hf
DIFFUSERS_PIPE_DICT = OrderedDict([
@@ -83,7 +94,9 @@ DIFFUSERS_PIPE_DICT = OrderedDict([
("Zero123PlusPipeline", Zero123PlusPipeline),
("DiffusionPipeline", DiffusionPipeline),
("StableDiffusionPipeline", StableDiffusionPipeline),
("Era3DPipeline", StableUnCLIPImg2ImgPipeline)
("Era3DPipeline", StableUnCLIPImg2ImgPipeline),
("Unique3DImage2MVCustomPipeline", StableDiffusionImage2MVCustomPipeline),
("Unique3DImageCustomPipeline", StableDiffusionImageCustomPipeline),
])
DIFFUSERS_SCHEDULER_DICT = OrderedDict([
@@ -332,7 +345,205 @@ class Save_3DGS:
gs_ply.write(save_path)
return (save_path, )
class Image_Add_Pure_Color_Background:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"masks": ("MASK",),
"R": ("INT", {"default": 255, "min": 0, "max": 255}),
"G": ("INT", {"default": 255, "min": 0, "max": 255}),
"B": ("INT", {"default": 255, "min": 0, "max": 255}),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"images",
)
FUNCTION = "image_add_bg"
CATEGORY = "Comfy3D/Preprocessor"
def image_add_bg(self, images, masks, R, G, B):
"""
bg_mask = bg_mask.unsqueeze(3)
inv_bg_mask = torch.ones_like(bg_mask) - bg_mask
color = torch.tensor([R, G, B]).to(image.dtype) / 255
color_bg = color.repeat(bg_mask.shape)
image = inv_bg_mask * image + bg_mask * color_bg
"""
image_pils = torch_imgs_to_pils(images, masks)
image_pils = pils_rgba_to_rgb(image_pils, (R, G, B))
images = pils_to_torch_imgs(image_pils, images.device)
return (images,)
class Make_Image_Grid:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"grid_side_num": ("INT", {"default": 1, "min": 1, "max": 8192}),
"use_rows": ("BOOLEAN", {"default": True},),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"image_grid",
)
FUNCTION = "make_image_grid"
CATEGORY = "Comfy3D/Preprocessor"
def make_image_grid(self, images, grid_side_num, use_rows):
pil_image_list = torch_imgs_to_pils(images)
if use_rows:
rows = grid_side_num
clos = None
else:
clos = grid_side_num
rows = None
image_grid = pil_make_image_grid(pil_image_list, rows, clos)
image_grid = TF.to_tensor(image_grid).permute(1, 2, 0).unsqueeze(0) # [1, H, W, 3]
return (image_grid,)
class Split_Image_Grid:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"grid_side_num": ("INT", {"default": 1, "min": 1, "max": 8192}),
"use_rows": ("BOOLEAN", {"default": True},),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"images",
)
FUNCTION = "split_image_grid"
CATEGORY = "Comfy3D/Preprocessor"
def split_image_grid(self, image, grid_side_num, use_rows):
image_pil = torch_imgs_to_pils(image)[0]
if use_rows:
rows = grid_side_num
clos = None
else:
clos = grid_side_num
rows = None
image_pils = pil_split_image(image_pil, rows, clos)
images = pils_to_torch_imgs(image_pils, image.device)
return (images,)
class Get_Masks_From_Normal_Maps:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"normal_maps": ("IMAGE",),
},
}
RETURN_TYPES = (
"MASK",
)
RETURN_NAMES = (
"normal_masks",
)
FUNCTION = "make_image_grid"
CATEGORY = "Comfy3D/Preprocessor"
def make_image_grid(self, normal_maps):
from Unique3D.scripts.utils import get_normal_map_masks
pil_normal_list = torch_imgs_to_pils(normal_maps)
normal_masks = get_normal_map_masks(pil_normal_list)
normal_masks = torch.stack(normal_masks, dim=0).to(normal_maps.dtype).to(normal_maps.device)
return (normal_masks,)
class Rotate_Normal_Maps_Horizontally:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"normal_maps": ("IMAGE",),
"normal_masks": ("MASK",),
"clockwise": ("BOOLEAN", {"default": True},),
},
}
RETURN_TYPES = (
"IMAGE",
)
RETURN_NAMES = (
"normal_maps",
)
FUNCTION = "make_image_grid"
CATEGORY = "Comfy3D/Preprocessor"
def make_image_grid(self, normal_maps, normal_masks, clockwise):
if normal_maps.shape[0] > 1:
from Unique3D.scripts.utils import rotate_normals_torch
pil_image_list = torch_imgs_to_pils(normal_maps, normal_masks)
pil_image_list = rotate_normals_torch(pil_image_list, return_types='pil', rotate_direction=int(clockwise))
normal_maps = pils_to_torch_imgs(pil_image_list, normal_maps.device)
return (normal_maps,)
class Fast_Clean_Mesh:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
"apply_smooth": ("BOOLEAN", {"default": True},),
"smooth_step": ("INT", {"default": 1, "min": 0, "max": 0xffffffffffffffff}),
"apply_sub_divide": ("BOOLEAN", {"default": True},),
"sub_divide_threshold": ("FLOAT", {"default": 0.25, 'step': 0.001}),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "clean_mesh"
CATEGORY = "Comfy3D/Preprocessor"
def clean_mesh(self, mesh, apply_smooth, smooth_step, apply_sub_divide, sub_divide_threshold):
meshes = simple_clean_mesh(to_pyml_mesh(mesh.v, mesh.f), apply_smooth=apply_smooth, stepsmoothnum=smooth_step, apply_sub_divide=apply_sub_divide, sub_divide_threshold=sub_divide_threshold).to(DEVICE)
vertices, faces, _ = from_py3d_mesh(meshes)
mesh = Mesh(v=vertices, f=faces, device=DEVICE)
return (mesh,)
class Switch_3DGS_Axis:
@classmethod
def INPUT_TYPES(cls):
@@ -454,8 +665,8 @@ class Stack_Orbit_Camera_Poses:
def INPUT_TYPES(cls):
return {
"required": {
"orbit_radius_start": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_radius_stop": ("FLOAT", {"default": 0.0, 'step': 0.0001}),
"orbit_radius_start": ("FLOAT", {"default": 1.75, 'step': 0.0001}),
"orbit_radius_stop": ("FLOAT", {"default": 1.75, 'step': 0.0001}),
"orbit_radius_step": ("FLOAT", {"default": 0.1, 'step': 0.0001}),
"elevation_start": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
"elevation_stop": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
@@ -635,84 +846,7 @@ class Stack_Orbit_Camera_Poses:
orbit_center_Z_list.append(campose[5])
return (orbit_camposes, orbit_radius_list, elevation_list, azimuth_list, orbit_center_X_list, orbit_center_Y_list, orbit_center_Z_list, )
class Generate_Orbit_Camera_Poses:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_images": ("IMAGE",),
"generate_pose_command": ("STRING", {
"default": "#([start_reference_image_index : end_reference_image_index], orbit_radius, elevation_angle [-90, 90], start_azimuth_angle [0, 360], end_azimuth_angle [0, 360])\n([0:30], 1.75, 0, 0, 360)",
"multiline": True
}),
},
}
RETURN_TYPES = (
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
)
RETURN_NAMES = (
"orbit_camposes",
)
FUNCTION = "get_camposes"
CATEGORY = "Comfy3D/Preprocessor"
class Slice_Camposes:
def __init__(self, start_reference_image_index, end_reference_image_index, camposes_start_to_end):
self.start_reference_image_index = start_reference_image_index
self.end_reference_image_index = end_reference_image_index
self.camposes_start_to_end = camposes_start_to_end
def get_camposes(self, reference_images, generate_pose_command):
orbit_camposes = []
self.ref_imgs_num_minus_1 = len(reference_images) - 1
# To match pattern ( [ start_reference_image_index : end_reference_image_index ] , orbit_radius, elevation_angle , start_azimuth_angle , end_azimuth_angle )
pattern = re.compile(r"\([ \t]*\[[ \t]*(\d+)[ \t]*:[ \t]*(\d+)[ \t]*\][ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*,[ \t]*([\d]+\.?[\d]*)[ \t]*\)")
all_matches = pattern.findall(generate_pose_command)
all_slice_camposes = []
for match in all_matches:
start_reference_image_index, end_reference_image_index, orbit_radius, elevation_angle, start_azimuth_angle, end_azimuth_angle = (int(s) if i < 2 else float(s) for i, s in enumerate(match))
end_reference_image_index = min(end_reference_image_index, self.ref_imgs_num_minus_1)
if start_reference_image_index <= end_reference_image_index:
azimuth_imgs_num = end_reference_image_index - start_reference_image_index + 1
# calculate all the reference camera azimuth angles
camposes_start_to_end = []
if start_azimuth_angle > end_azimuth_angle:
azimuth_angle_interval = -(end_azimuth_angle + 360 - start_azimuth_angle) / azimuth_imgs_num
else:
azimuth_angle_interval = (end_azimuth_angle - start_azimuth_angle) / azimuth_imgs_num
now_azimuth_angle = start_azimuth_angle
for _ in range(azimuth_imgs_num):
camposes_start_to_end.append((orbit_radius, elevation_angle, now_azimuth_angle, 0.0, 0.0, 0.0))
now_azimuth_angle = (now_azimuth_angle + azimuth_angle_interval) % 360
all_slice_camposes.append(Generate_Orbit_Camera_Poses.Slice_Camposes(start_reference_image_index, end_reference_image_index, camposes_start_to_end))
else:
cstr(f"[{self.__class__.__name__}] start_reference_image_index: {start_reference_image_index} must smaller than or equal to end_reference_image_index: {end_reference_image_index}").error.print()
all_slice_camposes = sorted(all_slice_camposes, key=lambda slice_camposes:slice_camposes.start_reference_image_index)
last_end_index_plus_1 = 0
for slice_camposes in all_slice_camposes:
if last_end_index_plus_1 == slice_camposes.start_reference_image_index:
orbit_camposes.extend(slice_camposes.camposes_start_to_end)
last_end_index_plus_1 = slice_camposes.end_reference_image_index + 1
else:
orbit_camposes = []
cstr(f"[{self.__class__.__name__}] Last end_reference_image_index: {end_reference_image_index} plus 1 must equal to current start_reference_image_index: {start_reference_image_index}").error.print()
return (orbit_camposes, )
class Get_Camposes_From_List_Indexed:
RETURN_TYPES = ("ORBIT_CAMPOSES",)
@@ -1276,6 +1410,7 @@ class Load_Diffusers_Pipeline:
"diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),),
"model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
"custom_pipeline": ("STRING", {"default": "", "multiline": False}),
"download_from_remote": ("BOOLEAN", {"default": True}),
},
}
@@ -1288,13 +1423,14 @@ class Load_Diffusers_Pipeline:
FUNCTION = "load_diffusers_pipe"
CATEGORY = "Comfy3D/Import|Export"
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name, custom_pipeline):
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name, custom_pipeline, download_from_remote):
# resume pretrained checkpoint
ckpt_path = os.path.join(ROOT_PATH, "checkpoints", model_name)
from huggingface_hub import snapshot_download
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model", ignore_patterns=["*.json"])
if download_from_remote:
from huggingface_hub import snapshot_download
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model", ignore_patterns=["*.json"])
diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name]
@@ -1306,7 +1442,6 @@ class Load_Diffusers_Pipeline:
ckpt_path,
torch_dtype=WEIGHT_DTYPE,
custom_pipeline=custom_pipeline,
trust_remote_code=True,
).to(DEVICE)
pipe.enable_xformers_memory_efficient_attention()
@@ -1372,7 +1507,7 @@ class Set_Diffusers_Pipeline_State_Dict:
state_dict = torch.load(ckpt_path, map_location='cpu')
pipe.unet.load_state_dict(state_dict, strict=True)
pipe.enable_xformers_memory_efficient_attention()
pipe = pipe.to(DEVICE)
return (pipe, )
@@ -1459,7 +1594,7 @@ class Wonder3D_MVDiffusion_Model:
return (mv_images, mv_normals, )
def prepare_data(self, ref_image, ref_mask):
single_image = torch_img_to_pil_rgba(ref_image, ref_mask)
single_image = torch_imgs_to_pils(ref_image, ref_mask)[0]
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
dataset = MVSingleImageDataset(fix_cam_pose_dir=abs_fix_cam_pose_dir, num_views=6, img_wh=[256, 256], bg_color='white', single_image=single_image)
return dataset[0]
@@ -1959,7 +2094,7 @@ class CRM_Images_MVDiffusion_Model:
mv_guidance_scale,
num_inference_steps,
):
pixel_img = torch_img_to_pil_rgba(reference_image, reference_mask)
pixel_img = torch_imgs_to_pils(reference_image, reference_mask)[0]
pixel_img = CRMSampler.process_pixel_img(pixel_img)
multiview_images = CRMSampler.stage1_sample(
@@ -2027,7 +2162,7 @@ class CRM_CCMs_MVDiffusion_Model:
mv_guidance_scale,
num_inference_steps,
):
pixel_img = torch_img_to_pil_rgba(reference_image, reference_mask)
pixel_img = torch_imgs_to_pils(reference_image, reference_mask)[0]
pixel_img = CRMSampler.process_pixel_img(pixel_img)
multiview_CCMs = CRMSampler.stage2_sample(
@@ -2152,7 +2287,7 @@ class Zero123Plus_Diffusion_Model:
num_inference_steps,
):
single_image = torch_img_to_pil_rgba(reference_image, reference_mask)
single_image = torch_imgs_to_pils(reference_image, reference_mask)[0]
seed = int(seed)
generator = torch.Generator(device=zero123plus_pipe.unet.device).manual_seed(seed)
@@ -2330,7 +2465,7 @@ class Era3D_Diffusion_Model:
config_path = os.path.join(ROOT_PATH, self.era3d_config_path)
cfg = load_config_era3d(config_path)
single_image = torch_img_to_pil_rgba(reference_image, reference_mask)
single_image = torch_imgs_to_pils(reference_image, reference_mask)[0]
# Get the dataset
cfg.dataset.prompt_embeds_path = os.path.join(ROOT_PATH, cfg.dataset.prompt_embeds_path)
@@ -2368,8 +2503,8 @@ class Era3D_Diffusion_Model:
images_pred = out[bsz:]
# [N, 3, H, W] -> [N, H, W, 3]
multiview_images = images_pred.permute(0, 2, 3, 1)
multiview_normals = normals_pred.permute(0, 2, 3, 1)
multiview_images = images_pred.permute(0, 2, 3, 1).to(reference_image.dtype)
multiview_normals = normals_pred.permute(0, 2, 3, 1).to(reference_image.dtype)
azimuths = [0, 45, 90, 180, -90, -45]
elevations = [0.0] * 6
@@ -2527,4 +2662,266 @@ class FlexiCubes_MVS:
mesh = fc_trainer.get_mesh()
return (mesh, )
return (mesh, )
class Load_Unique3D_Custom_UNet:
checkpoints_dir = "checkpoints/Wuvin/Unique3D"
config_root_dir = "configs/Unique3D_configs"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"pipe": ("DIFFUSERS_PIPE",),
"config_name": (["image2mvimage", "image2normal"],),
},
}
RETURN_TYPES = (
"DIFFUSERS_PIPE",
)
RETURN_NAMES = (
"pipe",
)
FUNCTION = "load_diffusers_unet"
CATEGORY = "Comfy3D/Import|Export"
def load_diffusers_unet(self, pipe, config_name):
from Unique3D.custum_3d_diffusion.trainings.config_classes import ExprimentConfig
from Unique3D.custum_3d_diffusion.custum_modules.unifield_processor import AttnConfig, ConfigurableUNet2DConditionModel
from Unique3D.custum_3d_diffusion.trainings.utils import load_config
cfg_path = os.path.join(ROOT_PATH, self.config_root_dir, config_name + ".yaml")
checkpoint_dir_path = os.path.join(ROOT_PATH, self.checkpoints_dir, config_name)
checkpoint_path = os.path.join(checkpoint_dir_path, "unet_state_dict.pth")
cfg: ExprimentConfig = load_config(ExprimentConfig, cfg_path)
if cfg.init_config.init_unet_path == "":
cfg.init_config.init_unet_path = checkpoint_dir_path
init_config: AttnConfig = load_config(AttnConfig, cfg.init_config)
configurable_unet = ConfigurableUNet2DConditionModel(init_config, WEIGHT_DTYPE)
configurable_unet.enable_xformers_memory_efficient_attention()
state_dict = torch.load(checkpoint_path)
configurable_unet.unet.load_state_dict(state_dict, strict=False)
# Move unet, vae and text_encoder to device and cast to weight_dtype
configurable_unet.unet.to(DEVICE, dtype=WEIGHT_DTYPE)
pipe.unet = configurable_unet.unet
return (pipe, )
class Unique3D_Diffusion_Model:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"unique3d_pipe": ("DIFFUSERS_PIPE",),
"reference_image": ("IMAGE",),
"seed": ("INT", {"default": 1145, "min": 0, "max": 0xffffffffffffffff}),
"guidance_scale": ("FLOAT", {"default": 1.5, "min": 0.0, "step": 0.01}),
"num_inference_steps": ("INT", {"default": 30, "min": 1}),
"image_resolution": ([256, 512],),
"radius": ("FLOAT", {"default": 4.0, "min": 0.1, "step": 0.01}),
"preprocess_images": ("BOOLEAN", {"default": True},),
},
}
RETURN_TYPES = (
"IMAGE",
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
)
RETURN_NAMES = (
"multiviews",
"orbit_camposes",
)
FUNCTION = "run_model"
CATEGORY = "Comfy3D/Algorithm"
@torch.no_grad()
def run_model(
self,
unique3d_pipe,
reference_image, # Need to have white background
seed,
guidance_scale,
num_inference_steps,
image_resolution,
radius,
preprocess_images,
):
from Unique3D.scripts.utils import simple_image_preprocess
pil_image_list = torch_imgs_to_pils(reference_image)
for i in range(len(pil_image_list)):
if preprocess_images:
pil_image_list[i] = simple_image_preprocess(pil_image_list[i])
pil_image_list = pils_rgba_to_rgb(pil_image_list, bkgd="WHITE")
generator = torch.Generator(device=unique3d_pipe.unet.device).manual_seed(seed)
image_pils = unique3d_pipe(
image=pil_image_list,
generator=generator,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=image_resolution,
height=image_resolution,
height_cond=image_resolution,
width_cond=image_resolution,
).images
# [N, H, W, 3]
multiview_images = pils_to_torch_imgs(image_pils, reference_image.device)
azimuths = [0, 90, 180, -90]
elevations = [0.0] * 4
radius = [radius] * 4
center = [0.0] * 4
orbit_camposes = [azimuths, elevations, radius, center, center, center]
return (multiview_images, orbit_camposes)
class Fast_Normal_Maps_To_Mesh:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"front_side_back_normal_maps": ("IMAGE",),
"front_side_back_normal_masks": ("MASK",),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_fast_recon"
CATEGORY = "Comfy3D/Algorithm"
def run_fast_recon(self, front_side_back_normal_maps, front_side_back_normal_masks):
pil_normal_list = torch_imgs_to_pils(front_side_back_normal_maps, front_side_back_normal_masks)
meshes = fast_geo(pil_normal_list[0], pil_normal_list[2], pil_normal_list[1])
vertices, faces, _ = from_py3d_mesh(meshes)
mesh = Mesh(v=vertices, f=faces, device=DEVICE)
return (mesh,)
class ExplicitTarget_Mesh_Optimization:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
"normal_maps": ("IMAGE",),
"normal_masks": ("MASK",),
"reconstruction_steps": ("INT", {"default": 200, "min": 1, "max": 0xffffffffffffffff}),
"coarse_reconstruct_resolution": ("INT", {"default": 512, "min": 128, "max": 8192}),
"loss_expansion_weight": ("FLOAT", {"default": 0.1, "min": 0.01, "step": 0.01}),
"refinement_steps": ("INT", {"default": 100, "min": 1, "max": 0xffffffffffffffff}),
"target_warmup_update_num": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}),
"target_update_interval": ("INT", {"default": 20, "min": 1, "max": 0xffffffffffffffff}),
},
"optional": {
"normal_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
}
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_ET_mesh_optimization"
CATEGORY = "Comfy3D/Algorithm"
def run_ET_mesh_optimization(
self,
mesh,
normal_maps,
normal_masks,
reconstruction_steps,
coarse_reconstruct_resolution,
loss_expansion_weight,
refinement_steps,
target_warmup_update_num,
target_update_interval,
normal_orbit_camera_poses=None,
):
#TODO For now only support four orthographic view with elevation equals zero
#azimuths, elevations, radius = normal_orbit_camera_poses[0], normal_orbit_camera_poses[1], normal_orbit_camera_poses[2]
pil_normal_list = torch_imgs_to_pils(normal_maps, normal_masks)
normal_stg1 = [img.resize((coarse_reconstruct_resolution, coarse_reconstruct_resolution)) for img in pil_normal_list]
with torch.inference_mode(False):
vertices, faces = reconstruct_stage1(normal_stg1, steps=reconstruction_steps, vertices=mesh.v, faces=mesh.f, loss_expansion_weight=loss_expansion_weight)
vertices, faces = run_mesh_refine(vertices, faces, pil_normal_list, steps=refinement_steps, update_normal_interval=target_update_interval, update_warmup=target_warmup_update_num, )
mesh = Mesh(v=vertices, f=faces, device=DEVICE)
return (mesh,)
class ExplicitTarget_Color_Projection:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mesh": ("MESH",),
"reference_images": ("IMAGE",),
"reference_masks": ("MASK",),
"projection_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"complete_unseen_rgb": ("BOOLEAN", {"default": True},),
},
"optional": {
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
}
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_color_projection"
CATEGORY = "Comfy3D/Algorithm"
def run_color_projection(
self,
mesh,
reference_images,
reference_masks,
projection_resolution,
complete_unseen_rgb,
reference_orbit_camera_poses=None,
):
pil_image_list = torch_imgs_to_pils(reference_images, reference_masks)
meshes = to_py3d_mesh(mesh.v, mesh.f)
#TODO Convert camera format, currently only support elevation equal to zero
if reference_orbit_camera_poses is None:
img_num = len(reference_images)
interval = 360 / img_num
angle = 0
azimuths = []
for i in range(0, img_num):
azimuths.append(angle)
angle += interval
else:
azimuths = [360 + angle if angle < 0 else angle for angle in reference_orbit_camera_poses[0]]
cam_list = get_cameras_list(azimuths, DEVICE, focal=1)
new_meshes = multiview_color_projection(meshes, pil_image_list, resolution=projection_resolution, device=DEVICE, complete_unseen=complete_unseen_rgb, confidence_threshold=0.2, cameras_list=cam_list)
vertices, faces, vertex_colors = from_py3d_mesh(new_meshes)
vertices = vertices / 2 * 1.35
mesh = Mesh(v=vertices, f=faces, vc=vertex_colors, device=DEVICE)
mesh.auto_normal()
return (mesh,)
shared_utils/image_utils.py
+116 -12
View File
@@ -1,7 +1,9 @@
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
import numpy as np
from PIL import Image
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False):
# [N, H, W, C] -> [N, C, H, W]
@@ -11,20 +13,27 @@ def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False):
img_new = img_new.permute(0, 2, 3, 1)
return img_new
def torch_img_to_pil_rgba(img, mask):
def torch_imgs_to_pils(images, masks=None):
"""
img (torch): [1, H, W, C] or [H, W, C]
mask (torch): [1, H, W] or [H, W]
images (torch): [N, H, W, C] or [H, W, C]
masks (torch): [N, H, W] or [H, W]
"""
if len(img.shape) == 4:
img = img.squeeze(0)
if len(mask.shape) == 3:
mask = mask.squeeze(0).unsqueeze(2)
if len(images.shape) == 3:
images = images.unsqueeze(0)
single_image = torch.cat((img, mask), dim=2).detach().cpu().numpy()
single_image = Image.fromarray((single_image * 255).astype(np.uint8), mode="RGBA")
return single_image
if masks is not None:
if len(masks.shape) == 2:
masks = masks.unsqueeze(0)
masks = masks.unsqueeze(3)
images = torch.cat((images, masks), dim=3)
mode="RGBA"
else:
mode="RGB"
pil_image_list = [Image.fromarray((images[i].detach().cpu().numpy() * 255).astype(np.uint8), mode=mode) for i in range(images.shape[0])]
return pil_image_list
def troch_image_dilate(img):
"""
@@ -41,4 +50,99 @@ def troch_image_dilate(img):
dilate_img = cv2.dilate(img, kernel, iterations=1)
img = img * (1 - mask) + dilate_img * mask
img = (img.clip(0, 255) / 255).astype(np.float32)
return torch.from_numpy(img)
return torch.from_numpy(img)
def pils_to_torch_imgs(pils: Union[Image.Image, List[Image.Image]], device="cuda"):
if isinstance(pils, Image.Image):
pils = [pils]
images = []
for pil in pils:
if pil.mode == "RGBA":
pil = pil.convert('RGB')
images.append(TF.to_tensor(pil).permute(1, 2, 0))
images = torch.stack(images, dim=0)
return images
def pils_rgba_to_rgb(pils: Union[Image.Image, List[Image.Image]], bkgd="WHITE"):
if isinstance(pils, Image.Image):
pils = [pils]
rgbs = []
for pil in pils:
if pil.mode == 'RGBA':
new_image = Image.new("RGBA", pil.size, bkgd)
new_image.paste(pil, (0, 0), pil)
rgbs.append(new_image.convert('RGB'))
else:
rgbs.append(pil)
return rgbs
def pil_split_image(image, rows=None, cols=None):
"""
inverse function of make_image_grid
"""
# image is in square
if rows is None and cols is None:
# image.size [W, H]
rows = 1
cols = image.size[0] // image.size[1]
assert cols * image.size[1] == image.size[0]
subimg_size = image.size[1]
elif rows is None:
subimg_size = image.size[0] // cols
rows = image.size[1] // subimg_size
assert rows * subimg_size == image.size[1]
elif cols is None:
subimg_size = image.size[1] // rows
cols = image.size[0] // subimg_size
assert cols * subimg_size == image.size[0]
else:
subimg_size = image.size[1] // rows
assert cols * subimg_size == image.size[0]
subimgs = []
for i in range(rows):
for j in range(cols):
subimg = image.crop((j*subimg_size, i*subimg_size, (j+1)*subimg_size, (i+1)*subimg_size))
subimgs.append(subimg)
return subimgs
def pil_make_image_grid(images, rows=None, cols=None):
if rows is None and cols is None:
rows = 1
cols = len(images)
if rows is None:
rows = len(images) // cols
if len(images) % cols != 0:
rows += 1
if cols is None:
cols = len(images) // rows
if len(images) % rows != 0:
cols += 1
total_imgs = rows * cols
if total_imgs > len(images):
images += [Image.new(images[0].mode, images[0].size) for _ in range(total_imgs - len(images))]
w, h = images[0].size
grid = Image.new(images[0].mode, size=(cols * w, rows * h))
for i, img in enumerate(images):
grid.paste(img, box=(i % cols * w, i // cols * h))
return grid
def pils_erode_masks(mask_list):
out_mask_list = []
for idx, mask in enumerate(mask_list):
arr = np.array(mask)
alpha = (arr[:, :, 3] > 127).astype(np.uint8)
# erode 1px
import cv2
alpha = cv2.erode(alpha, np.ones((3, 3), np.uint8), iterations=1)
alpha = (alpha * 255).astype(np.uint8)
out_mask_list.append(Image.fromarray(alpha[:, :, None]))
return out_mask_list
utils.py
+26
View File
@@ -0,0 +1,26 @@
import logging
import sys
class WarningFilter(logging.Filter):
def filter(self, record):
if record.levelno == logging.WARNING:
record.msg = f"Warn!: {record.msg}"
return True
def create_handler(stream, levels, formatter):
handler = logging.StreamHandler(stream)
handler.setLevel(min(levels))
handler.addFilter(lambda record: record.levelno in levels)
handler.addFilter(WarningFilter()) # Apply the custom filter
handler.setFormatter(formatter)
return handler
def setup_logger(logger_name, level, stdout_levels, stderr_levels, formatter):
logger = logging.getLogger(logger_name)
for handler in logger.handlers[:]:
logger.removeHandler(handler)
logger.setLevel(level)
stdout_handler = create_handler(sys.stdout, stdout_levels, formatter)
stderr_handler = create_handler(sys.stderr, stderr_levels, formatter)
logger.addHandler(stdout_handler)
logger.addHandler(stderr_handler)
web/js/threeVisualizer.js
+1 -1
View File
@@ -25,7 +25,7 @@ const pmremGenerator = new THREE.PMREMGenerator( renderer );
// scene
const scene = new THREE.Scene();
scene.background = new THREE.Color( 0x000000 );
scene.background = new THREE.Color( 0xffffff );
scene.environment = pmremGenerator.fromScene( new RoomEnvironment( renderer ), 0.04 ).texture;
const ambientLight = new THREE.AmbientLight( 0xffffff , 3.0 );