Add Era3D multi-view diffusion model
This commit is contained in:
Mr. For Example
@@ -13,10 +13,15 @@ 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: 23/05/2024**]
|
||||
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 02/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_)
|
||||
|
||||
- **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*
|
||||
|
||||
|
||||
- **InstantMesh Reconstruction Model**: [TencentARC/InstantMesh](https://github.com/TencentARC/InstantMesh)
|
||||
- Sparse multi-view images with white background to 3D Mesh with RGB texture
|
||||
- Works with arbitrary MVDiffusion models (Probably works best with Zero123++, but also works with CRM MVDiffusion model)
|
||||
@@ -25,13 +30,13 @@ This is an extensive node suite that enables ComfyUI to process 3D inputs (Mesh
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/33aecedb-f595-4c12-90dd-89d5f718598e"></video>
|
||||
|
||||
- **Zero123++**: [SUDO-AI-3D/zero123plus](https://github.com/SUDO-AI-3D/zero123plus)
|
||||
- Single image to 6 multi-view images
|
||||
- Single image to 6 view images with resulution: 320X320
|
||||
|
||||
- **CRM**: [thu-ml/CRM](https://github.com/thu-ml/CRM)
|
||||
- Three stages pipeline:
|
||||
1. Single image to 6 multi-view images (Front, Back, Left, Right, Top & Down)
|
||||
2. Single image & 6 multi-view images to 6 same views CCMs (Canonical Coordinate Maps)
|
||||
3. 6 multi-view images & CCMs to 3D mesh
|
||||
1. Single image to 6 view images (Front, Back, Left, Right, Top & Down)
|
||||
2. Single image & 6 view images to 6 same views CCMs (Canonical Coordinate Maps)
|
||||
3. 6 view images & CCMs to 3D mesh
|
||||
- *Note: For low vram pc, if you can't fit all three models for each stages into your GPU memory, then you can divide those three stages into different comfy workflow and run them separately*
|
||||
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/cf68bb83-9244-44df-9db8-f80eb3fdc29e"></video>
|
||||
|
||||
@@ -0,0 +1,547 @@
|
||||
{
|
||||
"last_node_id": 28,
|
||||
"last_link_id": 48,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 25,
|
||||
"type": "[Comfy3D] InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1627.019611278981,
|
||||
108.94257514767159
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 122
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"link": 43
|
||||
},
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"link": 48
|
||||
},
|
||||
{
|
||||
"name": "orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 47
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
46
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
30,
|
||||
1024
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 28,
|
||||
"type": "[Comfy3D] Switch Mesh Axis",
|
||||
"pos": [
|
||||
1988.0196112789795,
|
||||
107.94257514767159
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 130
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 46
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "switched_mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
44
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Switch Mesh Axis"
|
||||
},
|
||||
"widgets_values": [
|
||||
"+y",
|
||||
"+z",
|
||||
"+x",
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 24,
|
||||
"type": "[Comfy3D] Load InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1248,
|
||||
109
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"links": [
|
||||
43
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"instant_mesh_large.ckpt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 22,
|
||||
"type": "[Comfy3D] Era3D Diffusion Model",
|
||||
"pos": [
|
||||
1118,
|
||||
232
|
||||
],
|
||||
"size": {
|
||||
"0": 418.1999816894531,
|
||||
"1": 242
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "era3d_pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 38
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 40
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 39
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiviews",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
41,
|
||||
48
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "multiview_normals",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
42
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "orbit_camposes",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"links": [
|
||||
47
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 2
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Era3D Diffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
420,
|
||||
600,
|
||||
"fixed",
|
||||
3,
|
||||
40,
|
||||
1,
|
||||
4
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 26,
|
||||
"type": "[Comfy3D] Save 3D Mesh",
|
||||
"pos": [
|
||||
2354.019611278979,
|
||||
107.94257514767159
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 44
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
45
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"InstantMesh_test\\Priest_Cat.obj"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "[Comfy3D] Load Diffusers Pipeline",
|
||||
"pos": [
|
||||
417,
|
||||
195
|
||||
],
|
||||
"size": {
|
||||
"0": 383.9324645996094,
|
||||
"1": 109.8599853515625
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
38
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Era3DPipeline",
|
||||
"pengHTYX/MacLab-Era3D-512-6view",
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
483,
|
||||
365
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
40
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
5
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"PriestCat.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "InvertMask",
|
||||
"pos": [
|
||||
878,
|
||||
417
|
||||
],
|
||||
"size": {
|
||||
"0": 140,
|
||||
"1": 26
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 5
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
39
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "InvertMask"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 27,
|
||||
"type": "[Comfy3D] Preview 3DMesh",
|
||||
"pos": [
|
||||
2702,
|
||||
109
|
||||
],
|
||||
"size": [
|
||||
627.5795474783936,
|
||||
527.5795474783936
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh_file_path",
|
||||
"type": "STRING",
|
||||
"link": 45,
|
||||
"widget": {
|
||||
"name": "mesh_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 11,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1627,
|
||||
354
|
||||
],
|
||||
"size": {
|
||||
"0": 1043.55126953125,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 41
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 23,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1629,
|
||||
676
|
||||
],
|
||||
"size": {
|
||||
"0": 1043.55126953125,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 42
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
5,
|
||||
6,
|
||||
1,
|
||||
7,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
38,
|
||||
1,
|
||||
0,
|
||||
22,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
39,
|
||||
7,
|
||||
0,
|
||||
22,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
40,
|
||||
6,
|
||||
0,
|
||||
22,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
41,
|
||||
22,
|
||||
0,
|
||||
11,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
42,
|
||||
22,
|
||||
1,
|
||||
23,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
43,
|
||||
24,
|
||||
0,
|
||||
25,
|
||||
0,
|
||||
"LRM_MODEL"
|
||||
],
|
||||
[
|
||||
44,
|
||||
28,
|
||||
0,
|
||||
26,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
45,
|
||||
26,
|
||||
0,
|
||||
27,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
46,
|
||||
25,
|
||||
0,
|
||||
28,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
47,
|
||||
22,
|
||||
2,
|
||||
25,
|
||||
2,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
48,
|
||||
22,
|
||||
0,
|
||||
25,
|
||||
1,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.620921323059155,
|
||||
"offset": [
|
||||
-581.8092565730503,
|
||||
135.36562552877655
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||
*.model filter=lfs diff=lfs merge=lfs -text
|
||||
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||
*.zst filter=lfs diff=lfs merge=lfs -text
|
||||
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1,3 @@
|
||||
---
|
||||
license: apache-2.0
|
||||
---
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_name_or_path": "stabilityai/stable-diffusion-2-1-unclip",
|
||||
"architectures": [
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1280,
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 5120,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"model_type": "clip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 32,
|
||||
"patch_size": 14,
|
||||
"projection_dim": 1024,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.37.2"
|
||||
}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"_class_name": "DDPMScheduler",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"beta_end": 0.02,
|
||||
"beta_schedule": "squaredcos_cap_v2",
|
||||
"beta_start": 0.0001,
|
||||
"clip_sample": true,
|
||||
"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,
|
||||
"steps_offset": 0,
|
||||
"thresholding": false,
|
||||
"timestep_spacing": "leading",
|
||||
"trained_betas": null,
|
||||
"variance_type": "fixed_small"
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"_class_name": "StableUnCLIPImageNormalizer",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"_name_or_path": "stabilityai/stable-diffusion-2-1-unclip",
|
||||
"embedding_dim": 1024
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
{
|
||||
"_class_name": "StableUnCLIPImg2ImgPipeline",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"feature_extractor": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"image_encoder": [
|
||||
"transformers",
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"image_noising_scheduler": [
|
||||
"diffusers",
|
||||
"DDPMScheduler"
|
||||
],
|
||||
"image_normalizer": [
|
||||
"stable_diffusion",
|
||||
"StableUnCLIPImageNormalizer"
|
||||
],
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"DDIMScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"CLIPTokenizer"
|
||||
],
|
||||
"unet": [
|
||||
"era3d.mvdiffusion.models.unet_mv2d_condition",
|
||||
"UNetMV2DConditionModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
{
|
||||
"_class_name": "DDIMScheduler",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"beta_end": 0.012,
|
||||
"beta_schedule": "linear",
|
||||
"beta_start": 0.00085,
|
||||
"clip_sample": false,
|
||||
"clip_sample_range": 1.0,
|
||||
"dynamic_thresholding_ratio": 0.995,
|
||||
"num_train_timesteps": 1000,
|
||||
"prediction_type": "v_prediction",
|
||||
"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
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_name_or_path": "stabilityai/stable-diffusion-2-1-unclip",
|
||||
"architectures": [
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 0,
|
||||
"dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1024,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4096,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_hidden_layers": 23,
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 512,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.37.2",
|
||||
"vocab_size": 49408
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": "!",
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "!",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49406": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49407": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"bos_token": "<|startoftext|>",
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"do_lower_case": true,
|
||||
"eos_token": "<|endoftext|>",
|
||||
"errors": "replace",
|
||||
"model_max_length": 77,
|
||||
"pad_token": "!",
|
||||
"tokenizer_class": "CLIPTokenizer",
|
||||
"unk_token": "<|endoftext|>"
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,95 @@
|
||||
{
|
||||
"_class_name": "UNetMV2DConditionModel",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"_name_or_path": "../checkpoint_backup/output/unit-unclip-512-6view-randomele-self+row-8w-selfcd-rowmv-2block-linear-depthfilter-step-removewh-bs256-three/checkpoint-40000/unet_ema",
|
||||
"act_fn": "silu",
|
||||
"addition_channels": [
|
||||
1280,
|
||||
1280,
|
||||
1280
|
||||
],
|
||||
"addition_downsample": false,
|
||||
"addition_embed_type": null,
|
||||
"addition_embed_type_num_heads": 64,
|
||||
"addition_time_embed_dim": null,
|
||||
"attention_head_dim": [
|
||||
5,
|
||||
10,
|
||||
20,
|
||||
20
|
||||
],
|
||||
"block_out_channels": [
|
||||
320,
|
||||
640,
|
||||
1280,
|
||||
1280
|
||||
],
|
||||
"cd_attention_last": false,
|
||||
"cd_attention_mid": false,
|
||||
"center_input_sample": false,
|
||||
"class_embed_type": "projection",
|
||||
"class_embeddings_concat": false,
|
||||
"conv_in_kernel": 3,
|
||||
"conv_out_kernel": 3,
|
||||
"cross_attention_dim": 1024,
|
||||
"cross_attention_norm": null,
|
||||
"decay": 0.9999,
|
||||
"down_block_types": [
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"CrossAttnDownBlockMV2D",
|
||||
"DownBlock2D"
|
||||
],
|
||||
"downsample_padding": 1,
|
||||
"dual_cross_attention": false,
|
||||
"encoder_hid_dim": null,
|
||||
"encoder_hid_dim_type": null,
|
||||
"flip_sin_to_cos": true,
|
||||
"freq_shift": 0,
|
||||
"in_channels": 8,
|
||||
"inv_gamma": 1.0,
|
||||
"layers_per_block": 2,
|
||||
"mid_block_only_cross_attention": null,
|
||||
"mid_block_scale_factor": 1,
|
||||
"mid_block_type": "UNetMidBlockMV2DCrossAttn",
|
||||
"min_decay": 0.0,
|
||||
"multiview_attention": true,
|
||||
"mvcd_attention": true,
|
||||
"norm_eps": 1e-05,
|
||||
"norm_num_groups": 32,
|
||||
"num_attention_heads": null,
|
||||
"num_class_embeds": null,
|
||||
"num_regress_blocks": 3,
|
||||
"num_views": 6,
|
||||
"only_cross_attention": false,
|
||||
"optimization_step": 40000,
|
||||
"out_channels": 4,
|
||||
"power": 0.6666666666666666,
|
||||
"projection_camera_embeddings_input_dim": 4,
|
||||
"projection_class_embeddings_input_dim": 2048,
|
||||
"regress_elevation": true,
|
||||
"regress_focal_length": true,
|
||||
"resnet_out_scale_factor": 1.0,
|
||||
"resnet_skip_time_act": false,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"sample_size": 64,
|
||||
"selfattn_block": "self_rowwise",
|
||||
"sparse_mv_attention": true,
|
||||
"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",
|
||||
"CrossAttnUpBlockMV2D",
|
||||
"CrossAttnUpBlockMV2D",
|
||||
"CrossAttnUpBlockMV2D"
|
||||
],
|
||||
"upcast_attention": true,
|
||||
"update_after_step": 0,
|
||||
"use_dino": false,
|
||||
"use_ema_warmup": false,
|
||||
"use_linear_projection": true
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.26.0.dev0",
|
||||
"_name_or_path": "stabilityai/stable-diffusion-2-1-unclip",
|
||||
"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,
|
||||
"layers_per_block": 2,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 768,
|
||||
"scaling_factor": 0.18215,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
num_views: 6
|
||||
|
||||
dataset:
|
||||
prompt_embeds_path: era3d/mvdiffusion/data/fixed_prompt_embeds_6view
|
||||
num_views: ${num_views}
|
||||
bg_color: 'white'
|
||||
img_wh: [512, 512]
|
||||
num_validation_samples: 1000
|
||||
@@ -0,0 +1,138 @@
|
||||
# import decord
|
||||
# decord.bridge.set_bridge('torch')
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
from einops import rearrange
|
||||
from typing import Literal, Tuple, Optional, Any
|
||||
import glob
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class MVDiffusionDatasetV1(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
root_dir: str,
|
||||
num_views: int,
|
||||
bg_color: Any,
|
||||
img_wh: Tuple[int, int],
|
||||
validation: bool = False,
|
||||
num_validation_samples: int = 64,
|
||||
num_samples: Optional[int] = None,
|
||||
caption_path: Optional[str] = None,
|
||||
elevation_range_deg: Tuple[float,float] = (-90, 90),
|
||||
azimuth_range_deg: Tuple[float, float] = (0, 360),
|
||||
):
|
||||
self.all_obj_paths = sorted(glob.glob(os.path.join(root_dir, "*/*")))
|
||||
if not validation:
|
||||
self.all_obj_paths = self.all_obj_paths[:-num_validation_samples]
|
||||
else:
|
||||
self.all_obj_paths = self.all_obj_paths[-num_validation_samples:]
|
||||
if num_samples is not None:
|
||||
self.all_obj_paths = self.all_obj_paths[:num_samples]
|
||||
self.all_obj_ids = [os.path.basename(path) for path in self.all_obj_paths]
|
||||
self.num_views = num_views
|
||||
self.bg_color = bg_color
|
||||
self.img_wh = img_wh
|
||||
|
||||
def get_bg_color(self):
|
||||
if self.bg_color == 'white':
|
||||
bg_color = np.array([1., 1., 1.], dtype=np.float32)
|
||||
elif self.bg_color == 'black':
|
||||
bg_color = np.array([0., 0., 0.], dtype=np.float32)
|
||||
elif self.bg_color == 'gray':
|
||||
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
elif self.bg_color == 'random':
|
||||
bg_color = np.random.rand(3)
|
||||
elif isinstance(self.bg_color, float):
|
||||
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return bg_color
|
||||
|
||||
def load_image(self, img_path, bg_color, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
img = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
assert img.shape[-1] == 4 # RGBA
|
||||
|
||||
alpha = img[...,3:4]
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_obj_ids)
|
||||
|
||||
def __getitem__(self, index):
|
||||
obj_path = self.all_obj_paths[index]
|
||||
obj_id = self.all_obj_ids[index]
|
||||
with open(os.path.join(obj_path, 'meta.json')) as f:
|
||||
meta = json.loads(f.read())
|
||||
|
||||
num_views_all = len(meta['locations'])
|
||||
num_groups = num_views_all // self.num_views
|
||||
|
||||
# random a set of 4 views
|
||||
# the data is arranged in ascending order of the azimuth angle
|
||||
group_ids = random.sample(range(num_groups), k=2)
|
||||
cond_group_id, tgt_group_id = group_ids
|
||||
cond_location = meta['locations'][cond_group_id * self.num_views + random.randint(0, self.num_views - 1)]
|
||||
tgt_locations = meta['locations'][tgt_group_id * self.num_views : tgt_group_id * self.num_views + self.num_views]
|
||||
# random an order
|
||||
start_id = random.randint(0, self.num_views - 1)
|
||||
tgt_locations = tgt_locations[start_id:] + tgt_locations[:start_id]
|
||||
|
||||
cond_elevation = cond_location['elevation']
|
||||
cond_azimuth = cond_location['azimuth']
|
||||
tgt_elevations = [loc['elevation'] for loc in tgt_locations]
|
||||
tgt_azimuths = [loc['azimuth'] for loc in tgt_locations]
|
||||
|
||||
elevations = [ele - cond_elevation for ele in tgt_elevations]
|
||||
azimuths = [(azi - cond_azimuth) % (math.pi * 2) for azi in tgt_azimuths]
|
||||
elevations = torch.as_tensor(elevations).float()
|
||||
azimuths = torch.as_tensor(azimuths).float()
|
||||
elevations_cond = torch.as_tensor([cond_elevation] * self.num_views).float()
|
||||
|
||||
bg_color = self.get_bg_color()
|
||||
img_tensors_in = [
|
||||
self.load_image(os.path.join(obj_path, cond_location['frames'][0]['name']), bg_color, return_type='pt').permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
img_tensors_out = []
|
||||
for loc in tgt_locations:
|
||||
img_path = os.path.join(obj_path, loc['frames'][0]['name'])
|
||||
img_tensor = self.load_image(img_path, bg_color, return_type="pt").permute(2, 0, 1)
|
||||
img_tensors_out.append(img_tensor)
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
|
||||
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
|
||||
|
||||
return {
|
||||
'elevations_cond': elevations_cond,
|
||||
'elevations_cond_deg': torch.rad2deg(elevations_cond),
|
||||
'elevations': elevations,
|
||||
'azimuths': azimuths,
|
||||
'elevations_deg': torch.rad2deg(elevations),
|
||||
'azimuths_deg': torch.rad2deg(azimuths),
|
||||
'imgs_in': img_tensors_in,
|
||||
'imgs_out': img_tensors_out,
|
||||
'camera_embeddings': camera_embeddings
|
||||
}
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
# import decord
|
||||
# decord.bridge.set_bridge('torch')
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
from einops import rearrange
|
||||
from typing import Literal, Tuple, Optional, Any
|
||||
import glob
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from .normal_utils import trans_normal, img2normal, normal2img
|
||||
|
||||
"""
|
||||
load normal and color images together
|
||||
"""
|
||||
class MVDiffusionDatasetV2(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
root_dir: str,
|
||||
num_views: int,
|
||||
bg_color: Any,
|
||||
img_wh: Tuple[int, int],
|
||||
validation: bool = False,
|
||||
num_validation_samples: int = 64,
|
||||
num_samples: Optional[int] = None,
|
||||
caption_path: Optional[str] = None,
|
||||
elevation_range_deg: Tuple[float,float] = (-90, 90),
|
||||
azimuth_range_deg: Tuple[float, float] = (0, 360),
|
||||
):
|
||||
self.all_obj_paths = sorted(glob.glob(os.path.join(root_dir, "*/*")))
|
||||
if not validation:
|
||||
self.all_obj_paths = self.all_obj_paths[:-num_validation_samples]
|
||||
else:
|
||||
self.all_obj_paths = self.all_obj_paths[-num_validation_samples:]
|
||||
if num_samples is not None:
|
||||
self.all_obj_paths = self.all_obj_paths[:num_samples]
|
||||
self.all_obj_ids = [os.path.basename(path) for path in self.all_obj_paths]
|
||||
self.num_views = num_views
|
||||
self.bg_color = bg_color
|
||||
self.img_wh = img_wh
|
||||
|
||||
def get_bg_color(self):
|
||||
if self.bg_color == 'white':
|
||||
bg_color = np.array([1., 1., 1.], dtype=np.float32)
|
||||
elif self.bg_color == 'black':
|
||||
bg_color = np.array([0., 0., 0.], dtype=np.float32)
|
||||
elif self.bg_color == 'gray':
|
||||
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
elif self.bg_color == 'random':
|
||||
bg_color = np.random.rand(3)
|
||||
elif isinstance(self.bg_color, float):
|
||||
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return bg_color
|
||||
|
||||
def load_image(self, img_path, bg_color, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
img = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
assert img.shape[-1] == 4 # RGBA
|
||||
|
||||
alpha = img[...,3:4]
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img, alpha
|
||||
|
||||
def load_normal(self, img_path, bg_color, alpha, RT_w2c=None, RT_w2c_cond=None, return_type='np'):
|
||||
# not using cv2 as may load in uint16 format
|
||||
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
|
||||
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
|
||||
# pil always returns uint8
|
||||
normal = np.array(Image.open(img_path).resize(self.img_wh))
|
||||
|
||||
assert normal.shape[-1] == 3 # RGB
|
||||
|
||||
normal = trans_normal(img2normal(normal), RT_w2c, RT_w2c_cond)
|
||||
img = normal2img(normal)
|
||||
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_obj_ids)
|
||||
|
||||
def __getitem__(self, index):
|
||||
obj_path = self.all_obj_paths[index]
|
||||
obj_id = self.all_obj_ids[index]
|
||||
with open(os.path.join(obj_path, 'meta.json')) as f:
|
||||
meta = json.loads(f.read())
|
||||
|
||||
num_views_all = len(meta['locations'])
|
||||
num_groups = num_views_all // self.num_views
|
||||
|
||||
# random a set of 4 views
|
||||
# the data is arranged in ascending order of the azimuth angle
|
||||
group_ids = random.sample(range(num_groups), k=2)
|
||||
cond_group_id, tgt_group_id = group_ids
|
||||
cond_location = meta['locations'][cond_group_id * self.num_views + random.randint(0, self.num_views - 1)]
|
||||
tgt_locations = meta['locations'][tgt_group_id * self.num_views : tgt_group_id * self.num_views + self.num_views]
|
||||
# random an order
|
||||
start_id = random.randint(0, self.num_views - 1)
|
||||
tgt_locations = tgt_locations[start_id:] + tgt_locations[:start_id]
|
||||
|
||||
cond_elevation = cond_location['elevation']
|
||||
cond_azimuth = cond_location['azimuth']
|
||||
cond_c2w = cond_location['transform_matrix']
|
||||
cond_w2c = np.linalg.inv(cond_c2w)
|
||||
tgt_elevations = [loc['elevation'] for loc in tgt_locations]
|
||||
tgt_azimuths = [loc['azimuth'] for loc in tgt_locations]
|
||||
tgt_c2ws = [loc['transform_matrix'] for loc in tgt_locations]
|
||||
tgt_w2cs = [np.linalg.inv(loc['transform_matrix']) for loc in tgt_locations]
|
||||
|
||||
elevations = [ele - cond_elevation for ele in tgt_elevations]
|
||||
azimuths = [(azi - cond_azimuth) % (math.pi * 2) for azi in tgt_azimuths]
|
||||
elevations = torch.as_tensor(elevations).float()
|
||||
azimuths = torch.as_tensor(azimuths).float()
|
||||
elevations_cond = torch.as_tensor([cond_elevation] * self.num_views).float()
|
||||
|
||||
bg_color = self.get_bg_color()
|
||||
img_tensors_in = [
|
||||
self.load_image(os.path.join(obj_path, cond_location['frames'][0]['name']), bg_color, return_type='pt')[0].permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
img_tensors_out = []
|
||||
normal_tensors_out = []
|
||||
for loc, tgt_w2c in zip(tgt_locations, tgt_w2cs):
|
||||
img_path = os.path.join(obj_path, loc['frames'][0]['name'])
|
||||
img_tensor, alpha = self.load_image(img_path, bg_color, return_type="pt")
|
||||
img_tensor = img_tensor.permute(2, 0, 1)
|
||||
img_tensors_out.append(img_tensor)
|
||||
|
||||
normal_path = os.path.join(obj_path, loc['frames'][1]['name'])
|
||||
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt").permute(2, 0, 1)
|
||||
normal_tensors_out.append(normal_tensor)
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
|
||||
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
|
||||
|
||||
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
|
||||
|
||||
return {
|
||||
'elevations_cond': elevations_cond,
|
||||
'elevations_cond_deg': torch.rad2deg(elevations_cond),
|
||||
'elevations': elevations,
|
||||
'azimuths': azimuths,
|
||||
'elevations_deg': torch.rad2deg(elevations),
|
||||
'azimuths_deg': torch.rad2deg(azimuths),
|
||||
'imgs_in': img_tensors_in,
|
||||
'imgs_out': img_tensors_out,
|
||||
'normals_out': normal_tensors_out,
|
||||
'camera_embeddings': camera_embeddings
|
||||
}
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,78 @@
|
||||
from transformers import CLIPTokenizer, CLIPTextModel
|
||||
import torch
|
||||
import os
|
||||
|
||||
root = '/mnt/data/lipeng/'
|
||||
pretrained_model_name_or_path = 'stabilityai/stable-diffusion-2-1-unclip'
|
||||
|
||||
|
||||
weight_dtype = torch.float16
|
||||
device = torch.device("cuda:0")
|
||||
tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
|
||||
text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
|
||||
text_encoder = text_encoder.to(device, dtype=weight_dtype)
|
||||
|
||||
def generate_mv_embeds():
|
||||
path = './fixed_prompt_embeds_8view'
|
||||
os.makedirs(path, exist_ok=True)
|
||||
views = ["front", "front_right", "right", "back_right", "back", " back_left", "left", "front_left"]
|
||||
# views = ["front", "front_right", "right", "back", "left", "front_left"]
|
||||
# views = ["front", "right", "back", "left"]
|
||||
clr_prompt = [f"a rendering image of 3D models, {view} view, color map." for view in views]
|
||||
normal_prompt = [f"a rendering image of 3D models, {view} view, normal map." for view in views]
|
||||
|
||||
|
||||
for id, text_prompt in enumerate([clr_prompt, normal_prompt]):
|
||||
print(text_prompt)
|
||||
text_inputs = tokenizer(text_prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt").to(device)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = tokenizer(text_prompt, padding="longest", return_tensors="pt").input_ids
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
||||
text_input_ids, untruncated_ids):
|
||||
removed_text = tokenizer.batch_decode(
|
||||
untruncated_ids[:, tokenizer.model_max_length - 1 : -1]
|
||||
)
|
||||
if hasattr(text_encoder.config, "use_attention_mask") and text_encoder.config.use_attention_mask:
|
||||
attention_mask = text_inputs.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
prompt_embeds = text_encoder(text_input_ids.to(device), attention_mask=attention_mask,)
|
||||
prompt_embeds = prompt_embeds[0].detach().cpu()
|
||||
print(prompt_embeds.shape)
|
||||
|
||||
|
||||
# print(prompt_embeds.dtype)
|
||||
if id == 0:
|
||||
torch.save(prompt_embeds, f'./{path}/clr_embeds.pt')
|
||||
else:
|
||||
torch.save(prompt_embeds, f'./{path}/normal_embeds.pt')
|
||||
print('done')
|
||||
|
||||
|
||||
def generate_img_embeds():
|
||||
path = './fixed_prompt_embeds_persp2ortho'
|
||||
os.makedirs(path, exist_ok=True)
|
||||
text_prompt = ["a orthogonal renderining image of 3D models"]
|
||||
print(text_prompt)
|
||||
text_inputs = tokenizer(text_prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt").to(device)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = tokenizer(text_prompt, padding="longest", return_tensors="pt").input_ids
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
||||
text_input_ids, untruncated_ids):
|
||||
removed_text = tokenizer.batch_decode(
|
||||
untruncated_ids[:, tokenizer.model_max_length - 1 : -1]
|
||||
)
|
||||
if hasattr(text_encoder.config, "use_attention_mask") and text_encoder.config.use_attention_mask:
|
||||
attention_mask = text_inputs.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
prompt_embeds = text_encoder(text_input_ids.to(device), attention_mask=attention_mask,)
|
||||
prompt_embeds = prompt_embeds[0].detach().cpu()
|
||||
print(prompt_embeds.shape)
|
||||
|
||||
# print(prompt_embeds.dtype)
|
||||
|
||||
torch.save(prompt_embeds, f'./{path}/embeds.pt')
|
||||
print('done')
|
||||
|
||||
generate_img_embeds()
|
||||
@@ -0,0 +1,78 @@
|
||||
import numpy as np
|
||||
def deg2rad(deg):
|
||||
return deg*np.pi/180
|
||||
|
||||
def inv_RT(RT):
|
||||
# RT_h = np.concatenate([RT, np.array([[0,0,0,1]])], axis=0)
|
||||
RT_inv = np.linalg.inv(RT)
|
||||
|
||||
return RT_inv[:3, :]
|
||||
def camNormal2worldNormal(rot_c2w, camNormal):
|
||||
H,W,_ = camNormal.shape
|
||||
normal_img = np.matmul(rot_c2w[None, :, :], camNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
return normal_img
|
||||
|
||||
def worldNormal2camNormal(rot_w2c, normal_map_world):
|
||||
H,W,_ = normal_map_world.shape
|
||||
# normal_img = np.matmul(rot_w2c[None, :, :], worldNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
|
||||
|
||||
# faster version
|
||||
# Reshape the normal map into a 2D array where each row represents a normal vector
|
||||
normal_map_flat = normal_map_world.reshape(-1, 3)
|
||||
|
||||
# Transform the normal vectors using the transformation matrix
|
||||
normal_map_camera_flat = np.dot(normal_map_flat, rot_w2c.T)
|
||||
|
||||
# Reshape the transformed normal map back to its original shape
|
||||
normal_map_camera = normal_map_camera_flat.reshape(normal_map_world.shape)
|
||||
|
||||
return normal_map_camera
|
||||
|
||||
def trans_normal(normal, RT_w2c, RT_w2c_target):
|
||||
|
||||
# normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
|
||||
# normal_target_cam = worldNormal2camNormal(RT_w2c_target[:3,:3], normal_world)
|
||||
|
||||
relative_RT = np.matmul(RT_w2c_target[:3,:3], np.linalg.inv(RT_w2c[:3,:3]))
|
||||
return worldNormal2camNormal(relative_RT[:3,:3], normal)
|
||||
|
||||
def trans_normal_complex(normal, RT_w2c, RT_w2c_rela_to_cond):
|
||||
# camview -> world -> condview
|
||||
normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
|
||||
# debug_normal_world = normal2img(normal_world)
|
||||
|
||||
# relative_RT = np.matmul(RT_w2c_rela_to_cond[:3,:3], np.linalg.inv(RT_w2c[:3,:3]))
|
||||
normal_target_cam = worldNormal2camNormal(RT_w2c_rela_to_cond[:3,:3], normal_world)
|
||||
# normal_condview = normal2img(normal_target_cam)
|
||||
return normal_target_cam
|
||||
def img2normal(img):
|
||||
return (img/255.)*2-1
|
||||
|
||||
def normal2img(normal):
|
||||
return np.uint8((normal*0.5+0.5)*255)
|
||||
|
||||
def norm_normalize(normal, dim=-1):
|
||||
|
||||
normal = normal/(np.linalg.norm(normal, axis=dim, keepdims=True)+1e-6)
|
||||
|
||||
return normal
|
||||
|
||||
def plot_grid_images(images, row, col, path=None):
|
||||
import cv2
|
||||
"""
|
||||
Args:
|
||||
images: np.array [B, H, W, 3]
|
||||
row:
|
||||
col:
|
||||
save_path:
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
images = images.detach().cpu().numpy()
|
||||
assert row * col == images.shape[0]
|
||||
images = np.vstack([np.hstack(images[r * col:(r + 1) * col]) for r in range(row)])
|
||||
if path:
|
||||
cv2.imwrite(path, images[:,:,::-1] * 255)
|
||||
return images
|
||||
@@ -0,0 +1,248 @@
|
||||
from typing import Dict
|
||||
import numpy as np
|
||||
from omegaconf import DictConfig, ListConfig
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from pathlib import Path
|
||||
import json
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from einops import rearrange
|
||||
from typing import Literal, Tuple, Optional, Any
|
||||
import cv2
|
||||
import random
|
||||
|
||||
import json
|
||||
import os, sys
|
||||
import math
|
||||
|
||||
from glob import glob
|
||||
|
||||
import PIL.Image
|
||||
from .normal_utils import trans_normal, normal2img, img2normal
|
||||
import pdb
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def add_margin(pil_img, color=0, size=256):
|
||||
width, height = pil_img.size
|
||||
result = Image.new(pil_img.mode, (size, size), color)
|
||||
result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
|
||||
return result
|
||||
|
||||
def scale_and_place_object(image, scale_factor):
|
||||
assert np.shape(image)[-1]==4 # RGBA
|
||||
|
||||
# Extract the alpha channel (transparency) and the object (RGB channels)
|
||||
alpha_channel = image[:, :, 3]
|
||||
|
||||
# Find the bounding box coordinates of the object
|
||||
coords = cv2.findNonZero(alpha_channel)
|
||||
x, y, width, height = cv2.boundingRect(coords)
|
||||
|
||||
# Calculate the scale factor for resizing
|
||||
original_height, original_width = image.shape[:2]
|
||||
|
||||
if width > height:
|
||||
size = width
|
||||
original_size = original_width
|
||||
else:
|
||||
size = height
|
||||
original_size = original_height
|
||||
|
||||
scale_factor = min(scale_factor, size / (original_size+0.0))
|
||||
|
||||
new_size = scale_factor * original_size
|
||||
scale_factor = new_size / size
|
||||
|
||||
# Calculate the new size based on the scale factor
|
||||
new_width = int(width * scale_factor)
|
||||
new_height = int(height * scale_factor)
|
||||
|
||||
center_x = original_width // 2
|
||||
center_y = original_height // 2
|
||||
|
||||
paste_x = center_x - (new_width // 2)
|
||||
paste_y = center_y - (new_height // 2)
|
||||
|
||||
# Resize the object (RGB channels) to the new size
|
||||
rescaled_object = cv2.resize(image[y:y+height, x:x+width], (new_width, new_height))
|
||||
|
||||
# Create a new RGBA image with the resized image
|
||||
new_image = np.zeros((original_height, original_width, 4), dtype=np.uint8)
|
||||
|
||||
new_image[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = rescaled_object
|
||||
|
||||
return new_image
|
||||
|
||||
class SingleImageDataset(Dataset):
|
||||
def __init__(self,
|
||||
num_views: int,
|
||||
img_wh: Tuple[int, int],
|
||||
bg_color: str,
|
||||
crop_size: int = 224,
|
||||
single_image: Optional[PIL.Image.Image] = None,
|
||||
images_dir: Optional[str] = None,
|
||||
num_validation_samples: Optional[int] = None,
|
||||
filepaths: Optional[list] = None,
|
||||
cond_type: Optional[str] = None,
|
||||
prompt_embeds_path: Optional[str] = None,
|
||||
gt_path: Optional[str] = None,
|
||||
dtype: Optional[torch.dtype] = torch.float16,
|
||||
) -> None:
|
||||
"""Create a dataset from a folder of images.
|
||||
If you pass in a root directory it will be searched for images
|
||||
ending in ext (ext can be a list)
|
||||
"""
|
||||
self.images_dir = images_dir
|
||||
self.num_views = num_views
|
||||
self.img_wh = img_wh
|
||||
self.crop_size = crop_size
|
||||
self.bg_color = bg_color
|
||||
self.cond_type = cond_type
|
||||
self.gt_path = gt_path
|
||||
self.dtype = dtype
|
||||
|
||||
|
||||
if single_image is None:
|
||||
if filepaths is None:
|
||||
# Get a list of all files in the directory
|
||||
file_list = os.listdir(self.images_dir)
|
||||
else:
|
||||
file_list = filepaths
|
||||
|
||||
# Filter the files that end with .png or .jpg
|
||||
self.file_list = [file for file in file_list if file.endswith(('.png', '.jpg', '.webp'))]
|
||||
else:
|
||||
self.file_list = None
|
||||
|
||||
# load all images
|
||||
self.all_images = []
|
||||
self.all_alphas = []
|
||||
bg_color = self.get_bg_color()
|
||||
|
||||
if single_image is not None:
|
||||
image, alpha = self.load_image(None, bg_color, return_type='pt', Imagefile=single_image)
|
||||
self.all_images.append(image)
|
||||
self.all_alphas.append(alpha)
|
||||
else:
|
||||
for file in self.file_list:
|
||||
print(os.path.join(self.images_dir, file))
|
||||
image, alpha = self.load_image(os.path.join(self.images_dir, file), bg_color, return_type='pt')
|
||||
self.all_images.append(image)
|
||||
self.all_alphas.append(alpha)
|
||||
|
||||
|
||||
self.all_images = self.all_images[:num_validation_samples]
|
||||
self.all_alphas = self.all_alphas[:num_validation_samples]
|
||||
|
||||
try:
|
||||
self.normal_text_embeds = torch.load(f'{prompt_embeds_path}/normal_embeds.pt')
|
||||
self.color_text_embeds = torch.load(f'{prompt_embeds_path}/clr_embeds.pt') # 4view
|
||||
except:
|
||||
self.color_text_embeds = torch.load(f'{prompt_embeds_path}/embeds.pt')
|
||||
self.normal_text_embeds = None
|
||||
|
||||
def __len__(self):
|
||||
return len(self.all_images)
|
||||
|
||||
def get_bg_color(self):
|
||||
if self.bg_color == 'white':
|
||||
bg_color = np.array([1., 1., 1.], dtype=np.float32)
|
||||
elif self.bg_color == 'black':
|
||||
bg_color = np.array([0., 0., 0.], dtype=np.float32)
|
||||
elif self.bg_color == 'gray':
|
||||
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
||||
elif self.bg_color == 'random':
|
||||
bg_color = np.random.rand(3)
|
||||
elif isinstance(self.bg_color, float):
|
||||
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return bg_color
|
||||
|
||||
|
||||
def load_image(self, img_path, bg_color, return_type='np', Imagefile=None):
|
||||
# pil always returns uint8
|
||||
if Imagefile is None:
|
||||
image_input = Image.open(img_path)
|
||||
else:
|
||||
image_input = Imagefile
|
||||
image_size = self.img_wh[0]
|
||||
|
||||
if self.crop_size!=-1:
|
||||
alpha_np = np.asarray(image_input)[:, :, 3]
|
||||
coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)]
|
||||
min_x, min_y = np.min(coords, 0)
|
||||
max_x, max_y = np.max(coords, 0)
|
||||
ref_img_ = image_input.crop((min_x, min_y, max_x, max_y))
|
||||
h, w = ref_img_.height, ref_img_.width
|
||||
scale = self.crop_size / max(h, w)
|
||||
h_, w_ = int(scale * h), int(scale * w)
|
||||
ref_img_ = ref_img_.resize((w_, h_))
|
||||
image_input = add_margin(ref_img_, size=image_size)
|
||||
else:
|
||||
image_input = add_margin(image_input, size=max(image_input.height, image_input.width))
|
||||
image_input = image_input.resize((image_size, image_size))
|
||||
|
||||
# img = scale_and_place_object(img, self.scale_ratio)
|
||||
img = np.array(image_input)
|
||||
img = img.astype(np.float32) / 255. # [0, 1]
|
||||
assert img.shape[-1] == 4 # RGBA
|
||||
|
||||
alpha = img[...,3:4]
|
||||
img = img[...,:3] * alpha + bg_color * (1 - alpha)
|
||||
|
||||
if return_type == "np":
|
||||
pass
|
||||
elif return_type == "pt":
|
||||
img = torch.from_numpy(img)
|
||||
alpha = torch.from_numpy(alpha)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return img, alpha
|
||||
|
||||
|
||||
def __getitem__(self, index):
|
||||
image = self.all_images[index%len(self.all_images)]
|
||||
alpha = self.all_alphas[index%len(self.all_images)]
|
||||
if self.file_list is not None:
|
||||
filename = self.file_list[index%len(self.all_images)].replace(".png", "")
|
||||
else:
|
||||
filename = 'null'
|
||||
img_tensors_in = [
|
||||
image.permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
|
||||
alpha_tensors_in = [
|
||||
alpha.permute(2, 0, 1)
|
||||
] * self.num_views
|
||||
|
||||
img_tensors_in = torch.stack(img_tensors_in, dim=0).to(self.dtype) # (Nv, 3, H, W)
|
||||
alpha_tensors_in = torch.stack(alpha_tensors_in, dim=0).to(self.dtype) # (Nv, 3, H, W)
|
||||
|
||||
if self.gt_path is not None:
|
||||
gt_image = self.gt_images[index%len(self.all_images)]
|
||||
gt_alpha = self.gt_alpha[index%len(self.all_images)]
|
||||
gt_img_tensors_in = [gt_image.permute(2, 0, 1) ] * self.num_views
|
||||
gt_alpha_tensors_in = [gt_alpha.permute(2, 0, 1) ] * self.num_views
|
||||
gt_img_tensors_in = torch.stack(gt_img_tensors_in, dim=0).to(self.dtype)
|
||||
gt_alpha_tensors_in = torch.stack(gt_alpha_tensors_in, dim=0).to(self.dtype)
|
||||
|
||||
normal_prompt_embeddings = self.normal_text_embeds if hasattr(self, 'normal_text_embeds') else None
|
||||
color_prompt_embeddings = self.color_text_embeds if hasattr(self, 'color_text_embeds') else None
|
||||
|
||||
out = {
|
||||
'imgs_in': img_tensors_in,
|
||||
'alphas': alpha_tensors_in,
|
||||
'normal_prompt_embeddings': normal_prompt_embeddings,
|
||||
'color_prompt_embeddings': color_prompt_embeddings,
|
||||
'filename': filename,
|
||||
}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,978 @@
|
||||
# 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.
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import ImagePositionalEmbeddings
|
||||
from diffusers.utils import BaseOutput, deprecate
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
from diffusers.models.attention import FeedForward, AdaLayerNorm, AdaLayerNormZero, Attention
|
||||
from diffusers.models.embeddings import PatchEmbed
|
||||
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
from einops import rearrange
|
||||
import pdb
|
||||
import random
|
||||
import math
|
||||
|
||||
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerMV2DModelOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`Transformer2DModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
|
||||
The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
|
||||
distributions for the unnoised latent pixels.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class TransformerMV2DModel(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
A 2D Transformer model for image-like data.
|
||||
|
||||
Parameters:
|
||||
num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
|
||||
in_channels (`int`, *optional*):
|
||||
The number of channels in the input and output (specify if the input is **continuous**).
|
||||
num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
|
||||
sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
|
||||
This is fixed during training since it is used to learn a number of position embeddings.
|
||||
num_vector_embeds (`int`, *optional*):
|
||||
The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
|
||||
Includes the class for the masked latent pixel.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
|
||||
num_embeds_ada_norm ( `int`, *optional*):
|
||||
The number of diffusion steps used during training. Pass if at least one of the norm_layers is
|
||||
`AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
|
||||
added to the hidden states.
|
||||
|
||||
During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
|
||||
attention_bias (`bool`, *optional*):
|
||||
Configure if the `TransformerBlocks` attention should contain a bias parameter.
|
||||
"""
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
out_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
sample_size: Optional[int] = None,
|
||||
num_vector_embeds: Optional[int] = None,
|
||||
patch_size: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_type: str = "layer_norm",
|
||||
norm_elementwise_affine: bool = True,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool=False,
|
||||
cd_attention_mid: bool=False,
|
||||
multiview_attention: bool=True,
|
||||
sparse_mv_attention: bool = True, # not used
|
||||
mvcd_attention: bool=False
|
||||
):
|
||||
super().__init__()
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
|
||||
# Define whether input is continuous or discrete depending on configuration
|
||||
self.is_input_continuous = (in_channels is not None) and (patch_size is None)
|
||||
self.is_input_vectorized = num_vector_embeds is not None
|
||||
self.is_input_patches = in_channels is not None and patch_size is not None
|
||||
|
||||
if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
|
||||
deprecation_message = (
|
||||
f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
|
||||
" incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
|
||||
" Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
|
||||
" results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
|
||||
" would be very nice if you could open a Pull request for the `transformer/config.json` file"
|
||||
)
|
||||
deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
|
||||
norm_type = "ada_norm"
|
||||
|
||||
if self.is_input_continuous and self.is_input_vectorized:
|
||||
raise ValueError(
|
||||
f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
|
||||
" sure that either `in_channels` or `num_vector_embeds` is None."
|
||||
)
|
||||
elif self.is_input_vectorized and self.is_input_patches:
|
||||
raise ValueError(
|
||||
f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
|
||||
" sure that either `num_vector_embeds` or `num_patches` is None."
|
||||
)
|
||||
elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
|
||||
raise ValueError(
|
||||
f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
|
||||
f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
|
||||
)
|
||||
|
||||
# 2. Define input layers
|
||||
if self.is_input_continuous:
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
if use_linear_projection:
|
||||
self.proj_in = LoRACompatibleLinear(in_channels, inner_dim)
|
||||
else:
|
||||
self.proj_in = LoRACompatibleConv(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
|
||||
elif self.is_input_vectorized:
|
||||
assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
|
||||
assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
|
||||
|
||||
self.height = sample_size
|
||||
self.width = sample_size
|
||||
self.num_vector_embeds = num_vector_embeds
|
||||
self.num_latent_pixels = self.height * self.width
|
||||
|
||||
self.latent_image_embedding = ImagePositionalEmbeddings(
|
||||
num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
|
||||
)
|
||||
elif self.is_input_patches:
|
||||
assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
|
||||
|
||||
self.height = sample_size
|
||||
self.width = sample_size
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.pos_embed = PatchEmbed(
|
||||
height=sample_size,
|
||||
width=sample_size,
|
||||
patch_size=patch_size,
|
||||
in_channels=in_channels,
|
||||
embed_dim=inner_dim,
|
||||
)
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicMVTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
norm_type=norm_type,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
mvcd_attention=mvcd_attention
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
if self.is_input_continuous:
|
||||
# TODO: should use out_channels for continuous projections
|
||||
if use_linear_projection:
|
||||
self.proj_out = LoRACompatibleLinear(inner_dim, in_channels)
|
||||
else:
|
||||
self.proj_out = LoRACompatibleConv(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
|
||||
elif self.is_input_vectorized:
|
||||
self.norm_out = nn.LayerNorm(inner_dim)
|
||||
self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
|
||||
elif self.is_input_patches:
|
||||
self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
|
||||
self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
|
||||
self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
|
||||
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
|
||||
self-attention.
|
||||
timestep ( `torch.LongTensor`, *optional*):
|
||||
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
|
||||
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
|
||||
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
|
||||
`AdaLayerZeroNorm`.
|
||||
encoder_attention_mask ( `torch.Tensor`, *optional*):
|
||||
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
|
||||
|
||||
* Mask `(batch, sequence_length)` True = keep, False = discard.
|
||||
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
|
||||
|
||||
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
|
||||
above. This bias will be added to the cross-attention scores.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
|
||||
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
|
||||
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
|
||||
# expects mask of shape:
|
||||
# [batch, key_tokens]
|
||||
# adds singleton query_tokens dimension:
|
||||
# [batch, 1, key_tokens]
|
||||
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
|
||||
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
|
||||
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
|
||||
if attention_mask is not None and attention_mask.ndim == 2:
|
||||
# assume that mask is expressed as:
|
||||
# (1 = keep, 0 = discard)
|
||||
# convert mask into a bias that can be added to attention scores:
|
||||
# (keep = +0, discard = -10000.0)
|
||||
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
|
||||
encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# 1. Input
|
||||
if self.is_input_continuous:
|
||||
batch, _, height, width = hidden_states.shape
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
|
||||
else:
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
elif self.is_input_vectorized:
|
||||
hidden_states = self.latent_image_embedding(hidden_states)
|
||||
elif self.is_input_patches:
|
||||
hidden_states = self.pos_embed(hidden_states)
|
||||
|
||||
# 2. Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
timestep=timestep,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
if self.is_input_continuous:
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
else:
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
output = hidden_states + residual
|
||||
elif self.is_input_vectorized:
|
||||
hidden_states = self.norm_out(hidden_states)
|
||||
logits = self.out(hidden_states)
|
||||
# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
|
||||
logits = logits.permute(0, 2, 1)
|
||||
|
||||
# log(p(x_0))
|
||||
output = F.log_softmax(logits.double(), dim=1).float()
|
||||
elif self.is_input_patches:
|
||||
# TODO: cleanup!
|
||||
conditioning = self.transformer_blocks[0].norm1.emb(
|
||||
timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
|
||||
hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
|
||||
hidden_states = self.proj_out_2(hidden_states)
|
||||
|
||||
# unpatchify
|
||||
height = width = int(hidden_states.shape[1] ** 0.5)
|
||||
hidden_states = hidden_states.reshape(
|
||||
shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
|
||||
)
|
||||
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
|
||||
output = hidden_states.reshape(
|
||||
shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return TransformerMV2DModelOutput(sample=output)
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class BasicMVTransformerBlock(nn.Module):
|
||||
r"""
|
||||
A basic Transformer block.
|
||||
|
||||
Parameters:
|
||||
dim (`int`): The number of channels in the input and output.
|
||||
num_attention_heads (`int`): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`): The number of channels in each head.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
|
||||
only_cross_attention (`bool`, *optional*):
|
||||
Whether to use only cross-attention layers. In this case two cross attention layers are used.
|
||||
double_self_attention (`bool`, *optional*):
|
||||
Whether to use two self-attention layers. In this case no cross attention layers are used.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
||||
num_embeds_ada_norm (:
|
||||
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
|
||||
attention_bias (:
|
||||
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
dropout=0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_type: str = "layer_norm",
|
||||
final_dropout: bool = False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
mvcd_attention: bool = False,
|
||||
rowwise_attention: bool = True
|
||||
):
|
||||
super().__init__()
|
||||
self.only_cross_attention = only_cross_attention
|
||||
|
||||
self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
|
||||
self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
|
||||
|
||||
if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
|
||||
raise ValueError(
|
||||
f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
|
||||
f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
|
||||
)
|
||||
|
||||
# Define 3 blocks. Each block has its own normalization layer.
|
||||
# 1. Self-Attn
|
||||
if self.use_ada_layer_norm:
|
||||
self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
|
||||
else:
|
||||
self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
|
||||
self.multiview_attention = multiview_attention
|
||||
self.mvcd_attention = mvcd_attention
|
||||
self.rowwise_attention = multiview_attention and rowwise_attention
|
||||
|
||||
# rowwise multiview attention
|
||||
|
||||
print('INFO: using row wise attention...')
|
||||
|
||||
self.attn1 = CustomAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=MVAttnProcessor()
|
||||
)
|
||||
|
||||
# 2. Cross-Attn
|
||||
if cross_attention_dim is not None or double_self_attention:
|
||||
# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
|
||||
# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
|
||||
# the second cross attention block.
|
||||
self.norm2 = (
|
||||
AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
if self.use_ada_layer_norm
|
||||
else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
)
|
||||
self.attn2 = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=cross_attention_dim if not double_self_attention else None,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
) # is self-attn if encoder_hidden_states is none
|
||||
else:
|
||||
self.norm2 = None
|
||||
self.attn2 = None
|
||||
|
||||
# 3. Feed-forward
|
||||
self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
|
||||
self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)
|
||||
|
||||
# let chunk size default to None
|
||||
self._chunk_size = None
|
||||
self._chunk_dim = 0
|
||||
|
||||
self.num_views = num_views
|
||||
|
||||
self.cd_attention_last = cd_attention_last
|
||||
|
||||
if self.cd_attention_last:
|
||||
# Joint task -Attn
|
||||
self.attn_joint = CustomJointAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=JointAttnProcessor()
|
||||
)
|
||||
nn.init.zeros_(self.attn_joint.to_out[0].weight.data)
|
||||
self.norm_joint = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
|
||||
|
||||
|
||||
self.cd_attention_mid = cd_attention_mid
|
||||
|
||||
if self.cd_attention_mid:
|
||||
print("joint twice")
|
||||
# Joint task -Attn
|
||||
self.attn_joint_twice = CustomJointAttention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
|
||||
upcast_attention=upcast_attention,
|
||||
processor=JointAttnProcessor()
|
||||
)
|
||||
nn.init.zeros_(self.attn_joint_twice.to_out[0].weight.data)
|
||||
self.norm_joint_twice = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
|
||||
|
||||
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
|
||||
# Sets chunk feed-forward
|
||||
self._chunk_size = chunk_size
|
||||
self._chunk_dim = dim
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
assert attention_mask is None # not supported yet
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
# 1. Self-Attention
|
||||
if self.use_ada_layer_norm:
|
||||
norm_hidden_states = self.norm1(hidden_states, timestep)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
|
||||
hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
else:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
|
||||
cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
|
||||
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
|
||||
attention_mask=attention_mask,
|
||||
multiview_attention=self.multiview_attention,
|
||||
mvcd_attention=self.mvcd_attention,
|
||||
num_views=self.num_views,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# joint attention twice
|
||||
if self.cd_attention_mid:
|
||||
norm_hidden_states = (
|
||||
self.norm_joint_twice(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_joint_twice(hidden_states)
|
||||
)
|
||||
hidden_states = self.attn_joint_twice(norm_hidden_states) + hidden_states
|
||||
|
||||
# 2. Cross-Attention
|
||||
if self.attn2 is not None:
|
||||
norm_hidden_states = (
|
||||
self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
|
||||
)
|
||||
|
||||
attn_output = self.attn2(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=encoder_attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# 3. Feed-forward
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
|
||||
if self._chunk_size is not None:
|
||||
# "feed_forward_chunk_size" can be used to save memory
|
||||
if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
|
||||
raise ValueError(
|
||||
f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
|
||||
)
|
||||
|
||||
num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
|
||||
ff_output = torch.cat(
|
||||
[self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
|
||||
dim=self._chunk_dim,
|
||||
)
|
||||
else:
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
|
||||
if self.cd_attention_last:
|
||||
norm_hidden_states = (
|
||||
self.norm_joint(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_joint(hidden_states)
|
||||
)
|
||||
hidden_states = self.attn_joint(norm_hidden_states) + hidden_states
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CustomAttention(Attention):
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, use_memory_efficient_attention_xformers: bool, *args, **kwargs
|
||||
):
|
||||
processor = XFormersMVAttnProcessor()
|
||||
self.set_processor(processor)
|
||||
# print("using xformers attention processor")
|
||||
|
||||
|
||||
class CustomJointAttention(Attention):
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, use_memory_efficient_attention_xformers: bool, *args, **kwargs
|
||||
):
|
||||
processor = XFormersJointAttnProcessor()
|
||||
self.set_processor(processor)
|
||||
# print("using xformers attention processor")
|
||||
|
||||
class MVAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_views=1,
|
||||
multiview_attention=True
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
height = int(math.sqrt(sequence_length))
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# print('query', query.shape, 'key', key.shape, 'value', value.shape)
|
||||
#([bx4, 1024, 320]) key torch.Size([bx4, 1024, 320]) value torch.Size([bx4, 1024, 320])
|
||||
# pdb.set_trace()
|
||||
# multi-view self-attention
|
||||
key = rearrange(key, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height)
|
||||
value = rearrange(value, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height)
|
||||
query = rearrange(query, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height) # torch.Size([192, 384, 320])
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
hidden_states = rearrange(hidden_states, "(b h) (v w) c -> (b v) (h w) c", v=num_views, h=height)
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class XFormersMVAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_views=1,
|
||||
multiview_attention=True,
|
||||
mvcd_attention=False,
|
||||
):
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
height = int(math.sqrt(sequence_length))
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# from yuancheng; here attention_mask is None
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
print('Warning: using group norm, pay attention to use it in row-wise attention')
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key_raw = attn.to_k(encoder_hidden_states)
|
||||
value_raw = attn.to_v(encoder_hidden_states)
|
||||
|
||||
# print('query', query.shape, 'key', key.shape, 'value', value.shape)
|
||||
# pdb.set_trace()
|
||||
|
||||
key = rearrange(key_raw, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height)
|
||||
value = rearrange(value_raw, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height)
|
||||
query = rearrange(query, "(b v) (h w) c -> (b h) (v w) c", v=num_views, h=height) # torch.Size([192, 384, 320])
|
||||
if mvcd_attention:
|
||||
# memory efficient, cross domain attention
|
||||
key_0, key_1 = torch.chunk(key_raw, dim=0, chunks=2) # keys shape (b t) d c
|
||||
value_0, value_1 = torch.chunk(value_raw, dim=0, chunks=2)
|
||||
key_cross = torch.concat([key_1, key_0], dim=0)
|
||||
value_cross = torch.concat([value_1, value_0], dim=0) # shape (b t) d c
|
||||
key = torch.cat([key, key_cross], dim=1)
|
||||
value = torch.cat([value, value_cross], dim=1) # shape (b t) (t+1 d) c
|
||||
|
||||
|
||||
query = attn.head_to_batch_dim(query) # torch.Size([960, 384, 64])
|
||||
key = attn.head_to_batch_dim(key)
|
||||
value = attn.head_to_batch_dim(value)
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
# print(hidden_states.shape)
|
||||
hidden_states = rearrange(hidden_states, "(b h) (v w) c -> (b v) (h w) c", v=num_views, h=height)
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class XFormersJointAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_tasks=2
|
||||
):
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
# from yuancheng; here attention_mask is None
|
||||
if attention_mask is not None:
|
||||
# expand our mask's singleton query_tokens dimension:
|
||||
# [batch*heads, 1, key_tokens] ->
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# so that it can be added as a bias onto the attention scores that xformers computes:
|
||||
# [batch*heads, query_tokens, key_tokens]
|
||||
# we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
|
||||
_, query_tokens, _ = hidden_states.shape
|
||||
attention_mask = attention_mask.expand(-1, query_tokens, -1)
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
assert num_tasks == 2 # only support two tasks now
|
||||
|
||||
key_0, key_1 = torch.chunk(key, dim=0, chunks=2) # keys shape (b t) d c
|
||||
value_0, value_1 = torch.chunk(value, dim=0, chunks=2)
|
||||
key = torch.cat([key_0, key_1], dim=1) # (b t) 2d c
|
||||
value = torch.cat([value_0, value_1], dim=1) # (b t) 2d c
|
||||
key = torch.cat([key]*2, dim=0) # ( 2 b t) 2d c
|
||||
value = torch.cat([value]*2, dim=0) # (2 b t) 2d c
|
||||
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class JointAttnProcessor:
|
||||
r"""
|
||||
Default processor for performing attention-related computations.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
temb=None,
|
||||
num_tasks=2
|
||||
):
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
if attn.spatial_norm is not None:
|
||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
assert num_tasks == 2 # only support two tasks now
|
||||
|
||||
key_0, key_1 = torch.chunk(key, dim=0, chunks=2) # keys shape (b t) d c
|
||||
value_0, value_1 = torch.chunk(value, dim=0, chunks=2)
|
||||
key = torch.cat([key_0, key_1], dim=1) # (b t) 2d c
|
||||
value = torch.cat([value_0, value_1], dim=1) # (b t) 2d c
|
||||
key = torch.cat([key]*2, dim=0) # ( 2 b t) 2d c
|
||||
value = torch.cat([value]*2, dim=0) # (2 b t) 2d c
|
||||
|
||||
|
||||
query = attn.head_to_batch_dim(query).contiguous()
|
||||
key = attn.head_to_batch_dim(key).contiguous()
|
||||
value = attn.head_to_batch_dim(value).contiguous()
|
||||
|
||||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||||
hidden_states = torch.bmm(attention_probs, value)
|
||||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,971 @@
|
||||
# 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.
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.utils import is_torch_version, logging
|
||||
from diffusers.models.normalization import AdaGroupNorm
|
||||
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
|
||||
from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel
|
||||
from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D
|
||||
|
||||
from diffusers.models.unets.unet_2d_blocks import DownBlock2D, ResnetDownsampleBlock2D, AttnDownBlock2D, CrossAttnDownBlock2D, SimpleCrossAttnDownBlock2D, SkipDownBlock2D, AttnSkipDownBlock2D, DownEncoderBlock2D, AttnDownEncoderBlock2D, KDownBlock2D, KCrossAttnDownBlock2D
|
||||
from diffusers.models.unets.unet_2d_blocks import UpBlock2D, ResnetUpsampleBlock2D, CrossAttnUpBlock2D, SimpleCrossAttnUpBlock2D, AttnUpBlock2D, SkipUpBlock2D, AttnSkipUpBlock2D, UpDecoderBlock2D, AttnUpDecoderBlock2D, KUpBlock2D, KCrossAttnUpBlock2D
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def get_down_block(
|
||||
down_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
temb_channels,
|
||||
add_downsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
transformer_layers_per_block=1,
|
||||
num_attention_heads=None,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
downsample_padding=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
resnet_skip_time_act=False,
|
||||
resnet_out_scale_factor=1.0,
|
||||
cross_attention_norm=None,
|
||||
attention_head_dim=None,
|
||||
downsample_type=None,
|
||||
num_views=1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
selfattn_block: str = "custom",
|
||||
mvcd_attention: bool=False,
|
||||
use_dino: bool = False
|
||||
):
|
||||
# If attn head dim is not defined, we default it to the number of heads
|
||||
if attention_head_dim is None:
|
||||
logger.warn(
|
||||
f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
|
||||
)
|
||||
attention_head_dim = num_attention_heads
|
||||
|
||||
down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
|
||||
if down_block_type == "DownBlock2D":
|
||||
return DownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "ResnetDownsampleBlock2D":
|
||||
return ResnetDownsampleBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
)
|
||||
elif down_block_type == "AttnDownBlock2D":
|
||||
if add_downsample is False:
|
||||
downsample_type = None
|
||||
else:
|
||||
downsample_type = downsample_type or "conv" # default to 'conv'
|
||||
return AttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
downsample_type=downsample_type,
|
||||
)
|
||||
elif down_block_type == "CrossAttnDownBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
|
||||
return CrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
# custom MV2D attention block
|
||||
elif down_block_type == "CrossAttnDownBlockMV2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockMV2D")
|
||||
return CrossAttnDownBlockMV2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
selfattn_block=selfattn_block,
|
||||
mvcd_attention=mvcd_attention,
|
||||
use_dino=use_dino
|
||||
)
|
||||
elif down_block_type == "SimpleCrossAttnDownBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")
|
||||
return SimpleCrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
only_cross_attention=only_cross_attention,
|
||||
cross_attention_norm=cross_attention_norm,
|
||||
)
|
||||
elif down_block_type == "SkipDownBlock2D":
|
||||
return SkipDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "AttnSkipDownBlock2D":
|
||||
return AttnSkipDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "DownEncoderBlock2D":
|
||||
return DownEncoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "AttnDownEncoderBlock2D":
|
||||
return AttnDownEncoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif down_block_type == "KDownBlock2D":
|
||||
return KDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
)
|
||||
elif down_block_type == "KCrossAttnDownBlock2D":
|
||||
return KCrossAttnDownBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
add_self_attention=True if not add_downsample else False,
|
||||
)
|
||||
raise ValueError(f"{down_block_type} does not exist.")
|
||||
|
||||
|
||||
def get_up_block(
|
||||
up_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
prev_output_channel,
|
||||
temb_channels,
|
||||
add_upsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
transformer_layers_per_block=1,
|
||||
num_attention_heads=None,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
resnet_skip_time_act=False,
|
||||
resnet_out_scale_factor=1.0,
|
||||
cross_attention_norm=None,
|
||||
attention_head_dim=None,
|
||||
upsample_type=None,
|
||||
num_views=1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
selfattn_block: str = "custom",
|
||||
mvcd_attention: bool=False,
|
||||
use_dino: bool = False
|
||||
):
|
||||
# If attn head dim is not defined, we default it to the number of heads
|
||||
if attention_head_dim is None:
|
||||
logger.warn(
|
||||
f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
|
||||
)
|
||||
attention_head_dim = num_attention_heads
|
||||
|
||||
up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
|
||||
if up_block_type == "UpBlock2D":
|
||||
return UpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "ResnetUpsampleBlock2D":
|
||||
return ResnetUpsampleBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
)
|
||||
elif up_block_type == "CrossAttnUpBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
|
||||
return CrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
# custom MV2D attention block
|
||||
elif up_block_type == "CrossAttnUpBlockMV2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockMV2D")
|
||||
return CrossAttnUpBlockMV2D(
|
||||
num_layers=num_layers,
|
||||
transformer_layers_per_block=transformer_layers_per_block,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
selfattn_block=selfattn_block,
|
||||
mvcd_attention=mvcd_attention,
|
||||
use_dino=use_dino
|
||||
)
|
||||
elif up_block_type == "SimpleCrossAttnUpBlock2D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
|
||||
return SimpleCrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
skip_time_act=resnet_skip_time_act,
|
||||
output_scale_factor=resnet_out_scale_factor,
|
||||
only_cross_attention=only_cross_attention,
|
||||
cross_attention_norm=cross_attention_norm,
|
||||
)
|
||||
elif up_block_type == "AttnUpBlock2D":
|
||||
if add_upsample is False:
|
||||
upsample_type = None
|
||||
else:
|
||||
upsample_type = upsample_type or "conv" # default to 'conv'
|
||||
|
||||
return AttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
upsample_type=upsample_type,
|
||||
)
|
||||
elif up_block_type == "SkipUpBlock2D":
|
||||
return SkipUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "AttnSkipUpBlock2D":
|
||||
return AttnSkipUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif up_block_type == "UpDecoderBlock2D":
|
||||
return UpDecoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
elif up_block_type == "AttnUpDecoderBlock2D":
|
||||
return AttnUpDecoderBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
elif up_block_type == "KUpBlock2D":
|
||||
return KUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
)
|
||||
elif up_block_type == "KCrossAttnUpBlock2D":
|
||||
return KCrossAttnUpBlock2D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_head_dim=attention_head_dim,
|
||||
)
|
||||
|
||||
raise ValueError(f"{up_block_type} does not exist.")
|
||||
|
||||
|
||||
class UNetMidBlockMV2DCrossAttn(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
output_scale_factor=1.0,
|
||||
cross_attention_dim=1280,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
selfattn_block: str = "custom",
|
||||
mvcd_attention: bool=False,
|
||||
use_dino: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
||||
if selfattn_block == "custom":
|
||||
from .transformer_mv2d import TransformerMV2DModel
|
||||
elif selfattn_block == "rowwise":
|
||||
from .transformer_mv2d_rowwise import TransformerMV2DModel
|
||||
elif selfattn_block == "self_rowwise":
|
||||
from .transformer_mv2d_self_rowwise import TransformerMV2DModel
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
# there is always at least one resnet
|
||||
resnets = [
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
]
|
||||
attentions = []
|
||||
|
||||
for _ in range(num_layers):
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
in_channels // num_attention_heads,
|
||||
in_channels=in_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention,
|
||||
use_dino=use_dino
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
dino_feature: Optional[torch.FloatTensor] = None
|
||||
) -> torch.FloatTensor:
|
||||
hidden_states = self.resnets[0](hidden_states, temb)
|
||||
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
dino_feature=dino_feature,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossAttnUpBlockMV2D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
prev_output_channel: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
add_upsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
selfattn_block: str = "custom",
|
||||
mvcd_attention: bool=False,
|
||||
use_dino: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
|
||||
if selfattn_block == "custom":
|
||||
from .transformer_mv2d import TransformerMV2DModel
|
||||
elif selfattn_block == "rowwise":
|
||||
from .transformer_mv2d_rowwise import TransformerMV2DModel
|
||||
elif selfattn_block == "self_rowwise":
|
||||
from .transformer_mv2d_self_rowwise import TransformerMV2DModel
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
for i in range(num_layers):
|
||||
res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
|
||||
resnet_in_channels = prev_output_channel if i == 0 else out_channels
|
||||
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=resnet_in_channels + res_skip_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
out_channels // num_attention_heads,
|
||||
in_channels=out_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention,
|
||||
use_dino=use_dino
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
if add_upsample:
|
||||
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
||||
else:
|
||||
self.upsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
upsample_size: Optional[int] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
dino_feature: Optional[torch.FloatTensor] = None
|
||||
):
|
||||
for resnet, attn in zip(self.resnets, self.attentions):
|
||||
# pop res hidden states
|
||||
res_hidden_states = res_hidden_states_tuple[-1]
|
||||
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
||||
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
||||
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet),
|
||||
hidden_states,
|
||||
temb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(attn, return_dict=False),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
dino_feature,
|
||||
None, # timestep
|
||||
None, # class_labels
|
||||
cross_attention_kwargs,
|
||||
attention_mask,
|
||||
encoder_attention_mask,
|
||||
**ckpt_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
dino_feature=dino_feature,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
if self.upsamplers is not None:
|
||||
for upsampler in self.upsamplers:
|
||||
hidden_states = upsampler(hidden_states, upsample_size)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossAttnDownBlockMV2D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
transformer_layers_per_block: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
num_attention_heads=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
downsample_padding=1,
|
||||
add_downsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
num_views: int = 1,
|
||||
cd_attention_last: bool = False,
|
||||
cd_attention_mid: bool = False,
|
||||
multiview_attention: bool = True,
|
||||
sparse_mv_attention: bool = False,
|
||||
selfattn_block: str = "custom",
|
||||
mvcd_attention: bool=False,
|
||||
use_dino: bool = False
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.num_attention_heads = num_attention_heads
|
||||
if selfattn_block == "custom":
|
||||
from .transformer_mv2d import TransformerMV2DModel
|
||||
elif selfattn_block == "rowwise":
|
||||
from .transformer_mv2d_rowwise import TransformerMV2DModel
|
||||
elif selfattn_block == "self_rowwise":
|
||||
from .transformer_mv2d_self_rowwise import TransformerMV2DModel
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
for i in range(num_layers):
|
||||
in_channels = in_channels if i == 0 else out_channels
|
||||
resnets.append(
|
||||
ResnetBlock2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
)
|
||||
)
|
||||
if not dual_cross_attention:
|
||||
attentions.append(
|
||||
TransformerMV2DModel(
|
||||
num_attention_heads,
|
||||
out_channels // num_attention_heads,
|
||||
in_channels=out_channels,
|
||||
num_layers=transformer_layers_per_block,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
num_views=num_views,
|
||||
cd_attention_last=cd_attention_last,
|
||||
cd_attention_mid=cd_attention_mid,
|
||||
multiview_attention=multiview_attention,
|
||||
sparse_mv_attention=sparse_mv_attention,
|
||||
mvcd_attention=mvcd_attention,
|
||||
use_dino=use_dino
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
|
||||
if add_downsample:
|
||||
self.downsamplers = nn.ModuleList(
|
||||
[
|
||||
Downsample2D(
|
||||
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
||||
)
|
||||
]
|
||||
)
|
||||
else:
|
||||
self.downsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
temb: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
dino_feature: Optional[torch.FloatTensor] = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
additional_residuals=None,
|
||||
):
|
||||
output_states = ()
|
||||
|
||||
blocks = list(zip(self.resnets, self.attentions))
|
||||
|
||||
for i, (resnet, attn) in enumerate(blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet),
|
||||
hidden_states,
|
||||
temb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(attn, return_dict=False),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
dino_feature,
|
||||
None, # timestep
|
||||
None, # class_labels
|
||||
cross_attention_kwargs,
|
||||
attention_mask,
|
||||
encoder_attention_mask,
|
||||
**ckpt_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
dino_feature=dino_feature,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
attention_mask=attention_mask,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# apply additional residuals to the output of the last pair of resnet and attention blocks
|
||||
if i == len(blocks) - 1 and additional_residuals is not None:
|
||||
hidden_states = hidden_states + additional_residuals
|
||||
|
||||
output_states = output_states + (hidden_states,)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
for downsampler in self.downsamplers:
|
||||
hidden_states = downsampler(hidden_states)
|
||||
|
||||
output_states = output_states + (hidden_states,)
|
||||
|
||||
return hidden_states, output_states
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,637 @@
|
||||
import inspect
|
||||
import warnings
|
||||
from typing import Callable, List, Optional, Union, Dict, Any
|
||||
import PIL
|
||||
import torch
|
||||
from packaging import version
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, CLIPFeatureExtractor, CLIPTokenizer, CLIPTextModel
|
||||
from diffusers.utils.import_utils import is_accelerate_available
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.models.embeddings import get_timestep_embedding
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import deprecate, logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
||||
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
|
||||
import os
|
||||
import torchvision.transforms.functional as TF
|
||||
from einops import rearrange
|
||||
import comfy.utils
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
class StableUnCLIPImg2ImgPipeline(DiffusionPipeline):
|
||||
"""
|
||||
Pipeline for text-guided image to image generation using stable unCLIP.
|
||||
|
||||
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
|
||||
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
|
||||
|
||||
Args:
|
||||
feature_extractor ([`CLIPFeatureExtractor`]):
|
||||
Feature extractor for image pre-processing before being encoded.
|
||||
image_encoder ([`CLIPVisionModelWithProjection`]):
|
||||
CLIP vision model for encoding images.
|
||||
image_normalizer ([`StableUnCLIPImageNormalizer`]):
|
||||
Used to normalize the predicted image embeddings before the noise is applied and un-normalize the image
|
||||
embeddings after the noise has been applied.
|
||||
image_noising_scheduler ([`KarrasDiffusionSchedulers`]):
|
||||
Noise schedule for adding noise to the predicted image embeddings. The amount of noise to add is determined
|
||||
by `noise_level` in `StableUnCLIPPipeline.__call__`.
|
||||
tokenizer (`CLIPTokenizer`):
|
||||
Tokenizer of class
|
||||
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
|
||||
text_encoder ([`CLIPTextModel`]):
|
||||
Frozen text-encoder.
|
||||
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
|
||||
scheduler ([`KarrasDiffusionSchedulers`]):
|
||||
A scheduler to be used in combination with `unet` to denoise the encoded image latents.
|
||||
vae ([`AutoencoderKL`]):
|
||||
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
|
||||
"""
|
||||
# image encoding components
|
||||
feature_extractor: CLIPFeatureExtractor
|
||||
image_encoder: CLIPVisionModelWithProjection
|
||||
# image noising components
|
||||
image_normalizer: StableUnCLIPImageNormalizer
|
||||
image_noising_scheduler: KarrasDiffusionSchedulers
|
||||
# regular denoising components
|
||||
tokenizer: CLIPTokenizer
|
||||
text_encoder: CLIPTextModel
|
||||
unet: UNet2DConditionModel
|
||||
scheduler: KarrasDiffusionSchedulers
|
||||
vae: AutoencoderKL
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# image encoding components
|
||||
feature_extractor: CLIPFeatureExtractor,
|
||||
image_encoder: CLIPVisionModelWithProjection,
|
||||
# image noising components
|
||||
image_normalizer: StableUnCLIPImageNormalizer,
|
||||
image_noising_scheduler: KarrasDiffusionSchedulers,
|
||||
# regular denoising components
|
||||
tokenizer: CLIPTokenizer,
|
||||
text_encoder: CLIPTextModel,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
# vae
|
||||
vae: AutoencoderKL,
|
||||
num_views: int = 4,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
feature_extractor=feature_extractor,
|
||||
image_encoder=image_encoder,
|
||||
image_normalizer=image_normalizer,
|
||||
image_noising_scheduler=image_noising_scheduler,
|
||||
tokenizer=tokenizer,
|
||||
text_encoder=text_encoder,
|
||||
unet=unet,
|
||||
scheduler=scheduler,
|
||||
vae=vae,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
self.num_views: int = num_views
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing
|
||||
def enable_vae_slicing(self):
|
||||
r"""
|
||||
Enable sliced VAE decoding.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
|
||||
steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.vae.enable_slicing()
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing
|
||||
def disable_vae_slicing(self):
|
||||
r"""
|
||||
Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
|
||||
computing decoding in one step.
|
||||
"""
|
||||
self.vae.disable_slicing()
|
||||
|
||||
def enable_sequential_cpu_offload(self, gpu_id=0):
|
||||
r"""
|
||||
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, the pipeline's
|
||||
models have their state dicts saved to CPU and then are moved to a `torch.device('meta') and loaded to GPU only
|
||||
when their specific submodule has its `forward` method called.
|
||||
"""
|
||||
if is_accelerate_available():
|
||||
from accelerate import cpu_offload
|
||||
else:
|
||||
raise ImportError("Please install accelerate via `pip install accelerate`")
|
||||
|
||||
device = torch.device(f"cuda:{gpu_id}")
|
||||
|
||||
# TODO: self.image_normalizer.{scale,unscale} are not covered by the offload hooks, so they fails if added to the list
|
||||
models = [
|
||||
self.image_encoder,
|
||||
self.text_encoder,
|
||||
self.unet,
|
||||
self.vae,
|
||||
]
|
||||
for cpu_offloaded_model in models:
|
||||
if cpu_offloaded_model is not None:
|
||||
cpu_offload(cpu_offloaded_model, device)
|
||||
|
||||
@property
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
|
||||
def _execution_device(self):
|
||||
r"""
|
||||
Returns the device on which the pipeline's models will be executed. After calling
|
||||
`pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
|
||||
hooks.
|
||||
"""
|
||||
if not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
|
||||
def _encode_prompt(
|
||||
self,
|
||||
prompt,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
negative_prompt=None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
lora_scale: Optional[float] = None,
|
||||
):
|
||||
r"""
|
||||
Encodes the prompt into text encoder hidden states.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
prompt to be encoded
|
||||
device: (`torch.device`):
|
||||
torch device
|
||||
num_images_per_prompt (`int`):
|
||||
number of images that should be generated per prompt
|
||||
do_classifier_free_guidance (`bool`):
|
||||
whether to use classifier free guidance or not
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
|
||||
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
"""
|
||||
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
# 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
|
||||
normal_prompt_embeds, color_prompt_embeds = torch.chunk(prompt_embeds, 2, dim=0)
|
||||
|
||||
prompt_embeds = torch.cat([normal_prompt_embeds, normal_prompt_embeds, color_prompt_embeds, color_prompt_embeds], 0)
|
||||
|
||||
return prompt_embeds
|
||||
|
||||
def _encode_image(
|
||||
self,
|
||||
image_pil,
|
||||
device,
|
||||
num_images_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
noise_level: int=0,
|
||||
generator: Optional[torch.Generator] = None
|
||||
):
|
||||
dtype = next(self.image_encoder.parameters()).dtype
|
||||
# ______________________________clip image embedding______________________________
|
||||
image = self.feature_extractor(images=image_pil, return_tensors="pt").pixel_values
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
image_embeds = self.image_encoder(image).image_embeds
|
||||
|
||||
image_embeds = self.noise_image_embeddings(
|
||||
image_embeds=image_embeds,
|
||||
noise_level=noise_level,
|
||||
generator=generator,
|
||||
)
|
||||
# duplicate image embeddings for each generation per prompt, using mps friendly method
|
||||
# image_embeds = image_embeds.unsqueeze(1)
|
||||
# note: the condition input is same
|
||||
image_embeds = image_embeds.repeat(num_images_per_prompt, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
normal_image_embeds, color_image_embeds = torch.chunk(image_embeds, 2, dim=0)
|
||||
negative_prompt_embeds = torch.zeros_like(normal_image_embeds)
|
||||
|
||||
# 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_embeds = torch.cat([negative_prompt_embeds, normal_image_embeds, negative_prompt_embeds, color_image_embeds], 0)
|
||||
|
||||
# _____________________________vae input latents__________________________________________________
|
||||
image_pt = torch.stack([TF.to_tensor(img) for img in image_pil], dim=0).to(device)
|
||||
image_pt = image_pt * 2.0 - 1.0
|
||||
image_latents = self.vae.encode(image_pt).latent_dist.mode() * self.vae.config.scaling_factor
|
||||
# Note: repeat differently from official pipelines
|
||||
image_latents = image_latents.repeat(num_images_per_prompt, 1, 1, 1)
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
normal_image_latents, color_image_latents = torch.chunk(image_latents, 2, dim=0)
|
||||
image_latents = torch.cat([torch.zeros_like(normal_image_latents), normal_image_latents,
|
||||
torch.zeros_like(color_image_latents), color_image_latents], 0)
|
||||
|
||||
return image_embeds, image_latents
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
|
||||
def decode_latents(self, latents):
|
||||
latents = 1 / self.vae.config.scaling_factor * latents
|
||||
image = self.vae.decode(latents).sample
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
|
||||
return image
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def check_inputs(
|
||||
self,
|
||||
prompt,
|
||||
image,
|
||||
height,
|
||||
width,
|
||||
callback_steps,
|
||||
noise_level,
|
||||
):
|
||||
if height % 8 != 0 or width % 8 != 0:
|
||||
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
|
||||
|
||||
if (callback_steps is None) or (
|
||||
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
|
||||
f" {type(callback_steps)}."
|
||||
)
|
||||
|
||||
if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
||||
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
||||
|
||||
|
||||
if noise_level < 0 or noise_level >= self.image_noising_scheduler.config.num_train_timesteps:
|
||||
raise ValueError(
|
||||
f"`noise_level` must be between 0 and {self.image_noising_scheduler.config.num_train_timesteps - 1}, inclusive."
|
||||
)
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
|
||||
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
|
||||
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_unclip.StableUnCLIPPipeline.noise_image_embeddings
|
||||
def noise_image_embeddings(
|
||||
self,
|
||||
image_embeds: torch.Tensor,
|
||||
noise_level: int,
|
||||
noise: Optional[torch.FloatTensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
):
|
||||
"""
|
||||
Add noise to the image embeddings. The amount of noise is controlled by a `noise_level` input. A higher
|
||||
`noise_level` increases the variance in the final un-noised images.
|
||||
|
||||
The noise is applied in two ways
|
||||
1. A noise schedule is applied directly to the embeddings
|
||||
2. A vector of sinusoidal time embeddings are appended to the output.
|
||||
|
||||
In both cases, the amount of noise is controlled by the same `noise_level`.
|
||||
|
||||
The embeddings are normalized before the noise is applied and un-normalized after the noise is applied.
|
||||
"""
|
||||
if noise is None:
|
||||
noise = randn_tensor(
|
||||
image_embeds.shape, generator=generator, device=image_embeds.device, dtype=image_embeds.dtype
|
||||
)
|
||||
|
||||
noise_level = torch.tensor([noise_level] * image_embeds.shape[0], device=image_embeds.device)
|
||||
|
||||
image_embeds = self.image_normalizer.scale(image_embeds)
|
||||
|
||||
image_embeds = self.image_noising_scheduler.add_noise(image_embeds, timesteps=noise_level, noise=noise)
|
||||
|
||||
image_embeds = self.image_normalizer.unscale(image_embeds)
|
||||
|
||||
noise_level = get_timestep_embedding(
|
||||
timesteps=noise_level, embedding_dim=image_embeds.shape[-1], flip_sin_to_cos=True, downscale_freq_shift=0
|
||||
)
|
||||
|
||||
# `get_timestep_embeddings` does not contain any weights and will always return f32 tensors,
|
||||
# but we might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
noise_level = noise_level.to(image_embeds.dtype)
|
||||
|
||||
image_embeds = torch.cat((image_embeds, noise_level), 1)
|
||||
|
||||
return image_embeds
|
||||
|
||||
@torch.no_grad()
|
||||
# @replace_example_docstring(EXAMPLE_DOC_STRING)
|
||||
def __call__(
|
||||
self,
|
||||
image: Union[torch.FloatTensor, PIL.Image.Image],
|
||||
prompt: Union[str, List[str]],
|
||||
prompt_embeds: torch.FloatTensor = None,
|
||||
dino_feature: torch.FloatTensor = None,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float = 10,
|
||||
negative_prompt: Optional[Union[str, List[str]]] = None,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: 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,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
noise_level: int = 0,
|
||||
image_embeds: Optional[torch.FloatTensor] = None,
|
||||
return_elevation_focal: Optional[bool] = False,
|
||||
gt_img_in: Optional[torch.FloatTensor] = None,
|
||||
):
|
||||
r"""
|
||||
Function invoked when calling the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
||||
instead.
|
||||
image (`torch.FloatTensor` or `PIL.Image.Image`):
|
||||
`Image`, or tensor representing an image batch. The image will be encoded to its CLIP embedding which
|
||||
the unet will be conditioned on. Note that the image is _not_ encoded by the vae and then used as the
|
||||
latents in the denoising process such as in the standard stable diffusion text guided image variation
|
||||
process.
|
||||
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 20):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
guidance_scale (`float`, *optional*, defaults to 10.0):
|
||||
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
||||
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
||||
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
||||
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
||||
usually at the expense of lower image quality.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
|
||||
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
eta (`float`, *optional*, defaults to 0.0):
|
||||
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
||||
[`schedulers.DDIMScheduler`], will be ignored for others.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
||||
to make generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor will ge generated by sampling using the supplied random `generator`.
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generate image. Choose between
|
||||
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
||||
plain tuple.
|
||||
callback (`Callable`, *optional*):
|
||||
A function that will be called every `callback_steps` steps during inference. The function will be
|
||||
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 will be called. If not specified, the callback will be
|
||||
called at every step.
|
||||
cross_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
|
||||
noise_level (`int`, *optional*, defaults to `0`):
|
||||
The amount of noise to add to the image embeddings. A higher `noise_level` increases the variance in
|
||||
the final un-noised images. See `StableUnCLIPPipeline.noise_image_embeddings` for details.
|
||||
image_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated CLIP embeddings to condition the unet on. Note that these are not latents to be used in
|
||||
the denoising process. If you want to provide pre-generated latents, pass them to `__call__` as
|
||||
`latents`.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~pipelines.ImagePipelineOutput`] or `tuple`: [`~ pipeline_utils.ImagePipelineOutput`] if `return_dict` is
|
||||
True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated images.
|
||||
"""
|
||||
# 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(
|
||||
prompt=prompt,
|
||||
image=image,
|
||||
height=height,
|
||||
width=width,
|
||||
callback_steps=callback_steps,
|
||||
noise_level=noise_level
|
||||
)
|
||||
|
||||
# 2. Define call parameters
|
||||
if isinstance(image, list):
|
||||
batch_size = len(image)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
batch_size = image.shape[0]
|
||||
assert batch_size >= self.num_views and batch_size % self.num_views == 0
|
||||
elif isinstance(image, PIL.Image.Image):
|
||||
image = [image]*self.num_views*2
|
||||
batch_size = self.num_views*2
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt] * self.num_views * 2
|
||||
|
||||
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 prompt
|
||||
text_encoder_lora_scale = (
|
||||
cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
|
||||
)
|
||||
prompt_embeds = self._encode_prompt(
|
||||
prompt=prompt,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
negative_prompt=negative_prompt,
|
||||
prompt_embeds=prompt_embeds,
|
||||
negative_prompt_embeds=negative_prompt_embeds,
|
||||
lora_scale=text_encoder_lora_scale,
|
||||
)
|
||||
|
||||
|
||||
# 4. Encoder input image
|
||||
if isinstance(image, list):
|
||||
image_pil = image
|
||||
elif isinstance(image, torch.Tensor):
|
||||
image_pil = [TF.to_pil_image(image[i]) for i in range(image.shape[0])]
|
||||
noise_level = torch.tensor([noise_level], device=device)
|
||||
image_embeds, image_latents = self._encode_image(
|
||||
image_pil=image_pil,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
noise_level=noise_level,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
# 6. Prepare latent variables
|
||||
num_channels_latents = self.unet.config.out_channels
|
||||
if gt_img_in is not None:
|
||||
latents = gt_img_in * self.scheduler.init_noise_sigma
|
||||
else:
|
||||
latents = self.prepare_latents(
|
||||
batch_size=batch_size,
|
||||
num_channels_latents=num_channels_latents,
|
||||
height=height,
|
||||
width=width,
|
||||
dtype=prompt_embeds.dtype,
|
||||
device=device,
|
||||
generator=generator,
|
||||
latents=latents,
|
||||
)
|
||||
|
||||
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
eles, focals = [], []
|
||||
# 8. Denoising loop
|
||||
comfy_pbar = comfy.utils.ProgressBar(num_inference_steps)
|
||||
for i, t in enumerate(self.progress_bar(timesteps)):
|
||||
if do_classifier_free_guidance:
|
||||
normal_latents, color_latents = torch.chunk(latents, 2, dim=0)
|
||||
latent_model_input = torch.cat([normal_latents, normal_latents, color_latents, color_latents], 0)
|
||||
else:
|
||||
latent_model_input = latents
|
||||
latent_model_input = torch.cat([
|
||||
latent_model_input, image_latents
|
||||
], dim=1)
|
||||
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
||||
|
||||
# predict the noise residual
|
||||
unet_out = self.unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
dino_feature=dino_feature,
|
||||
class_labels=image_embeds,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
return_dict=False)
|
||||
|
||||
noise_pred = unet_out[0]
|
||||
if return_elevation_focal:
|
||||
uncond_pose, pose = torch.chunk(unet_out[1], 2, 0)
|
||||
pose = uncond_pose + guidance_scale * (pose - uncond_pose)
|
||||
ele = pose[:, 0].detach().cpu().numpy() # b
|
||||
eles.append(ele)
|
||||
focal = pose[:, 1].detach().cpu().numpy()
|
||||
focals.append(focal)
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
normal_noise_pred_uncond, normal_noise_pred_text, color_noise_pred_uncond, color_noise_pred_text = torch.chunk(noise_pred, 4, dim=0)
|
||||
|
||||
noise_pred_uncond, noise_pred_text = torch.cat([normal_noise_pred_uncond, color_noise_pred_uncond], 0), torch.cat([normal_noise_pred_text, color_noise_pred_text], 0)
|
||||
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, return_dict=False)[0]
|
||||
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
callback(i, t, latents)
|
||||
|
||||
comfy_pbar.update_absolute(i + 1)
|
||||
|
||||
# 9. Post-processing
|
||||
if not output_type == "latent":
|
||||
if num_channels_latents == 8:
|
||||
latents = torch.cat([latents[:, :4], latents[:, 4:]], dim=0)
|
||||
with torch.no_grad():
|
||||
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
|
||||
else:
|
||||
image = latents
|
||||
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
|
||||
# Offload last model to CPU
|
||||
# if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
|
||||
# self.final_offload_hook.offload()
|
||||
if not return_dict:
|
||||
return (image, )
|
||||
if return_elevation_focal:
|
||||
return ImagePipelineOutput(images=image), eles, focals
|
||||
else:
|
||||
return ImagePipelineOutput(images=image)
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
from omegaconf import OmegaConf
|
||||
from packaging import version
|
||||
from typing import Dict, Optional, List
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# ============ Register OmegaConf Recolvers ============= #
|
||||
OmegaConf.register_new_resolver('calc_exp_lr_decay_rate', lambda factor, n: factor**(1./n), replace=True)
|
||||
OmegaConf.register_new_resolver('add', lambda a, b: a + b, replace=True)
|
||||
OmegaConf.register_new_resolver('sub', lambda a, b: a - b, replace=True)
|
||||
OmegaConf.register_new_resolver('mul', lambda a, b: a * b, replace=True)
|
||||
OmegaConf.register_new_resolver('div', lambda a, b: a / b, replace=True)
|
||||
OmegaConf.register_new_resolver('idiv', lambda a, b: a // b, replace=True)
|
||||
OmegaConf.register_new_resolver('basename', lambda p: os.path.basename(p), replace=True)
|
||||
# ======================================================= #
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
num_views: int
|
||||
dataset: Dict
|
||||
|
||||
|
||||
def load_config(*yaml_files):
|
||||
yaml_confs = [OmegaConf.load(f) for f in yaml_files]
|
||||
conf = OmegaConf.merge(*yaml_confs)
|
||||
OmegaConf.resolve(conf)
|
||||
|
||||
schema = OmegaConf.structured(TestConfig)
|
||||
conf = OmegaConf.merge(schema, conf)
|
||||
return conf
|
||||
|
||||
|
||||
def config_to_primitive(config, resolve=True):
|
||||
return OmegaConf.to_container(config, resolve=resolve)
|
||||
|
||||
|
||||
def dump_config(path, config):
|
||||
with open(path, 'w') as fp:
|
||||
OmegaConf.save(config=config, f=fp)
|
||||
|
||||
def parse_version(ver):
|
||||
return version.parse(ver)
|
||||
@@ -44,7 +44,6 @@ from .mesh_processer.mesh_utils import (
|
||||
calculate_max_sh_degree_from_gs_ply,
|
||||
marching_cubes_density_to_mesh,
|
||||
color_func_to_albedo,
|
||||
|
||||
)
|
||||
from .algorithms.main_3DGS import GaussianSplatting, GaussianSplattingCameraController, GSParams
|
||||
from .algorithms.main_3DGS_renderer import GaussianSplattingRenderer
|
||||
@@ -65,12 +64,15 @@ from .tgs.utils.misc import todevice, get_device
|
||||
from .lgm.core.options import config_defaults
|
||||
from .lgm.mvdream.pipeline_mvdream import MVDreamPipeline
|
||||
from .mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePipeline
|
||||
from .mvdiffusion.data.single_image_dataset import SingleImageDataset
|
||||
from .mvdiffusion.data.single_image_dataset import SingleImageDataset as MVSingleImageDataset
|
||||
from .mvdiffusion.utils.misc import load_config as load_config_wonder3d
|
||||
from .tsr.system import TSR
|
||||
from .crm.model.crm.model import CRM
|
||||
from .zero123plus.pipeline import Zero123PlusPipeline
|
||||
from .instant_mesh.utils.camera_util import oribt_camera_poses_to_input_cameras
|
||||
from .era3d.mvdiffusion.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline
|
||||
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 .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames, resume_or_download_model_from_hf
|
||||
@@ -81,6 +83,7 @@ DIFFUSERS_PIPE_DICT = OrderedDict([
|
||||
("Zero123PlusPipeline", Zero123PlusPipeline),
|
||||
("DiffusionPipeline", DiffusionPipeline),
|
||||
("StableDiffusionPipeline", StableDiffusionPipeline),
|
||||
("Era3DPipeline", StableUnCLIPImg2ImgPipeline)
|
||||
])
|
||||
|
||||
DIFFUSERS_SCHEDULER_DICT = OrderedDict([
|
||||
@@ -127,7 +130,8 @@ AZIMUTH_MAX = 180.0
|
||||
|
||||
WEIGHT_DTYPE = torch.float16
|
||||
|
||||
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
DEVICE = torch.device(DEVICE_STR)
|
||||
|
||||
class Preview_3DGS:
|
||||
|
||||
@@ -1288,11 +1292,9 @@ class Load_Diffusers_Pipeline:
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(ROOT_PATH, "checkpoints", model_name)
|
||||
if not os.path.exists(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {model_name}").warning.print()
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model")
|
||||
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]
|
||||
|
||||
@@ -1450,15 +1452,16 @@ class Wonder3D_MVDiffusion_Model:
|
||||
).images
|
||||
|
||||
num_views = out.shape[0] // 2
|
||||
mv_images = out[num_views:].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
|
||||
mv_normals = out[:num_views].permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
|
||||
# [N, 3, H, W] -> [N, H, W, 3]
|
||||
mv_images = out[num_views:].permute(0, 2, 3, 1)
|
||||
mv_normals = out[:num_views].permute(0, 2, 3, 1)
|
||||
|
||||
return (mv_images, mv_normals, )
|
||||
|
||||
def prepare_data(self, ref_image, ref_mask):
|
||||
single_image = torch_img_to_pil_rgba(ref_image, ref_mask)
|
||||
abs_fix_cam_pose_dir = os.path.join(ROOT_PATH, self.fix_cam_pose_dir)
|
||||
dataset = SingleImageDataset(fix_cam_pose_dir=abs_fix_cam_pose_dir, num_views=6, img_wh=[256, 256], bg_color='white', single_image=single_image)
|
||||
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]
|
||||
|
||||
class MVDream_Model:
|
||||
@@ -1606,13 +1609,12 @@ class Large_Multiview_Gaussian_Model:
|
||||
|
||||
@torch.no_grad()
|
||||
def run_LGM(self, multiview_images, lgm_model):
|
||||
device = "cuda"
|
||||
ref_image_torch = prepare_torch_img(multiview_images, lgm_model.opt.input_size, lgm_model.opt.input_size, device) # [4, 3, 256, 256]
|
||||
ref_image_torch = prepare_torch_img(multiview_images, lgm_model.opt.input_size, lgm_model.opt.input_size, DEVICE_STR) # [4, 3, 256, 256]
|
||||
ref_image_torch = TF.normalize(ref_image_torch, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
|
||||
rays_embeddings = lgm_model.prepare_default_rays(device)
|
||||
rays_embeddings = lgm_model.prepare_default_rays(DEVICE_STR)
|
||||
ref_image_torch = torch.cat([ref_image_torch, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, 256, 256]
|
||||
|
||||
with torch.autocast(device_type=device, dtype=WEIGHT_DTYPE):
|
||||
with torch.autocast(device_type=DEVICE_STR, dtype=WEIGHT_DTYPE):
|
||||
# generate gaussians
|
||||
gaussians = lgm_model.forward_gaussians(ref_image_torch)
|
||||
|
||||
@@ -1789,10 +1791,11 @@ class Load_TripoSR_Model:
|
||||
def load_TSR(self, model_name, chunk_size):
|
||||
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
config_path = os.path.join(ROOT_PATH, self.tsr_config_path)
|
||||
|
||||
tsr_model = TSR.from_pretrained(
|
||||
weight_path = ckpt_path,
|
||||
config_path = os.path.join(ROOT_PATH, self.tsr_config_path)
|
||||
weight_path=ckpt_path,
|
||||
config_path=config_path
|
||||
)
|
||||
|
||||
tsr_model.renderer.set_chunk_size(chunk_size)
|
||||
@@ -2278,6 +2281,104 @@ class InstantMesh_Reconstruction_Model:
|
||||
mesh = Mesh(v=vertices, f=faces, vt=uvs, ft=mesh_tex_idx, albedo=tex_map, device=DEVICE)
|
||||
mesh.auto_normal()
|
||||
return (mesh,)
|
||||
|
||||
class Era3D_Diffusion_Model:
|
||||
|
||||
era3d_config_path = "configs/era3d_config.yaml"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"era3d_pipe": ("DIFFUSERS_PIPE",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"reference_mask": ("MASK",),
|
||||
"image_crop_size": ("INT", {"default": 420, "min": 400, "max": 8192}),
|
||||
"seed": ("INT", {"default": 600, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "step": 0.01}),
|
||||
"num_inference_steps": ("INT", {"default": 40, "min": 1}),
|
||||
"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.01}),
|
||||
"radius": ("FLOAT", {"default": 4.0, "min": 0.1, "step": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"IMAGE",
|
||||
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"multiviews",
|
||||
"multiview_normals",
|
||||
"orbit_camposes",
|
||||
)
|
||||
FUNCTION = "run_model"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
@torch.no_grad()
|
||||
def run_model(
|
||||
self,
|
||||
era3d_pipe,
|
||||
reference_image,
|
||||
reference_mask,
|
||||
image_crop_size,
|
||||
seed,
|
||||
guidance_scale,
|
||||
num_inference_steps,
|
||||
eta,
|
||||
radius,
|
||||
):
|
||||
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)
|
||||
|
||||
# Get the dataset
|
||||
cfg.dataset.prompt_embeds_path = os.path.join(ROOT_PATH, cfg.dataset.prompt_embeds_path)
|
||||
dataset = Era3DSingleImageDataset(
|
||||
single_image=single_image,
|
||||
crop_size=image_crop_size,
|
||||
dtype=WEIGHT_DTYPE,
|
||||
**cfg.dataset
|
||||
)
|
||||
|
||||
# Get input data
|
||||
img_batch = dataset.__getitem__(0)
|
||||
|
||||
imgs_in = torch.cat([img_batch['imgs_in']]*2, dim=0).to(DEVICE).to(WEIGHT_DTYPE) # (B*Nv, 3, H, W) B==1
|
||||
#num_views = imgs_in.shape[1]
|
||||
|
||||
normal_prompt_embeddings, clr_prompt_embeddings = img_batch['normal_prompt_embeddings'], img_batch['color_prompt_embeddings']
|
||||
prompt_embeddings = torch.cat([normal_prompt_embeddings, clr_prompt_embeddings], dim=0).to(DEVICE).to(WEIGHT_DTYPE) # (B*Nv, N, C) B==1
|
||||
|
||||
generator = torch.Generator(device=era3d_pipe.unet.device).manual_seed(seed)
|
||||
|
||||
# sampling
|
||||
with torch.autocast(DEVICE_STR):
|
||||
unet_out = era3d_pipe(
|
||||
imgs_in, None, prompt_embeds=prompt_embeddings,
|
||||
generator=generator, guidance_scale=guidance_scale, output_type='pt', num_images_per_prompt=1,
|
||||
num_inference_steps=num_inference_steps, eta=eta
|
||||
)
|
||||
|
||||
out = unet_out.images
|
||||
bsz = out.shape[0] // 2
|
||||
|
||||
# (1, 3, 512, 512)
|
||||
normals_pred = out[:bsz]
|
||||
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)
|
||||
|
||||
azimuths = [0, 45, 90, 180, -90, -45]
|
||||
elevations = [0.0] * 6
|
||||
radius = [radius] * 6
|
||||
center = [0.0] * 6
|
||||
|
||||
orbit_camposes = [azimuths, elevations, radius, center, center, center]
|
||||
|
||||
return (multiview_images, multiview_normals, orbit_camposes)
|
||||
|
||||
class Instant_NGP:
|
||||
|
||||
|
||||
Reference in New Issue
Block a user