diff --git a/.gitignore b/.gitignore index 2d06e3e..8f12504 100644 --- a/.gitignore +++ b/.gitignore @@ -156,6 +156,7 @@ cython_debug/ *.bin *.pth *.safetensors +*.pkl # Ignore test outputs outputs/ diff --git a/Checkpoints/CharacterGen/.gitattributes b/Checkpoints/CharacterGen/.gitattributes new file mode 100644 index 0000000..a6344aa --- /dev/null +++ b/Checkpoints/CharacterGen/.gitattributes @@ -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 diff --git a/Checkpoints/CharacterGen/2D_Stage/models/checkpoint/scaler.pt b/Checkpoints/CharacterGen/2D_Stage/models/checkpoint/scaler.pt new file mode 100644 index 0000000..38ac882 --- /dev/null +++ b/Checkpoints/CharacterGen/2D_Stage/models/checkpoint/scaler.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3c63c3d2e7b7f18788aca8a5a0e766ad4633f631d49192602d8532b6b0df7615 +size 559 diff --git a/Checkpoints/CharacterGen/2D_Stage/models/image_encoder/config.json b/Checkpoints/CharacterGen/2D_Stage/models/image_encoder/config.json new file mode 100644 index 0000000..9c841dd --- /dev/null +++ b/Checkpoints/CharacterGen/2D_Stage/models/image_encoder/config.json @@ -0,0 +1,23 @@ +{ + "_name_or_path": "./image_encoder", + "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.28.0.dev0" +} diff --git a/Checkpoints/CharacterGen/3D_Stage/models/base/README.md b/Checkpoints/CharacterGen/3D_Stage/models/base/README.md new file mode 100644 index 0000000..41acf20 --- /dev/null +++ b/Checkpoints/CharacterGen/3D_Stage/models/base/README.md @@ -0,0 +1,60 @@ +--- +license: apache-2.0 +tags: +- dino +- vision +--- + +# Vision Transformer (base-sized model) trained using DINOv2 + +Vision Transformer (ViT) model trained using the DINOv2 method. It was introduced in the paper [DINOv2: Learning Robust Visual Features without Supervision](https://arxiv.org/abs/2304.07193) by Oquab et al. and first released in [this repository](https://github.com/facebookresearch/dinov2). + +Disclaimer: The team releasing DINOv2 did not write a model card for this model so this model card has been written by the Hugging Face team. + +## Model description + +The Vision Transformer (ViT) is a transformer encoder model (BERT-like) pretrained on a large collection of images in a self-supervised fashion. + +Images are presented to the model as a sequence of fixed-size patches, which are linearly embedded. One also adds a [CLS] token to the beginning of a sequence to use it for classification tasks. One also adds absolute position embeddings before feeding the sequence to the layers of the Transformer encoder. + +Note that this model does not include any fine-tuned heads. + +By pre-training the model, it learns an inner representation of images that can then be used to extract features useful for downstream tasks: if you have a dataset of labeled images for instance, you can train a standard classifier by placing a linear layer on top of the pre-trained encoder. One typically places a linear layer on top of the [CLS] token, as the last hidden state of this token can be seen as a representation of an entire image. + +## Intended uses & limitations + +You can use the raw model for feature extraction. See the [model hub](https://huggingface.co/models?search=facebook/dinov2) to look for +fine-tuned versions on a task that interests you. + +### How to use + +Here is how to use this model: + +```python +from transformers import AutoImageProcessor, AutoModel +from PIL import Image +import requests + +url = 'http://images.cocodataset.org/val2017/000000039769.jpg' +image = Image.open(requests.get(url, stream=True).raw) + +processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base') +model = AutoModel.from_pretrained('facebook/dinov2-base') + +inputs = processor(images=image, return_tensors="pt") +outputs = model(**inputs) +last_hidden_states = outputs.last_hidden_state +``` + +### BibTeX entry and citation info + +```bibtex +misc{oquab2023dinov2, + title={DINOv2: Learning Robust Visual Features without Supervision}, + author={Maxime Oquab and Timothée Darcet and Théo Moutakanni and Huy Vo and Marc Szafraniec and Vasil Khalidov and Pierre Fernandez and Daniel Haziza and Francisco Massa and Alaaeldin El-Nouby and Mahmoud Assran and Nicolas Ballas and Wojciech Galuba and Russell Howes and Po-Yao Huang and Shang-Wen Li and Ishan Misra and Michael Rabbat and Vasu Sharma and Gabriel Synnaeve and Hu Xu and Hervé Jegou and Julien Mairal and Patrick Labatut and Armand Joulin and Piotr Bojanowski}, + year={2023}, + eprint={2304.07193}, + archivePrefix={arXiv}, + primaryClass={cs.CV} +} +``` \ No newline at end of file diff --git a/Checkpoints/CharacterGen/3D_Stage/models/base/config.json b/Checkpoints/CharacterGen/3D_Stage/models/base/config.json new file mode 100644 index 0000000..5df0129 --- /dev/null +++ b/Checkpoints/CharacterGen/3D_Stage/models/base/config.json @@ -0,0 +1,24 @@ +{ + "architectures": [ + "Dinov2Model" + ], + "attention_probs_dropout_prob": 0.0, + "drop_path_rate": 0.0, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.0, + "hidden_size": 768, + "image_size": 518, + "initializer_range": 0.02, + "layer_norm_eps": 1e-06, + "layerscale_value": 1.0, + "mlp_ratio": 4, + "model_type": "dinov2", + "num_attention_heads": 12, + "num_channels": 3, + "num_hidden_layers": 12, + "patch_size": 14, + "qkv_bias": true, + "torch_dtype": "float32", + "transformers_version": "4.31.0.dev0", + "use_swiglu_ffn": false +} diff --git a/Checkpoints/CharacterGen/3D_Stage/models/base/preprocessor_config.json b/Checkpoints/CharacterGen/3D_Stage/models/base/preprocessor_config.json new file mode 100644 index 0000000..bae170e --- /dev/null +++ b/Checkpoints/CharacterGen/3D_Stage/models/base/preprocessor_config.json @@ -0,0 +1,27 @@ +{ + "crop_size": { + "height": 256, + "width": 256 + }, + "do_center_crop": false, + "do_convert_rgb": false, + "do_normalize": true, + "do_rescale": true, + "do_resize": false, + "image_mean": [ + 0.485, + 0.456, + 0.406 + ], + "image_processor_type": "BitImageProcessor", + "image_std": [ + 0.229, + 0.224, + 0.225 + ], + "resample": 3, + "rescale_factor": 0.00392156862745098, + "size": { + "shortest_edge": 256 + } +} diff --git a/Checkpoints/CharacterGen/README.md b/Checkpoints/CharacterGen/README.md new file mode 100644 index 0000000..1317488 --- /dev/null +++ b/Checkpoints/CharacterGen/README.md @@ -0,0 +1,22 @@ +--- +license: apache-2.0 +language: +- en +pipeline_tag: image-to-3d +--- + +This is the model parameters of [CharacterGen](https://github.com/zjp-shadow/CharacterGen). +Thanks for downloading the checkpoint~ +If you find our work helpful, please consider citing + +```bibtex +@article +{peng2024charactergen, +title ={CharacterGen: Efficient 3D Character Generation from Single Images with Multi-View Pose Canonicalization}, +author ={Hao-Yang Peng and Jia-Peng Zhang and Meng-Hao Guo and Yan-Pei Cao and Shi-Min Hu}, +journal ={ACM Transactions on Graphics (TOG)}, +year ={2024}, +volume ={43}, +number ={4}, +doi ={10.1145/3658217} +} \ No newline at end of file diff --git a/Checkpoints/Diffusers/Wuvin/Unique3D/.gitattributes b/Checkpoints/Diffusers/Wuvin/Unique3D/.gitattributes new file mode 100644 index 0000000..a6344aa --- /dev/null +++ b/Checkpoints/Diffusers/Wuvin/Unique3D/.gitattributes @@ -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 diff --git a/Checkpoints/Diffusers/Wuvin/Unique3D/README.md b/Checkpoints/Diffusers/Wuvin/Unique3D/README.md new file mode 100644 index 0000000..b7f59e5 --- /dev/null +++ b/Checkpoints/Diffusers/Wuvin/Unique3D/README.md @@ -0,0 +1,5 @@ +--- +license: mit +--- +To enable download model using huggingface_hub package +
Copied from https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt \ No newline at end of file diff --git a/Configs/CharacterGen_configs/Stage_2D_infer.yaml b/Configs/CharacterGen_configs/Stage_2D_infer.yaml new file mode 100644 index 0000000..c437f31 --- /dev/null +++ b/Configs/CharacterGen_configs/Stage_2D_infer.yaml @@ -0,0 +1,23 @@ +pretrained_model_path: "stabilityai/stable-diffusion-2-1" +image_encoder_path: "./2D_Stage/models/image_encoder" +ckpt_dir: "./2D_Stage/models/checkpoint" + +validation: + use_inv_latent: False + video_length: 4 + +use_pose_guider: True +use_noise: False +use_shifted_noise: False +unet_condition_type: image + +unet_from_pretrained_kwargs: + camera_embedding_type: 'e_de_da_sincos' + projection_class_embeddings_input_dim: 10 # modify + joint_attention: false # modify + num_views: 4 + sample_size: 96 + zero_init_conv_in: false + zero_init_camera_projection: false + in_channels: 4 + use_safetensors: true \ No newline at end of file diff --git a/Configs/CharacterGen_configs/Stage_3D_infer.yaml b/Configs/CharacterGen_configs/Stage_3D_infer.yaml new file mode 100644 index 0000000..2a386a8 --- /dev/null +++ b/Configs/CharacterGen_configs/Stage_3D_infer.yaml @@ -0,0 +1,104 @@ +system_cls: CharacterGen.Stage_3D.lrm.systems.multiview_lrm.MultiviewLRM +data: + cond_width: 504 + cond_height: 504 + +system: + weights: 3D_Stage/models/lrm.ckpt + + weights_ignore_modules: + - decoder.heads.density + + check_train_every_n_steps: 100 + + camera_embedder_cls: CharacterGen.Stage_3D.lrm.models.camera.LinearCameraEmbedder + camera_embedder: + in_channels: 16 + out_channels: 768 + conditions: + - c2w_cond + + # image tokenizer transforms input images to tokens + image_tokenizer_cls: CharacterGen.Stage_3D.lrm.models.tokenizers.image.DINOV2SingleImageTokenizer + image_tokenizer: + pretrained_model_name_or_path: "3D_Stage/models/base" + freeze_backbone_params: false + enable_memory_efficient_attention: true + enable_gradient_checkpointing: true + # camera modulation to the DINO transformer layers + modulation: true + modulation_zero_init: true + modulation_single_layer: true + modulation_cond_dim: ${system.camera_embedder.out_channels} + + # tokenizer gives a tokenized representation for the 3D scene + # triplane tokens in this case + tokenizer_cls: CharacterGen.Stage_3D.lrm.models.tokenizers.triplane.TriplaneLearnablePositionalEmbedding + tokenizer: + plane_size: 32 + num_channels: 512 + + # backbone network is a transformer that takes scene tokens (potentially with conditional image tokens) + # and outputs scene tokens of the same size + backbone_cls: CharacterGen.Stage_3D.lrm.models.transformers.transformer_1d.Transformer1D + backbone: + in_channels: ${system.tokenizer.num_channels} + num_attention_heads: 16 + attention_head_dim: 64 + num_layers: 12 + cross_attention_dim: 768 # hard-code, =DINO feature dim + # camera modulation to the transformer layers + # if not needed, set norm_type=layer_norm and do not specify cond_dim_ada_norm_continuous + norm_type: "layer_norm" + enable_memory_efficient_attention: true + gradient_checkpointing: true + + # post processor takes scene tokens and outputs the final scene parameters that will be used for rendering + # in this case, triplanes are upsampled and the features are condensed + post_processor_cls: CharacterGen.Stage_3D.lrm.models.networks.TriplaneUpsampleNetwork + post_processor: + in_channels: 512 + out_channels: 80 + + renderer_cls: CharacterGen.Stage_3D.lrm.models.renderers.triplane_dmtet.TriplaneDMTetRenderer + renderer: + radius: 0.6 # slightly larger than 0.5 + feature_reduction: concat + sdf_bias: -2. + tet_dir: "tets" + isosurface_resolution: 256 + enable_isosurface_grid_deformation: false + sdf_activation: negative + + decoder_cls: CharacterGen.Stage_3D.lrm.models.networks.MultiHeadMLP + decoder: + in_channels: 240 # 3 * 80 + n_neurons: 64 + n_hidden_layers_share: 8 + heads: + - name: sdf + out_channels: 1 + n_hidden_layers: 1 + output_activation: null + - name: features + out_channels: 3 + n_hidden_layers: 1 + output_activation: null # activate in material + activation: silu + chunk_mode: deferred + chunk_size: 131072 + + exporter: + fmt: "obj" + #visual: "vertex" + visual: "uv" + save_uv: True + save_texture: True + uv_unwrap_method: "open3d" + output_path: "3D_Stage/outputs" + + material_cls: CharacterGen.Stage_3D.lrm.models.materials.no_material.NoMaterial + + background_cls: CharacterGen.Stage_3D.lrm.models.background.solid_color_background.SolidColorBackground + background: + color: [0.5, 0.5, 0.5] \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose.json b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose.json new file mode 100644 index 0000000..40ed1ce --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose.json @@ -0,0 +1,38 @@ +[ + [ + [ + 0, 0, -1, 0, + 0, 1, 0, 0, + 1, 0, 0, 0, + 1.5, 0, 0, 1 + ], + "pose0.png" + ], + [ + [ + 0, 0, 1, 0, + 0, 1, 0, 0, + -1, 0, 0, 0, + -1.5, 0, 0, 1 + ], + "pose1.png" + ], + [ + [ + 0, 0, 1, 0, + 0, 1, 0, 0, + -1, 0, 0, 0, + -1.5, 0, 0, 1 + ], + "pose2.png" + ], + [ + [ + -1, 0, 0, 0, + 0, 1, 0, 0, + 0, 0, -1, 0, + 0, 0, -1.5, 1 + ], + "pose3.png" + ] +] \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose0.png b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose0.png new file mode 100644 index 0000000..7e044b3 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose0.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose1.png b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose1.png new file mode 100644 index 0000000..fd4093c Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose1.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose2.png b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose2.png new file mode 100644 index 0000000..537d77b Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose2.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose3.png b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose3.png new file mode 100644 index 0000000..c562d5f Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_2D/material/pose3.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/PoseGuider.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/PoseGuider.py new file mode 100644 index 0000000..7d0e1e2 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/PoseGuider.py @@ -0,0 +1,59 @@ +import os +import torch +import torch.nn as nn +import torch.nn.init as init +from einops import rearrange + +class PoseGuider(nn.Module): + def __init__(self, noise_latent_channels=4): + super(PoseGuider, self).__init__() + + self.conv_layers = nn.Sequential( + nn.Conv2d(in_channels=3, out_channels=16, kernel_size=4, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), + nn.ReLU() + ) + + # Final projection layer + self.final_proj = nn.Conv2d(in_channels=128, out_channels=noise_latent_channels, kernel_size=1) + + # Initialize layers + self._initialize_weights() + + def _initialize_weights(self): + # Initialize weights with Gaussian distribution and zero out the final layer + for m in self.conv_layers: + if isinstance(m, nn.Conv2d): + init.normal_(m.weight, mean=0.0, std=0.02) + if m.bias is not None: + init.zeros_(m.bias) + + init.zeros_(self.final_proj.weight) + if self.final_proj.bias is not None: + init.zeros_(self.final_proj.bias) + + def forward(self, pose_image): + x = self.conv_layers(pose_image) + x = self.final_proj(x) + + return x + + @classmethod + def from_pretrained(pretrained_model_path): + if not os.path.exists(pretrained_model_path): + print(f"There is no model file in {pretrained_model_path}") + print(f"loaded PoseGuider's pretrained weights from {pretrained_model_path} ...") + + state_dict = torch.load(pretrained_model_path, map_location="cpu") + model = PoseGuider(noise_latent_channels=4) + m, u = model.load_state_dict(state_dict, strict=False) + print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") + params = [p.numel() if "temporal" in n else 0 for n, p in model.named_parameters()] + print(f"### PoseGuider's Parameters: {sum(params) / 1e6} M") + + return model diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/attention.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/attention.py new file mode 100644 index 0000000..e4be16c --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/attention.py @@ -0,0 +1,344 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py + +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers import ModelMixin +from diffusers.utils import BaseOutput +from diffusers.utils.import_utils import is_xformers_available +from diffusers.models.attention import CrossAttention, FeedForward, AdaLayerNorm + +from einops import rearrange, repeat + + +@dataclass +class Transformer3DModelOutput(BaseOutput): + sample: torch.FloatTensor + + +if is_xformers_available(): + import xformers + import xformers.ops +else: + xformers = None + + +class Transformer3DModel(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_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, + 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, + use_attn_temp: 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 + + # Define input layers + 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 = nn.Linear(in_channels, inner_dim) + else: + self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0) + + # Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + 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, + use_attn_temp = use_attn_temp, + ) + for d in range(num_layers) + ] + ) + + # 4. Define output layers + if use_linear_projection: + self.proj_out = nn.Linear(in_channels, inner_dim) + else: + self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True): + # Input + assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + video_length = hidden_states.shape[2] + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=video_length) + + batch, channel, height, weight = 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 * weight, inner_dim) + else: + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim) + hidden_states = self.proj_in(hidden_states) + + # Blocks + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + video_length=video_length + ) + + # Output + if not self.use_linear_projection: + hidden_states = ( + hidden_states.reshape(batch, height, weight, 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, weight, inner_dim).permute(0, 3, 1, 2).contiguous() + ) + + output = hidden_states + residual + + output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length) + if not return_dict: + return (output,) + + return Transformer3DModelOutput(sample=output) + + +class BasicTransformerBlock(nn.Module): + 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, + upcast_attention: bool = False, + use_attn_temp: bool = False + ): + super().__init__() + self.only_cross_attention = only_cross_attention + self.use_ada_layer_norm = num_embeds_ada_norm is not None + self.use_attn_temp = use_attn_temp + # SC-Attn + self.attn1 = SparseCausalAttention( + 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, + ) + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim) + + # Cross-Attn + if cross_attention_dim is not None: + self.attn2 = CrossAttention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + else: + self.attn2 = None + + if cross_attention_dim is not None: + self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim) + else: + self.norm2 = None + + # Feed-forward + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn) + self.norm3 = nn.LayerNorm(dim) + + # Temp-Attn + if self.use_attn_temp: + self.attn_temp = CrossAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + nn.init.zeros_(self.attn_temp.to_out[0].weight.data) + self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim) + + def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool): + if not is_xformers_available(): + print("Here is how to install it") + raise ModuleNotFoundError( + "Refer to https://github.com/facebookresearch/xformers for more information on how to install" + " xformers", + name="xformers", + ) + elif not torch.cuda.is_available(): + raise ValueError( + "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only" + " available for GPU " + ) + else: + try: + # Make sure we can run the memory efficient attention + _ = xformers.ops.memory_efficient_attention( + torch.randn((1, 2, 40), device="cuda"), + torch.randn((1, 2, 40), device="cuda"), + torch.randn((1, 2, 40), device="cuda"), + ) + except Exception as e: + raise e + self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers + if self.attn2 is not None: + self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers + #self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers + + def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None): + # SparseCausal-Attention + norm_hidden_states = ( + self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states) + ) + + if self.only_cross_attention: + hidden_states = ( + self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states + ) + else: + hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states + + if self.attn2 is not None: + # Cross-Attention + norm_hidden_states = ( + self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) + ) + hidden_states = ( + self.attn2( + norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask + ) + + hidden_states + ) + + # Feed-forward + hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states + + # Temporal-Attention + if self.use_attn_temp: + d = hidden_states.shape[1] + hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length) + norm_hidden_states = ( + self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states) + ) + hidden_states = self.attn_temp(norm_hidden_states) + hidden_states + hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) + + return hidden_states + + +class SparseCausalAttention(CrossAttention): + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None, use_full_attn=True): + batch_size, sequence_length, _ = hidden_states.shape + + encoder_hidden_states = encoder_hidden_states + + if self.group_norm is not None: + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = self.to_q(hidden_states) + # query = rearrange(query, "(b f) d c -> b (f d) c", f=video_length) + dim = query.shape[-1] + query = self.reshape_heads_to_batch_dim(query) + + if self.added_kv_proj_dim is not None: + raise NotImplementedError + + encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + former_frame_index = torch.arange(video_length) - 1 + former_frame_index[0] = 0 + + key = rearrange(key, "(b f) d c -> b f d c", f=video_length) + if not use_full_attn: + key = torch.cat([key[:, [0] * video_length], key[:, former_frame_index]], dim=2) + else: + # key = torch.cat([key[:, [0] * video_length], key[:, [1] * video_length], key[:, [2] * video_length], key[:, [3] * video_length]], dim=2) + key_video_length = [key[:, [i] * video_length] for i in range(video_length)] + key = torch.cat(key_video_length, dim=2) + key = rearrange(key, "b f d c -> (b f) d c") + + value = rearrange(value, "(b f) d c -> b f d c", f=video_length) + if not use_full_attn: + value = torch.cat([value[:, [0] * video_length], value[:, former_frame_index]], dim=2) + else: + # value = torch.cat([value[:, [0] * video_length], value[:, [1] * video_length], value[:, [2] * video_length], value[:, [3] * video_length]], dim=2) + value_video_length = [value[:, [i] * video_length] for i in range(video_length)] + value = torch.cat(value_video_length, dim=2) + value = rearrange(value, "b f d c -> (b f) d c") + + key = self.reshape_heads_to_batch_dim(key) + value = self.reshape_heads_to_batch_dim(value) + + if attention_mask is not None: + if attention_mask.shape[-1] != query.shape[1]: + target_length = query.shape[1] + attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) + attention_mask = attention_mask.repeat_interleave(self.heads, dim=0) + + # attention, what we cannot get enough of + if self._use_memory_efficient_attention_xformers: + hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask) + # Some versions of xformers return output in fp32, cast it back to the dtype of the input + hidden_states = hidden_states.to(query.dtype) + else: + if self._slice_size is None or query.shape[0] // self._slice_size == 1: + hidden_states = self._attention(query, key, value, attention_mask) + else: + hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask) + + # linear proj + hidden_states = self.to_out[0](hidden_states) + + # dropout + hidden_states = self.to_out[1](hidden_states) + return hidden_states diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/imageproj.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/imageproj.py new file mode 100644 index 0000000..63e2052 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/imageproj.py @@ -0,0 +1,118 @@ +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +import math + +import torch +import torch.nn as nn + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/refunet.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/refunet.py new file mode 100644 index 0000000..efe4e7b --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/refunet.py @@ -0,0 +1,125 @@ +import torch +from einops import rearrange +from typing import Any, Dict, Optional +from diffusers.utils.import_utils import is_xformers_available +from .transformer_mv2d import XFormersMVAttnProcessor, MVAttnProcessor +class ReferenceOnlyAttnProc(torch.nn.Module): + def __init__( + self, + chained_proc, + enabled=False, + name=None + ) -> None: + super().__init__() + self.enabled = enabled + self.chained_proc = chained_proc + self.name = name + + def __call__( + self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, + mode="w", ref_dict: dict = None, is_cfg_guidance = False,num_views=4, + multiview_attention=True, + cross_domain_attention=False, + ) -> Any: + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + # print(self.enabled) + if self.enabled: + if mode == 'w': + ref_dict[self.name] = encoder_hidden_states + res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, num_views=1, + multiview_attention=False, + cross_domain_attention=False,) + elif mode == 'r': + encoder_hidden_states = rearrange(encoder_hidden_states, '(b t) d c-> b (t d) c', t=num_views) + if self.name in ref_dict: + encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1).unsqueeze(1).repeat(1,num_views,1,1).flatten(0,1) + res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, num_views=num_views, + multiview_attention=False, + cross_domain_attention=False,) + elif mode == 'm': + encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict[self.name]], dim=1) + elif mode == 'n': + encoder_hidden_states = rearrange(encoder_hidden_states, '(b t) d c-> b (t d) c', t=num_views) + encoder_hidden_states = torch.cat([encoder_hidden_states], dim=1).unsqueeze(1).repeat(1,num_views,1,1).flatten(0,1) + res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask, num_views=num_views, + multiview_attention=False, + cross_domain_attention=False,) + else: + assert False, mode + else: + res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask) + return res + +class RefOnlyNoisedUNet(torch.nn.Module): + def __init__(self, unet, train_sched, val_sched) -> None: + super().__init__() + self.unet = unet + self.train_sched = train_sched + self.val_sched = val_sched + + unet_lora_attn_procs = dict() + for name, _ in unet.attn_processors.items(): + if is_xformers_available(): + default_attn_proc = XFormersMVAttnProcessor() + else: + default_attn_proc = MVAttnProcessor() + unet_lora_attn_procs[name] = ReferenceOnlyAttnProc( + default_attn_proc, enabled=name.endswith("attn1.processor"), name=name) + + self.unet.set_attn_processor(unet_lora_attn_procs) + + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.unet, name) + + def forward_cond(self, noisy_cond_lat, timestep, encoder_hidden_states, class_labels, ref_dict, is_cfg_guidance, **kwargs): + if is_cfg_guidance: + encoder_hidden_states = encoder_hidden_states[1:] + class_labels = class_labels[1:] + self.unet( + noisy_cond_lat, timestep, + encoder_hidden_states=encoder_hidden_states, + class_labels=class_labels, + cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict), + **kwargs + ) + + def forward( + self, sample, timestep, encoder_hidden_states, class_labels=None, + *args, cross_attention_kwargs, + down_block_res_samples=None, mid_block_res_sample=None, + **kwargs + ): + cond_lat = cross_attention_kwargs['cond_lat'] + is_cfg_guidance = cross_attention_kwargs.get('is_cfg_guidance', False) + noise = torch.randn_like(cond_lat) + if self.training: + noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep) + noisy_cond_lat = self.train_sched.scale_model_input(noisy_cond_lat, timestep) + else: + noisy_cond_lat = self.val_sched.add_noise(cond_lat, noise, timestep.reshape(-1)) + noisy_cond_lat = self.val_sched.scale_model_input(noisy_cond_lat, timestep.reshape(-1)) + ref_dict = {} + self.forward_cond( + noisy_cond_lat, timestep, + encoder_hidden_states, class_labels, + ref_dict, is_cfg_guidance, **kwargs + ) + weight_dtype = self.unet.dtype + return self.unet( + sample, timestep, + encoder_hidden_states, *args, + class_labels=class_labels, + cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=is_cfg_guidance), + down_block_additional_residuals=[ + sample.to(dtype=weight_dtype) for sample in down_block_res_samples + ] if down_block_res_samples is not None else None, + mid_block_additional_residual=( + mid_block_res_sample.to(dtype=weight_dtype) + if mid_block_res_sample is not None else None + ), + **kwargs + ) \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/resnet.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/resnet.py new file mode 100644 index 0000000..9b8b0fc --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/resnet.py @@ -0,0 +1,210 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from einops import rearrange + + +class InflatedConv3d(nn.Conv2d): + def forward(self, x): + video_length = x.shape[2] + + x = rearrange(x, "b c f h w -> (b f) c h w") + x = super().forward(x) + x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length) + + return x + + +class Upsample3D(nn.Module): + def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + + conv = None + if use_conv_transpose: + raise NotImplementedError + elif use_conv: + conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1) + + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward(self, hidden_states, output_size=None): + assert hidden_states.shape[1] == self.channels + + if self.use_conv_transpose: + raise NotImplementedError + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if output_size is None: + hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest") + else: + hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest") + + # If the input is bfloat16, we cast back to bfloat16 + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + + +class Downsample3D(nn.Module): + def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + self.name = name + + if use_conv: + conv = InflatedConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding) + else: + raise NotImplementedError + + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, hidden_states): + assert hidden_states.shape[1] == self.channels + if self.use_conv and self.padding == 0: + raise NotImplementedError + + assert hidden_states.shape[1] == self.channels + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class ResnetBlock3D(nn.Module): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout=0.0, + temb_channels=512, + groups=32, + groups_out=None, + pre_norm=True, + eps=1e-6, + non_linearity="swish", + time_embedding_norm="default", + output_scale_factor=1.0, + use_in_shortcut=None, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.time_embedding_norm = time_embedding_norm + self.output_scale_factor = output_scale_factor + + if groups_out is None: + groups_out = groups + + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + time_emb_proj_out_channels = out_channels + elif self.time_embedding_norm == "scale_shift": + time_emb_proj_out_channels = out_channels * 2 + else: + raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ") + + self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels) + else: + self.time_emb_proj = None + + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + + if non_linearity == "swish": + self.nonlinearity = lambda x: F.silu(x) + elif non_linearity == "mish": + self.nonlinearity = Mish() + elif non_linearity == "silu": + self.nonlinearity = nn.SiLU() + + self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = InflatedConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, input_tensor, temb): + hidden_states = input_tensor + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv1(hidden_states) + + if temb is not None: + # temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None] + temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, :, None, None].permute(0,2,1,3,4) + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + + +class Mish(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states)) \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/transformer_mv2d.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/transformer_mv2d.py new file mode 100644 index 0000000..bf9ca74 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/transformer_mv2d.py @@ -0,0 +1,1010 @@ +# 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 +try: + from diffusers.utils import maybe_allow_in_graph +except: + 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 + + +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, + joint_attention: bool=False, + joint_attention_twice: bool=False, + multiview_attention: bool=True, + cross_domain_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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: bool = False + ): + 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.cross_domain_attention = cross_domain_attention + # import pdb;pdb.set_trace() + 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.joint_attention = joint_attention + + if self.joint_attention: + # 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.joint_attention_twice = joint_attention_twice + + if self.joint_attention_twice: + 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, + num_views=self.num_views, + multiview_attention=self.multiview_attention, + cross_domain_attention=self.cross_domain_attention, + **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.joint_attention_twice: + 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.joint_attention: + 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 + ) + 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 + if multiview_attention: + if num_views <= 6: + # after use xformer; possible to train with 6 views + # key = rearrange(key, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0) + # value = rearrange(value, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0) + key = rearrange(key, '(b t) d c-> b (t d) c', t=num_views).unsqueeze(1).repeat(1,num_views,1,1).flatten(0,1) + value = rearrange(value, '(b t) d c-> b (t d) c', t=num_views).unsqueeze(1).repeat(1,num_views,1,1).flatten(0,1) + + else:# apply sparse attention + pass + # print("use sparse attention") + # # seems that the sparse random sampling cause problems + # # don't use random sampling, just fix the indexes + # onekey = rearrange(key, "(b t) d c -> b t d c", t=num_views) + # onevalue = rearrange(value, "(b t) d c -> b t d c", t=num_views) + # allkeys = [] + # allvalues = [] + # all_indexes = { + # 0 : [0, 2, 3, 4], + # 1: [0, 1, 3, 5], + # 2: [0, 2, 3, 4], + # 3: [0, 2, 3, 4], + # 4: [0, 2, 3, 4], + # 5: [0, 1, 3, 5] + # } + # for jj in range(num_views): + # # valid_index = [x for x in range(0, num_views) if x!= jj] + # # indexes = random.sample(valid_index, 3) + [jj] + [0] + # indexes = all_indexes[jj] + + # indexes = torch.tensor(indexes).long().to(key.device) + # allkeys.append(onekey[:, indexes]) + # allvalues.append(onevalue[:, indexes]) + # keys = torch.stack(allkeys, dim=1) # checked, should be dim=1 + # values = torch.stack(allvalues, dim=1) + # key = rearrange(keys, 'b t f d c -> (b t) (f d) c') + # value = rearrange(values, 'b t f d c -> (b t) (f d) 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 + + +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, + cross_domain_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 + ) + 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_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) + #([bx4, 1024, 320]) key torch.Size([bx4, 1024, 320]) value torch.Size([bx4, 1024, 320]) + # pdb.set_trace() + # multi-view self-attention + if multiview_attention: + key = rearrange(key_raw, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0) + value = rearrange(value_raw, "(b t) d c -> b (t d) c", t=num_views).repeat_interleave(num_views, dim=0) + + if cross_domain_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 + else: + # print("don't use multiview attention.") + key = key_raw + value = value_raw + + query = attn.head_to_batch_dim(query) + 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) + + 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 \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet.py new file mode 100644 index 0000000..fa156af --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet.py @@ -0,0 +1,497 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_condition.py + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import os +import json + +import torch +import torch.nn as nn +import torch.utils.checkpoint + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers import ModelMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from .unet_blocks import ( + CrossAttnDownBlock3D, + CrossAttnUpBlock3D, + DownBlock3D, + UNetMidBlock3DCrossAttn, + UpBlock3D, + get_down_block, + get_up_block, +) +from .resnet import InflatedConv3d + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNet3DConditionOutput(BaseOutput): + sample: torch.FloatTensor + + +class UNet3DConditionModel(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D", + ), + mid_block_type: str = "UNetMidBlock3DCrossAttn", + up_block_types: Tuple[str] = ( + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D" + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-5, + cross_attention_dim: int = 1280, + attention_head_dim: Union[int, Tuple[int]] = 8, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + use_attn_temp: bool = False, + camera_input_dim: int = 12, + camera_hidden_dim: int = 320, + camera_output_dim: int = 1280, + ): + super().__init__() + + self.sample_size = sample_size + time_embed_dim = block_out_channels[0] * 4 + + # input + self.conv_in = InflatedConv3d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)) + + # time + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + + self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + else: + self.class_embedding = None + + # camera metrix + # def init_linear(l, stddev): + # nn.init.normal_(l.weight, std=stddev) + # if l.bias is not None: + # nn.init.constant_(l.bias, 0.0) + # self.camera_embedding_1 = nn.Linear(camera_input_dim, camera_hidden_dim) + # self.camera_embedding_2 = nn.Linear(camera_hidden_dim, camera_output_dim) + # init_linear(self.camera_embedding_1, 0.25) + # init_linear(self.camera_embedding_2, 0.25) + + self.camera_embedding = nn.Sequential( + nn.Linear(camera_input_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + self.down_blocks = nn.ModuleList([]) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + temb_channels=time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + use_attn_temp=use_attn_temp + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock3DCrossAttn": + self.mid_block = UNetMidBlock3DCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the videos + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_attention_head_dim = list(reversed(attention_head_dim)) + only_cross_attention = list(reversed(only_cross_attention)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=reversed_attention_head_dim[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + use_attn_temp=use_attn_temp, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps) + self.conv_act = nn.SiLU() + self.conv_out = InflatedConv3d(block_out_channels[0], out_channels, kernel_size=3, padding=1) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module will split the input tensor in slices, to compute attention + in several steps. This is useful to save some memory in exchange for a small speed decrease. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_slicable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_slicable_dims(module) + + num_slicable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_slicable_layers * [1] + + slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + camera_matrixs: Optional[torch.Tensor] = None, + class_labels: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[UNet3DConditionOutput, Tuple]: + r""" + Args: + sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor + timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps + encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # prepare attention_mask + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + # time + timesteps = timestep + if not torch.is_tensor(timesteps): + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # timesteps does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=self.dtype) + emb = self.time_embedding(t_emb) #torch.Size([32, 1280]) + emb = torch.unsqueeze(emb, 1) + if camera_matrixs is not None: + # came emb + cam_emb = self.camera_embedding(camera_matrixs) + # cam_emb = self.camera_embedding_2(cam_emb) + emb = emb.repeat(1,cam_emb.shape[1],1) #torch.Size([32, 4, 1280]) + emb = emb + cam_emb + + # import pdb;pdb.set_trace() + if self.class_embedding is not None: + # if class_labels is None: + # raise ValueError("class_labels should be provided when num_class_embeds > 0") + if class_labels is not None: + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + class_emb = self.class_embedding(class_labels) + emb = emb + class_emb + + # pre-process + sample = self.conv_in(sample) + + # down + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb) + + down_block_res_samples += res_samples + + # mid + sample = self.mid_block( + sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask + ) + + # up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + upsample_size=upsample_size, + attention_mask=attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size + ) + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if not return_dict: + return (sample,) + + return UNet3DConditionOutput(sample=sample) + + @classmethod + def from_pretrained_2d(cls, pretrained_model_path, subfolder=None): + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, subfolder) + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f"{config_file} does not exist") + with open(config_file, "r") as f: + config = json.load(f) + config["_class_name"] = cls.__name__ + config["down_block_types"] = [ + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D" + ] + config["up_block_types"] = [ + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D" + ] + + from diffusers.utils import WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME + # model = cls.from_config(config) + # model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + # if not os.path.isfile(model_file): + # raise RuntimeError(f"{model_file} does not exist") + # state_dict = torch.load(model_file, map_location="cpu") + + import safetensors + model = cls.from_config(config) + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + if not os.path.isfile(model_file): + model_file = os.path.join(pretrained_model_path, SAFETENSORS_WEIGHTS_NAME) + if not os.path.isfile(model_file): + raise RuntimeError(f"{model_file} does not exist") + else: + state_dict = safetensors.torch.load_file(model_file, device="cpu") + else: + state_dict = torch.load(model_file, map_location="cpu") + + for k, v in model.state_dict().items(): + if '_temp.' in k or 'camera_embedding' in k or 'class_embedding' in k: + state_dict.update({k: v}) + for k in list(state_dict.keys()): + if 'camera_embedding_' in k: + v = state_dict.pop(k) + model.load_state_dict(state_dict) + + return model \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_blocks.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_blocks.py new file mode 100644 index 0000000..ac9b78d --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_blocks.py @@ -0,0 +1,596 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py + +import torch +from torch import nn + +# from .attention import Transformer3DModel +from .resnet import Downsample3D, ResnetBlock3D, Upsample3D + + +def get_down_block( + down_block_type, + num_layers, + in_channels, + out_channels, + temb_channels, + add_downsample, + resnet_eps, + resnet_act_fn, + attn_num_head_channels, + 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", + use_attn_temp=False, +): + down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type + if down_block_type == "DownBlock3D": + return DownBlock3D( + 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 == "CrossAttnDownBlock3D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D") + return CrossAttnDownBlock3D( + 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, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + 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, + use_attn_temp=use_attn_temp, + ) + 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, + attn_num_head_channels, + 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", + use_attn_temp=False, +): + up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + if up_block_type == "UpBlock3D": + return UpBlock3D( + 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 == "CrossAttnUpBlock3D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D") + return CrossAttnUpBlock3D( + 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, + attn_num_head_channels=attn_num_head_channels, + 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, + use_attn_temp=use_attn_temp, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlock3DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: 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, + attn_num_head_channels=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + dual_cross_attention=False, + use_linear_projection=False, + upcast_attention=False, + ): + super().__init__() + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock3D( + 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 dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + in_channels // attn_num_head_channels, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + ) + resnets.append( + ResnetBlock3D( + 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, temb=None, encoder_hidden_states=None, attention_mask=None): + 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).sample + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class CrossAttnDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: 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, + attn_num_head_channels=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, + use_attn_temp=False, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + 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 dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + 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, + use_attn_temp=use_attn_temp, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + 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, temb=None, encoder_hidden_states=None, attention_mask=None): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + 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 + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class DownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: 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, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + 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, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + 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, temb=None): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnUpBlock3D(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, + 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, + attn_num_head_channels=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, + use_attn_temp=False, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + 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( + ResnetBlock3D( + 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 dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + 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, + use_attn_temp=use_attn_temp, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + encoder_hidden_states=None, + upsample_size=None, + attention_mask=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 + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class UpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: 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, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + + 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( + ResnetBlock3D( + 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, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample3D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None): + for resnet in self.resnets: + # 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): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_blocks.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_blocks.py new file mode 100644 index 0000000..0072567 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_blocks.py @@ -0,0 +1,926 @@ +# 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.attention 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 .transformer_mv2d import TransformerMV2DModel + +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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: 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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + 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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: 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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + 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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: 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) + + # 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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + ) + 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, + ) -> 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, + 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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: bool=False + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + 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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + ) + 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 + if num_views == 4: + self.gradient_checkpointing = False + else: + 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, + ): + 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, + None, # timestep + None, # class_labels + cross_attention_kwargs, + attention_mask, + encoder_attention_mask, + **ckpt_kwargs, + )[0] + # 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, + # return_dict=False, + # )[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, + 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, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: bool=False + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + 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, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + ) + 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 + if num_views == 4: + self.gradient_checkpointing = False + else: + self.gradient_checkpointing = False + + 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, + 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, + None, # timestep + None, # class_labels + cross_attention_kwargs, + attention_mask, + encoder_attention_mask, + **ckpt_kwargs, + )[0] + else: + # import ipdb + # ipdb.set_trace() + 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, + 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 + diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_condition.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_condition.py new file mode 100644 index 0000000..be58de7 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_condition.py @@ -0,0 +1,1506 @@ +# 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, List, Optional, Tuple, Union +import os + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from einops import rearrange + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor +from diffusers.models.embeddings import ( + GaussianFourierProjection, + ImageHintTimeEmbedding, + ImageProjection, + ImageTimeEmbedding, + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.modeling_utils import ModelMixin, load_state_dict, _load_state_dict_into_model +from diffusers.models.unets.unet_2d_blocks import ( + CrossAttnDownBlock2D, + CrossAttnUpBlock2D, + DownBlock2D, + UNetMidBlock2DCrossAttn, + UNetMidBlock2DSimpleCrossAttn, + UpBlock2D, +) +from huggingface_hub.constants import HF_HUB_CACHE, HF_HUB_OFFLINE +from diffusers.utils import ( + SAFETENSORS_WEIGHTS_NAME, + WEIGHTS_NAME, + _add_variant, + _get_model_file, + deprecate, + is_accelerate_available, + is_torch_version, + logging, +) +from diffusers import __version__ +from .unet_mv2d_blocks import ( + CrossAttnDownBlockMV2D, + CrossAttnUpBlockMV2D, + UNetMidBlockMV2DCrossAttn, + get_down_block, + get_up_block, +) +from diffusers.models.attention_processor import Attention, AttnProcessor +from diffusers.utils.import_utils import is_xformers_available +from .transformer_mv2d import XFormersMVAttnProcessor, MVAttnProcessor +from .refunet import ReferenceOnlyAttnProc + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNetMV2DConditionOutput(BaseOutput): + """ + The output of [`UNet2DConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + +class UNetMV2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or + `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, normalization and activation layers is skipped in post-processing. + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + num_attention_heads (`int`, *optional*): + The number of attention heads. If not defined, defaults to `attention_head_dim` + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + addition_time_embed_dim: (`int`, *optional*, defaults to `None`): + Dimension for the timestep embeddings. + num_class_embeds (`int`, *optional*, defaults to `None`): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, defaults to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, defaults to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, defaults to `None`): + Optional activation function to use only once on the time embeddings before they are passed to the rest of + the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str`, *optional*, defaults to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, defaults to `None`): + The dimension of `cond_proj` layer in the timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. + conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. + projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when + `class_embed_type="projection"`. Required when `class_embed_type="projection"`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the + `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` + otherwise. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlockMV2DCrossAttn", + up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D"), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = "positional", + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + num_views: int = 1, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: bool = False, + camera_input_dim: int = 12, + camera_hidden_dim: int = 320, + camera_output_dim: int = 1280, + + ): + super().__init__() + + self.sample_size = sample_size + + if num_attention_heads is not None: + raise ValueError( + "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." + ) + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding + ) + + # time + if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos + ) + timestep_input_dim = time_embed_dim + elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 + self.encoder_hid_proj = ImageProjection( + image_embed_dim=encoder_hid_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif class_embed_type == "simple_projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) + self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif addition_embed_type == "image": + # Kandinsky 2.2 + self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type == "image_hint": + # Kandinsky 2.2 ControlNet + self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type is not None: + raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") + + if time_embedding_act_fn is None: + self.time_embed_act = None + else: + self.time_embed_act = get_activation(time_embedding_act_fn) + + self.camera_embedding = nn.Sequential( + nn.Linear(camera_input_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock2DCrossAttn": + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + # custom MV2D attention block + elif mid_block_type == "UNetMidBlockMV2DCrossAttn": + self.mid_block = UNetMidBlockMV2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": + self.mid_block = UNetMidBlock2DSimpleCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + cross_attention_dim=cross_attention_dim[-1], + attention_head_dim=attention_head_dim[-1], + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + only_cross_attention=mid_block_only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps + ) + + self.conv_act = get_activation(act_fn) + + else: + self.conv_norm_out = None + self.conv_act = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + self.conv_out = nn.Conv2d( + block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "set_processor"): + processors[f"{name}.processor"] = module.processor + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + self.set_attn_processor(AttnProcessor()) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + # def _set_gradient_checkpointing(self, module, value=False): + # if isinstance(module, (CrossAttnDownBlock2D, CrossAttnDownBlockMV2D, DownBlock2D, CrossAttnUpBlock2D, CrossAttnUpBlockMV2D, UpBlock2D)): + # module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + camera_matrixs: Optional[torch.Tensor] = None, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[UNetMV2DConditionOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # 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: + # 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(sample.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: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + emb = self.time_embedding(t_emb, timestep_cond) + + # import pdb; pdb.set_trace() + if camera_matrixs is not None: + emb = torch.unsqueeze(emb, 1) + # came emb + cam_emb = self.camera_embedding(camera_matrixs) + # cam_emb = self.camera_embedding_2(cam_emb) + # import ipdb + # ipdb.set_trace() + emb = emb.repeat(1,cam_emb.shape[1],1) #torch.Size([32, 4, 1280]) + emb = emb + cam_emb + emb = rearrange(emb, "b f c -> (b f) c", f=emb.shape[1]) + + aug_emb = None + + if self.class_embedding is not None and class_labels is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + # 2. pre-process + sample = rearrange(sample, "b c f h w -> (b f) c h w", f=sample.shape[2]) + sample = self.conv_in(sample) + # 3. down + + is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None + is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + # print("after down: ", sample.mean(), emb.mean()) + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if is_controlnet: + sample = sample + mid_block_additional_residual + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size + ) + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if not return_dict: + return (sample,) + + return UNetMV2DConditionOutput(sample=sample) + + @classmethod + def from_pretrained_2d( + cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + camera_embedding_type: str, num_views: int, sample_size: int, + zero_init_conv_in: bool = True, zero_init_camera_projection: bool = False, + projection_class_embeddings_input_dim: int=6, joint_attention: bool = False, + joint_attention_twice: bool = False, multiview_attention: bool = True, + cross_domain_attention: bool = False, + in_channels: int = 8, out_channels: int = 4, local_crossattn=False, + **kwargs + ): + r""" + Instantiate a pretrained PyTorch model from a pretrained model configuration. + + The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To + train the model, set it back in training mode with `model.train()`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on + the Hub. + - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved + with [`~ModelMixin.save_pretrained`]. + + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the + dtype is automatically derived from the model's weights. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info (`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + from_flax (`bool`, *optional*, defaults to `False`): + Load the model weights from a Flax checkpoint save file. + subfolder (`str`, *optional*, defaults to `""`): + The subfolder location of a model file within a larger model repository on the Hub or locally. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you're downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be defined for each + parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the + same device. + + Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier for the maximum memory. Will default to the maximum memory available for + each GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + The path to offload weights if `device_map` contains the value `"disk"`. + offload_state_dict (`bool`, *optional*): + If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if + the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` + when there is some disk offload. + low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): + Speed up model loading only loading the pretrained weights and not initializing the weights. This also + tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this + argument to `True` will raise an error. + variant (`str`, *optional*): + Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the + `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` + weights. If set to `False`, `safetensors` weights are not loaded. + + + + To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with + `huggingface-cli login`. You can also activate the special + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + firewalled environment. + + + + Example: + + ```py + from diffusers import UNet2DConditionModel + + unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") + ``` + + If you get the error message below, you need to finetune the weights for your downstream task: + + ```bash + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated + You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. + ``` + """ + cache_dir = kwargs.pop("cache_dir", HF_HUB_CACHE) + ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False) + force_download = kwargs.pop("force_download", False) + from_flax = kwargs.pop("from_flax", False) + resume_download = kwargs.pop("resume_download", False) + proxies = kwargs.pop("proxies", None) + output_loading_info = kwargs.pop("output_loading_info", False) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + torch_dtype = kwargs.pop("torch_dtype", None) + subfolder = kwargs.pop("subfolder", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + + # if use_safetensors and not is_safetensors_available(): + # raise ValueError( + # "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetensors" + # ) + + allow_pickle = False + if use_safetensors is None: + # use_safetensors = is_safetensors_available() + use_safetensors = False + allow_pickle = True + + if device_map is not None and not is_accelerate_available(): + raise NotImplementedError( + "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set" + " `device_map=None`. You can install accelerate with `pip install accelerate`." + ) + + # Check if we can handle device_map and dispatching the weights + if device_map is not None and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `device_map=None`." + ) + + # Load config if we don't provide a configuration + config_path = pretrained_model_name_or_path + + user_agent = { + "diffusers": __version__, + "file_type": "model", + "framework": "pytorch", + } + + # load config + config, unused_kwargs, commit_hash = cls.load_config( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + return_commit_hash=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + user_agent=user_agent, + **kwargs, + ) + + # modify config + config["_class_name"] = cls.__name__ + config['in_channels'] = in_channels + config['out_channels'] = out_channels + config['sample_size'] = sample_size # training resolution + config['num_views'] = num_views + config['joint_attention'] = joint_attention + config['joint_attention_twice'] = joint_attention_twice + config['multiview_attention'] = multiview_attention + config['cross_domain_attention'] = cross_domain_attention + config["down_block_types"] = [ + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "DownBlock2D" + ] + config['mid_block_type'] = "UNetMidBlockMV2DCrossAttn" + config["up_block_types"] = [ + "UpBlock2D", + "CrossAttnUpBlockMV2D", + "CrossAttnUpBlockMV2D", + "CrossAttnUpBlockMV2D" + ] + config['class_embed_type'] = 'projection' + if camera_embedding_type == 'e_de_da_sincos': + config['projection_class_embeddings_input_dim'] = projection_class_embeddings_input_dim # default 6 + else: + raise NotImplementedError + + # load model + model_file = None + if from_flax: + raise NotImplementedError + else: + if use_safetensors: + try: + model_file = _get_model_file( + pretrained_model_name_or_path, + weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant), + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + commit_hash=commit_hash, + ) + except IOError as e: + if not allow_pickle: + raise e + pass + if model_file is None: + model_file = _get_model_file( + pretrained_model_name_or_path, + weights_name=_add_variant(WEIGHTS_NAME, variant), + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + commit_hash=commit_hash, + ) + + model = cls.from_config(config, **unused_kwargs) + if local_crossattn: + unet_lora_attn_procs = dict() + for name, _ in model.attn_processors.items(): + if not name.endswith("attn1.processor"): + default_attn_proc = AttnProcessor() + elif is_xformers_available(): + default_attn_proc = XFormersMVAttnProcessor() + else: + default_attn_proc = MVAttnProcessor() + unet_lora_attn_procs[name] = ReferenceOnlyAttnProc( + default_attn_proc, enabled=name.endswith("attn1.processor"), name=name + ) + model.set_attn_processor(unet_lora_attn_procs) + state_dict = load_state_dict(model_file, variant=variant) + model._convert_deprecated_attention_blocks(state_dict) + + conv_in_weight = state_dict['conv_in.weight'] + conv_out_weight = state_dict['conv_out.weight'] + model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model_2d( + model, + state_dict, + model_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=True, + ) + if any([key == 'conv_in.weight' for key, _, _ in mismatched_keys]): + # initialize from the original SD structure + model.conv_in.weight.data[:,:4] = conv_in_weight + + # whether to place all zero to new layers? + if zero_init_conv_in: + model.conv_in.weight.data[:,4:] = 0. + + if any([key == 'conv_out.weight' for key, _, _ in mismatched_keys]): + # initialize from the original SD structure + model.conv_out.weight.data[:,:4] = conv_out_weight + if out_channels == 8: # copy for the last 4 channels + model.conv_out.weight.data[:, 4:] = conv_out_weight + + if zero_init_camera_projection: + for p in model.class_embedding.parameters(): + torch.nn.init.zeros_(p) + + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + "error_msgs": error_msgs, + } + + if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}." + ) + elif torch_dtype is not None: + model = model.to(torch_dtype) + + model.register_to_config(_name_or_path=pretrained_model_name_or_path) + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + if output_loading_info: + return model, loading_info + + return model + + @classmethod + def _load_pretrained_model_2d( + cls, + model, + state_dict, + resolved_archive_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=False, + ): + # Retrieve missing & unexpected_keys + model_state_dict = model.state_dict() + loaded_keys = list(state_dict.keys()) + + expected_keys = list(model_state_dict.keys()) + + original_loaded_keys = loaded_keys + + missing_keys = list(set(expected_keys) - set(loaded_keys)) + unexpected_keys = list(set(loaded_keys) - set(expected_keys)) + + # Make sure we are able to load base models as well as derived models (with heads) + model_to_load = model + + def _find_mismatched_keys( + state_dict, + model_state_dict, + loaded_keys, + ignore_mismatched_sizes, + ): + mismatched_keys = [] + if ignore_mismatched_sizes: + for checkpoint_key in loaded_keys: + model_key = checkpoint_key + + if ( + model_key in model_state_dict + and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape + ): + mismatched_keys.append( + (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape) + ) + del state_dict[checkpoint_key] + return mismatched_keys + + if state_dict is not None: + # Whole checkpoint + mismatched_keys = _find_mismatched_keys( + state_dict, + model_state_dict, + original_loaded_keys, + ignore_mismatched_sizes, + ) + error_msgs = _load_state_dict_into_model(model_to_load, state_dict) + + if len(error_msgs) > 0: + error_msg = "\n\t".join(error_msgs) + if "size mismatch" in error_msg: + error_msg += ( + "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." + ) + raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") + + if len(unexpected_keys) > 0: + logger.warning( + f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task" + " or with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly" + " identical (initializing a BertForSequenceClassification model from a" + " BertForSequenceClassification model)." + ) + else: + logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.info( + f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the" + f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions" + " without further training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be" + " able to use it for predictions and inference." + ) + + return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs + diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_ref.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_ref.py new file mode 100644 index 0000000..bef070c --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/models/unet_mv2d_ref.py @@ -0,0 +1,1567 @@ +# 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, List, Optional, Tuple, Union +import os + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from einops import rearrange + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor +from diffusers.models.embeddings import ( + GaussianFourierProjection, + ImageHintTimeEmbedding, + ImageProjection, + ImageTimeEmbedding, + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.lora import LoRALinearLayer + +from diffusers.models.modeling_utils import ModelMixin, load_state_dict, _load_state_dict_into_model +from diffusers.models.unets.unet_2d_blocks import ( + CrossAttnDownBlock2D, + CrossAttnUpBlock2D, + DownBlock2D, + UNetMidBlock2DCrossAttn, + UNetMidBlock2DSimpleCrossAttn, + UpBlock2D, +) +from huggingface_hub.constants import HF_HUB_CACHE, HF_HUB_OFFLINE +from diffusers.utils import ( + SAFETENSORS_WEIGHTS_NAME, + WEIGHTS_NAME, + _add_variant, + _get_model_file, + deprecate, + is_accelerate_available, + is_torch_version, + logging, +) +from diffusers import __version__ +from .unet_mv2d_blocks import ( + CrossAttnDownBlockMV2D, + CrossAttnUpBlockMV2D, + UNetMidBlockMV2DCrossAttn, + get_down_block, + get_up_block, +) +from diffusers.models.attention_processor import Attention, AttnProcessor +from diffusers.utils.import_utils import is_xformers_available +from .transformer_mv2d import XFormersMVAttnProcessor, MVAttnProcessor +from .refunet import ReferenceOnlyAttnProc + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNetMV2DRefOutput(BaseOutput): + """ + The output of [`UNet2DConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + +class Identity(torch.nn.Module): + r"""A placeholder identity operator that is argument-insensitive. + + Args: + args: any argument (unused) + kwargs: any keyword argument (unused) + + Shape: + - Input: :math:`(*)`, where :math:`*` means any number of dimensions. + - Output: :math:`(*)`, same shape as the input. + + Examples:: + + >>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False) + >>> input = torch.randn(128, 20) + >>> output = m(input) + >>> print(output.size()) + torch.Size([128, 20]) + + """ + def __init__(self, scale=None, *args, **kwargs) -> None: + super(Identity, self).__init__() + + def forward(self, input, *args, **kwargs): + return input + + + +class _LoRACompatibleLinear(nn.Module): + """ + A Linear layer that can be used with LoRA. + """ + + def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = None, **kwargs): + super().__init__(*args, **kwargs) + self.lora_layer = lora_layer + + def set_lora_layer(self, lora_layer: Optional[LoRALinearLayer]): + self.lora_layer = lora_layer + + def _fuse_lora(self): + pass + + def _unfuse_lora(self): + pass + + def forward(self, hidden_states, scale=None, lora_scale: int = 1): + return hidden_states + +class UNetMV2DRefModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or + `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, normalization and activation layers is skipped in post-processing. + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + num_attention_heads (`int`, *optional*): + The number of attention heads. If not defined, defaults to `attention_head_dim` + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + addition_time_embed_dim: (`int`, *optional*, defaults to `None`): + Dimension for the timestep embeddings. + num_class_embeds (`int`, *optional*, defaults to `None`): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, defaults to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, defaults to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, defaults to `None`): + Optional activation function to use only once on the time embeddings before they are passed to the rest of + the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str`, *optional*, defaults to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, defaults to `None`): + The dimension of `cond_proj` layer in the timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. + conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. + projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when + `class_embed_type="projection"`. Required when `class_embed_type="projection"`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the + `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` + otherwise. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlockMV2DCrossAttn", + up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D"), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = "positional", + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + num_views: int = 1, + joint_attention: bool = False, + joint_attention_twice: bool = False, + multiview_attention: bool = True, + cross_domain_attention: bool = False, + camera_input_dim: int = 12, + camera_hidden_dim: int = 320, + camera_output_dim: int = 1280, + + ): + super().__init__() + + self.sample_size = sample_size + + if num_attention_heads is not None: + raise ValueError( + "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." + ) + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding + ) + + # time + if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos + ) + timestep_input_dim = time_embed_dim + elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 + self.encoder_hid_proj = ImageProjection( + image_embed_dim=encoder_hid_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif class_embed_type == "simple_projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) + self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif addition_embed_type == "image": + # Kandinsky 2.2 + self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type == "image_hint": + # Kandinsky 2.2 ControlNet + self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type is not None: + raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") + + if time_embedding_act_fn is None: + self.time_embed_act = None + else: + self.time_embed_act = get_activation(time_embedding_act_fn) + + self.camera_embedding = nn.Sequential( + nn.Linear(camera_input_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock2DCrossAttn": + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + # custom MV2D attention block + elif mid_block_type == "UNetMidBlockMV2DCrossAttn": + self.mid_block = UNetMidBlockMV2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": + self.mid_block = UNetMidBlock2DSimpleCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + cross_attention_dim=cross_attention_dim[-1], + attention_head_dim=attention_head_dim[-1], + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + only_cross_attention=mid_block_only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + num_views=num_views, + joint_attention=joint_attention, + joint_attention_twice=joint_attention_twice, + multiview_attention=multiview_attention, + cross_domain_attention=cross_domain_attention + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + # if norm_num_groups is not None: + # self.conv_norm_out = nn.GroupNorm( + # num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps + # ) + + # self.conv_act = get_activation(act_fn) + + # else: + # self.conv_norm_out = None + # self.conv_act = None + + # conv_out_padding = (conv_out_kernel - 1) // 2 + # self.conv_out = nn.Conv2d( + # block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding + # ) + + self.up_blocks[3].attentions[2].transformer_blocks[0].attn1.to_q = _LoRACompatibleLinear() + self.up_blocks[3].attentions[2].transformer_blocks[0].attn1.to_k = _LoRACompatibleLinear() + self.up_blocks[3].attentions[2].transformer_blocks[0].attn1.to_v = _LoRACompatibleLinear() + self.up_blocks[3].attentions[2].transformer_blocks[0].attn1.to_out = nn.ModuleList([Identity(), Identity()]) + self.up_blocks[3].attentions[2].transformer_blocks[0].norm2 = Identity() + self.up_blocks[3].attentions[2].transformer_blocks[0].attn2 = None + self.up_blocks[3].attentions[2].transformer_blocks[0].norm3 = Identity() + self.up_blocks[3].attentions[2].transformer_blocks[0].ff = Identity() + self.up_blocks[3].attentions[2].proj_out = Identity() + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "set_processor"): + processors[f"{name}.processor"] = module.processor + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + self.set_attn_processor(AttnProcessor()) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (CrossAttnDownBlock2D, CrossAttnDownBlockMV2D, DownBlock2D, CrossAttnUpBlock2D, CrossAttnUpBlockMV2D, UpBlock2D)): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + camera_matrixs: Optional[torch.Tensor] = None, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[UNetMV2DRefOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # 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: + # 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(sample.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: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + emb = self.time_embedding(t_emb, timestep_cond) + + # import pdb; pdb.set_trace() + if camera_matrixs is not None: + emb = torch.unsqueeze(emb, 1) + # came emb + cam_emb = self.camera_embedding(camera_matrixs) + # cam_emb = self.camera_embedding_2(cam_emb) + emb = emb.repeat(1,cam_emb.shape[1],1) #torch.Size([32, 4, 1280]) + emb = emb + cam_emb + emb = rearrange(emb, "b f c -> (b f) c", f=emb.shape[1]) + + aug_emb = None + + if self.class_embedding is not None and class_labels is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + # 2. pre-process + sample = rearrange(sample, "b c f h w -> (b f) c h w", f=sample.shape[2]) + sample = self.conv_in(sample) + # 3. down + + is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None + is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + # print("after down: ", sample.mean(), emb.mean()) + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if is_controlnet: + sample = sample + mid_block_additional_residual + + # print("after mid: ", sample.mean()) + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size + ) + + # 6. post-process + # if self.conv_norm_out: + # sample = self.conv_norm_out(sample) + # sample = self.conv_act(sample) + # sample = self.conv_out(sample) + + if not return_dict: + return (sample,) + + return UNetMV2DRefOutput(sample=sample) + + @classmethod + def from_pretrained_2d( + cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + camera_embedding_type: str, num_views: int, sample_size: int, + zero_init_conv_in: bool = True, zero_init_camera_projection: bool = False, + projection_class_embeddings_input_dim: int=6, joint_attention: bool = False, + joint_attention_twice: bool = False, multiview_attention: bool = True, + cross_domain_attention: bool = False, + in_channels: int = 8, out_channels: int = 4, local_crossattn=False, + **kwargs + ): + r""" + Instantiate a pretrained PyTorch model from a pretrained model configuration. + + The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To + train the model, set it back in training mode with `model.train()`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on + the Hub. + - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved + with [`~ModelMixin.save_pretrained`]. + + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the + dtype is automatically derived from the model's weights. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info (`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + from_flax (`bool`, *optional*, defaults to `False`): + Load the model weights from a Flax checkpoint save file. + subfolder (`str`, *optional*, defaults to `""`): + The subfolder location of a model file within a larger model repository on the Hub or locally. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you're downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be defined for each + parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the + same device. + + Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier for the maximum memory. Will default to the maximum memory available for + each GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + The path to offload weights if `device_map` contains the value `"disk"`. + offload_state_dict (`bool`, *optional*): + If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if + the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` + when there is some disk offload. + low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): + Speed up model loading only loading the pretrained weights and not initializing the weights. This also + tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this + argument to `True` will raise an error. + variant (`str`, *optional*): + Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the + `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors` + weights. If set to `False`, `safetensors` weights are not loaded. + + + + To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with + `huggingface-cli login`. You can also activate the special + ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a + firewalled environment. + + + + Example: + + ```py + from diffusers import UNet2DConditionModel + + unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet") + ``` + + If you get the error message below, you need to finetune the weights for your downstream task: + + ```bash + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated + You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. + ``` + """ + cache_dir = kwargs.pop("cache_dir", HF_HUB_CACHE) + ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False) + force_download = kwargs.pop("force_download", False) + from_flax = kwargs.pop("from_flax", False) + resume_download = kwargs.pop("resume_download", False) + proxies = kwargs.pop("proxies", None) + output_loading_info = kwargs.pop("output_loading_info", False) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + torch_dtype = kwargs.pop("torch_dtype", None) + subfolder = kwargs.pop("subfolder", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + + # if use_safetensors and not is_safetensors_available(): + # raise ValueError( + # "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetensors" + # ) + + allow_pickle = False + if use_safetensors is None: + # use_safetensors = is_safetensors_available() + use_safetensors = False + allow_pickle = True + + if device_map is not None and not is_accelerate_available(): + raise NotImplementedError( + "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set" + " `device_map=None`. You can install accelerate with `pip install accelerate`." + ) + + # Check if we can handle device_map and dispatching the weights + if device_map is not None and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `device_map=None`." + ) + + # Load config if we don't provide a configuration + config_path = pretrained_model_name_or_path + + user_agent = { + "diffusers": __version__, + "file_type": "model", + "framework": "pytorch", + } + + # load config + config, unused_kwargs, commit_hash = cls.load_config( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + return_commit_hash=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + user_agent=user_agent, + **kwargs, + ) + + # modify config + config["_class_name"] = cls.__name__ + config['in_channels'] = in_channels + config['out_channels'] = out_channels + config['sample_size'] = sample_size # training resolution + config['num_views'] = num_views + config['joint_attention'] = joint_attention + config['joint_attention_twice'] = joint_attention_twice + config['multiview_attention'] = multiview_attention + config['cross_domain_attention'] = cross_domain_attention + config["down_block_types"] = [ + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "CrossAttnDownBlockMV2D", + "DownBlock2D" + ] + config['mid_block_type'] = "UNetMidBlockMV2DCrossAttn" + config["up_block_types"] = [ + "UpBlock2D", + "CrossAttnUpBlockMV2D", + "CrossAttnUpBlockMV2D", + "CrossAttnUpBlockMV2D" + ] + config['class_embed_type'] = 'projection' + if camera_embedding_type == 'e_de_da_sincos': + config['projection_class_embeddings_input_dim'] = projection_class_embeddings_input_dim # default 6 + else: + raise NotImplementedError + + # load model + model_file = None + if from_flax: + raise NotImplementedError + else: + if use_safetensors: + try: + model_file = _get_model_file( + pretrained_model_name_or_path, + weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant), + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + commit_hash=commit_hash, + ) + except IOError as e: + if not allow_pickle: + raise e + pass + if model_file is None: + model_file = _get_model_file( + pretrained_model_name_or_path, + weights_name=_add_variant(WEIGHTS_NAME, variant), + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + commit_hash=commit_hash, + ) + + model = cls.from_config(config, **unused_kwargs) + if local_crossattn: + unet_lora_attn_procs = dict() + for name, _ in model.attn_processors.items(): + if not name.endswith("attn1.processor"): + default_attn_proc = AttnProcessor() + elif is_xformers_available(): + default_attn_proc = XFormersMVAttnProcessor() + else: + default_attn_proc = MVAttnProcessor() + unet_lora_attn_procs[name] = ReferenceOnlyAttnProc( + default_attn_proc, enabled=name.endswith("attn1.processor"), name=name + ) + model.set_attn_processor(unet_lora_attn_procs) + state_dict = load_state_dict(model_file, variant=variant) + model._convert_deprecated_attention_blocks(state_dict) + + conv_in_weight = state_dict['conv_in.weight'] + conv_out_weight = state_dict['conv_out.weight'] + model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model_2d( + model, + state_dict, + model_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=True, + ) + if any([key == 'conv_in.weight' for key, _, _ in mismatched_keys]): + # initialize from the original SD structure + model.conv_in.weight.data[:,:4] = conv_in_weight + + # whether to place all zero to new layers? + if zero_init_conv_in: + model.conv_in.weight.data[:,4:] = 0. + + if any([key == 'conv_out.weight' for key, _, _ in mismatched_keys]): + # initialize from the original SD structure + model.conv_out.weight.data[:,:4] = conv_out_weight + if out_channels == 8: # copy for the last 4 channels + model.conv_out.weight.data[:, 4:] = conv_out_weight + + if zero_init_camera_projection: + for p in model.class_embedding.parameters(): + torch.nn.init.zeros_(p) + + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + "error_msgs": error_msgs, + } + + if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}." + ) + elif torch_dtype is not None: + model = model.to(torch_dtype) + + model.register_to_config(_name_or_path=pretrained_model_name_or_path) + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + if output_loading_info: + return model, loading_info + + return model + + @classmethod + def _load_pretrained_model_2d( + cls, + model, + state_dict, + resolved_archive_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=False, + ): + # Retrieve missing & unexpected_keys + model_state_dict = model.state_dict() + loaded_keys = list(state_dict.keys()) + + expected_keys = list(model_state_dict.keys()) + + original_loaded_keys = loaded_keys + + missing_keys = list(set(expected_keys) - set(loaded_keys)) + unexpected_keys = list(set(loaded_keys) - set(expected_keys)) + + # Make sure we are able to load base models as well as derived models (with heads) + model_to_load = model + + def _find_mismatched_keys( + state_dict, + model_state_dict, + loaded_keys, + ignore_mismatched_sizes, + ): + mismatched_keys = [] + if ignore_mismatched_sizes: + for checkpoint_key in loaded_keys: + model_key = checkpoint_key + + if ( + model_key in model_state_dict + and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape + ): + mismatched_keys.append( + (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape) + ) + del state_dict[checkpoint_key] + return mismatched_keys + + if state_dict is not None: + # Whole checkpoint + mismatched_keys = _find_mismatched_keys( + state_dict, + model_state_dict, + original_loaded_keys, + ignore_mismatched_sizes, + ) + error_msgs = _load_state_dict_into_model(model_to_load, state_dict) + + if len(error_msgs) > 0: + error_msg = "\n\t".join(error_msgs) + if "size mismatch" in error_msg: + error_msg += ( + "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." + ) + raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") + + if len(unexpected_keys) > 0: + logger.warning( + f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task" + " or with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly" + " identical (initializing a BertForSequenceClassification model from a" + " BertForSequenceClassification model)." + ) + else: + logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.info( + f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the" + f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions" + " without further training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be" + " able to use it for predictions and inference." + ) + + return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs + diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/pipelines/pipeline_tuneavideo.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/pipelines/pipeline_tuneavideo.py new file mode 100644 index 0000000..5934929 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/pipelines/pipeline_tuneavideo.py @@ -0,0 +1,589 @@ +# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py + +import tqdm + +import inspect +from typing import Callable, List, Optional, Union +from dataclasses import dataclass + +import numpy as np +import torch + +from diffusers.utils import is_accelerate_available +from packaging import version +from transformers import CLIPTextModel, CLIPTokenizer +import torchvision.transforms.functional as TF + +from diffusers.configuration_utils import FrozenDict +from diffusers.models import AutoencoderKL +from diffusers import DiffusionPipeline +from diffusers.schedulers import ( + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) +from diffusers.utils import deprecate, logging, BaseOutput + +from einops import rearrange + +from ..models.unet import UNet3DConditionModel +from torchvision.transforms import InterpolationMode + +import comfy.utils + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class TuneAVideoPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class TuneAVideoPipeline(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet3DConditionModel, + + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ref_unet = None, + feature_extractor=None, + image_encoder=None + ): + super().__init__() + self.ref_unet = ref_unet + self.feature_extractor = feature_extractor + self.image_encoder = image_encoder + + if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" + f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " + "to update the config accordingly as leaving `steps_offset` 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 `scheduler/scheduler_config.json`" + " file" + ) + deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["steps_offset"] = 1 + scheduler._internal_dict = FrozenDict(new_config) + + if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." + " `clip_sample` should be set to False in the configuration file. Please make sure to update the" + " config accordingly as not setting `clip_sample` in the config might lead 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 `scheduler/scheduler_config.json` file" + ) + deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["clip_sample"] = False + scheduler._internal_dict = FrozenDict(new_config) + + is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse( + version.parse(unet.config._diffusers_version).base_version + ) < version.parse("0.9.0.dev0") + is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 + if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: + deprecation_message = ( + "The configuration file of the unet has set the default `sample_size` to smaller than" + " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" + " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" + " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" + " in the config might lead 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 `unet/config.json` file" + ) + deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(unet.config) + new_config["sample_size"] = 64 + unet._internal_dict = FrozenDict(new_config) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + 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}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + + @property + def _execution_device(self): + if self.device != torch.device("meta") or 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 + + def _encode_image(self, image_pil, device, num_images_per_prompt, do_classifier_free_guidance, img_proj=None): + dtype = next(self.image_encoder.parameters()).dtype + + # image_pt = self.feature_extractor(images=image_pil, return_tensors="pt").pixel_values + # image_pt = image_pt.to(device=device, dtype=dtype) + # image_embeddings = self.image_encoder(image_pt).image_embeds + # image_embeddings = image_embeddings.unsqueeze(1) + + # # image encoding + clip_image_mean = torch.as_tensor(self.feature_extractor.image_mean)[:,None,None].to(device, dtype=torch.float32) + clip_image_std = torch.as_tensor(self.feature_extractor.image_std)[:,None,None].to(device, dtype=torch.float32) + imgs_in_proc = TF.resize(image_pil, (self.feature_extractor.crop_size['height'], self.feature_extractor.crop_size['width']), interpolation=InterpolationMode.BICUBIC) + # do the normalization in float32 to preserve precision + imgs_in_proc = ((imgs_in_proc.float() - clip_image_mean) / clip_image_std).to(dtype) + if img_proj is None: + # (B*Nv, 1, 768) + image_embeddings = self.image_encoder(imgs_in_proc).image_embeds.unsqueeze(1) + # duplicate image embeddings for each generation per prompt, using mps friendly method + # Note: repeat differently from official pipelines + # B1B2B3B4 -> B1B2B3B4B1B2B3B4 + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(num_images_per_prompt, 1, 1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + else: + if do_classifier_free_guidance: + negative_image_proc = torch.zeros_like(imgs_in_proc) + + # 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 + imgs_in_proc = torch.cat([negative_image_proc, imgs_in_proc]) + + image_embeds = self.image_encoder(imgs_in_proc, output_hidden_states=True).hidden_states[-2] + image_embeddings = img_proj(image_embeds) + + # image_embeddings_unet = rearrange(image_embeddings_unet, 'B Nv d c -> (B Nv) d c') + + # image_pt = torch.stack([TF.to_tensor(img) for img in image_pil], dim=0).to(device) + # image_pil = image_pil * 2.0 - 1.0 + image_latents = self.vae.encode(image_pil* 2.0 - 1.0).latent_dist.mode() * self.vae.config.scaling_factor + + # Note: repeat differently from official pipelines + # B1B2B3B4 -> B1B2B3B4B1B2B3B4 + image_latents = image_latents.repeat(num_images_per_prompt, 1, 1, 1) + + # if do_classifier_free_guidance: + # image_latents = torch.cat([torch.zeros_like(image_latents), image_latents]) + + return image_embeddings, image_latents + + def _encode_prompt(self, prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt): + batch_size = len(prompt) if isinstance(prompt, list) else 1 + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(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 = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1) + text_embeddings = text_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1) + uncond_embeddings = uncond_embeddings.view(batch_size * num_videos_per_prompt, seq_len, -1) + + # 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 + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + video_length = latents.shape[2] + latents = 1 / 0.18215 * latents + latents = rearrange(latents, "b c f h w -> (b f) c h w") + video = self.vae.decode(latents).sample + video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length) + video = (video / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + video = video.cpu().float().numpy() + return video + + 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, height, width, callback_steps): + if 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 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)}." + ) + + def prepare_latents(self, batch_size, num_channels_latents, video_length, height, width, dtype, device, generator, latents=None): + shape = (batch_size, num_channels_latents, video_length, 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: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + shape = (1,) + shape[1:] + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + 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 + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + image: Union[str, List[str]], + video_length: Optional[int], + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "tensor", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + camera_matrixs = None, + class_labels = None, + prompt_ids = None, + unet_condition_type = None, + pose_guider = None, + pose_image = None, + img_proj=None, + use_noise=True, + use_shifted_noise=False, + rescale = 0.7, + **kwargs, + ): + # 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 + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + if isinstance(image, list): + batch_size = len(image) + else: + batch_size = image.shape[0] + # assert batch_size >= video_length and batch_size % video_length == 0 + # Define call parameters + batch_size = 1 if isinstance(prompt, str) else len(prompt) + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input image + # if isinstance(image, list): + # image_pil = image + # elif isinstance(image, torch.Tensor): + # image_pil = [TF.to_pil_image(image[i]) for i in range(image.shape[0])] + # encode input reference image + image_embeddings, image_latents = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance, img_proj=img_proj) #torch.Size([64, 1, 768]) torch.Size([64, 4, 32, 32]) + image_latents = rearrange(image_latents, "(b f) c h w -> b c f h w", f=1) #torch.Size([64, 4, 1, 32, 32]) + + # Encode input prompt_id + # encoder_hidden_states = self.text_encoder(prompt_ids)[0] #torch.Size([32, 77, 768]) + + # Encode input prompt + text_embeddings = self._encode_prompt( #torch.Size([64, 77, 768]) + prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt + ) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # Prepare latent variables + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( #torch.Size([32, 4, 4, 32, 32]) + batch_size * num_videos_per_prompt, + num_channels_latents, + video_length, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + latents_dtype = latents.dtype + # import ipdb + # ipdb.set_trace() + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + # prepare camera_matrix + if camera_matrixs is not None: + camera_matrixs = torch.cat([camera_matrixs] * 2) if do_classifier_free_guidance else camera_matrixs #(64, 4, 12) + # Denoising loop + comfy_pbar = comfy.utils.ProgressBar(num_inference_steps) + + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + if pose_guider is not None: + if len(pose_image.shape) == 5: + pose_embeds = pose_guider(rearrange(pose_image, "b f c h w -> (b f) c h w")) + pose_embeds = rearrange(pose_embeds, "(b f) c h w-> b c f h w ", f=video_length) + else: + pose_embeds = pose_guider(pose_image).unsqueeze(0) + pose_embeds = torch.cat([pose_embeds]*2, dim=0) + # import ipdb + # ipdb.set_trace() + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(tqdm.tqdm(timesteps)): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + if pose_guider is not None: + latent_model_input = latent_model_input + pose_embeds + + noise_cond = torch.randn_like(image_latents) + if use_noise: + cond_latents = self.scheduler.add_noise(image_latents, noise_cond, t) + else: + cond_latents = image_latents + cond_latent_model_input = torch.cat([cond_latents] * 2) if do_classifier_free_guidance else cond_latents + cond_latent_model_input = self.scheduler.scale_model_input(cond_latent_model_input, t) + + # predict the noise residual + # ref text condition + ref_dict = {} + if self.ref_unet is not None: + noise_pred_cond = self.ref_unet( + cond_latent_model_input, #torch.Size([64, 4, 1, 32, 32]) + t, #torch.Size([32]) + encoder_hidden_states=text_embeddings.to(torch.float32), #torch.Size([64, 77, 768]) + cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict) + ).sample.to(dtype=latents_dtype) + + # if torch.isnan(noise_pred_cond).any(): + # ipdb.set_trace() + # Predict the noise residual and compute loss + # model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, camera_matrixs).sample + # unet + #text condition for unet + text_embeddings_unet = text_embeddings.unsqueeze(1).repeat(1,latents.shape[2],1,1) + text_embeddings_unet = rearrange(text_embeddings_unet, 'B Nv d c -> (B Nv) d c') + #image condition for unet + image_embeddings_unet = image_embeddings.unsqueeze(1).repeat(1,latents.shape[2],1, 1) + image_embeddings_unet = rearrange(image_embeddings_unet, 'B Nv d c -> (B Nv) d c') + + if unet_condition_type == 'text': + encoder_hidden_states_unet_cond = text_embeddings_unet + elif unet_condition_type == 'image': + encoder_hidden_states_unet_cond = image_embeddings_unet + else: + raise('need unet_condition_type') + + if self.ref_unet is not None: + noise_pred = self.unet( + latent_model_input.to(torch.float32), #torch.Size([64, 4, 4, 32, 32]) + t, + encoder_hidden_states=encoder_hidden_states_unet_cond.to(torch.float32), + camera_matrixs=camera_matrixs.to(torch.float32), #torch.Size([64, 4, 12]) + cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance) + # cross_attention_kwargs=dict(mode="n", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance) + ).sample.to(dtype=latents_dtype) + else: + noise_pred = self.unet( + latent_model_input.to(torch.float32), #torch.Size([64, 4, 4, 32, 32]) + t, + encoder_hidden_states=encoder_hidden_states_unet_cond.to(torch.float32), + camera_matrixs=camera_matrixs.to(torch.float32), #torch.Size([64, 4, 12]) + # cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance) + cross_attention_kwargs=dict(mode="n", ref_dict=ref_dict, is_cfg_guidance=do_classifier_free_guidance) + ).sample.to(dtype=latents_dtype) + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + if use_shifted_noise: + # Apply regular classifier-free guidance. + cfg = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + # Calculate standard deviations. + std_pos = noise_pred_text.std([1,2,3], keepdim=True) + std_cfg = cfg.std([1,2,3], keepdim=True) + # Apply guidance rescale with fused operations. + factor = std_pos / std_cfg + factor = rescale * factor + (1 - rescale) + noise_pred = cfg * factor + else: + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + # noise_pred_uncond_, noise_pred_text_ = noise_pred_cond.chunk(2) + # noise_pred_cond = noise_pred_uncond_ + guidance_scale * (noise_pred_text_ - noise_pred_uncond_) + + # compute the previous noisy sample x_t -> x_t-1 + noise_pred = rearrange(noise_pred, "(b f) c h w -> b c f h w", f=video_length) + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + # noise_pred_cond = rearrange(noise_pred_cond, "(b f) c h w -> b c f h w", f=1) + # cond_latents = self.scheduler.step(noise_pred_cond, t, cond_latents, **extra_step_kwargs).prev_sample + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + comfy_pbar.update_absolute(i + 1) + + # Post-processing + video = self.decode_latents(latents) + + # Convert to tensor + if output_type == "tensor": + video = torch.from_numpy(video) + + if not return_dict: + return video + + return TuneAVideoPipelineOutput(videos=video) diff --git a/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/util.py b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/util.py new file mode 100644 index 0000000..3c4daa3 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_2D/tuneavideo/util.py @@ -0,0 +1,128 @@ +import os +import imageio +import numpy as np +from typing import Union +import cv2 +import torch +import torchvision + +from tqdm import tqdm +from einops import rearrange + +def shifted_noise(betas, image_d=512, noise_d=256, shifted_noise=True): + alphas = 1 - betas + alphas_bar = torch.cumprod(alphas, dim=0) + d = (image_d / noise_d) ** 2 + if shifted_noise: + alphas_bar = alphas_bar / (d - (d - 1) * alphas_bar) + alphas_bar_sqrt = torch.sqrt(alphas_bar) + alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() + alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() + # Shift so last timestep is zero. + alphas_bar_sqrt -= alphas_bar_sqrt_T + # Scale so first timestep is back to old value. + alphas_bar_sqrt *= alphas_bar_sqrt_0 / ( + alphas_bar_sqrt_0 - alphas_bar_sqrt_T) + + # Convert alphas_bar_sqrt to betas + alphas_bar = alphas_bar_sqrt ** 2 + alphas = alphas_bar[1:] / alphas_bar[:-1] + alphas = torch.cat([alphas_bar[0:1], alphas]) + betas = 1 - alphas + return betas + +def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=4, fps=8): + videos = rearrange(videos, "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + outputs.append(x) + + os.makedirs(os.path.dirname(path), exist_ok=True) + imageio.mimsave(path, outputs, duration=1000/fps) + +def save_imgs_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=4, fps=8): + videos = rearrange(videos, "b c t h w -> t b c h w") + for i, x in enumerate(videos): + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + os.makedirs(os.path.dirname(path), exist_ok=True) + cv2.imwrite(os.path.join(path, f'view_{i}.png'), x[:,:,::-1]) + +def imgs_grid(videos: torch.Tensor, rescale=False, n_rows=4, fps=8): + videos = rearrange(videos, "b c t h w -> t b c h w") + image_list = [] + for i, x in enumerate(videos): + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + # image_list.append(x[:,:,::-1]) + image_list.append(x) + return image_list + +# DDIM Inversion +@torch.no_grad() +def init_prompt(prompt, pipeline): + uncond_input = pipeline.tokenizer( + [""], padding="max_length", max_length=pipeline.tokenizer.model_max_length, + return_tensors="pt" + ) + uncond_embeddings = pipeline.text_encoder(uncond_input.input_ids.to(pipeline.device))[0] + text_input = pipeline.tokenizer( + [prompt], + padding="max_length", + max_length=pipeline.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_embeddings = pipeline.text_encoder(text_input.input_ids.to(pipeline.device))[0] + context = torch.cat([uncond_embeddings, text_embeddings]) + + return context + + +def next_step(model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, + sample: Union[torch.FloatTensor, np.ndarray], ddim_scheduler): + timestep, next_timestep = min( + timestep - ddim_scheduler.config.num_train_timesteps // ddim_scheduler.num_inference_steps, 999), timestep + alpha_prod_t = ddim_scheduler.alphas_cumprod[timestep] if timestep >= 0 else ddim_scheduler.final_alpha_cumprod + alpha_prod_t_next = ddim_scheduler.alphas_cumprod[next_timestep] + beta_prod_t = 1 - alpha_prod_t + next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5 + next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output + next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction + return next_sample + + +def get_noise_pred_single(latents, t, context, unet): + noise_pred = unet(latents, t, encoder_hidden_states=context)["sample"] + return noise_pred + + +@torch.no_grad() +def ddim_loop(pipeline, ddim_scheduler, latent, num_inv_steps, prompt): + context = init_prompt(prompt, pipeline) + uncond_embeddings, cond_embeddings = context.chunk(2) + all_latent = [latent] + latent = latent.clone().detach() + for i in tqdm(range(num_inv_steps)): + t = ddim_scheduler.timesteps[len(ddim_scheduler.timesteps) - i - 1] + noise_pred = get_noise_pred_single(latent.to(torch.float32), t, cond_embeddings.to(torch.float32), pipeline.unet) + latent = next_step(noise_pred, t, latent, ddim_scheduler) + all_latent.append(latent) + return all_latent + + +@torch.no_grad() +def ddim_inversion(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt=""): + ddim_latents = ddim_loop(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt) + return ddim_latents diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/character_gen_texture_refine.py b/Gen_3D_Modules/CharacterGen/Stage_3D/character_gen_texture_refine.py new file mode 100644 index 0000000..7d87c44 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/character_gen_texture_refine.py @@ -0,0 +1,185 @@ +import os +import numpy as np +import torch +import nvdiffrast.torch as dr +import json +import torch.nn.functional as F +from PIL import Image +import pymeshlab +import cv2 + +def back_to_texture(glctx, look_at, pos, tri, tex, uv, uv_idx, idx, vn): + rast_out, rast_out_db = dr.rasterize(glctx, pos, tri, resolution=[tex.shape[0],tex.shape[1]]) + gb_normal, _ = dr.interpolate(vn[None], rast_out, tri) + gb_normal = F.normalize(gb_normal, dim=-1) + if idx == 2 or idx == 0: + filter_camera = [torch.tensor([[[[1,0.,0.]]]]).cuda(), torch.tensor([[[[-1,0.,0.]]]]).cuda()] + else: + filter_camera = [torch.tensor([[[[0,-1.,0.]]]]).cuda(), torch.tensor([[[[0,1.,0.]]]]).cuda()] + nmasks = [] + for fc in filter_camera: + nmasks.append(((gb_normal * fc) > 0.75).int().sum(keepdim=True, dim=-1)) + gb_normal_mask = 1 - (nmasks[0] | nmasks[1]) + #Image.fromarray(np.clip(gb_normal_mask[0,...,0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"mask_normal_{idx}.png") + gb_mask = rast_out[...,3:4] > 0 + tri_list = torch.unique(rast_out[...,3:4].reshape(-1)) + tri_list = (tri_list[1:] - 1).to(torch.int32) + pos = pos[0] + + depth_map = rast_out[...,3:4].clone() + depth_map[depth_map > 0] = 1 + depth_map = depth_map.to(torch.float32) + dmax = (rast_out[...,2:3] * gb_mask).max() + uv = torch.cat([uv * 2 - 1, torch.zeros(uv.shape[0], 1).cuda(), torch.ones(uv.shape[0], 1).cuda()], dim=1).unsqueeze(0) + uv_idx = uv_idx[tri_list.to(torch.long)] + rast_uv, rast_uv_db = dr.rasterize(glctx, uv, uv_idx, resolution=(1024, 1024)) + pos_clip = torch.cat([pos[...,:2], pos[...,3:]], -1) + pos_2d, _ = dr.interpolate(pos_clip, rast_uv, tri[tri_list.to(torch.long)]) # pos (x, y, z, w) + pos_coord = (pos_2d[...,:2] / (pos_2d[...,2:3] + 1e-6) + 1) / 2. + texture_mask = (rast_uv[...,3:4] > 0).int() + color = dr.texture(tex[None, ...] * gb_normal_mask, pos_coord, filter_mode='linear') + color_mask = dr.texture(gb_normal_mask.to(torch.float32), pos_coord, filter_mode='linear') + color_mask[color_mask > 0.82] = 1 + color_mask[color_mask <= 0.82] = 0 + color_mask = color_mask.to(torch.int32) + #Image.fromarray(np.clip(color_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"depth_{idx}.png") + texture_mask = texture_mask * color_mask + #Image.fromarray(np.clip(color[0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"{idx}.png") + #Image.fromarray(np.clip(texture_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).convert("RGB").save(f"mask-{idx}.png") + return color, texture_mask, rast_uv + +def perspective(fovy=0.6913, aspect=1.0, n=0.1, f=1000.0, device=None): + y = np.tan(fovy / 2) + return torch.tensor([[1/(y*aspect), 0, 0, 0], + [ 0, 1/-y, 0, 0], + [ 0, 0, -(f+n)/(f-n), -(2*f*n)/(f-n)], + [ 0, 0, -1, 0]]).to(torch.float32).cuda() + +def rec_mvp(trans, h, w): + mv = trans + fov = 40. / 180. * np.pi + proj = perspective(fov, h / w, n=0.1, f=1000) + mvp = proj @ mv + return mvp + +def aggregate_texture(kd_map, textures, texture_masks, rast_uvs): + texture = torch.zeros_like(textures[0]) + texture_mask = torch.zeros_like(texture_masks[0]) + ctex = [] + for idx in range(len(textures)): + ctex.append(textures[idx] * texture_masks[idx] + 10 * (1 - texture_masks[idx])) + cat_textures = torch.stack(ctex, dim=-2) + dis_measure = (cat_textures - kd_map.unsqueeze(-2)).abs().sum(-1) + _, choose_idx = dis_measure.min(-1) + + choose_idx = choose_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 1, 1, 3) + final_texture_map = torch.gather(cat_textures, 3, choose_idx).squeeze(-2) + #cv2.imwrite("final_texture_map.png", cv2.cvtColor((final_texture_map[0].cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB)) + #cv2.imwrite("final_texture_mask.png", (texture_mask[0].cpu().numpy() * 255).astype(np.uint8)) + zero_mask = (final_texture_map.max(dim=-1, keepdim=True)[0] > 0.1) + close_mask = ((final_texture_map[0] - kd_map).abs().sum(dim=-1, keepdim=True) < 1.0).int() + for idx in range(len(textures)): + texture += textures[idx] * texture_masks[idx] + texture_mask |= texture_masks[idx] + texture_mask = texture_mask * zero_mask * close_mask[None] + optimize_mask = (texture_mask == 0).int() + + #import pdb; pdb.set_trace() + #mask = (texture_mask[0].cpu().numpy() * 255).astype(np.uint8) + #cv2.imwrite("mask.png", mask) + #kernel = np.ones((5,5), np.uint8) + #dilated = cv2.dilate(mask, kernel, iterations=1) + #cv2.imwrite("di_mask.png", dilated) + #texture_mask[0] = torch.from_numpy(dilated).unsqueeze(-1).to(torch.float32) / 255. + + final_texture_map = final_texture_map[0] * texture_mask[0] + Image.fromarray(np.rint(final_texture_map.cpu().numpy() * 255).astype(np.uint8)).save(f"final_texture.png") + + #cv2.imwrite("kd_map.png", cv2.cvtColor((kd_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB)) + #cv2.imwrite("texture_map.png", cv2.cvtColor((final_texture_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB)) + #result = cv2.seamlessClone((final_texture_map.cpu().numpy() * 255).astype(np.uint8), (kd_map.cpu().numpy() * 255).astype(np.uint8), mask, (mask.shape[1]//2, mask.shape[0]//2), cv2.NORMAL_CLONE) + #cv2.imwrite("result.png", cv2.cvtColor(result * 255, cv2.COLOR_BGR2RGB)) + + kd_map = kd_map * (1 - texture_mask[0]) + final_texture_map + return kd_map, optimize_mask + +def refine(save_path, front_image, back_image, left_image, right_image): + ms = pymeshlab.MeshSet() + mesh_path = f"{save_path}/model-00.obj" + ms.load_new_mesh(mesh_path) + ms.apply_coord_laplacian_smoothing(stepsmoothnum=10) + tl = open(mesh_path, "r").readlines() + tex_uv = [] + uv_idx = [] + for line in tl: + if line.startswith("vt"): + uvs = line.split(" ")[1:3] + tex_uv += [float(uvs[0]), 1.0-float(uvs[1])] + tex_uv = torch.from_numpy(np.array(tex_uv)).to(torch.float32).cuda().reshape(-1, 2) + m = ms.current_mesh() + v_matrix = m.vertex_matrix() + f_matrix = m.face_matrix() + vn = m.vertex_normal_matrix() + uv_idx = torch.arange(f_matrix.shape[0] * 3).reshape(-1, 3).to(torch.int32).cuda() + vn = torch.tensor(vn).contiguous().cuda().to(torch.float32) + + frames = [] + front_camera = torch.tensor([[ + 1,0,0,0, + 0,0,1,0, + 0,-1,0,-1.5, + 0,0,0,1, + ]]).to(torch.float32).reshape(4,4).cuda() + back_camera = torch.tensor([[ + 1,0,0,0, + 0,0,1,0, + 0,1,0,-1.5, + 0,0,0,1, + ]]).to(torch.float32).reshape(4,4).cuda() + right_camera = torch.tensor([[ + 0,-1,0,0, + 0,0,1,0, + 1,0,0,-1.5, + 0,0,0,1, + ]]).to(torch.float32).reshape(4,4).cuda() + left_camera = torch.tensor([[ + 0,1,0,0, + 0,0,1,0, + -1,0,0,-1.5, + 0,0,0,1, + ]]).to(torch.float32).reshape(4,4).cuda() + frames = [front_camera, left_camera, back_camera, right_camera] + + target_images = [] + for target_image in [front_image, left_image, back_image, right_image]: + target_images.append(torch.from_numpy(np.asarray(target_image.convert("RGB"))).to(torch.float32).cuda() / 255.) + + pos = torch.tensor(v_matrix, dtype=torch.float32).contiguous().cuda() + tri = torch.tensor(f_matrix, dtype=torch.int32).contiguous().cuda() + + kd_map = (torch.tensor(np.asarray(Image.open(f"{save_path}/texture_kd.jpg"))) / 255.).cuda() + translate_tensor = torch.zeros((1,1,3)).cuda() + pos = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()],-1).unsqueeze(0) + glctx = dr.RasterizeCudaContext() + target_texture = [] + target_mask = [] + rast_uvs = [] + with torch.no_grad(): + for idx, trans in enumerate(frames): + target_image = target_images[idx] + look_at = -torch.linalg.inv(trans)[:3,2] + mvp = rec_mvp(trans, h=target_images[0].shape[0], w=target_images[0].shape[1]) + trans_pos = pos.clone() + trans_pos[...,:3] += translate_tensor + view_pos = torch.matmul(mvp, trans_pos.unsqueeze(-1)).squeeze(-1) + texture, mask, rast_uv = back_to_texture(glctx, look_at, view_pos, tri, target_image, tex_uv, uv_idx, idx, vn) + target_texture.append(texture) + target_mask.append(mask) + rast_uvs.append(rast_uv) + kd_map, opt_mask = aggregate_texture(kd_map, target_texture, target_mask, rast_uvs) + opt_mask = opt_mask[0] + Image.fromarray((np.clip(kd_map.detach().cpu().numpy() * 255, 0, 255)).astype(np.uint8)).save(f"{save_path}/refined_texture_kd.jpg") + + #ms.save_current_mesh(f"{save_path}/model-00.obj") + with open(f"{save_path}/model-00.mtl", "w") as f: + f.write(f"newmtl default\nKa 0.0 0.0 0.0\nmap_Kd refined_texture_kd.jpg\nKs 0.0 0.0 0.0") \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/__init__.py new file mode 100644 index 0000000..69f5951 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/__init__.py @@ -0,0 +1,29 @@ +import importlib + + +def find(cls_string): + module_string = ".".join(cls_string.split(".")[:-1]) + cls_name = cls_string.split(".")[-1] + module = importlib.import_module(module_string, package=None) + cls = getattr(module, cls_name) + return cls + + +### grammar sugar for logging utilities ### +import logging + +logger = logging.getLogger("pytorch_lightning") + +from pytorch_lightning.utilities.rank_zero import ( + rank_zero_debug, + rank_zero_info, + rank_zero_only, +) + +debug = rank_zero_debug +info = rank_zero_info + + +@rank_zero_only +def warn(*args, **kwargs): + logger.warn(*args, **kwargs) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/base.py new file mode 100644 index 0000000..0aadefd --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/base.py @@ -0,0 +1,18 @@ +from dataclasses import dataclass + +from ...utils.base import BaseModule +from ...utils.typing import * + + +class BaseBackground(BaseModule): + @dataclass + class Config(BaseModule.Config): + pass + + cfg: Config + + def configure(self): + pass + + def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]: + raise NotImplementedError diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/solid_color_background.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/solid_color_background.py new file mode 100644 index 0000000..304eb62 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/background/solid_color_background.py @@ -0,0 +1,56 @@ +import random +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from .base import BaseBackground +from ...utils.typing import * + + +class SolidColorBackground(BaseBackground): + @dataclass + class Config(BaseBackground.Config): + n_output_dims: int = 3 + color: Tuple = (1.0, 1.0, 1.0) + learned: bool = False + random_aug: bool = False + random_aug_prob: float = 0.5 + + cfg: Config + + def configure(self) -> None: + self.env_color: Float[Tensor, "Nc"] + if self.cfg.learned: + self.env_color = nn.Parameter( + torch.as_tensor(self.cfg.color, dtype=torch.float32) + ) + else: + self.register_buffer( + "env_color", torch.as_tensor(self.cfg.color, dtype=torch.float32) + ) + + def forward( + self, + dirs: Float[Tensor, "B H W Nc"], + color_spec: Optional[Float[Tensor, "Nc"]] = None, + ) -> Float[Tensor, "B H W Nc"]: + color = torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(dirs) * ( + color_spec if color_spec is not None else self.env_color + ) + if ( + self.training + and self.cfg.random_aug + and random.random() < self.cfg.random_aug_prob + ): + # use random background color with probability random_aug_prob + # color = color * 0 + ( + # torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(dirs) * + # torch.rand(self.cfg.n_output_dims).to(dirs) + # ) + color = color * 0 + ( # prevent checking for unused parameters in DDP + torch.rand(dirs.shape[0], 1, 1, self.cfg.n_output_dims) + .to(dirs) + .expand(*dirs.shape[:-1], -1) + ) + return color diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/camera.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/camera.py new file mode 100644 index 0000000..62461c3 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/camera.py @@ -0,0 +1,33 @@ +from dataclasses import dataclass, field + +import torch +import torch.nn as nn + +from ..utils.base import BaseModule +from ..utils.typing import * + + +class LinearCameraEmbedder(BaseModule): + @dataclass + class Config(BaseModule.Config): + in_channels: int = 0 + out_channels: int = 0 + conditions: List[str] = field(default_factory=list) + + cfg: Config + + def configure(self) -> None: + super().configure() + self.linear = nn.Linear(self.cfg.in_channels, self.cfg.out_channels) + + def forward(self, **kwargs): + cond_tensors = [] + for cond_name in self.cfg.conditions: + assert cond_name in kwargs + cond = kwargs[cond_name] + # cond in shape (B, Nv, ...) + cond_tensors.append(cond.view(*cond.shape[:2], -1)) + cond_tensor = torch.cat(cond_tensors, dim=-1) + assert cond_tensor.shape[-1] == self.cfg.in_channels + embedding = self.linear(cond_tensor) + return embedding diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/base.py new file mode 100644 index 0000000..9b65247 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/base.py @@ -0,0 +1,32 @@ +from dataclasses import dataclass + +from ..renderers.base import BaseRenderer +from ...utils.base import BaseObject +from ...utils.typing import * + + +@dataclass +class ExporterOutput: + save_name: str + save_type: str + params: Dict[str, Any] + + +class Exporter(BaseObject): + @dataclass + class Config(BaseObject.Config): + save_video: bool = False + + cfg: Config + + def configure(self, renderer: BaseRenderer) -> None: + self.renderer = renderer + + def __call__(self, *args, **kwargs) -> List[ExporterOutput]: + raise NotImplementedError + + +class DummyExporter(Exporter): + def __call__(self, *args, **kwargs) -> List[ExporterOutput]: + # DummyExporter does not export anything + return [] diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/mesh_exporter.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/mesh_exporter.py new file mode 100644 index 0000000..0352aef --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/exporters/mesh_exporter.py @@ -0,0 +1,262 @@ +from dataclasses import dataclass, field +import os + +import cv2 +import numpy as np +import torch + +from CharacterGen.Stage_3D import lrm +from ..renderers.base import BaseRenderer +from .base import Exporter, ExporterOutput +from ..mesh import Mesh +from ...utils.rasterize import NVDiffRasterizerContext +from ...utils.typing import * +from ...utils.misc import time_recorder as tr, time_recorder_enabled + + +def uv_padding_cpu(image, hole_mask, padding): + uv_padding_size = padding + inpaint_image = ( + cv2.inpaint( + (image.detach().cpu().numpy() * 255).astype(np.uint8), + (hole_mask.detach().cpu().numpy() * 255).astype(np.uint8), + uv_padding_size, + cv2.INPAINT_TELEA, + ) + / 255.0 + ) + return torch.from_numpy(inpaint_image).to(image) + + +def uv_padding_cvc(image, hole_mask, padding): + import cvcuda + + torch_to_cvc = lambda x, layout: cvcuda.as_tensor(x, layout) + cvc_to_torch = lambda x: torch.as_tensor(x.cuda()) + + uv_padding_size = padding + image_cvc = torch_to_cvc((image.detach() * 255).to(torch.uint8), "HWC") + hole_mask_cvc = torch_to_cvc((hole_mask.detach() * 255).to(torch.uint8), "HW") + inpaint_image = cvcuda.inpaint(image_cvc, hole_mask_cvc, uv_padding_size) + inpaint_image = cvc_to_torch(inpaint_image) / 255.0 + return inpaint_image.to(image) + + +def uv_padding(image, hole_mask, padding): + try: + inpaint_image = uv_padding_cvc(image, hole_mask, padding) + except: + lrm.info(f"CVCUDA not available, fallback to CPU UV padding.") + inpaint_image = uv_padding_cpu(image, hole_mask, padding) + return inpaint_image + + +class MeshExporter(Exporter): + @dataclass + class Config(Exporter.Config): + fmt: str = "obj" # in ['obj', 'glb'] + visual: str = "uv" # in ['uv', 'vertex'] + save_name: str = "model" + save_normal: bool = False + save_uv: bool = True + save_texture: bool = True + texture_size: int = 1024 + texture_format: str = "jpg" + uv_unwrap_method: str = "xatlas" + xatlas_chart_options: dict = field(default_factory=dict) + xatlas_pack_options: dict = field(default_factory=dict) + smartuv_options: dict = field(default_factory=dict) + uv_padding_size: int = 2 + subdivide: bool = False + post_process: bool = False + post_process_options: dict = field(default_factory=dict) + context_type: str = "gl" + output_path: str = "outputs" + + cfg: Config + + def configure(self, renderer: BaseRenderer) -> None: + super().configure(renderer) + self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device) + if self.cfg.fmt == "obj-mtl": + lrm.warn( + f"fmt=obj-mtl is deprecated, please us fmt=obj and visual=uv instead." + ) + self.cfg.fmt = "obj" + self.cfg.visual = "uv" + + if self.cfg.fmt == "glb": + assert self.cfg.visual in [ + "vertex", + "uv-blender", + ], "GLB format only supports visual=vertex and visual=uv-blender!" + + def get_geometry(self, scene_code: torch.Tensor) -> Mesh: + tr.start("Surface extraction") + mesh: Mesh = self.renderer.isosurface(scene_code) + tr.end("Surface extraction") + return mesh + + def get_texture_maps( + self, scene_code: torch.Tensor, mesh: Mesh + ) -> Dict[str, torch.Tensor]: + assert mesh.has_uv + # clip space transform + uv_clip = mesh.v_tex * 2.0 - 1.0 + # pad to four component coordinate + uv_clip4 = torch.cat( + ( + uv_clip, + torch.zeros_like(uv_clip[..., 0:1]), + torch.ones_like(uv_clip[..., 0:1]), + ), + dim=-1, + ) + # rasterize + rast, _ = self.ctx.rasterize_one( + uv_clip4, + mesh.t_tex_idx, + (self.cfg.texture_size, self.cfg.texture_size), + ) + + hole_mask = ~(rast[:, :, 3] > 0) + + # Interpolate world space position + gb_pos, _ = self.ctx.interpolate_one( + mesh.v_pos, rast[None, ...], mesh.t_pos_idx + ) + gb_pos = gb_pos[0] + + # Sample out textures from MLP + tr.start("Query color") + geo_out = self.renderer.query(scene_code, points=gb_pos) + tr.end("Query color") + mat_out = self.renderer.material.export(points=gb_pos, **geo_out) + + textures = {} + tr.start("UV padding") + if "albedo" in mat_out: + textures["map_Kd"] = uv_padding( + mat_out["albedo"], hole_mask, self.cfg.uv_padding_size + ) + else: + lrm.warn( + "save_texture is True but no albedo texture found, using default white texture" + ) + if "metallic" in mat_out: + textures["map_Pm"] = uv_padding( + mat_out["metallic"], hole_mask, self.cfg.uv_padding_size + ) + if "roughness" in mat_out: + textures["map_Pr"] = uv_padding( + mat_out["roughness"], hole_mask, self.cfg.uv_padding_size + ) + if "bump" in mat_out: + textures["map_Bump"] = uv_padding( + mat_out["bump"], hole_mask, self.cfg.uv_padding_size + ) + tr.end("UV padding") + return textures + + def __call__(self, names, scene_codes) -> List[ExporterOutput]: + outputs = [] + for name, scene_code in zip(names, scene_codes): + mesh = self.get_geometry(scene_code) + if self.cfg.post_process: + tr.start("Mesh post-processing") + mesh = mesh.post_process(self.cfg.post_process_options) + tr.end("Mesh post-processing") + if self.cfg.visual == "uv": + output = self.export_model_with_mtl( + name, self.cfg.fmt, scene_code, mesh + ) + elif self.cfg.visual == "vertex": + output = self.export_model(name, self.cfg.fmt, scene_code, mesh) + elif self.cfg.visual == "uv-blender": + output = self.export_model_blender(name, self.cfg.fmt, scene_code, mesh) + else: + raise ValueError(f"Unsupported visual format: {self.cfg.visual}") + outputs.append(output) + return outputs + + def export_model_with_mtl( + self, name: str, fmt: str, scene_code: torch.Tensor, mesh: Mesh + ) -> ExporterOutput: + params = { + "mesh": mesh, + "save_mat": True, + "save_normal": self.cfg.save_normal, + "save_uv": self.cfg.save_uv, + "save_vertex_color": False, + "map_Kd": None, # Base Color + "map_Ks": None, # Specular + "map_Bump": None, # Normal + # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering + "map_Pm": None, # Metallic + "map_Pr": None, # Roughness + "map_format": self.cfg.texture_format, + } + + if self.cfg.save_uv: + mesh.unwrap_uv( + self.cfg.uv_unwrap_method, + self.cfg.xatlas_chart_options, + self.cfg.xatlas_pack_options, + self.cfg.smartuv_options, + ) + + if self.cfg.save_texture: + lrm.info("Exporting textures ...") + assert self.cfg.save_uv, "save_uv must be True when save_texture is True" + + with time_recorder_enabled(): + textures = self.get_texture_maps(scene_code, mesh) + params.update(textures) + os.makedirs(self.cfg.output_path, exist_ok=True) + np.savez(f"{self.cfg.output_path}/tex_info.npz", v_tex=mesh.v_tex.cpu().numpy(), t_tex_idx=mesh.t_tex_idx.cpu().numpy()) + return ExporterOutput( + save_name=f"{self.cfg.save_name}-{name}.{fmt}", save_type=fmt, params=params + ) + + def export_model( + self, name: str, fmt: str, scene_code, mesh: Mesh + ) -> ExporterOutput: + params = { + "mesh": mesh, + "save_mat": False, + "save_normal": self.cfg.save_normal, + "save_uv": self.cfg.save_uv, + "save_vertex_color": False, + "map_Kd": None, # Base Color + "map_Ks": None, # Specular + "map_Bump": None, # Normal + # ref: https://en.wikipedia.org/wiki/Wavefront_.obj_file#Physically-based_Rendering + "map_Pm": None, # Metallic + "map_Pr": None, # Roughness + "map_format": self.cfg.texture_format, + } + + if self.cfg.save_uv: + mesh.unwrap_uv( + self.cfg.uv_unwrap_method, + self.cfg.xatlas_chart_options, + self.cfg.xatlas_pack_options, + self.cfg.smartuv_options, + ) + + if self.cfg.save_texture: + lrm.info("Exporting textures ...") + geo_out = self.renderer.query(scene_code, points=mesh.v_pos) + mat_out = self.renderer.material.export(points=mesh.v_pos, **geo_out) + + if "albedo" in mat_out: + mesh.set_vertex_color(mat_out["albedo"]) + params["save_vertex_color"] = True + else: + lrm.warn( + "save_texture is True but no albedo texture found, not saving vertex color" + ) + + return ExporterOutput( + save_name=f"{self.cfg.save_name}-{name}.{fmt}", save_type=fmt, params=params + ) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/isosurface.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/isosurface.py new file mode 100644 index 0000000..c5dc45e --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/isosurface.py @@ -0,0 +1,271 @@ +import numpy as np +import torch +import torch.nn as nn + +from CharacterGen.Stage_3D import lrm +from ..models.mesh import Mesh +from ..utils.typing import * +from ..utils.ops import scale_tensor + + +class IsosurfaceHelper(nn.Module): + points_range: Tuple[float, float] = (0, 1) + + @property + def grid_vertices(self) -> Float[Tensor, "N 3"]: + raise NotImplementedError + + +class MarchingCubeCPUHelper(IsosurfaceHelper): + def __init__(self, resolution: int) -> None: + super().__init__() + self.resolution = resolution + import mcubes + + self.mc_func: Callable = mcubes.marching_cubes + self._grid_vertices: Optional[Float[Tensor, "N3 3"]] = None + self._dummy: Float[Tensor, "..."] + self.register_buffer( + "_dummy", torch.zeros(0, dtype=torch.float32), persistent=False + ) + + @property + def grid_vertices(self) -> Float[Tensor, "N3 3"]: + if self._grid_vertices is None: + # keep the vertices on CPU so that we can support very large resolution + x, y, z = ( + torch.linspace(*self.points_range, self.resolution), + torch.linspace(*self.points_range, self.resolution), + torch.linspace(*self.points_range, self.resolution), + ) + x, y, z = torch.meshgrid(x, y, z, indexing="ij") + verts = torch.cat( + [x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1 + ).reshape(-1, 3) + self._grid_vertices = verts + return self._grid_vertices + + def forward( + self, + level: Float[Tensor, "N3 1"], + deformation: Optional[Float[Tensor, "N3 3"]] = None, + ) -> Mesh: + if deformation is not None: + lrm.warn( + f"{self.__class__.__name__} does not support deformation. Ignoring." + ) + level = -level.view(self.resolution, self.resolution, self.resolution) + print(level.shape, level.min(), level.max()) + v_pos, t_pos_idx = self.mc_func( + level.detach().cpu().numpy(), 0.0 + ) # transform to numpy + # test + v_pos, t_pos_idx = ( + torch.from_numpy(v_pos).float().to(self._dummy.device), + torch.from_numpy(t_pos_idx.astype(np.int64)).long().to(self._dummy.device), + ) # transform back to torch tensor on CUDA + v_pos = v_pos / (self.resolution - 1.0) + return Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx) + + +def get_center_boundary_index(verts): + # Assuming the verts are in range [-1.0, 1.0] + length_ = torch.linalg.norm(verts**2, ord=2, dim=-1, keepdim=False) + center_idx = torch.argmin(length_) + # center_idx = torch.where(length_ < 0.1)[0] + boundary_neg = verts == verts.max() + boundary_pos = verts == verts.min() + boundary = torch.bitwise_or(boundary_pos, boundary_neg) + boundary = torch.sum(boundary.float(), dim=-1) + boundary_idx = torch.nonzero(boundary) + return center_idx.unsqueeze(0), boundary_idx.squeeze(dim=-1) + + +class MarchingTetrahedraHelper(IsosurfaceHelper): + def __init__(self, resolution: int, tets_path: str): + super().__init__() + self.resolution = resolution + self.tets_path = tets_path + + self.triangle_table: Float[Tensor, "..."] + self.register_buffer( + "triangle_table", + torch.as_tensor( + [ + [-1, -1, -1, -1, -1, -1], + [1, 0, 2, -1, -1, -1], + [4, 0, 3, -1, -1, -1], + [1, 4, 2, 1, 3, 4], + [3, 1, 5, -1, -1, -1], + [2, 3, 0, 2, 5, 3], + [1, 4, 0, 1, 5, 4], + [4, 2, 5, -1, -1, -1], + [4, 5, 2, -1, -1, -1], + [4, 1, 0, 4, 5, 1], + [3, 2, 0, 3, 5, 2], + [1, 3, 5, -1, -1, -1], + [4, 1, 2, 4, 3, 1], + [3, 0, 4, -1, -1, -1], + [2, 0, 1, -1, -1, -1], + [-1, -1, -1, -1, -1, -1], + ], + dtype=torch.long, + ), + persistent=False, + ) + self.num_triangles_table: Integer[Tensor, "..."] + self.register_buffer( + "num_triangles_table", + torch.as_tensor( + [0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long + ), + persistent=False, + ) + self.base_tet_edges: Integer[Tensor, "..."] + self.register_buffer( + "base_tet_edges", + torch.as_tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long), + persistent=False, + ) + + tets = np.load(self.tets_path) + self._grid_vertices: Float[Tensor, "..."] + self.register_buffer( + "_grid_vertices", + torch.from_numpy(tets["vertices"]).float(), + persistent=False, + ) + self.indices: Integer[Tensor, "..."] + self.register_buffer( + "indices", torch.from_numpy(tets["indices"]).long(), persistent=False + ) + + self._all_edges: Optional[Integer[Tensor, "Ne 2"]] = None + self.center_indices, self.boundary_indices = get_center_boundary_index( + scale_tensor(self.grid_vertices, self.points_range, (-1.0, 1.0)) + ) + + def normalize_grid_deformation( + self, grid_vertex_offsets: Float[Tensor, "Nv 3"] + ) -> Float[Tensor, "Nv 3"]: + return ( + (self.points_range[1] - self.points_range[0]) + / (self.resolution) # half tet size is approximately 1 / self.resolution + * torch.tanh(grid_vertex_offsets) + ) # FIXME: hard-coded activation + + @property + def grid_vertices(self) -> Float[Tensor, "Nv 3"]: + return self._grid_vertices + + @property + def all_edges(self) -> Integer[Tensor, "Ne 2"]: + if self._all_edges is None: + # compute edges on GPU, or it would be VERY SLOW (basically due to the unique operation) + edges = torch.tensor( + [0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], + dtype=torch.long, + device=self.indices.device, + ) + _all_edges = self.indices[:, edges].reshape(-1, 2) + _all_edges_sorted = torch.sort(_all_edges, dim=1)[0] + _all_edges = torch.unique(_all_edges_sorted, dim=0) + self._all_edges = _all_edges + return self._all_edges + + def sort_edges(self, edges_ex2): + with torch.no_grad(): + order = (edges_ex2[:, 0] > edges_ex2[:, 1]).long() + order = order.unsqueeze(dim=1) + + a = torch.gather(input=edges_ex2, index=order, dim=1) + b = torch.gather(input=edges_ex2, index=1 - order, dim=1) + + return torch.stack([a, b], -1) + + def _forward(self, pos_nx3, sdf_n, tet_fx4): + with torch.no_grad(): + occ_n = sdf_n > 0 + occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4) + occ_sum = torch.sum(occ_fx4, -1) + valid_tets = (occ_sum > 0) & (occ_sum < 4) + occ_sum = occ_sum[valid_tets] + + # find all vertices + all_edges = tet_fx4[valid_tets][:, self.base_tet_edges].reshape(-1, 2) + all_edges = self.sort_edges(all_edges) + unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True) + + unique_edges = unique_edges.long() + mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1 + mapping = ( + torch.ones( + (unique_edges.shape[0]), dtype=torch.long, device=pos_nx3.device + ) + * -1 + ) + mapping[mask_edges] = torch.arange( + mask_edges.sum(), dtype=torch.long, device=pos_nx3.device + ) + idx_map = mapping[idx_map] # map edges to verts + + interp_v = unique_edges[mask_edges] + edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3) + edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1) + edges_to_interp_sdf[:, -1] *= -1 + + denominator = edges_to_interp_sdf.sum(1, keepdim=True) + + edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator + verts = (edges_to_interp * edges_to_interp_sdf).sum(1) + + idx_map = idx_map.reshape(-1, 6) + + v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=pos_nx3.device)) + tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1) + num_triangles = self.num_triangles_table[tetindex] + + # Generate triangle indices + faces = torch.cat( + ( + torch.gather( + input=idx_map[num_triangles == 1], + dim=1, + index=self.triangle_table[tetindex[num_triangles == 1]][:, :3], + ).reshape(-1, 3), + torch.gather( + input=idx_map[num_triangles == 2], + dim=1, + index=self.triangle_table[tetindex[num_triangles == 2]][:, :6], + ).reshape(-1, 3), + ), + dim=0, + ) + + return verts, faces + + def forward( + self, + level: Float[Tensor, "N3 1"], + deformation: Optional[Float[Tensor, "N3 3"]] = None, + ) -> Mesh: + if deformation is not None: + grid_vertices = self.grid_vertices + self.normalize_grid_deformation( + deformation + ) + else: + grid_vertices = self.grid_vertices + + v_pos, t_pos_idx = self._forward(grid_vertices, level, self.indices) + + mesh = Mesh( + v_pos=v_pos, + t_pos_idx=t_pos_idx, + # extras + grid_vertices=grid_vertices, + tet_edges=self.all_edges, + grid_level=level, + grid_deformation=deformation, + ) + + return mesh diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/lpips.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/lpips.py new file mode 100644 index 0000000..de43bc6 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/lpips.py @@ -0,0 +1,17 @@ +from ..utils.ops import scale_tensor +from ..utils.misc import get_device + + +class LPIPS: + def __init__(self): + import lpips + self.model = lpips.LPIPS(net="vgg").to(get_device()) + self.model.eval() + for params in self.model.parameters(): + params.requires_grad = False + self.model_input_range = (-1, 1) + + def __call__(self, x1, x2, return_layers=False, input_range=(0, 1)): + x1 = scale_tensor(x1, input_range, self.model_input_range) + x2 = scale_tensor(x2, input_range, self.model_input_range) + return self.model(x1, x2, retPerLayer=return_layers, normalize=False) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/base.py new file mode 100644 index 0000000..ce5b395 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/base.py @@ -0,0 +1,23 @@ +from dataclasses import dataclass + +from ...utils.base import BaseModule +from ...utils.typing import * + + +class BaseMaterial(BaseModule): + @dataclass + class Config(BaseModule.Config): + pass + + cfg: Config + requires_normal: bool = False + requires_tangent: bool = False + + def configure(self): + pass + + def forward(self, *args, **kwargs) -> Float[Tensor, "*B 3"]: + raise NotImplementedError + + def export(self, *args, **kwargs) -> Dict[str, Any]: + return {} diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/no_material.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/no_material.py new file mode 100644 index 0000000..61ba251 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/materials/no_material.py @@ -0,0 +1,60 @@ +import random +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from CharacterGen.Stage_3D import lrm +from .base import BaseMaterial +from ..networks import get_encoding, get_mlp +from ...utils.ops import dot, get_activation +from ...utils.typing import * + + +class NoMaterial(BaseMaterial): + @dataclass + class Config(BaseMaterial.Config): + n_output_dims: int = 3 + color_activation: str = "sigmoid" + input_feature_dims: Optional[int] = None + mlp_network_config: Optional[dict] = None + requires_normal: bool = False + + cfg: Config + + def configure(self) -> None: + self.use_network = False + if ( + self.cfg.input_feature_dims is not None + and self.cfg.mlp_network_config is not None + ): + self.network = get_mlp( + self.cfg.input_feature_dims, + self.cfg.n_output_dims, + self.cfg.mlp_network_config, + ) + self.use_network = True + self.requires_normal = self.cfg.requires_normal + + def forward( + self, features: Float[Tensor, "B ... Nf"], **kwargs + ) -> Float[Tensor, "B ... Nc"]: + if not self.use_network: + assert ( + features.shape[-1] == self.cfg.n_output_dims + ), f"Expected {self.cfg.n_output_dims} output dims, only got {features.shape[-1]} dims input." + color = get_activation(self.cfg.color_activation)(features) + else: + color = self.network(features.view(-1, features.shape[-1])).view( + *features.shape[:-1], self.cfg.n_output_dims + ) + color = get_activation(self.cfg.color_activation)(color) + return color + + def export(self, features: Float[Tensor, "*N Nf"], **kwargs) -> Dict[str, Any]: + color = self(features, **kwargs).clamp(0, 1) + assert color.shape[-1] >= 3, "Output color must have at least 3 channels" + if color.shape[-1] > 3: + lrm.warn("Output color has >3 channels, treating the first 3 as RGB") + return {"albedo": color[..., :3]} diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py new file mode 100644 index 0000000..771e436 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py @@ -0,0 +1,470 @@ +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F + +from CharacterGen.Stage_3D import lrm +from ..utils.ops import dot +from ..utils.typing import * +from ..utils.misc import time_recorder as tr, time_recorder_enabled + + +class Mesh: + def __init__( + self, v_pos: Float[Tensor, "Nv 3"], t_pos_idx: Integer[Tensor, "Nf 3"], **kwargs + ) -> None: + self.v_pos: Float[Tensor, "Nv 3"] = v_pos + self.t_pos_idx: Integer[Tensor, "Nf 3"] = t_pos_idx + self._v_nrm: Optional[Float[Tensor, "Nv 3"]] = None + self._v_tng: Optional[Float[Tensor, "Nv 3"]] = None + self._v_tex: Optional[Float[Tensor, "Nt 3"]] = None + self._t_tex_idx: Optional[Float[Tensor, "Nf 3"]] = None + self._v_rgb: Optional[Float[Tensor, "Nv 3"]] = None + self._edges: Optional[Integer[Tensor, "Ne 2"]] = None + self.extras: Dict[str, Any] = {} + for k, v in kwargs.items(): + self.add_extra(k, v) + + def add_extra(self, k, v) -> None: + self.extras[k] = v + + def remove_outlier(self, outlier_n_faces_threshold: Union[int, float]) -> Mesh: + if self.requires_grad: + lrm.debug("Mesh is differentiable, not removing outliers") + return self + + # use trimesh to first split the mesh into connected components + # then remove the components with less than n_face_threshold faces + import trimesh + + # construct a trimesh object + mesh = trimesh.Trimesh( + vertices=self.v_pos.detach().cpu().numpy(), + faces=self.t_pos_idx.detach().cpu().numpy(), + ) + + # split the mesh into connected components + components = mesh.split(only_watertight=False) + # log the number of faces in each component + lrm.debug( + "Mesh has {} components, with faces: {}".format( + len(components), [c.faces.shape[0] for c in components] + ) + ) + + n_faces_threshold: int + if isinstance(outlier_n_faces_threshold, float): + # set the threshold to the number of faces in the largest component multiplied by outlier_n_faces_threshold + n_faces_threshold = int( + max([c.faces.shape[0] for c in components]) * outlier_n_faces_threshold + ) + else: + # set the threshold directly to outlier_n_faces_threshold + n_faces_threshold = outlier_n_faces_threshold + + # log the threshold + lrm.debug( + "Removing components with less than {} faces".format(n_faces_threshold) + ) + + # remove the components with less than n_face_threshold faces + components = [c for c in components if c.faces.shape[0] >= n_faces_threshold] + + # log the number of faces in each component after removing outliers + lrm.debug( + "Mesh has {} components after removing outliers, with faces: {}".format( + len(components), [c.faces.shape[0] for c in components] + ) + ) + # merge the components + mesh = trimesh.util.concatenate(components) + + # convert back to our mesh format + v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos) + t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx) + + clean_mesh = Mesh(v_pos, t_pos_idx) + # keep the extras unchanged + + if len(self.extras) > 0: + clean_mesh.extras = self.extras + lrm.debug( + f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}" + ) + return clean_mesh + + def subdivide(self): + if self.requires_grad: + lrm.debug("Mesh is differentiable, not performing subdivision") + return self + + import trimesh + + mesh = trimesh.Trimesh( + vertices=self.v_pos.detach().cpu().numpy(), + faces=self.t_pos_idx.detach().cpu().numpy(), + ) + + mesh.subdivide_loop() + + v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos) + t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx) + + subdivided_mesh = Mesh(v_pos, t_pos_idx) + + if len(self.extras) > 0: + subdivided_mesh.extras = self.extras + lrm.debug( + f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}" + ) + + return subdivided_mesh + + def post_process(self, options): + if self.requires_grad: + lrm.debug("Mesh is differentiable, not performing post processing") + return self + + from extern.mesh_process.MeshProcess import process_mesh + + v_pos, t_pos_idx = process_mesh( + vertices=self.v_pos.detach().cpu().numpy(), + faces=self.t_pos_idx.detach().cpu().numpy(), + **options, + ) + + v_pos = torch.from_numpy(v_pos).to(self.v_pos).contiguous() + t_pos_idx = torch.from_numpy(t_pos_idx).to(self.t_pos_idx).contiguous() + + processed_mesh = Mesh(v_pos, t_pos_idx) + + if len(self.extras) > 0: + processed_mesh.extras = self.extras + lrm.debug( + f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}" + ) + + return processed_mesh + + @property + def requires_grad(self): + return self.v_pos.requires_grad + + @property + def v_nrm(self): + if self._v_nrm is None: + self._v_nrm = self._compute_vertex_normal() + return self._v_nrm + + @property + def v_tng(self): + if self._v_tng is None: + self._v_tng = self._compute_vertex_tangent() + return self._v_tng + + @property + def v_tex(self): + if self._v_tex is None: + self._v_tex, self._t_tex_idx = self._unwrap_uv() + return self._v_tex + + @property + def t_tex_idx(self): + if self._t_tex_idx is None: + self._v_tex, self._t_tex_idx = self._unwrap_uv() + return self._t_tex_idx + + @property + def v_rgb(self): + return self._v_rgb + + @property + def edges(self): + if self._edges is None: + self._edges = self._compute_edges() + return self._edges + + def _compute_vertex_normal(self): + i0 = self.t_pos_idx[:, 0] + i1 = self.t_pos_idx[:, 1] + i2 = self.t_pos_idx[:, 2] + + v0 = self.v_pos[i0, :] + v1 = self.v_pos[i1, :] + v2 = self.v_pos[i2, :] + + face_normals = torch.cross(v1 - v0, v2 - v0) + + # Splat face normals to vertices + v_nrm = torch.zeros_like(self.v_pos) + v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals) + v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals) + v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals) + + # Normalize, replace zero (degenerated) normals with some default value + v_nrm = torch.where( + dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm) + ) + v_nrm = F.normalize(v_nrm, dim=1) + + if torch.is_anomaly_enabled(): + assert torch.all(torch.isfinite(v_nrm)) + + return v_nrm + + def _compute_vertex_tangent(self): + vn_idx = [None] * 3 + pos = [None] * 3 + tex = [None] * 3 + for i in range(0, 3): + pos[i] = self.v_pos[self.t_pos_idx[:, i]] + tex[i] = self.v_tex[self.t_tex_idx[:, i]] + # t_nrm_idx is always the same as t_pos_idx + vn_idx[i] = self.t_pos_idx[:, i] + + tangents = torch.zeros_like(self.v_nrm) + tansum = torch.zeros_like(self.v_nrm) + + # Compute tangent space for each triangle + uve1 = tex[1] - tex[0] + uve2 = tex[2] - tex[0] + pe1 = pos[1] - pos[0] + pe2 = pos[2] - pos[0] + + nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2] + denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1] + + # Avoid division by zero for degenerated texture coordinates + tang = nom / torch.where( + denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6) + ) + + # Update all 3 vertices + for i in range(0, 3): + idx = vn_idx[i][:, None].repeat(1, 3) + tangents.scatter_add_(0, idx, tang) # tangents[n_i] = tangents[n_i] + tang + tansum.scatter_add_( + 0, idx, torch.ones_like(tang) + ) # tansum[n_i] = tansum[n_i] + 1 + tangents = tangents / tansum + + # Normalize and make sure tangent is perpendicular to normal + tangents = F.normalize(tangents, dim=1) + tangents = F.normalize(tangents - dot(tangents, self.v_nrm) * self.v_nrm) + + if torch.is_anomaly_enabled(): + assert torch.all(torch.isfinite(tangents)) + + return tangents + + def _unwrap_uv_open3d( + self + ): + import open3d as o3d + mesh = o3d.t.geometry.TriangleMesh() + mesh.vertex.positions = o3d.core.Tensor(self.v_pos.detach().cpu().numpy()) + mesh.triangle.indices = o3d.core.Tensor(self.t_pos_idx.cpu().numpy()) + mesh.compute_uvatlas(size=1024) + texture_uvs = torch.from_numpy(mesh.triangle.texture_uvs.numpy()).reshape(-1, 2).cuda() + indices = torch.arange(self.t_pos_idx.shape[0] * 3).reshape(-1, 3).to(torch.int64).cuda() + # Add a wood texture and visualize + return texture_uvs, indices + + def _unwrap_uv_xatlas( + self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {} + ): + lrm.info("Using xatlas to perform UV unwrapping, may take a while ...") + + import xatlas + + atlas = xatlas.Atlas() + atlas.add_mesh( + self.v_pos.detach().cpu().numpy(), + self.t_pos_idx.cpu().numpy(), + ) + co = xatlas.ChartOptions() + po = xatlas.PackOptions() + for k, v in xatlas_chart_options.items(): + setattr(co, k, v) + for k, v in xatlas_pack_options.items(): + setattr(po, k, v) + atlas.generate(co, po) + vmapping, indices, uvs = atlas.get_mesh(0) + vmapping = ( + torch.from_numpy( + vmapping.astype(np.uint64, casting="same_kind").view(np.int64) + ) + .to(self.v_pos.device) + .long() + ) + uvs = torch.from_numpy(uvs).to(self.v_pos.device).float() + indices = ( + torch.from_numpy( + indices.astype(np.uint64, casting="same_kind").view(np.int64) + ) + .to(self.v_pos.device) + .long() + ) + return uvs, indices + + def _unwrap_uv_smartuv(self, options: dict = {}): + from extern.mesh_process.MeshProcess import ( + mesh_to_bpy, + get_uv_from_bpy, + bpy_context, + bpy_export, + ) + from CharacterGen.Stage_3D.lrm.utils.misc import time_recorder as tr + + v_pos, t_pos_idx = self.v_pos.cpu().numpy(), self.t_pos_idx.cpu().numpy() + with bpy_context(): + mesh_bpy = mesh_to_bpy("_", v_pos, t_pos_idx) + v_tex = get_uv_from_bpy(mesh_bpy, **options).astype(np.float32) + + assert v_tex.shape[0] == self.t_pos_idx.shape[0] * 3 + + t_tex_idx = torch.arange( + self.t_pos_idx.shape[0] * 3, device=self.t_pos_idx.device, dtype=torch.long + ).reshape(-1, 3) + + """ + # super efficient de-duplication + v_tex_u_uint32 = v_tex[..., 0].view(np.uint32) + v_tex_v_uint32 = v_tex[..., 1].view(np.uint32) + v_hashed = (v_tex_u_uint32.astype(np.uint64) << 32) | v_tex_v_uint32 + v_hashed = torch.from_numpy(v_hashed.view(np.int64)).to(self.v_pos.device) + + v_tex = torch.from_numpy(v_tex).to( + device=self.v_pos.device, dtype=torch.float32 + ) + t_pos_idx_f3 = torch.arange( + self.t_pos_idx.shape[0] * 3, device=self.t_pos_idx.device, dtype=torch.long + ).reshape(-1, 3) + v_pos_f3 = self.v_pos[self.t_pos_idx].reshape(-1, 3) + + # super efficient de-duplication + v_hashed_dedup, inverse_indices = torch.unique(v_hashed, return_inverse=True) + dedup_size, full_size = v_hashed_dedup.shape[0], inverse_indices.shape[0] + indices = torch.scatter_reduce( + torch.full( + [dedup_size], + fill_value=full_size, + device=inverse_indices.device, + dtype=torch.long, + ), + index=inverse_indices, + src=torch.arange( + full_size, device=inverse_indices.device, dtype=torch.int64 + ), + dim=0, + reduce="amin", + ) + v_tex = v_tex[indices] + t_tex_idx = inverse_indices.reshape(-1, 3) + + v_pos = v_pos_f3[indices] + t_pos_idx = inverse_indices[t_pos_idx_f3] + """ + + return self.v_pos, self.t_pos_idx, v_tex, t_tex_idx + + def unwrap_uv( + self, + method: str, + xatlas_chart_options: dict = {}, + xatlas_pack_options: dict = {}, + smartuv_options: dict = {}, + ): + if method == "xatlas": + with time_recorder_enabled(): + tr.start("UV unwrapping xatlas") + self._v_tex, self._t_tex_idx = self._unwrap_uv_xatlas( + xatlas_chart_options, xatlas_pack_options + ) + tr.end("UV unwrapping xatlas") + elif method == "open3d": + with time_recorder_enabled(): + tr.start("UV unwrapping o3d") + self._v_tex, self._t_tex_idx = self._unwrap_uv_open3d() + tr.end("UV unwrapping o3d") + elif method == "smartuv": + with time_recorder_enabled(): + tr.start("UV unwrapping smartuv") + ( + self.v_pos, + self.t_pos_idx, + self._v_tex, + self._t_tex_idx, + ) = self._unwrap_uv_smartuv(smartuv_options) + tr.end("UV unwrapping smartuv") + else: + raise NotImplementedError + + def set_vertex_color(self, v_rgb): + assert v_rgb.shape[0] == self.v_pos.shape[0] + self._v_rgb = v_rgb + + def set_uv(self, v_tex, t_tex_idx): + self._v_tex = v_tex + self._t_tex_idx = t_tex_idx + + @property + def has_uv(self): + return self._v_tex is not None and self._t_tex_idx is not None + + def _compute_edges(self): + # Compute edges + edges = torch.cat( + [ + self.t_pos_idx[:, [0, 1]], + self.t_pos_idx[:, [1, 2]], + self.t_pos_idx[:, [2, 0]], + ], + dim=0, + ) + edges = edges.sort()[0] + edges = torch.unique(edges, dim=0) + return edges + + def normal_consistency(self) -> Float[Tensor, ""]: + edge_nrm: Float[Tensor, "Ne 2 3"] = self.v_nrm[self.edges] + nc = ( + 1.0 - torch.cosine_similarity(edge_nrm[:, 0], edge_nrm[:, 1], dim=-1) + ).mean() + return nc + + def _laplacian_uniform(self): + # from stable-dreamfusion + # https://github.com/ashawkey/stable-dreamfusion/blob/8fb3613e9e4cd1ded1066b46e80ca801dfb9fd06/nerf/renderer.py#L224 + verts, faces = self.v_pos, self.t_pos_idx + + V = verts.shape[0] + F = faces.shape[0] + + # Neighbor indices + ii = faces[:, [1, 2, 0]].flatten() + jj = faces[:, [2, 0, 1]].flatten() + adj = torch.stack([torch.cat([ii, jj]), torch.cat([jj, ii])], dim=0).unique( + dim=1 + ) + adj_values = torch.ones(adj.shape[1]).to(verts) + + # Diagonal indices + diag_idx = adj[0] + + # Build the sparse matrix + idx = torch.cat((adj, torch.stack((diag_idx, diag_idx), dim=0)), dim=1) + values = torch.cat((-adj_values, adj_values)) + + # The coalesce operation sums the duplicate indices, resulting in the + # correct diagonal + return torch.sparse_coo_tensor(idx, values, (V, V)).coalesce() + + def laplacian(self) -> Float[Tensor, ""]: + with torch.no_grad(): + L = self._laplacian_uniform() + loss = L.mm(self.v_pos) + loss = loss.norm(dim=1) + loss = loss.mean() + return loss diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/networks.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/networks.py new file mode 100644 index 0000000..44ffc40 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/networks.py @@ -0,0 +1,390 @@ +from dataclasses import dataclass, field +from copy import deepcopy + +import torch +import torch.nn as nn +from einops import rearrange + +from ..utils.base import BaseModule +from ..utils.ops import get_activation +from ..utils.typing import * + + +class TriplaneUpsampleNetwork(BaseModule): + @dataclass + class Config(BaseModule.Config): + in_channels: int = 1024 + out_channels: int = 80 + + cfg: Config + + def configure(self) -> None: + super().configure() + self.upsample = nn.ConvTranspose2d( + self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2 + ) + + def forward( + self, triplanes: Float[Tensor, "B 3 Ci Hp Wp"] + ) -> Float[Tensor, "B 3 Co Hp2 Wp2"]: + triplanes_up = rearrange( + self.upsample( + rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3) + ), + "(B Np) Co Hp Wp -> B Np Co Hp Wp", + Np=3, + ) + return triplanes_up + + +class MLP(nn.Module): + def __init__( + self, + dim_in: int, + dim_out: int, + n_neurons: int, + n_hidden_layers: int, + activation: str = "relu", + output_activation: Optional[str] = None, + bias: bool = True, + weight_init: Optional[str] = "kaiming_uniform", + bias_init: Optional[str] = None, + ): + super().__init__() + layers = [ + self.make_linear( + dim_in, + n_neurons, + is_first=True, + is_last=False, + bias=bias, + weight_init=weight_init, + bias_init=bias_init, + ), + self.make_activation(activation), + ] + for i in range(n_hidden_layers - 1): + layers += [ + self.make_linear( + n_neurons, + n_neurons, + is_first=False, + is_last=False, + bias=bias, + weight_init=weight_init, + bias_init=bias_init, + ), + self.make_activation(activation), + ] + layers += [ + self.make_linear( + n_neurons, + dim_out, + is_first=False, + is_last=True, + bias=bias, + weight_init=weight_init, + bias_init=bias_init, + ) + ] + self.layers = nn.Sequential(*layers) + self.output_activation = get_activation(output_activation) + + def forward(self, x): + x = self.layers(x) + x = self.output_activation(x) + return x + + def make_linear( + self, + dim_in, + dim_out, + is_first, + is_last, + bias=True, + weight_init=None, + bias_init=None, + ): + layer = nn.Linear(dim_in, dim_out, bias=bias) + + if weight_init is None: + pass + elif weight_init == "kaiming_uniform": + torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu") + else: + raise NotImplementedError + + if bias: + if bias_init is None: + pass + elif bias_init == "zero": + torch.nn.init.zeros_(layer.bias) + else: + raise NotImplementedError + + return layer + + def make_activation(self, activation): + if activation == "relu": + return nn.ReLU(inplace=True) + elif activation == "silu": + return nn.SiLU(inplace=True) + else: + raise NotImplementedError + + +@dataclass +class HeadSpec: + name: str + out_channels: int + n_hidden_layers: int + output_activation: Optional[str] = None + + +class MultiHeadMLP(BaseModule): + @dataclass + class Config(BaseModule.Config): + in_channels: int = 0 + n_neurons: int = 0 + n_hidden_layers_share: int = 0 + heads: List[HeadSpec] = field(default_factory=lambda: []) + activation: str = "relu" + bias: bool = True + weight_init: Optional[str] = "kaiming_uniform" + bias_init: Optional[str] = None + chunk_mode: Optional[str] = None + chunk_size: int = -1 + + cfg: Config + + def configure(self) -> None: + super().configure() + shared_layers = [ + self.make_linear( + self.cfg.in_channels, + self.cfg.n_neurons, + bias=self.cfg.bias, + weight_init=self.cfg.weight_init, + bias_init=self.cfg.bias_init, + ), + self.make_activation(self.cfg.activation), + ] + for i in range(self.cfg.n_hidden_layers_share - 1): + shared_layers += [ + self.make_linear( + self.cfg.n_neurons, + self.cfg.n_neurons, + bias=self.cfg.bias, + weight_init=self.cfg.weight_init, + bias_init=self.cfg.bias_init, + ), + self.make_activation(self.cfg.activation), + ] + self.shared_layers = nn.Sequential(*shared_layers) + + assert len(self.cfg.heads) > 0 + heads = {} + for head in self.cfg.heads: + head_layers = [] + for i in range(head.n_hidden_layers): + head_layers += [ + self.make_linear( + self.cfg.n_neurons, + self.cfg.n_neurons, + bias=self.cfg.bias, + weight_init=self.cfg.weight_init, + bias_init=self.cfg.bias_init, + ), + self.make_activation(self.cfg.activation), + ] + head_layers += [ + self.make_linear( + self.cfg.n_neurons, + head.out_channels, + bias=self.cfg.bias, + weight_init=self.cfg.weight_init, + bias_init=self.cfg.bias_init, + ), + ] + heads[head.name] = nn.Sequential(*head_layers) + self.heads = nn.ModuleDict(heads) + + if self.cfg.chunk_mode is not None: + assert self.cfg.chunk_size > 0 + + def make_linear( + self, + dim_in, + dim_out, + bias=True, + weight_init=None, + bias_init=None, + ): + layer = nn.Linear(dim_in, dim_out, bias=bias) + + if weight_init is None: + pass + elif weight_init == "kaiming_uniform": + torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu") + else: + raise NotImplementedError + + if bias: + if bias_init is None: + pass + elif bias_init == "zero": + torch.nn.init.zeros_(layer.bias) + else: + raise NotImplementedError + + return layer + + def make_activation(self, activation): + if activation == "relu": + return nn.ReLU(inplace=True) + elif activation == "silu": + return nn.SiLU(inplace=True) + else: + raise NotImplementedError + + def forward( + self, x, include: Optional[List] = None, exclude: Optional[List] = None + ): + inp_shape = x.shape[:-1] + x = x.reshape(-1, x.shape[-1]) + + if self.cfg.chunk_mode is None: + shared_features = self.shared_layers(x) + elif self.cfg.chunk_mode == "deferred": + shared_features = DeferredFunc.apply( + self.shared_layers, x, self.cfg.chunk_size + ) + elif self.cfg.chunk_mode == "checkpointing": + shared_features = apply_batch_checkpointing( + self.shared_layers, x, self.cfg.chunk_size + ) + else: + raise NotImplementedError + + shared_features = shared_features.reshape(*inp_shape, -1) + + if include is not None and exclude is not None: + raise ValueError("Cannot specify both include and exclude.") + if include is not None: + heads = [h for h in self.cfg.heads if h.name in include] + elif exclude is not None: + heads = [h for h in self.cfg.heads if h.name not in exclude] + else: + heads = self.cfg.heads + + out = { + head.name: get_activation(head.output_activation)( + self.heads[head.name](shared_features) + ) + for head in heads + } + """ + # TypeError + if self.cfg.chunk_mode is None: + out = { + head.name: get_activation(head.output_activation)( + self.heads[head.name](shared_features) + ) + for head in heads + } + elif self.cfg.chunk_mode == "deferred": + out = { + head.name: get_activation(head.output_activation)( + DeferredFunc.apply(self.heads[head.name], shared_features, self.cfg.chunk_size) + ) + for head in heads + } + else: + raise NotImplementedError + """ + return out + + +class DeferredFunc(torch.autograd.Function): + # Note that forward, setup_context, and backward are @staticmethods + @staticmethod + def forward(ctx, model, x, chunk_size): + model_copy = deepcopy(model) + model_copy.requires_grad_(False) + + ret = [] + x_split = torch.split(x, chunk_size, dim=0) + + with torch.no_grad(): + for cur_x in x_split: + ret.append(model_copy(cur_x)) + + ctx.model = model + ctx.save_for_backward(x.detach(), torch.as_tensor(chunk_size)) + + ret = torch.cat(ret, dim=0) + + return ret + + # This function has only a single output, so it gets only one gradient + @staticmethod + def backward(ctx, grad_output): + model = ctx.model + x, chunk_size = ctx.saved_tensors + chunk_size = chunk_size.item() + + model_copy = deepcopy(model) + + x_split = torch.split(x, chunk_size, dim=0) + grad_output_split = torch.split(grad_output, chunk_size, 0) + grad_input_split = [] + + with torch.set_grad_enabled(True): + model_copy.requires_grad_(True) + model_copy.zero_grad() + for cur_x, cur_grad_output in zip(x_split, grad_output_split): + cur_x.requires_grad_(True) + cur_y = model_copy(cur_x) + cur_y.backward(cur_grad_output) + + grad_input_split.append(cur_x.grad.clone()) + + grad_input = torch.cat(grad_input_split, dim=0) + + model_copy_params = list(model_copy.parameters()) + model_params = list(model.parameters()) + + for param, param_copy in zip(model_params, model_copy_params): + if param.grad is None: + param.grad = param_copy.grad.clone() + else: + param.grad.add_(param_copy.grad) + + return None, grad_input, None + + +def apply_batch_checkpointing(func, x, chunk_size): + if chunk_size >= len(x): + # return func(x) + return torch.utils.checkpoint.checkpoint(func, x, use_reentrant=False) + + x_split = torch.split(x, chunk_size, dim=0) + + def cat_and_query(y_all, x): + return torch.cat([y_all, func(x)]) + + y_all = func(x_split[0]) + for cur_x in x_split[1:]: + y_all = torch.utils.checkpoint.checkpoint( + cat_and_query, y_all, cur_x, use_reentrant=False + ) + + return y_all + + +def get_encoding(n_input_dims: int, config) -> nn.Module: + raise NotImplementedError + + +def get_mlp(n_input_dims, n_output_dims, config) -> nn.Module: + raise NotImplementedError diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/base.py new file mode 100644 index 0000000..29e13b0 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/base.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass + +import torch + +from ..networks import MultiHeadMLP +from ..background.base import BaseBackground +from ..materials.base import BaseMaterial +from ...utils.base import BaseModule +from ...utils.typing import * + + +class BaseRenderer(BaseModule): + @dataclass + class Config(BaseModule.Config): + radius: float = 1.0 + + cfg: Config + + def configure( + self, + decoder: MultiHeadMLP, + material: BaseMaterial, + background: BaseBackground, + ) -> None: + super().configure() + + self.set_decoder(decoder) + self.set_material(material) + self.set_background(background) + + # set up bounding box + self.bbox: Float[Tensor, "2 3"] + self.register_buffer( + "bbox", + torch.as_tensor( + [ + [-self.cfg.radius, -self.cfg.radius, -self.cfg.radius], + [self.cfg.radius, self.cfg.radius, self.cfg.radius], + ], + dtype=torch.float32, + ), + ) + + def forward(self, *args, **kwargs) -> Dict[str, Any]: + raise NotImplementedError + + @property + def decoder(self) -> MultiHeadMLP: + return self.non_module("decoder") + + @property + def material(self) -> BaseMaterial: + return self.non_module("material") + + @property + def background(self) -> BaseBackground: + return self.non_module("background") + + def set_decoder(self, decoder: MultiHeadMLP) -> None: + self.register_non_module("decoder", decoder) + + def set_material(self, material: BaseMaterial) -> None: + self.register_non_module("material", material) + + def set_background(self, background: BaseBackground) -> None: + self.register_non_module("background", background) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/triplane_dmtet.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/triplane_dmtet.py new file mode 100644 index 0000000..72021db --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/renderers/triplane_dmtet.py @@ -0,0 +1,369 @@ +import os +from dataclasses import dataclass +from collections import defaultdict +from functools import partial + +import torch +import torch.nn.functional as F +from einops import rearrange, reduce + +from CharacterGen.Stage_3D import lrm +from ..renderers.base import BaseRenderer +from ..isosurface import MarchingTetrahedraHelper +from ...utils.ops import ( + get_activation, + scale_tensor, + dot, + chunk_batch +) +from ...utils.rasterize import NVDiffRasterizerContext +from ..mesh import Mesh +from ...utils.typing import * + + +class TriplaneDMTetRenderer(BaseRenderer): + @dataclass + class Config(BaseRenderer.Config): + feature_reduction: str = "concat" + sdf_activation: Optional[str] = None + sdf_bias: Union[str, float] = 0.0 + sdf_bias_params: Any = None + inside_out: bool = False + + isosurface_resolution: int = 256 + tet_dir: str = "tets/" + enable_isosurface_grid_deformation: bool = False + eval_chunk_size: int = 262144 + context_type: str = "gl" + + cfg: Config + + def configure(self, *args, **kwargs) -> None: + super().configure(*args, **kwargs) + + assert self.cfg.feature_reduction in ["concat", "mean"] + + self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device) + self.isosurface_helper = MarchingTetrahedraHelper( + self.cfg.isosurface_resolution, + os.path.join(self.cfg.tet_dir, f"{self.cfg.isosurface_resolution}_tets.npz"), + ).to(self.device) + + def query_triplane( + self, + positions: Float[Tensor, "*B N 3"], + triplanes: Float[Tensor, "*B 3 Cp Hp Wp"], + ) -> Dict[str, Tensor]: + batched = positions.ndim == 3 + if not batched: + # no batch dimension + triplanes = triplanes[None, ...] + positions = positions[None, ...] + # import ipdb + # ipdb.set_trace() + assert triplanes.ndim == 5 and positions.ndim == 3 + + # assume positions in [-1, 1] + # normalized to (-1, 1) for grid sample + positions = scale_tensor( + positions, (-self.cfg.radius, self.cfg.radius), (-1, 1) + ) + + indices2D: Float[Tensor, "B 3 N 2"] = torch.stack( + (positions[..., [0, 1]], positions[..., [0, 2]], positions[..., [1, 2]]), + dim=-3, + ) + out: Float[Tensor, "B3 Cp 1 N"] = F.grid_sample( + rearrange(triplanes, "B Np Cp Hp Wp -> (B Np) Cp Hp Wp", Np=3), + rearrange(indices2D, "B Np N Nd -> (B Np) () N Nd", Np=3), + align_corners=False, + mode="bilinear", + ) + if self.cfg.feature_reduction == "concat": + out = rearrange(out, "(B Np) Cp () N -> B N (Np Cp)", Np=3) + elif self.cfg.feature_reduction == "mean": + out = reduce(out, "(B Np) Cp () N -> B N Cp", Np=3, reduction="mean") + else: + raise NotImplementedError + + net_out: Dict[str, Float[Tensor, "B N ..."]] = self.decoder(out) + assert "sdf" in net_out + net_out["sdf"] = get_activation(self.cfg.sdf_activation)( + self.get_shifted_sdf(positions, net_out["sdf"]) + ) + + if not batched: + net_out = {k: v.squeeze(0) for k, v in net_out.items()} + + return net_out + + def get_shifted_sdf( + self, points: Float[Tensor, "*N Di"], sdf: Float[Tensor, "*N 1"] + ) -> Float[Tensor, "*N 1"]: + sdf_bias: Union[float, Float[Tensor, "*N 1"]] + if self.cfg.sdf_bias == "ellipsoid": + assert ( + isinstance(self.cfg.sdf_bias_params, Sized) + and len(self.cfg.sdf_bias_params) == 3 + ) + size = torch.as_tensor(self.cfg.sdf_bias_params).to(points) + sdf_bias = ((points / size) ** 2).sum( + dim=-1, keepdim=True + ).sqrt() - 1.0 # pseudo signed distance of an ellipsoid + elif self.cfg.sdf_bias == "sphere": + assert isinstance(self.cfg.sdf_bias_params, float) + radius = self.cfg.sdf_bias_params + sdf_bias = (points**2).sum(dim=-1, keepdim=True).sqrt() - radius + elif isinstance(self.cfg.sdf_bias, float): + sdf_bias = self.cfg.sdf_bias + else: + raise ValueError(f"Unknown sdf bias {self.cfg.sdf_bias}") + return sdf + sdf_bias + + def forward_single( + self, + triplane: Float[Tensor, "3 Cp Hp Wp"], + mvp_mtx: Float[Tensor, "Nv 4 4"], + camera_positions: Float[Tensor, "Nv 3"], + height: int, + width: int, + background_color: Optional[Float[Tensor, "3"]], + extra_sdf_query: Any = None, + ) -> Dict[str, Tensor]: + Nv = mvp_mtx.shape[0] + + out = {} + + query_vertices = [] + query_sizes = [] + + grid_vertices = scale_tensor( + self.isosurface_helper.grid_vertices, + self.isosurface_helper.points_range, + self.bbox, + ) + + query_vertices.append(grid_vertices) + query_sizes.append(len(grid_vertices)) + + if extra_sdf_query is not None: + query_vertices.append(extra_sdf_query) + query_sizes.append(len(extra_sdf_query)) + + query_vertices = torch.cat(query_vertices, dim=0) + triplane_out = self.query_triplane(query_vertices, triplane) + + all_sdf = triplane_out["sdf"] + if extra_sdf_query is not None: + sdf, sdf_ex_query = torch.split(all_sdf, query_sizes) + else: + sdf, sdf_ex_query = all_sdf, None + + out.update({"sdf_ex_query": sdf_ex_query}) + + if self.cfg.enable_isosurface_grid_deformation: + all_deformation = triplane_out["deformation"] + if extra_sdf_query is not None: + deformation, _ = torch.split(all_deformation, query_sizes) + else: + deformation, _ = all_deformation, None + else: + deformation = None + + # Fix some sdf if we observe empty shape (full positive or full negative) + pos_shape = torch.sum((sdf.squeeze(dim=-1) > 0).int(), dim=-1) + neg_shape = torch.sum((sdf.squeeze(dim=-1) < 0).int(), dim=-1) + zero_surface = torch.bitwise_or(pos_shape == 0, neg_shape == 0) + if torch.sum(zero_surface).item() > 0: + lrm.warn("Empty mesh! Fixing by adding fake faces.") + sdf = torch.nan_to_num(sdf, nan=0.0, posinf=1.0, neginf=-1.0) + update_sdf = torch.zeros_like(sdf) + max_sdf = sdf.max() + min_sdf = sdf.min() + update_sdf[self.isosurface_helper.center_indices] += ( + -1.0 - max_sdf + ) # greater than zero + update_sdf[self.isosurface_helper.boundary_indices] += ( + 1.0 - min_sdf + ) # smaller than zero + new_sdf = sdf.clone().detach() + if zero_surface: + new_sdf += update_sdf + update_mask = (update_sdf == 0).float() + # Regulraization here is used to push the sdf to be a different sign (make it not fully positive or fully negative) + sdf_reg_loss = sdf.abs().mean() + sdf_reg_loss = sdf_reg_loss * zero_surface.float() + sdf = sdf * update_mask + new_sdf * (1 - update_mask) + lrm.debug( + "max sdf: {}, min sdf: {}".format(sdf.max().item(), sdf.min().item()) + ) + out.update({"sdf_reg": sdf_reg_loss}) + + # Here we remove the gradient for the bad sdf (full positive or full negative) + if zero_surface: + sdf = sdf.detach() + + mesh: Mesh = self.isosurface_helper(sdf, deformation=deformation) + + mesh.v_pos = scale_tensor( + mesh.v_pos, self.isosurface_helper.points_range, self.bbox + ) # scale to bbox as the grid vertices are in [0, 1] + # import ipdb + # ipdb.set_trace() + v_pos_clip: Float[Tensor, "Nv V 4"] = self.ctx.vertex_transform( + mesh.v_pos, mvp_mtx + ) + rast, _ = self.ctx.rasterize(v_pos_clip, mesh.t_pos_idx, (height, width)) + mask = rast[..., 3:] > 0 + mask_aa = self.ctx.antialias(mask.float(), rast, v_pos_clip, mesh.t_pos_idx) + + out.update({"opacity": mask_aa, "mesh": mesh}) + + gb_normal, _ = self.ctx.interpolate_one(mesh.v_nrm, rast, mesh.t_pos_idx) + gb_normal = F.normalize(gb_normal, dim=-1) + gb_normal_aa = torch.lerp( + torch.zeros_like(gb_normal), (gb_normal + 1.0) / 2.0, mask.float() + ) + gb_normal_aa = self.ctx.antialias( + gb_normal_aa, rast, v_pos_clip, mesh.t_pos_idx + ) + out.update({"comp_normal": gb_normal_aa}) # in [0, 1] + + gb_pos, _ = self.ctx.interpolate_one(mesh.v_pos, rast, mesh.t_pos_idx) + + # FIXME: this depth corresponds to the one provided in the dataset, but assumes camera looking at scene center + gb_depth = dot( + gb_pos - camera_positions[:, None, None, :], + F.normalize(-camera_positions[:, None, None, :], dim=-1), + ) + + gb_depth = torch.lerp(torch.zeros_like(gb_depth), gb_depth, mask.float()) + out.update({"depth": gb_depth}) + + gb_viewdirs = F.normalize(gb_pos - camera_positions[:, None, None, :], dim=-1) + gb_rgb_fg = torch.zeros( + (Nv, height, width, 3), device=self.device, dtype=torch.float32 + ) + gb_rgb_bg = self.background(dirs=gb_viewdirs, color_spec=background_color) + + selector = mask[..., 0] + if selector.sum() > 0: + positions = gb_pos[selector] + geo_out = self.query_triplane(positions, triplane) + + extra_geo_info = {} + if self.material.requires_normal: + extra_geo_info["shading_normal"] = gb_normal[selector] + if self.material.requires_tangent: + gb_tangent, _ = self.ctx.interpolate_one( + mesh.v_tng, rast, mesh.t_pos_idx + ) + gb_tangent = F.normalize(gb_tangent, dim=-1) + extra_geo_info["tangent"] = gb_tangent[selector] + + rgb_fg = self.material( + viewdirs=gb_viewdirs[selector], + positions=positions, + **extra_geo_info, + **geo_out, + ) + + gb_rgb_fg[selector] = rgb_fg.to( + gb_rgb_fg.dtype + ) # TODO: don't know if this is correct + + gb_rgb = torch.lerp(gb_rgb_bg, gb_rgb_fg, mask.float()) + gb_rgb_aa = self.ctx.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx) + + out.update( + {"comp_rgb": gb_rgb_aa, "comp_rgb_fg": gb_rgb_fg, "comp_rgb_bg": gb_rgb_bg} + ) + + return out + + def forward( + self, + triplanes: Float[Tensor, "B 3 Cp Hp Wp"], + mvp_mtx: Float[Tensor, "B Nv 4 4"], + camera_positions: Float[Tensor, "B Nv 3"], + height: int, + width: int, + background_color: Optional[Float[Tensor, "B 3"]] = None, + extra_sdf_query: Optional[List[Float[Tensor, "N 3"]]] = None, + **kwargs, + ) -> Dict[str, Tensor]: + batch_size = triplanes.shape[0] + out_list = [] + for b in range(batch_size): + out_list.append( + self.forward_single( + triplanes[b], + mvp_mtx[b], + camera_positions[b], + height, + width, + background_color=background_color[b] + if background_color is not None + else None, + extra_sdf_query=extra_sdf_query[b] + if extra_sdf_query is not None + else None, + ) + ) + + out = defaultdict(list) + for out_ in out_list: + for k, v in out_.items(): + out[k].append(v) + + for k, v in out.items(): + # some properties cannot be batched + if isinstance(v[0], torch.Tensor) and ( + all([vv.ndim == 0 for vv in v]) + or all([vv.shape[0] == v[0].shape[0] for vv in v]) + ): + out[k] = torch.stack(v, dim=0) + else: + out[k] = v + + return out + + def isosurface(self, triplane: Float[Tensor, "3 Cp Hp Wp"]): + grid_vertices = scale_tensor( + self.isosurface_helper.grid_vertices, + self.isosurface_helper.points_range, + self.bbox, + ) + triplane_out = chunk_batch( + partial(self.query_triplane, triplanes=triplane), self.cfg.eval_chunk_size, grid_vertices, + ) + + sdf = triplane_out["sdf"] + + if self.cfg.inside_out: + sdf = -sdf + + if self.cfg.enable_isosurface_grid_deformation: + deformation = triplane_out["deformation"] + else: + deformation = None + + mesh: Mesh = self.isosurface_helper(sdf, deformation=deformation) + + mesh.v_pos = scale_tensor( + mesh.v_pos, self.isosurface_helper.points_range, self.bbox + ) + + return mesh + + def query( + self, triplane: Float[Tensor, "3 Cp Hp Wp"], points: Float[Tensor, "*N 3"] + ): + input_shape = points.shape[:-1] + triplane_out = chunk_batch( + partial(self.query_triplane, triplanes=triplane), self.cfg.eval_chunk_size, points.view(-1, 3) + ) + triplane_out = { + k: v.view(*input_shape, v.shape[-1]) for k, v in triplane_out.items() + } + return triplane_out diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/dinov2.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/dinov2.py new file mode 100644 index 0000000..d4d7eae --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/dinov2.py @@ -0,0 +1,1223 @@ +# coding=utf-8 +# Copyright 2023 Meta AI and The HuggingFace Inc. 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. +""" PyTorch DINOv2 model.""" + + +import collections.abc +import math +from typing import Dict, List, Optional, Set, Tuple, Union +from dataclasses import dataclass + +import torch +import torch.utils.checkpoint +from torch import nn +import torch.nn.functional as F +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.modeling_outputs import ( + BackboneOutput, + BaseModelOutput, + BaseModelOutputWithPooling, + ImageClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + find_pruneable_heads_and_indices, + prune_linear_layer, +) +from transformers.utils import ( + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from transformers.utils.backbone_utils import BackboneMixin +from transformers.models.dinov2.configuration_dinov2 import Dinov2Config +import xformers + +from ..transformers.attention import MemoryEfficientAttentionMixin +from ...utils.typing import * + + +logger = logging.get_logger(__name__) + +# General docstring +_CONFIG_FOR_DOC = "Dinov2Config" + +# Base docstring +_CHECKPOINT_FOR_DOC = "facebook/dinov2-base" +_EXPECTED_OUTPUT_SHAPE = [1, 257, 768] + +# Image classification docstring +_IMAGE_CLASS_CHECKPOINT = "facebook/dinov2-base" + + +DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "facebook/dinov2-base", + # See all DINOv2 models at https://huggingface.co/models?filter=dinov2 +] + + +class Dinov2Embeddings(nn.Module): + """ + Construct the CLS token, mask token, position and patch embeddings. + """ + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + + self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + # register as mask token as it's not used in optimization + # to avoid the use of find_unused_parameters_true + # self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size)) + self.register_buffer("mask_token", torch.zeros(1, config.hidden_size)) + self.patch_embeddings = Dinov2PatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter( + torch.randn(1, num_patches + 1, config.hidden_size) + ) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def interpolate_pos_encoding( + self, embeddings: torch.Tensor, height: int, width: int + ) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher + resolution images. + + Source: + https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + if num_patches == num_positions and height == width: + return self.position_embeddings + class_pos_embed = self.position_embeddings[:, 0] + patch_pos_embed = self.position_embeddings[:, 1:] + dim = embeddings.shape[-1] + height = height // self.config.patch_size + width = width // self.config.patch_size + # we add a small number to avoid floating point error in the interpolation + # see discussion at https://github.com/facebookresearch/dino/issues/8 + height, width = height + 0.1, width + 0.1 + patch_pos_embed = patch_pos_embed.reshape( + 1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim + ) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + scale_factor=( + height / math.sqrt(num_positions), + width / math.sqrt(num_positions), + ), + mode="bicubic", + align_corners=False, + ) + if ( + int(height) != patch_pos_embed.shape[-2] + or int(width) != patch_pos_embed.shape[-1] + ): + raise ValueError( + "Width or height does not match with the interpolated position embeddings" + ) + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) + + def forward( + self, + pixel_values: torch.Tensor, + bool_masked_pos: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + patch_embeddings = self.patch_embeddings(pixel_values) + embeddings = patch_embeddings + + if bool_masked_pos is not None: + embeddings = torch.where( + bool_masked_pos.unsqueeze(-1), + self.mask_token.to(embeddings.dtype).unsqueeze(0), + embeddings, + ) + + # add the [CLS] token to the embedded patch tokens + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, height, width + ) + + embeddings = self.dropout(embeddings) + + return embeddings + + +class Dinov2PatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = ( + image_size + if isinstance(image_size, collections.abc.Iterable) + else (image_size, image_size) + ) + patch_size = ( + patch_size + if isinstance(patch_size, collections.abc.Iterable) + else (patch_size, patch_size) + ) + num_patches = (image_size[1] // patch_size[1]) * ( + image_size[0] // patch_size[0] + ) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d( + num_channels, hidden_size, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """ + num_channels = pixel_values.shape[1] + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + f" Expected {self.num_channels} but got {num_channels}." + ) + """ + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Dinov2 +class Dinov2SelfAttention(nn.Module): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr( + config, "embedding_size" + ): + raise ValueError( + f"The hidden size {config.hidden_size,} is not a multiple of the number of attention " + f"heads {config.num_attention_heads}." + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.attention_probs_dropout_prob = config.attention_probs_dropout_prob + + self.query = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + self.key = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + self.value = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.use_memory_efficient_attention_xformers: bool = False + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward( + self, + hidden_states, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + mixed_query_layer = self.query(hidden_states) + + if self.use_memory_efficient_attention_xformers: + assert head_mask is None and not output_attentions + new_size = hidden_states.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + key_layer = self.key(hidden_states).view(new_size) + value_layer = self.value(hidden_states).view(new_size) + query_layer = mixed_query_layer.view(new_size) + context_layer = xformers.ops.memory_efficient_attention( + query_layer, key_layer, value_layer, p=self.attention_probs_dropout_prob + ) + context_layer = context_layer.view(*hidden_states.size()[:-1], -1) + elif hasattr(F, "scaled_dot_product_attention"): + assert head_mask is None and not output_attentions + new_size = hidden_states.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + key_layer = self.key(hidden_states).reshape(new_size).transpose(1, 2) + value_layer = self.value(hidden_states).reshape(new_size).transpose(1, 2) + query_layer = mixed_query_layer.reshape(new_size).transpose(1, 2) + context_layer = F.scaled_dot_product_attention( + query_layer, + key_layer, + value_layer, + dropout_p=self.attention_probs_dropout_prob, + is_causal=False, + ) + context_layer = context_layer.transpose(1, 2).reshape( + *hidden_states.size()[:-1], -1 + ) + else: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + query_layer = self.transpose_for_scores(mixed_query_layer) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + # Mask heads if we want to + if head_mask is not None: + attention_probs = attention_probs * head_mask + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = ( + (context_layer, attention_probs) if output_attentions else (context_layer,) + ) + + return outputs + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ): + self.use_memory_efficient_attention_xformers = valid + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Dinov2 +class Dinov2SelfOutput(nn.Module): + """ + The residual connection is defined in Dinov2Layer instead of here (as is the case with other models), due to the + layernorm applied before each block. + """ + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Dinov2 +class Dinov2Attention(nn.Module): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.attention = Dinov2SelfAttention(config) + self.output = Dinov2SelfOutput(config) + self.pruned_heads = set() + + def prune_heads(self, heads: Set[int]) -> None: + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, + self.attention.num_attention_heads, + self.attention.attention_head_size, + self.pruned_heads, + ) + + # Prune linear layers + self.attention.query = prune_linear_layer(self.attention.query, index) + self.attention.key = prune_linear_layer(self.attention.key, index) + self.attention.value = prune_linear_layer(self.attention.value, index) + self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) + + # Update hyper params and store pruned heads + self.attention.num_attention_heads = self.attention.num_attention_heads - len( + heads + ) + self.attention.all_head_size = ( + self.attention.attention_head_size * self.attention.num_attention_heads + ) + self.pruned_heads = self.pruned_heads.union(heads) + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + self_outputs = self.attention(hidden_states, head_mask, output_attentions) + + attention_output = self.output(self_outputs[0], hidden_states) + + outputs = (attention_output,) + self_outputs[ + 1: + ] # add attentions if we output them + return outputs + + +class Dinov2LayerScale(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.lambda1 = nn.Parameter( + config.layerscale_value * torch.ones(config.hidden_size) + ) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + return hidden_state * self.lambda1 + + +# Copied from transformers.models.beit.modeling_beit.drop_path +def drop_path( + input: torch.Tensor, drop_prob: float = 0.0, training: bool = False +) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks, + however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the + layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the + argument. + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * ( + input.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand( + shape, dtype=input.dtype, device=input.device + ) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +# Copied from transformers.models.beit.modeling_beit.BeitDropPath +class Dinov2DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: Optional[float] = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return "p={}".format(self.drop_prob) + + +class Dinov2MLP(nn.Module): + def __init__(self, config) -> None: + super().__init__() + in_features = out_features = config.hidden_size + hidden_features = int(config.hidden_size * config.mlp_ratio) + self.fc1 = nn.Linear(in_features, hidden_features, bias=True) + if isinstance(config.hidden_act, str): + self.activation = ACT2FN[config.hidden_act] + else: + self.activation = config.hidden_act + self.fc2 = nn.Linear(hidden_features, out_features, bias=True) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.fc1(hidden_state) + hidden_state = self.activation(hidden_state) + hidden_state = self.fc2(hidden_state) + return hidden_state + + +class Dinov2SwiGLUFFN(nn.Module): + def __init__(self, config) -> None: + super().__init__() + in_features = out_features = config.hidden_size + hidden_features = int(config.hidden_size * config.mlp_ratio) + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + + self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True) + self.weights_out = nn.Linear(hidden_features, out_features, bias=True) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.weights_in(hidden_state) + x1, x2 = hidden_state.chunk(2, dim=-1) + hidden = nn.functional.silu(x1) * x2 + return self.weights_out(hidden) + + +class Dinov2Layer(nn.Module, MemoryEfficientAttentionMixin): + """This corresponds to the Block class in the original implementation.""" + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm1_modulation = None + self.attention = Dinov2Attention(config) + self.layer_scale1 = Dinov2LayerScale(config) + self.drop_path1 = ( + Dinov2DropPath(config.drop_path_rate) + if config.drop_path_rate > 0.0 + else nn.Identity() + ) + + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm2_modulation = None + + if config.use_swiglu_ffn: + self.mlp = Dinov2SwiGLUFFN(config) + else: + self.mlp = Dinov2MLP(config) + self.layer_scale2 = Dinov2LayerScale(config) + self.drop_path2 = ( + Dinov2DropPath(config.drop_path_rate) + if config.drop_path_rate > 0.0 + else nn.Identity() + ) + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + hidden_states_norm = self.norm1(hidden_states) + if self.norm1_modulation is not None: + assert modulation_cond is not None + hidden_states_norm = self.norm1_modulation( + hidden_states_norm, modulation_cond + ) + self_attention_outputs = self.attention( + hidden_states_norm, # in Dinov2, layernorm is applied before self-attention + head_mask, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + + attention_output = self.layer_scale1(attention_output) + outputs = self_attention_outputs[ + 1: + ] # add self attentions if we output attention weights + + # first residual connection + hidden_states = attention_output + hidden_states + + # in Dinov2, layernorm is also applied after self-attention + layer_output = self.norm2(hidden_states) + if self.norm2_modulation is not None: + assert modulation_cond is not None + layer_output = self.norm2_modulation(layer_output, modulation_cond) + layer_output = self.mlp(layer_output) + layer_output = self.layer_scale2(layer_output) + + # second residual connection + layer_output = layer_output + hidden_states + + outputs = (layer_output,) + outputs + + return outputs + + def register_ada_norm_modulation(self, norm1_mod: nn.Module, norm2_mod: nn.Module): + self.norm1_modulation = norm1_mod + self.norm2_modulation = norm2_mod + + +# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->Dinov2 +class Dinov2Encoder(nn.Module, MemoryEfficientAttentionMixin): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.config = config + self.layer = nn.ModuleList( + [Dinov2Layer(config) for _ in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> Union[tuple, BaseModelOutput]: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_head_mask = head_mask[i] if head_mask is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + layer_head_mask, + modulation_cond, + use_reentrant=False, + ) + else: + layer_outputs = layer_module( + hidden_states, layer_head_mask, modulation_cond, output_attentions + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_self_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class Dinov2PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = Dinov2Config + base_model_prefix = "dinov2" + main_input_name = "pixel_values" + supports_gradient_checkpointing = True + + def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None: + """Initialize the weights""" + if isinstance(module, (nn.Linear, nn.Conv2d)): + # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid + # `trunc_normal_cpu` not implemented in `half` issues + module.weight.data = nn.init.trunc_normal_( + module.weight.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.weight.dtype) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + elif isinstance(module, Dinov2Embeddings): + module.position_embeddings.data = nn.init.trunc_normal_( + module.position_embeddings.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.position_embeddings.dtype) + + module.cls_token.data = nn.init.trunc_normal_( + module.cls_token.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.cls_token.dtype) + + def _set_gradient_checkpointing( + self, module: Dinov2Encoder, value: bool = False + ) -> None: + if isinstance(module, Dinov2Encoder): + module.gradient_checkpointing = value + + +DINOV2_START_DOCSTRING = r""" + This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it + as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and + behavior. + + Parameters: + config ([`Dinov2Config`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +DINOV2_BASE_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See + [`BitImageProcessor.preprocess`] for details. + + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for + pre-training. + + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +DINOV2_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See + [`BitImageProcessor.preprocess`] for details. + + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@dataclass +class CustomBaseModelOutputWithPooling(BaseModelOutputWithPooling): + patch_embeddings: Optional[torch.FloatTensor] = None + + +@add_start_docstrings( + "The bare DINOv2 Model transformer outputting raw hidden-states without any specific head on top.", + DINOV2_START_DOCSTRING, +) +class Dinov2Model(Dinov2PreTrainedModel, MemoryEfficientAttentionMixin): + def __init__(self, config: Dinov2Config): + super().__init__(config) + self.config = config + + self.embeddings = Dinov2Embeddings(config) + self.encoder = Dinov2Encoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> Dinov2PatchEmbeddings: + return self.embeddings.patch_embeddings + + def expand_input_channels(self, extra_input_channels: int) -> None: + if extra_input_channels == 0: + return + conv_old = self.embeddings.patch_embeddings.projection + conv_new = nn.Conv2d( + self.config.num_channels + extra_input_channels, + self.config.hidden_size, + kernel_size=self.config.patch_size, + stride=self.config.patch_size, + ).to(self.device) + with torch.no_grad(): + conv_new.weight[:, :3] = conv_old.weight + conv_new.bias = conv_old.bias + self.embeddings.patch_embeddings.projection = conv_new + del conv_old + + def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None: + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward(DINOV2_BASE_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPooling, + config_class=_CONFIG_FOR_DOC, + modality="vision", + expected_output=_EXPECTED_OUTPUT_SHAPE, + ) + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + bool_masked_pos: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + pixel_values, bool_masked_pos=bool_masked_pos + ) + + encoder_outputs = self.encoder( + embedding_output, + head_mask=head_mask, + modulation_cond=modulation_cond, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + pooled_output = sequence_output[:, 0, :] + + if not return_dict: + head_outputs = (sequence_output, pooled_output) + return head_outputs + encoder_outputs[1:] + + return CustomBaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + patch_embeddings=embedding_output, + ) + + def set_gradient_checkpointing(self, value: bool = False) -> None: + self._set_gradient_checkpointing(self.encoder, value) + + +@add_start_docstrings( + """ + Dinov2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state + of the [CLS] token) e.g. for ImageNet. + """, + DINOV2_START_DOCSTRING, +) +class Dinov2ForImageClassification(Dinov2PreTrainedModel): + def __init__(self, config: Dinov2Config) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.dinov2 = Dinov2Model(config) + + # Classifier head + self.classifier = ( + nn.Linear(config.hidden_size * 2, config.num_labels) + if config.num_labels > 0 + else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_IMAGE_CLASS_CHECKPOINT, + output_type=ImageClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.dinov2( + pixel_values, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] # batch_size, sequence_length, hidden_size + + cls_token = sequence_output[:, 0] + patch_tokens = sequence_output[:, 1:] + + linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1) + + logits = self.classifier(linear_input) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Dinov2 backbone, to be used with frameworks like DETR and MaskFormer. + """, + DINOV2_START_DOCSTRING, +) +class Dinov2Backbone(Dinov2PreTrainedModel, BackboneMixin): + def __init__(self, config): + super().__init__(config) + super()._init_backbone(config) + + self.num_features = [ + config.hidden_size for _ in range(config.num_hidden_layers + 1) + ] + self.embeddings = Dinov2Embeddings(config) + self.encoder = Dinov2Encoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> Dinov2PatchEmbeddings: + return self.embeddings.patch_embeddings + + @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + pixel_values: torch.Tensor, + output_hidden_states: Optional[bool] = None, + output_attentions: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> BackboneOutput: + """ + Returns: + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoBackbone + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base") + >>> model = AutoBackbone.from_pretrained( + ... "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"] + ... ) + + >>> inputs = processor(image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> feature_maps = outputs.feature_maps + >>> list(feature_maps[-1].shape) + [1, 768, 16, 16] + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + + embedding_output = self.embeddings(pixel_values) + + outputs = self.encoder( + embedding_output, + output_hidden_states=True, + output_attentions=output_attentions, + return_dict=return_dict, + ) + + hidden_states = outputs.hidden_states if return_dict else outputs[1] + + feature_maps = () + for stage, hidden_state in zip(self.stage_names, hidden_states): + if stage in self.out_features: + if self.config.apply_layernorm: + hidden_state = self.layernorm(hidden_state) + if self.config.reshape_hidden_states: + batch_size, _, height, width = pixel_values.shape + patch_size = self.config.patch_size + hidden_state = hidden_state[:, 1:, :].reshape( + batch_size, width // patch_size, height // patch_size, -1 + ) + hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() + feature_maps += (hidden_state,) + + if not return_dict: + if output_hidden_states: + output = (feature_maps,) + outputs[1:] + else: + output = (feature_maps,) + outputs[2:] + return output + + return BackboneOutput( + feature_maps=feature_maps, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=outputs.attentions if output_attentions else None, + ) + + +class CustomPatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__( + self, image_size: int, patch_size: int, num_channels: int, hidden_size: int + ): + super().__init__() + + image_size = ( + image_size + if isinstance(image_size, collections.abc.Iterable) + else (image_size, image_size) + ) + patch_size = ( + patch_size + if isinstance(patch_size, collections.abc.Iterable) + else (patch_size, patch_size) + ) + num_patches = (image_size[1] // patch_size[1]) * ( + image_size[0] // patch_size[0] + ) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d( + num_channels, hidden_size, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + num_channels = pixel_values.shape[1] + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + f" Expected {self.num_channels} but got {num_channels}." + ) + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +class CustomEmbeddings(nn.Module): + """ + Construct the CLS token, mask token, position and patch embeddings. + """ + + def __init__( + self, image_size: int, patch_size: int, num_channels: int, hidden_size: int + ) -> None: + super().__init__() + + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.hidden_size = hidden_size + + self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size)) + + self.patch_embeddings = CustomPatchEmbeddings( + image_size, patch_size, num_channels, hidden_size + ) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter( + torch.randn(1, num_patches + 1, self.hidden_size) + ) + + def interpolate_pos_encoding( + self, embeddings: torch.Tensor, height: int, width: int + ) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher + resolution images. + + Source: + https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + if num_patches == num_positions and height == width: + return self.position_embeddings + class_pos_embed = self.position_embeddings[:, 0] + patch_pos_embed = self.position_embeddings[:, 1:] + dim = embeddings.shape[-1] + height = height // self.patch_size + width = width // self.patch_size + # we add a small number to avoid floating point error in the interpolation + # see discussion at https://github.com/facebookresearch/dino/issues/8 + height, width = height + 0.1, width + 0.1 + patch_pos_embed = patch_pos_embed.reshape( + 1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim + ) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + scale_factor=( + height / math.sqrt(num_positions), + width / math.sqrt(num_positions), + ), + mode="bicubic", + align_corners=False, + ) + if ( + int(height) != patch_pos_embed.shape[-2] + or int(width) != patch_pos_embed.shape[-1] + ): + raise ValueError( + "Width or height does not match with the interpolated position embeddings" + ) + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) + + def forward( + self, + pixel_values: torch.Tensor, + ) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + patch_embeddings = self.patch_embeddings(pixel_values) + embeddings = patch_embeddings + + # add the [CLS] token to the embedded patch tokens + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, height, width + ) + + return embeddings diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/image.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/image.py new file mode 100644 index 0000000..1663470 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/image.py @@ -0,0 +1,268 @@ +from dataclasses import dataclass +import random + +import torch +import torch.nn as nn +from einops import rearrange + +from ...utils.base import BaseModule +from .dinov2 import Dinov2Model, CustomEmbeddings +from ..transformers.attention import Modulation +from ...utils.typing import * + + +class NaiveImageTokenizer(BaseModule): + @dataclass + class Config(BaseModule.Config): + num_tokens: int = 1024 + num_channels: int = 768 + + cfg: Config + + def configure(self) -> None: + super().configure() + + def forward(self, images: Float[Tensor, "B N C H W"]) -> Float[Tensor, "B Ct Nt"]: + return torch.rand( + ( + images.shape[0], + self.cfg.num_channels, + self.cfg.num_tokens, + ), + device=images.device, + dtype=images.dtype, + ) + + def detokenize(self, *args, **kwargs): + raise NotImplementedError + + +class DINOV2SingleImageTokenizer(BaseModule): + @dataclass + class Config(BaseModule.Config): + pretrained_model_name_or_path: str = "facebook/dinov2-base" + width: int = 224 + height: int = 224 + modulation: bool = False + modulation_zero_init: bool = False + modulation_single_layer: bool = False + modulation_cond_dim: int = 16 + freeze_backbone_params: bool = True + enable_memory_efficient_attention: bool = False + enable_gradient_checkpointing: bool = False + use_patch_embeddings: bool = False + patch_embeddings_aggr_method: str = "concat" + append_plucker_rays: bool = False + drop_rate: float = 0.0 + drop_type: str = "all_but_first" + + cfg: Config + + def configure(self) -> None: + super().configure() + model: Dinov2Model + + if self.cfg.freeze_backbone_params: + # freeze dino backbone parameters + self.register_non_module( + "model", + Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path).to( + self.device + ), + ) + + model = self.non_module("model") + for p in model.parameters(): + p.requires_grad_(False) + model.eval() + else: + self.model = Dinov2Model.from_pretrained( + self.cfg.pretrained_model_name_or_path + ).to(self.device) + model = self.model + + if self.cfg.append_plucker_rays: + model.expand_input_channels(6) + + model.set_use_memory_efficient_attention_xformers( + self.cfg.enable_memory_efficient_attention + ) + model.set_gradient_checkpointing(self.cfg.enable_gradient_checkpointing) + + # add modulation + if self.cfg.modulation: + modulations = [] + for layer in model.encoder.layer: + norm1_modulation = Modulation( + model.config.hidden_size, + self.cfg.modulation_cond_dim, + zero_init=self.cfg.modulation_zero_init, + single_layer=self.cfg.modulation_single_layer, + ) + norm2_modulation = Modulation( + model.config.hidden_size, + self.cfg.modulation_cond_dim, + zero_init=self.cfg.modulation_zero_init, + single_layer=self.cfg.modulation_single_layer, + ) + layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation) + modulations += [norm1_modulation, norm2_modulation] + self.modulations = nn.ModuleList(modulations) + + self.register_buffer( + "image_mean", + torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1), + persistent=False, + ) + self.register_buffer( + "image_std", + torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1), + persistent=False, + ) + + def forward( + self, + images: Float[Tensor, "B *N C H W"], + modulation_cond: Optional[Float[Tensor, "B *N Cc"]], + plucker_rays: Optional[Float[Tensor, "B *N 6 H W"]], + **kwargs + ) -> Float[Tensor, "B *N Ct Nt"]: + model: Dinov2Model + if self.cfg.freeze_backbone_params: + model = self.non_module("model") + else: + model = self.model + + packed = False + if images.ndim == 4: + packed = True + images = images.unsqueeze(1) + if modulation_cond is not None: + assert modulation_cond.ndim == 2 + modulation_cond = modulation_cond.unsqueeze(1) + if plucker_rays is not None: + assert plucker_rays.ndim == 4 + plucker_rays = plucker_rays.unsqueeze(1) + + if ( + self.training + and self.cfg.drop_rate > 0 + and random.random() < self.cfg.drop_rate + ): + if self.cfg.drop_type == "all_but_first": + drop_func = lambda x: x if x is None else x[:, 0:1] + images = drop_func(images) + modulation_cond = drop_func(modulation_cond) + plucker_rays = drop_func(plucker_rays) + else: + raise NotImplementedError + + batch_size, n_input_views = images.shape[:2] + images = (images - self.image_mean) / self.image_std + if self.cfg.append_plucker_rays and plucker_rays is not None: + images = torch.cat([images, plucker_rays], dim=2) + out = model( + rearrange(images, "B N C H W -> (B N) C H W"), + modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc") + if modulation_cond is not None + else None, + ) + local_features, global_features = out.last_hidden_state, out.pooler_output + if self.cfg.use_patch_embeddings: + patch_embeddings = out.patch_embeddings + if self.cfg.patch_embeddings_aggr_method == "concat": + local_features = torch.cat([local_features, patch_embeddings], dim=1) + elif self.cfg.patch_embeddings_aggr_method == "add": + local_features = local_features + patch_embeddings + else: + raise NotImplementedError + local_features = local_features.permute(0, 2, 1) + local_features = rearrange( + local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size + ) + if packed: + local_features = local_features.squeeze(1) + + return local_features + + def detokenize(self, *args, **kwargs): + raise NotImplementedError + + +class DINOV2CustomLowLevelSingleImageTokenizer(DINOV2SingleImageTokenizer): + @dataclass + class Config(DINOV2SingleImageTokenizer.Config): + custom_embeddings_aggr_method: str = "concat" + custom_embeddings_scale: float = 1.0 + + cfg: Config + + def configure(self) -> None: + super().configure() + self.custom_embeddings = CustomEmbeddings( + self.model.config.image_size, + self.model.config.patch_size, + self.model.config.num_channels, + self.model.config.hidden_size, + ) + self.custom_embeddings.load_state_dict( + self.model.embeddings.state_dict(), strict=False + ) + + def forward( + self, + images: Float[Tensor, "B *N C H W"], + modulation_cond: Optional[Float[Tensor, "B *N Cc"]], + **kwargs + ) -> Float[Tensor, "B *N Ct Nt"]: + model: Dinov2Model + if self.cfg.freeze_backbone_params: + model = self.non_module("model") + else: + model = self.model + + packed = False + if images.ndim == 4: + packed = True + images = images.unsqueeze(1) + if modulation_cond is not None: + assert modulation_cond.ndim == 2 + modulation_cond = modulation_cond.unsqueeze(1) + + batch_size, n_input_views = images.shape[:2] + images = (images - self.image_mean) / self.image_std + images = rearrange(images, "B N C H W -> (B N) C H W") + out = model( + images, + modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc") + if modulation_cond is not None + else None, + ) + local_features, global_features = out.last_hidden_state, out.pooler_output + if self.cfg.use_patch_embeddings: + patch_embeddings = out.patch_embeddings + if self.cfg.patch_embeddings_aggr_method == "concat": + local_features = torch.cat([local_features, patch_embeddings], dim=1) + elif self.cfg.patch_embeddings_aggr_method == "add": + local_features = local_features + patch_embeddings + else: + raise NotImplementedError + + custom_embeddings = ( + self.custom_embeddings(images) * self.cfg.custom_embeddings_scale + ) + if self.cfg.custom_embeddings_aggr_method == "concat": + local_features = torch.cat([local_features, custom_embeddings], dim=1) + elif self.cfg.custom_embeddings_aggr_method == "add": + local_features = local_features + custom_embeddings + else: + raise NotImplementedError + + local_features = local_features.permute(0, 2, 1) + local_features = rearrange( + local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size + ) + if packed: + local_features = local_features.squeeze(1) + + return local_features diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/triplane.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/triplane.py new file mode 100644 index 0000000..cfbbba5 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/tokenizers/triplane.py @@ -0,0 +1,49 @@ +from dataclasses import dataclass +import math + +import torch +import torch.nn as nn +from einops import rearrange, repeat + +from ...utils.base import BaseModule +from ...utils.typing import * + + +class TriplaneLearnablePositionalEmbedding(BaseModule): + @dataclass + class Config(BaseModule.Config): + plane_size: int = 32 + num_channels: int = 1024 + + cfg: Config + + def configure(self) -> None: + super().configure() + self.embeddings = nn.Parameter( + torch.randn( + (3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size), + dtype=torch.float32, + ) + * 1 + / math.sqrt(self.cfg.num_channels) + ) + + def forward(self, batch_size: int) -> Float[Tensor, "B Ct Nt"]: + return rearrange( + repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size), + "B Np Ct Hp Wp -> B Ct (Np Hp Wp)", + ) + + def detokenize( + self, tokens: Float[Tensor, "B Ct Nt"] + ) -> Float[Tensor, "B 3 Ct Hp Wp"]: + batch_size, Ct, Nt = tokens.shape + assert Nt == self.cfg.plane_size**2 * 3 + assert Ct == self.cfg.num_channels + return rearrange( + tokens, + "B Ct (Np Hp Wp) -> B Np Ct Hp Wp", + Np=3, + Hp=self.cfg.plane_size, + Wp=self.cfg.plane_size, + ) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/attention.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/attention.py new file mode 100644 index 0000000..1b5891f --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/attention.py @@ -0,0 +1,669 @@ +# 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. + +import torch +import torch.nn.functional as F +from torch import nn + +from diffusers.utils.torch_utils import maybe_allow_in_graph +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings + +from ...utils.typing import * + + +class MemoryEfficientAttentionMixin: + def enable_xformers_memory_efficient_attention( + self, attention_op: Optional[Callable] = None + ): + r""" + Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this + option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed + up during training is not guaranteed. + + + + ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes + precedent. + + + + Parameters: + attention_op (`Callable`, *optional*): + Override the default `None` operator for use as `op` argument to the + [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) + function of xFormers. + + Examples: + + ```py + >>> import torch + >>> from diffusers import DiffusionPipeline + >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp + + >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) + >>> pipe = pipe.to("cuda") + >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) + >>> # Workaround for not accepting attention shape using VAE for Flash Attention + >>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None) + ``` + """ + self.set_use_memory_efficient_attention_xformers(True, attention_op) + + def disable_xformers_memory_efficient_attention(self): + r""" + Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). + """ + self.set_use_memory_efficient_attention_xformers(False) + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ) -> None: + # Recursively walk through all the children. + # Any children which exposes the set_use_memory_efficient_attention_xformers method + # gets the message + def fn_recursive_set_mem_eff(module: torch.nn.Module): + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid, attention_op) + + for child in module.children(): + fn_recursive_set_mem_eff(child) + + for module in self.children(): + if isinstance(module, torch.nn.Module): + fn_recursive_set_mem_eff(module) + + +@maybe_allow_in_graph +class GatedSelfAttentionDense(nn.Module): + r""" + A gated self-attention dense layer that combines visual features and object features. + + Parameters: + query_dim (`int`): The number of channels in the query. + context_dim (`int`): The number of channels in the context. + n_heads (`int`): The number of heads to use for attention. + d_head (`int`): The number of channels in each head. + """ + + def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int): + super().__init__() + + # we need a linear projection since we need cat visual feature and obj feature + self.linear = nn.Linear(context_dim, query_dim) + + self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head) + self.ff = FeedForward(query_dim, activation_fn="geglu") + + self.norm1 = nn.LayerNorm(query_dim) + self.norm2 = nn.LayerNorm(query_dim) + + self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0))) + self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0))) + + self.enabled = True + + def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor: + if not self.enabled: + return x + + n_visual = x.shape[1] + objs = self.linear(objs) + + x = ( + x + + self.alpha_attn.tanh() + * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :] + ) + x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) + + return x + + +@maybe_allow_in_graph +class BasicTransformerBlock(nn.Module, MemoryEfficientAttentionMixin): + 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. + 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. + 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. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + """ + + 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, + cond_dim_ada_norm_continuous: 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, + attention_type: str = "default", + ): + 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" + self.use_ada_layer_norm_continuous = ( + cond_dim_ada_norm_continuous is not None + ) and norm_type == "ada_norm_continuous" + + assert ( + int(self.use_ada_layer_norm) + + int(self.use_ada_layer_norm_continuous) + + int(self.use_ada_layer_norm_zero) + <= 1 + ) + + 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_continuous: + self.norm1 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous) + 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.attn1 = Attention( + 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, + ) + + # 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. + if self.use_ada_layer_norm: + self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_continuous: + self.norm2 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous) + else: + self.norm2 = 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 + if self.use_ada_layer_norm_continuous: + self.norm3 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous) + else: + self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + ) + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense( + dim, cross_attention_dim, num_attention_heads, attention_head_dim + ) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + 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, + modulation_cond: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm1(hidden_states, modulation_cond) + 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) + + # 1. Retrieve lora scale. + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = ( + cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + ) + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if self.only_cross_attention + else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + # 2.5 ends + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm2(hidden_states, modulation_cond) + else: + norm_hidden_states = 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 + + # 4. Feed-forward + if self.use_ada_layer_norm_continuous: + norm_hidden_states = self.norm3(hidden_states, modulation_cond) + else: + 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, scale=lora_scale) + 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, scale=lora_scale) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + +class FeedForward(nn.Module): + r""" + A feed-forward layer. + + Parameters: + dim (`int`): The number of channels in the input. + dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. + mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + linear_cls = nn.Linear + + if activation_fn == "gelu": + act_fn = GELU(dim, inner_dim) + if activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh") + elif activation_fn == "geglu": + act_fn = GEGLU(dim, inner_dim) + elif activation_fn == "geglu-approximate": + act_fn = ApproximateGELU(dim, inner_dim) + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # project dropout + self.net.append(nn.Dropout(dropout)) + # project out + self.net.append(linear_cls(inner_dim, dim_out)) + # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +class GELU(nn.Module): + r""" + GELU activation function with tanh approximation support with `approximate="tanh"`. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. + """ + + def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + self.approximate = approximate + + def gelu(self, gate: torch.Tensor) -> torch.Tensor: + if gate.device.type != "mps": + return F.gelu(gate, approximate=self.approximate) + # mps: gelu is not implemented for float16 + return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to( + dtype=gate.dtype + ) + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + hidden_states = self.gelu(hidden_states) + return hidden_states + + +class GEGLU(nn.Module): + r""" + A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + """ + + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + linear_cls = nn.Linear + + self.proj = linear_cls(dim_in, dim_out * 2) + + def gelu(self, gate: torch.Tensor) -> torch.Tensor: + if gate.device.type != "mps": + return F.gelu(gate) + # mps: gelu is not implemented for float16 + return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype) + + def forward(self, hidden_states, scale: float = 1.0): + args = () + hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1) + return hidden_states * self.gelu(gate) + + +class ApproximateGELU(nn.Module): + r""" + The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2: + https://arxiv.org/abs/1606.08415. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + """ + + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + return x * torch.sigmoid(1.702 * x) + + +class AdaLayerNorm(nn.Module): + r""" + Norm layer modified to incorporate timestep embeddings. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the dictionary of embeddings. + """ + + def __init__(self, embedding_dim: int, num_embeddings: int): + super().__init__() + self.emb = nn.Embedding(num_embeddings, embedding_dim) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim * 2) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False) + + def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + emb = self.linear(self.silu(self.emb(timestep))) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + return x + + +class AdaLayerNormContinuous(nn.Module): + r""" + Norm layer modified to incorporate arbitrary continuous embeddings. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + """ + + def __init__(self, embedding_dim: int, condition_dim: int): + super().__init__() + self.silu = nn.SiLU() + self.linear1 = nn.Linear(condition_dim, condition_dim) + self.linear2 = nn.Linear(condition_dim, embedding_dim * 2) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False) + + def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: + emb = self.linear2(self.silu(self.linear1(condition))) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + return x + + +class Modulation(nn.Module): + def __init__(self, embedding_dim: int, condition_dim: int, zero_init: bool = False, single_layer: bool = False): + super().__init__() + self.silu = nn.SiLU() + if single_layer: + self.linear1 = nn.Identity() + else: + self.linear1 = nn.Linear(condition_dim, condition_dim) + + self.linear2 = nn.Linear(condition_dim, embedding_dim * 2) + + # Only zero init the last linear layer + if zero_init: + nn.init.zeros_(self.linear2.weight) + nn.init.zeros_(self.linear2.bias) + + def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor: + emb = self.linear2(self.silu(self.linear1(condition))) + scale, shift = torch.chunk(emb, 2, dim=1) + x = x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + return x + + +class AdaLayerNormZero(nn.Module): + r""" + Norm layer adaptive layer norm zero (adaLN-Zero). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the dictionary of embeddings. + """ + + def __init__(self, embedding_dim: int, num_embeddings: int): + super().__init__() + + self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, + x: torch.Tensor, + timestep: torch.Tensor, + class_labels: torch.LongTensor, + hidden_dtype: Optional[torch.dtype] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear( + self.silu(self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)) + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk( + 6, dim=1 + ) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaGroupNorm(nn.Module): + r""" + GroupNorm layer modified to incorporate timestep embeddings. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the dictionary of embeddings. + num_groups (`int`): The number of groups to separate the channels into. + act_fn (`str`, *optional*, defaults to `None`): The activation function to use. + eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability. + """ + + def __init__( + self, + embedding_dim: int, + out_dim: int, + num_groups: int, + act_fn: Optional[str] = None, + eps: float = 1e-5, + ): + super().__init__() + self.num_groups = num_groups + self.eps = eps + + if act_fn is None: + self.act = None + else: + self.act = get_activation(act_fn) + + self.linear = nn.Linear(embedding_dim, out_dim * 2) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + if self.act: + emb = self.act(emb) + emb = self.linear(emb) + emb = emb[:, :, None, None] + scale, shift = emb.chunk(2, dim=1) + + x = F.group_norm(x, self.num_groups, eps=self.eps) + x = x * (1 + scale) + shift + return x diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_1d.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_1d.py new file mode 100644 index 0000000..836ac96 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_1d.py @@ -0,0 +1,249 @@ +from dataclasses import dataclass + +import torch +from torch import nn + +from ...utils.base import BaseModule +from .attention import ( + BasicTransformerBlock, + MemoryEfficientAttentionMixin, +) +from ...utils.typing import * + + +class Transformer1D(BaseModule, MemoryEfficientAttentionMixin): + """ + A 1D Transformer model for sequence 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. + 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. + """ + + @dataclass + class Config(BaseModule.Config): + 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 + activation_fn: str = "geglu" + num_embeds_ada_norm: Optional[int] = None + cond_dim_ada_norm_continuous: Optional[int] = None + only_cross_attention: bool = False + double_self_attention: bool = False + upcast_attention: bool = False + norm_type: str = "layer_norm" + norm_elementwise_affine: bool = True + attention_type: str = "default" + enable_memory_efficient_attention: bool = False + gradient_checkpointing: bool = False + + cfg: Config + + def configure(self) -> None: + super().configure() + + self.num_attention_heads = self.cfg.num_attention_heads + self.attention_head_dim = self.cfg.attention_head_dim + inner_dim = self.num_attention_heads * self.attention_head_dim + + linear_cls = nn.Linear + + if self.cfg.norm_type == "layer_norm" and ( + self.cfg.num_embeds_ada_norm is not None + or self.cfg.cond_dim_ada_norm_continuous is not None + ): + raise ValueError("Incorrect norm_type.") + + # 2. Define input layers + self.in_channels = self.cfg.in_channels + + self.norm = torch.nn.GroupNorm( + num_groups=self.cfg.norm_num_groups, + num_channels=self.cfg.in_channels, + eps=1e-6, + affine=True, + ) + self.proj_in = linear_cls(self.cfg.in_channels, inner_dim) + + # 3. Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + self.num_attention_heads, + self.attention_head_dim, + dropout=self.cfg.dropout, + cross_attention_dim=self.cfg.cross_attention_dim, + activation_fn=self.cfg.activation_fn, + num_embeds_ada_norm=self.cfg.num_embeds_ada_norm, + cond_dim_ada_norm_continuous=self.cfg.cond_dim_ada_norm_continuous, + attention_bias=self.cfg.attention_bias, + only_cross_attention=self.cfg.only_cross_attention, + double_self_attention=self.cfg.double_self_attention, + upcast_attention=self.cfg.upcast_attention, + norm_type=self.cfg.norm_type, + norm_elementwise_affine=self.cfg.norm_elementwise_affine, + attention_type=self.cfg.attention_type, + ) + for d in range(self.cfg.num_layers) + ] + ) + + # 4. Define output layers + self.out_channels = ( + self.cfg.in_channels + if self.cfg.out_channels is None + else self.cfg.out_channels + ) + + self.proj_out = linear_cls(inner_dim, self.cfg.in_channels) + + self.gradient_checkpointing = self.cfg.gradient_checkpointing + + self.set_use_memory_efficient_attention_xformers( + self.cfg.enable_memory_efficient_attention + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + modulation_cond: Optional[torch.FloatTensor] = 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, + ): + """ + The [`Transformer1DModel`] 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`. + cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + 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 + batch, _, seq_len = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 1).reshape( + batch, seq_len, inner_dim + ) + hidden_states = self.proj_in(hidden_states) + + # 2. Blocks + for block in self.transformer_blocks: + if self.training and self.gradient_checkpointing: + hidden_states = torch.utils.checkpoint.checkpoint( + block, + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + modulation_cond, + cross_attention_kwargs, + class_labels, + use_reentrant=False, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + modulation_cond=modulation_cond, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + ) + + # 3. Output + hidden_states = self.proj_out(hidden_states) + hidden_states = ( + hidden_states.reshape(batch, seq_len, inner_dim) + .permute(0, 2, 1) + .contiguous() + ) + + output = hidden_states + residual + + return output diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_2d.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_2d.py new file mode 100644 index 0000000..c2e16d7 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/transformers/transformer_2d.py @@ -0,0 +1,414 @@ +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn +from diffusers.models.embeddings import ImagePositionalEmbeddings, PatchEmbed + +from CharacterGen.Stage_3D import lrm +from ...utils.base import BaseModule +from .attention import ( + BasicTransformerBlock, + MemoryEfficientAttentionMixin, +) +from ...utils.typing import * + + +class Transformer2D(BaseModule, MemoryEfficientAttentionMixin): + """ + 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. + """ + + @dataclass + class Config(BaseModule.Config): + 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 + double_self_attention: bool = False + upcast_attention: bool = False + norm_type: str = "layer_norm" + norm_elementwise_affine: bool = True + attention_type: str = "default" + enable_memory_efficient_attention: bool = False + gradient_checkpointing: bool = False + + cfg: Config + + def configure(self) -> None: + super().configure() + + self.use_linear_projection = self.cfg.use_linear_projection + self.num_attention_heads = self.cfg.num_attention_heads + self.attention_head_dim = self.cfg.attention_head_dim + inner_dim = self.num_attention_heads * self.attention_head_dim + + conv_cls = nn.Conv2d + linear_cls = nn.Linear + + # 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 = (self.cfg.in_channels is not None) and ( + self.cfg.patch_size is None + ) + self.is_input_vectorized = self.cfg.num_vector_embeds is not None + self.is_input_patches = ( + self.cfg.in_channels is not None and self.cfg.patch_size is not None + ) + + if ( + self.cfg.norm_type == "layer_norm" + and self.cfg.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" + ) + lrm.warn(deprecation_message) + self.cfg.norm_type = "ada_norm" + + if self.is_input_continuous and self.is_input_vectorized: + raise ValueError( + f"Cannot define both `in_channels`: {self.cfg.in_channels} and `num_vector_embeds`: {self.cfg.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`: {self.cfg.num_vector_embeds} and `patch_size`: {self.cfg.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`: {self.cfg.in_channels}, `num_vector_embeds`: {self.cfg.num_vector_embeds}, or patch_size:" + f" {self.cfg.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 = self.cfg.in_channels + + self.norm = torch.nn.GroupNorm( + num_groups=self.cfg.norm_num_groups, + num_channels=self.cfg.in_channels, + eps=1e-6, + affine=True, + ) + if self.cfg.use_linear_projection: + self.proj_in = linear_cls(self.cfg.in_channels, inner_dim) + else: + self.proj_in = conv_cls( + self.cfg.in_channels, inner_dim, kernel_size=1, stride=1, padding=0 + ) + elif self.is_input_vectorized: + assert ( + self.cfg.sample_size is not None + ), "Transformer2DModel over discrete input must provide sample_size" + assert ( + self.cfg.num_vector_embeds is not None + ), "Transformer2DModel over discrete input must provide num_embed" + + self.height = self.cfg.sample_size + self.width = self.cfg.sample_size + self.num_vector_embeds = self.cfg.num_vector_embeds + self.num_latent_pixels = self.height * self.width + + self.latent_image_embedding = ImagePositionalEmbeddings( + num_embed=self.cfg.num_vector_embeds, + embed_dim=inner_dim, + height=self.height, + width=self.width, + ) + elif self.is_input_patches: + assert ( + self.cfg.sample_size is not None + ), "Transformer2DModel over patched input must provide sample_size" + + self.height = self.cfg.sample_size + self.width = self.cfg.sample_size + + self.patch_size = self.cfg.patch_size + self.pos_embed = PatchEmbed( + height=self.cfg.sample_size, + width=self.cfg.sample_size, + patch_size=self.cfg.patch_size, + in_channels=self.cfg.in_channels, + embed_dim=inner_dim, + ) + + # 3. Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + self.num_attention_heads, + self.attention_head_dim, + dropout=self.cfg.dropout, + cross_attention_dim=self.cfg.cross_attention_dim, + activation_fn=self.cfg.activation_fn, + num_embeds_ada_norm=self.cfg.num_embeds_ada_norm, + attention_bias=self.cfg.attention_bias, + only_cross_attention=self.cfg.only_cross_attention, + double_self_attention=self.cfg.double_self_attention, + upcast_attention=self.cfg.upcast_attention, + norm_type=self.cfg.norm_type, + norm_elementwise_affine=self.cfg.norm_elementwise_affine, + attention_type=self.cfg.attention_type, + ) + for d in range(self.cfg.num_layers) + ] + ) + + # 4. Define output layers + self.out_channels = ( + self.cfg.in_channels + if self.cfg.out_channels is None + else self.cfg.out_channels + ) + if self.is_input_continuous: + # TODO: should use out_channels for continuous projections + if self.cfg.use_linear_projection: + self.proj_out = linear_cls(inner_dim, self.cfg.in_channels) + else: + self.proj_out = conv_cls( + inner_dim, self.cfg.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, self.cfg.patch_size * self.cfg.patch_size * self.out_channels + ) + + self.gradient_checkpointing = self.cfg.gradient_checkpointing + + self.set_use_memory_efficient_attention_xformers( + self.cfg.enable_memory_efficient_attention + ) + + 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, + ): + """ + 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`. + cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + 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: + if self.training and self.gradient_checkpointing: + hidden_states = torch.utils.checkpoint.checkpoint( + block, + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + cross_attention_kwargs, + class_labels, + use_reentrant=False, + ) + else: + 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, + ) + ) + + return output diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/base.py new file mode 100644 index 0000000..4026237 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/base.py @@ -0,0 +1,283 @@ +import os +from dataclasses import dataclass, field + +import pytorch_lightning as pl +import torch.nn.functional as F + +from CharacterGen.Stage_3D import lrm +from .utils import parse_optimizer, parse_scheduler +from ..utils.base import ( + Updateable, + update_end_if_possible, + update_if_possible, +) +from ..models.exporters.base import Exporter, ExporterOutput +from ..utils.config import parse_structured +from ..utils.misc import C, cleanup, get_device, load_module_weights +#from ..utils.saving import SaverMixin +from ..utils.typing import * + + +@dataclass +class BaseLossConfig: + pass + +#class BaseSystem(pl.LightningModule, Updateable, SaverMixin): +class BaseSystem(pl.LightningModule, Updateable): + @dataclass + class Config: + #loss: BaseLossConfig = BaseLossConfig() + optimizer: dict = field(default_factory=dict) + scheduler: Optional[dict] = None + weights: Optional[str] = None + weights_ignore_modules: Optional[List[str]] = None + weights_mapping: Optional[List[Dict[str, str]]] = None + check_train_every_n_steps: int = 0 + check_val_limit_rank: int = 8 + cleanup_after_validation_step: bool = False + cleanup_after_test_step: bool = False + + exporter_cls: str = "CharacterGen.Stage_3D.lrm.models.exporters.mesh_exporter.MeshExporter" + exporter: dict = field(default_factory=lambda: {"fmt": "obj", "save_uv": False}) + + cfg: Config + + def __init__(self, cfg, resumed=False) -> None: + super().__init__() + self.cfg = parse_structured(self.Config, cfg) + self._save_dir: Optional[str] = None + self._resumed: bool = resumed + self._resumed_eval: bool = False + self._resumed_eval_status: dict = {"global_step": 0, "current_epoch": 0} + + self.configure() + print(self.cfg.weights) + if self.cfg.weights is not None: + self.load_weights( + self.cfg.weights, + self.cfg.weights_ignore_modules, + self.cfg.weights_mapping, + ) + print("finish loading!!") + self.post_configure() + + def load_weights( + self, + weights: str, + ignore_modules: Optional[List[str]] = None, + mapping: Optional[List[Dict[str, str]]] = None, + ): + state_dict, epoch, global_step = load_module_weights( + weights, + ignore_modules=ignore_modules, + mapping=mapping, + map_location="cpu", + ) + self.load_state_dict(state_dict, strict=False) + # restore step-dependent states + self.do_update_step(epoch, global_step, on_load_weights=True) + + def set_resume_status(self, current_epoch: int, global_step: int): + # restore correct epoch and global step in eval + self._resumed_eval = True + self._resumed_eval_status["current_epoch"] = current_epoch + self._resumed_eval_status["global_step"] = global_step + + @property + def resumed(self): + # whether from resumed checkpoint + return self._resumed + + @property + def true_global_step(self): + if self._resumed_eval: + return self._resumed_eval_status["global_step"] + else: + return self.global_step + + @property + def true_current_epoch(self): + if self._resumed_eval: + return self._resumed_eval_status["current_epoch"] + else: + return self.current_epoch + + def configure(self) -> None: + pass + + def post_configure(self) -> None: + """ + executed after weights are loaded + """ + pass + + def C(self, value: Any) -> float: + return C(value, self.true_current_epoch, self.true_global_step) + + def configure_optimizers(self): + optim = parse_optimizer(self.cfg.optimizer, self) + ret = { + "optimizer": optim, + } + if self.cfg.scheduler is not None: + ret.update( + { + "lr_scheduler": parse_scheduler(self.cfg.scheduler, optim), + } + ) + return ret + + def on_fit_start(self) -> None: + if self._save_dir is not None: + lrm.info(f"Validation results will be saved to {self._save_dir}") + else: + lrm.warn( + f"Saving directory not set for the system, visualization results will not be saved" + ) + + def training_step(self, batch, batch_idx): + raise NotImplementedError + + def check_train(self, batch, outputs, **kwargs): + if ( + self.global_rank == 0 + and self.cfg.check_train_every_n_steps > 0 + and self.true_global_step % self.cfg.check_train_every_n_steps == 0 + ): + self.on_check_train(batch, outputs, **kwargs) + + def on_check_train(self, batch, outputs, **kwargs): + pass + + def validation_step(self, batch, batch_idx): + raise NotImplementedError + + def on_validation_epoch_end(self): + pass + + def test_step(self, batch, batch_idx): + raise NotImplementedError + + def on_test_epoch_end(self): + pass + + def on_test_end(self) -> None: + if self._save_dir is not None: + lrm.info(f"Test results saved to {self._save_dir}") + + def on_predict_start(self) -> None: + pass + + def predict_step(self, batch, batch_idx): + batch_size = batch["index"].shape[0] + scene_codes = self(batch) + for b in range(batch_size): + if batch["view_index"][b, 0] == 0: + exporter_output: List[ExporterOutput] = self.exporter( + batch["index"][b][None], scene_codes[b][None] + ) + for out in exporter_output: + save_func_name = f"save_{out.save_type}" + if not hasattr(self, save_func_name): + raise ValueError( + f"{save_func_name} not supported by the SaverMixin" + ) + save_func = getattr(self, save_func_name) + save_func( + f"it{self.true_global_step}-export/{out.save_name}", + **out.params, + ) + if self.exporter.cfg.save_video: + self.test_step(batch, batch_idx) + + def on_predict_epoch_end(self) -> None: + if self.exporter.cfg.save_video: + self.on_test_epoch_end() + + def on_predict_end(self) -> None: + if self._save_dir is not None: + lrm.info(f"Export assets saved to {self._save_dir}") + + def preprocess_data(self, batch, stage): + pass + + """ + Implementing on_after_batch_transfer of DataModule does the same. + But on_after_batch_transfer does not support DP. + """ + + def on_train_batch_start(self, batch, batch_idx, unused=0): + self.preprocess_data(batch, "train") + self.dataset = self.trainer.train_dataloader.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_validation_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "validation") + self.dataset = self.trainer.val_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_test_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "test") + self.dataset = self.trainer.test_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_predict_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "predict") + self.dataset = self.trainer.predict_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_train_batch_end(self, outputs, batch, batch_idx): + self.dataset = self.trainer.train_dataloader.dataset + update_end_if_possible( + self.dataset, self.true_current_epoch, self.true_global_step + ) + self.do_update_step_end(self.true_current_epoch, self.true_global_step) + + def on_validation_batch_end(self, outputs, batch, batch_idx): + self.dataset = self.trainer.val_dataloaders.dataset + update_end_if_possible( + self.dataset, self.true_current_epoch, self.true_global_step + ) + self.do_update_step_end(self.true_current_epoch, self.true_global_step) + if self.cfg.cleanup_after_validation_step: + # cleanup to save vram + cleanup() + + def on_test_batch_end(self, outputs, batch, batch_idx): + self.dataset = self.trainer.test_dataloaders.dataset + update_end_if_possible( + self.dataset, self.true_current_epoch, self.true_global_step + ) + self.do_update_step_end(self.true_current_epoch, self.true_global_step) + if self.cfg.cleanup_after_test_step: + # cleanup to save vram + cleanup() + + def on_predict_batch_end(self, outputs, batch, batch_idx): + self.dataset = self.trainer.predict_dataloaders.dataset + update_end_if_possible( + self.dataset, self.true_current_epoch, self.true_global_step + ) + self.do_update_step_end(self.true_current_epoch, self.true_global_step) + if self.cfg.cleanup_after_test_step: + # cleanup to save vram + cleanup() + + def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False): + pass + + def on_before_optimizer_step(self, optimizer): + """ + # some gradient-related debugging goes here, example: + from lightning.pytorch.utilities import grad_norm + norms = grad_norm(self.geometry, norm_type=2) + print(norms) + for name, p in self.named_parameters(): + if p.grad is None: + lrm.info(f"{name} does not receive gradients!") + """ + pass diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/multiview_lrm.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/multiview_lrm.py new file mode 100644 index 0000000..e543a5e --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/multiview_lrm.py @@ -0,0 +1,342 @@ +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F +from einops import rearrange + +from CharacterGen.Stage_3D import lrm +from CharacterGen.Stage_3D.lrm.models.mesh import Mesh +from CharacterGen.Stage_3D.lrm.systems.base import BaseLossConfig, BaseSystem +from CharacterGen.Stage_3D.lrm.utils.ops import binary_cross_entropy, get_plucker_rays +from CharacterGen.Stage_3D.lrm.utils.typing import * +from CharacterGen.Stage_3D.lrm.models.lpips import LPIPS +from CharacterGen.Stage_3D.lrm.utils.misc import time_recorder as tr + + +@dataclass +class MultiviewLRMLossConfig(BaseLossConfig): + lambda_mse: Any = 0.0 + lambda_mse_coarse: Any = 0.0 + lambda_smooth_l1: Any = 0.0 + lambda_smooth_l1_coarse: Any = 0.0 + lambda_lpips: Any = 0.0 + lambda_lpips_coarse: Any = 0.0 + lambda_mask: Any = 0.0 + lambda_mask_coarse: Any = 0.0 + + +class MultiviewLRM(BaseSystem): + @dataclass + class Config(BaseSystem.Config): + #loss: MultiviewLRMLossConfig = MultiviewLRMLossConfig() + + camera_embedder_cls: str = "" + camera_embedder: dict = field(default_factory=dict) + + image_tokenizer_cls: str = "" + image_tokenizer: dict = field(default_factory=dict) + + tokenizer_cls: str = "" + tokenizer: dict = field(default_factory=dict) + + backbone_cls: str = "" + backbone: dict = field(default_factory=dict) + + post_processor_cls: str = "" + post_processor: dict = field(default_factory=dict) + + decoder_cls: str = "" + decoder: dict = field(default_factory=dict) + + material_cls: str = "" + material: dict = field(default_factory=dict) + + background_cls: str = "" + background: dict = field(default_factory=dict) + + renderer_cls: str = "" + renderer: dict = field(default_factory=dict) + + resume_ckpt_path: str = "" + + cfg: Config + + def configure(self): + super().configure() + self.image_tokenizer = lrm.find(self.cfg.image_tokenizer_cls)( + self.cfg.image_tokenizer + ) + if self.cfg.image_tokenizer.modulation: + self.camera_embedder = lrm.find(self.cfg.camera_embedder_cls)( + self.cfg.camera_embedder + ) + self.tokenizer = lrm.find(self.cfg.tokenizer_cls)(self.cfg.tokenizer) + self.backbone = lrm.find(self.cfg.backbone_cls)(self.cfg.backbone) + self.post_processor = lrm.find(self.cfg.post_processor_cls)( + self.cfg.post_processor + ) + self.decoder = lrm.find(self.cfg.decoder_cls)(self.cfg.decoder) + self.material = lrm.find(self.cfg.material_cls)(self.cfg.material) + self.background = lrm.find(self.cfg.background_cls)(self.cfg.background) + self.renderer = lrm.find(self.cfg.renderer_cls)( + self.cfg.renderer, self.decoder, self.material, self.background + ) + + self.exporter = lrm.find(self.cfg.exporter_cls)( + self.cfg.exporter, self.renderer + ) + + def on_fit_start(self): + super().on_fit_start() + self.lpips_loss_fn = LPIPS() + + def forward(self, batch: Dict[str, Any]) -> Dict[str, Any]: + # batch["rgb_cond"]: B, N_cond, H, W, 3 + # batch["rgb"]: B, N_render, H, W, 3 + # batch["c2w_cond"]: B, N_cond, 4, 4 + # for single image input (like LRM), N_cond = 1 + + batch_size, n_input_views = batch["rgb_cond"].shape[:2] + + # Camera modulation + camera_embeds: Optional[Float[Tensor, "B Nv Cc"]] + if self.cfg.image_tokenizer.modulation: + camera_embeds = self.camera_embedder(**batch) + else: + camera_embeds = None + + tr.start("image tokenizer") + input_image_tokens: Float[Tensor, "B Nv Cit Nit"] = self.image_tokenizer( + rearrange(batch["rgb_cond"], "B Nv H W C -> B Nv C H W"), + modulation_cond=camera_embeds, + plucker_rays=rearrange( + get_plucker_rays(batch["rays_o_cond"], batch["rays_d_cond"]), + "B Nv H W C -> B Nv C H W", + ) + if "rays_o_cond" in batch + else None, + ) + tr.end("image tokenizer") + + input_image_tokens = rearrange( + input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=n_input_views + ) + + tokens: Float[Tensor, "B Ct Nt"] = self.tokenizer(batch_size) + + tr.start("backbone") + tokens = self.backbone( + tokens, + encoder_hidden_states=input_image_tokens, + modulation_cond=None, + ) + tr.end("backbone") + + scene_codes = self.post_processor(self.tokenizer.detokenize(tokens)) + return scene_codes + + def forward_renderer_nerf( + self, batch: Dict[str, Any], scene_codes + ) -> Dict[str, Any]: + tr.start("render") + render_out = self.renderer(scene_codes, **batch) + tr.end("render") + return render_out + + def training_step(self, batch, batch_idx): + scene_codes = self(batch) + out = self.forward_renderer_nerf(batch, scene_codes) + + loss = 0.0 + + for suffix in ["", "_coarse"]: + if not f"comp_rgb{suffix}" in out: + continue + + comp_rgb: Float[Tensor, "B Nv H W 3"] = out["comp_rgb{}".format(suffix)] + gt_rgb: Float[Tensor, "B Nv H W 3"] = batch["rgb"] + + self.log(f"train/comp_rgb_min{suffix}", comp_rgb.min()) + + loss_mse = F.mse_loss(comp_rgb, gt_rgb, reduction="mean") + self.log(f"train/loss_mse{suffix}", loss_mse) + loss += loss_mse * self.C(self.cfg.loss[f"lambda_mse{suffix}"]) + + loss_smooth_l1 = F.smooth_l1_loss( + comp_rgb, gt_rgb, beta=0.1, reduction="mean" + ) + self.log(f"train/loss_smooth_l1{suffix}", loss_smooth_l1) + loss += loss_smooth_l1 * self.C(self.cfg.loss[f"lambda_smooth_l1{suffix}"]) + + if self.C(self.cfg.loss[f"lambda_lpips{suffix}"]) > 0: + loss_lpips = self.lpips_loss_fn( + rearrange(comp_rgb, "B Nv H W C -> (B Nv) C H W"), + rearrange(gt_rgb, "B Nv H W C -> (B Nv) C H W"), + input_range=(0, 1), + ).mean() + self.log(f"train/loss_lpips{suffix}", loss_lpips) + loss += loss_lpips * self.C(self.cfg.loss[f"lambda_lpips{suffix}"]) + + loss_mask = binary_cross_entropy( + out[f"opacity{suffix}"].clamp(1e-5, 1 - 1e-5), batch["mask"] + ) + self.log(f"train/loss_mask{suffix}", loss_mask) + loss += loss_mask * self.C(self.cfg.loss[f"lambda_mask{suffix}"]) + + for name, value in self.cfg.loss.items(): + self.log(f"train_params/{name}", self.C(value)) + + # will execute self.on_check_train every self.cfg.check_train_every_n_steps steps + self.check_train( + batch, + out, + extra=f"m{loss_mse:.2f}_l{loss_smooth_l1:.2f}_p{loss_lpips:.2f}_ma{loss_mask:.2f}", + ) + + return {"loss": loss} + + def get_input_visualizations(self, batch): + return [ + { + "type": "rgb", + "img": rearrange(batch["rgb_cond"], "B N H W C -> (B H) (N W) C"), + "kwargs": {"data_format": "HWC"}, + } + ] + + def get_output_visualizations(self, batch, outputs): + out = outputs + images = [] + if "rgb" in batch: + images += [ + { + "type": "rgb", + "img": rearrange(batch["rgb"], "B N H W C -> (B H) (N W) C"), + "kwargs": {"data_format": "HWC"}, + }, + { + "type": "grayscale", + "img": rearrange(batch["mask"], "B N H W C -> (B H) (N W) C")[ + ..., 0 + ], + "kwargs": {"cmap": None, "data_range": None}, + }, + ] + for suffix in ["", "_coarse"]: + if not f"comp_rgb{suffix}" in out: + continue + images += [ + { + "type": "rgb", + "img": rearrange( + out[f"comp_rgb{suffix}"], "B N H W C -> (B H) (N W) C" + ), + "kwargs": {"data_format": "HWC"}, + }, + { + "type": "grayscale", + "img": rearrange( + out[f"opacity{suffix}"], "B N H W C -> (B H) (N W) C" + )[..., 0], + "kwargs": {"cmap": None, "data_range": None}, + }, + { + "type": "grayscale", + "img": rearrange( + out[f"depth{suffix}"], "B N H W C -> (B H) (N W) C" + )[..., 0], + "kwargs": {"cmap": None, "data_range": None}, + }, + ] + return images + + # def check_train(self, batch, outputs, **kwargs): + # self.on_check_train(batch, outputs, **kwargs) + + def on_check_train(self, batch, outputs, extra=""): + self.save_image_grid( + f"it{self.true_global_step}-train.jpg", + self.get_output_visualizations(batch, outputs), + name="train_step_output", + step=self.true_global_step, + ) + # self.save_image_grid( + # f"debug/it{self.true_global_step}-{self.global_rank}-{extra}.jpg", + # self.get_output_visualizations(batch, outputs), + # name="train_step_output", + # step=self.true_global_step, + # ) + # self.save_json( + # f"debug_list/it{self.true_global_step}-{self.global_rank}-ids.json", + # batch["scene_id"], + # ) + + def validation_step(self, batch, batch_idx): + scene_codes = self(batch) + out = self.forward_renderer_nerf(batch, scene_codes) + if ( + self.cfg.check_val_limit_rank > 0 + and self.global_rank < self.cfg.check_val_limit_rank + ): + self.save_image_grid( + f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}-input.jpg", + self.get_input_visualizations(batch), + name=f"validation_step_input_{self.global_rank}_{batch_idx}", + step=self.true_global_step, + ) + self.save_image_grid( + f"it{self.true_global_step}-validation-{self.global_rank}_{batch_idx}.jpg", + self.get_output_visualizations(batch, out), + name=f"validation_step_output_{self.global_rank}_{batch_idx}", + step=self.true_global_step, + ) + + def test_step(self, batch, batch_idx): + # not saved to wandb + scene_codes = self(batch) + out = self.forward_renderer_nerf(batch, scene_codes) + batch_size = batch["index"].shape[0] + for b in range(batch_size): + if batch["view_index"][b, 0] == 0: + self.save_image_grid( + f"it{self.true_global_step}-test/{batch['index'][b]}-input.jpg", + [ + { + "type": "rgb", + "img": rearrange( + batch["rgb_cond"][b], "N H W C -> H (N W) C" + ), + "kwargs": {"data_format": "HWC"}, + }, + ], + ) + self.save_image_grid( + f"it{self.true_global_step}-test/{batch['index'][b]}/{batch['view_index'][b,0]}.png", + [ + { + "type": "rgb", + "img": out["comp_rgb"][b][0], + "kwargs": {"data_format": "HWC"}, + }, + { + "type": "grayscale", + "img": out["depth"][b][0, ..., 0], + "kwargs": {"cmap": None, "data_range": None}, + }, + ], + ) + + def on_test_end(self): + if self.global_rank == 0: + self.save_img_sequences( + f"it{self.true_global_step}-test", + "(\d+)\.png", + save_format="mp4", + fps=30, + ) + + def get_geometry(self, scene_code: torch.Tensor) -> Mesh: + tr.start("Surface extraction") + mesh: Mesh = self.renderer.isosurface(scene_code) + tr.end("Surface extraction") + return mesh diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/utils.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/utils.py new file mode 100644 index 0000000..1884c57 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/systems/utils.py @@ -0,0 +1,100 @@ +import torch +import torch.nn as nn +from torch.optim import lr_scheduler + +from CharacterGen.Stage_3D import lrm + + +def get_scheduler(name): + if hasattr(lr_scheduler, name): + return getattr(lr_scheduler, name) + else: + raise NotImplementedError + + +def getattr_recursive(m, attr): + for name in attr.split("."): + m = getattr(m, name) + return m + + +def get_parameters(model, name): + module = getattr_recursive(model, name) + if isinstance(module, nn.Module): + return module.parameters() + elif isinstance(module, nn.Parameter): + return module + return [] + + +def parse_optimizer(config, model): + if hasattr(config, "params"): + params = [ + {"params": get_parameters(model, name), "name": name, **args} + for name, args in config.params.items() + ] + lrm.debug(f"Specify optimizer params: {config.params}") + else: + params = model.parameters() + if config.name in ["FusedAdam"]: + import apex + + optim = getattr(apex.optimizers, config.name)(params, **config.args) + elif config.name in ["Adam8bit", "AdamW8bit"]: + import bitsandbytes as bnb + + optim = bnb.optim.Adam8bit(params, **config.args) + else: + optim = getattr(torch.optim, config.name)(params, **config.args) + return optim + + +def parse_scheduler_to_instance(config, optimizer): + if config.name == "ChainedScheduler": + schedulers = [ + parse_scheduler_to_instance(conf, optimizer) for conf in config.schedulers + ] + scheduler = lr_scheduler.ChainedScheduler(schedulers) + elif config.name == "Sequential": + schedulers = [ + parse_scheduler_to_instance(conf, optimizer) for conf in config.schedulers + ] + scheduler = lr_scheduler.SequentialLR( + optimizer, schedulers, milestones=config.milestones + ) + else: + scheduler = getattr(lr_scheduler, config.name)(optimizer, **config.args) + return scheduler + + +def parse_scheduler(config, optimizer): + interval = config.get("interval", "epoch") + assert interval in ["epoch", "step"] + if config.name == "SequentialLR": + scheduler = { + "scheduler": lr_scheduler.SequentialLR( + optimizer, + [ + parse_scheduler(conf, optimizer)["scheduler"] + for conf in config.schedulers + ], + milestones=config.milestones, + ), + "interval": interval, + } + elif config.name == "ChainedScheduler": + scheduler = { + "scheduler": lr_scheduler.ChainedScheduler( + [ + parse_scheduler(conf, optimizer)["scheduler"] + for conf in config.schedulers + ] + ), + "interval": interval, + } + else: + scheduler = { + "scheduler": get_scheduler(config.name)(optimizer, **config.args), + "interval": interval, + } + return scheduler diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/__init__.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/base.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/base.py new file mode 100644 index 0000000..f78adda --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/base.py @@ -0,0 +1,123 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from .config import parse_structured +from .misc import get_device, load_module_weights +from .typing import * + + +class Configurable: + @dataclass + class Config: + pass + + def __init__(self, cfg: Optional[dict] = None) -> None: + super().__init__() + self.cfg = parse_structured(self.Config, cfg) + + +class Updateable: + def do_update_step( + self, epoch: int, global_step: int, on_load_weights: bool = False + ): + for attr in self.__dir__(): + if attr.startswith("_"): + continue + try: + module = getattr(self, attr) + except: + continue # ignore attributes like property, which can't be retrived using getattr? + if isinstance(module, Updateable): + module.do_update_step( + epoch, global_step, on_load_weights=on_load_weights + ) + self.update_step(epoch, global_step, on_load_weights=on_load_weights) + + def do_update_step_end(self, epoch: int, global_step: int): + for attr in self.__dir__(): + if attr.startswith("_"): + continue + try: + module = getattr(self, attr) + except: + continue # ignore attributes like property, which can't be retrived using getattr? + if isinstance(module, Updateable): + module.do_update_step_end(epoch, global_step) + self.update_step_end(epoch, global_step) + + def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False): + # override this method to implement custom update logic + # if on_load_weights is True, you should be careful doing things related to model evaluations, + # as the models and tensors are not guarenteed to be on the same device + pass + + def update_step_end(self, epoch: int, global_step: int): + pass + + +def update_if_possible(module: Any, epoch: int, global_step: int) -> None: + if isinstance(module, Updateable): + module.do_update_step(epoch, global_step) + + +def update_end_if_possible(module: Any, epoch: int, global_step: int) -> None: + if isinstance(module, Updateable): + module.do_update_step_end(epoch, global_step) + + +class BaseObject(Updateable): + @dataclass + class Config: + pass + + cfg: Config # add this to every subclass of BaseObject to enable static type checking + + def __init__( + self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs + ) -> None: + super().__init__() + self.cfg = parse_structured(self.Config, cfg) + self.device = get_device() + self.configure(*args, **kwargs) + + def configure(self, *args, **kwargs) -> None: + pass + + +class BaseModule(nn.Module, Updateable): + @dataclass + class Config: + weights: Optional[str] = None + + cfg: Config # add this to every subclass of BaseModule to enable static type checking + + def __init__( + self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs + ) -> None: + super().__init__() + self.cfg = parse_structured(self.Config, cfg) + self.device = get_device() + self._non_modules = {} + self.configure(*args, **kwargs) + if self.cfg.weights is not None: + # format: path/to/weights:module_name + weights_path, module_name = self.cfg.weights.split(":") + state_dict, epoch, global_step = load_module_weights( + weights_path, module_name=module_name, map_location="cpu" + ) + self.load_state_dict(state_dict) + self.do_update_step( + epoch, global_step, on_load_weights=True + ) # restore states + + def configure(self, *args, **kwargs) -> None: + pass + + def register_non_module(self, name: str, module: nn.Module) -> None: + # non-modules won't be treated as model parameters + self._non_modules[name] = module + + def non_module(self, name: str): + return self._non_modules.get(name, None) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/callbacks.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/callbacks.py new file mode 100644 index 0000000..1f4cbd6 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/callbacks.py @@ -0,0 +1,156 @@ +import os +import shutil +import subprocess + +import pytorch_lightning + +from .config import dump_config +from .misc import parse_version + +if parse_version(pytorch_lightning.__version__) > parse_version("1.8"): + from pytorch_lightning.callbacks import Callback +else: + from pytorch_lightning.callbacks.base import Callback + +from pytorch_lightning.callbacks.progress import TQDMProgressBar +from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn + + +class VersionedCallback(Callback): + def __init__(self, save_root, version=None, use_version=True): + self.save_root = save_root + self._version = version + self.use_version = use_version + + @property + def version(self) -> int: + """Get the experiment version. + + Returns: + The experiment version if specified else the next version. + """ + if self._version is None: + self._version = self._get_next_version() + return self._version + + def _get_next_version(self): + existing_versions = [] + if os.path.isdir(self.save_root): + for f in os.listdir(self.save_root): + bn = os.path.basename(f) + if bn.startswith("version_"): + dir_ver = os.path.splitext(bn)[0].split("_")[1].replace("/", "") + existing_versions.append(int(dir_ver)) + if len(existing_versions) == 0: + return 0 + return max(existing_versions) + 1 + + @property + def savedir(self): + if not self.use_version: + return self.save_root + return os.path.join( + self.save_root, + self.version + if isinstance(self.version, str) + else f"version_{self.version}", + ) + + +class CodeSnapshotCallback(VersionedCallback): + def __init__(self, save_root, version=None, use_version=True): + super().__init__(save_root, version, use_version) + + def get_file_list(self): + return [ + b.decode() + for b in set( + subprocess.check_output( + 'git ls-files -- ":!:load/*"', shell=True + ).splitlines() + ) + | set( # hard code, TODO: use config to exclude folders or files + subprocess.check_output( + "git ls-files --others --exclude-standard", shell=True + ).splitlines() + ) + ] + + @rank_zero_only + def save_code_snapshot(self): + os.makedirs(self.savedir, exist_ok=True) + for f in self.get_file_list(): + if not os.path.exists(f) or os.path.isdir(f): + continue + os.makedirs(os.path.join(self.savedir, os.path.dirname(f)), exist_ok=True) + shutil.copyfile(f, os.path.join(self.savedir, f)) + + def on_fit_start(self, trainer, pl_module): + try: + self.save_code_snapshot() + except: + rank_zero_warn( + "Code snapshot is not saved. Please make sure you have git installed and are in a git repository." + ) + + +class ConfigSnapshotCallback(VersionedCallback): + def __init__(self, config_path, config, save_root, version=None, use_version=True): + super().__init__(save_root, version, use_version) + self.config_path = config_path + self.config = config + + @rank_zero_only + def save_config_snapshot(self): + os.makedirs(self.savedir, exist_ok=True) + dump_config(os.path.join(self.savedir, "parsed.yaml"), self.config) + shutil.copyfile(self.config_path, os.path.join(self.savedir, "raw.yaml")) + + def on_fit_start(self, trainer, pl_module): + self.save_config_snapshot() + + +class CustomProgressBar(TQDMProgressBar): + def get_metrics(self, *args, **kwargs): + # don't show the version number + items = super().get_metrics(*args, **kwargs) + items.pop("v_num", None) + return items + + +class ProgressCallback(Callback): + def __init__(self, save_path): + super().__init__() + self.save_path = save_path + self._file_handle = None + + @property + def file_handle(self): + if self._file_handle is None: + self._file_handle = open(self.save_path, "w") + return self._file_handle + + @rank_zero_only + def write(self, msg: str) -> None: + self.file_handle.seek(0) + self.file_handle.truncate() + self.file_handle.write(msg) + self.file_handle.flush() + + @rank_zero_only + def on_train_batch_end(self, trainer, pl_module, *args, **kwargs): + self.write( + f"Generation progress: {pl_module.true_global_step / trainer.max_steps * 100:.2f}%" + ) + + @rank_zero_only + def on_validation_start(self, trainer, pl_module): + self.write(f"Rendering validation image ...") + + @rank_zero_only + def on_test_start(self, trainer, pl_module): + self.write(f"Rendering video ...") + + @rank_zero_only + def on_predict_start(self, trainer, pl_module): + self.write(f"Exporting mesh assets ...") diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/config.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/config.py new file mode 100644 index 0000000..0cb06ff --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/config.py @@ -0,0 +1,138 @@ +import os +from dataclasses import dataclass, field +from datetime import datetime + +from omegaconf import OmegaConf + +from CharacterGen.Stage_3D import lrm +from .typing import * + +# ============ Register OmegaConf Resolvers ============= # +OmegaConf.register_new_resolver("calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / 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) +OmegaConf.register_new_resolver("rmspace", lambda s, sub: s.replace(" ", sub), replace=True) +OmegaConf.register_new_resolver("tuple2", lambda s: [float(s), float(s)], replace=True) +OmegaConf.register_new_resolver("gt0", lambda s: s > 0, replace=True) +OmegaConf.register_new_resolver("not", lambda s: not s, replace=True) + + +def calc_num_train_steps(num_data, batch_size, max_epochs, num_nodes, num_cards=8): + return int(num_data / (num_nodes * num_cards * batch_size)) * max_epochs + + +OmegaConf.register_new_resolver("calc_num_train_steps", calc_num_train_steps) + +# ======================================================= # + + +# ============== Automatic Name Resolvers =============== # +def get_naming_convention(cfg): + # TODO + name = f"lrm_{cfg.system.backbone.num_layers}" + return name + + +# ======================================================= # + + +@dataclass +class ExperimentConfig: + name: str = "default" + description: str = "" + tag: str = "" + seed: int = 0 + use_timestamp: bool = True + timestamp: Optional[str] = None + exp_root_dir: str = "outputs" + + ### these shouldn't be set manually + exp_dir: str = "outputs/default" + trial_name: str = "exp" + trial_dir: str = "outputs/default/exp" + n_gpus: int = 1 + ### + + resume: Optional[str] = None + + data_cls: str = "" + data: dict = field(default_factory=dict) + + system_cls: str = "" + system: dict = field(default_factory=dict) + + # accept pytorch-lightning trainer parameters + # see https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api + trainer: dict = field(default_factory=dict) + + # accept pytorch-lightning checkpoint callback parameters + # see https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint + checkpoint: dict = field(default_factory=dict) + + +def load_config( + *yamls: str, cli_args: list = [], from_string=False, makedirs=False, **kwargs +) -> Any: + if from_string: + parse_func = OmegaConf.create + else: + parse_func = OmegaConf.load + yaml_confs = [] + for y in yamls: + conf = parse_func(y) + extends = conf.pop("extends", None) + if extends: + assert os.path.exists(extends), f"File {extends} does not exist." + yaml_confs.append(OmegaConf.load(extends)) + yaml_confs.append(conf) + cli_conf = OmegaConf.from_cli(cli_args) + cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs) + OmegaConf.resolve(cfg) + assert isinstance(cfg, DictConfig) + scfg: ExperimentConfig = parse_structured(ExperimentConfig, cfg) + + # post processing + # auto naming + if scfg.name == "auto": + scfg.name = get_naming_convention(scfg) + # add timestamp + if not scfg.tag and not scfg.use_timestamp: + raise ValueError("Either tag is specified or use_timestamp is True.") + scfg.trial_name = scfg.tag + # if resume from an existing config, scfg.timestamp should not be None + if scfg.timestamp is None: + scfg.timestamp = "" + if scfg.use_timestamp: + if scfg.n_gpus > 1: + lrm.warn( + "Timestamp is disabled when using multiple GPUs, please make sure you have a unique tag." + ) + else: + scfg.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S") + # make directories + scfg.trial_name += scfg.timestamp + scfg.exp_dir = os.path.join(scfg.exp_root_dir, scfg.name) + scfg.trial_dir = os.path.join(scfg.exp_dir, scfg.trial_name) + + if makedirs: + os.makedirs(scfg.trial_dir, exist_ok=True) + + return scfg + + +def config_to_primitive(config, resolve: bool = True) -> Any: + return OmegaConf.to_container(config, resolve=resolve) + + +def dump_config(path: str, config) -> None: + with open(path, "w") as fp: + OmegaConf.save(config=config, f=fp) + + +def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any: + scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg) + return scfg diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/misc.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/misc.py new file mode 100644 index 0000000..653b8a7 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/misc.py @@ -0,0 +1,211 @@ +import gc +import os +import re +import time +from collections import defaultdict +from contextlib import contextmanager + +import torch +from packaging import version + +from CharacterGen.Stage_3D import lrm +from .config import config_to_primitive +from .typing import * + + +def parse_version(ver: str): + return version.parse(ver) + + +def get_rank(): + # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing, + # therefore LOCAL_RANK needs to be checked first + rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK") + for key in rank_keys: + rank = os.environ.get(key) + if rank is not None: + return int(rank) + return 0 + + +def get_device(): + return torch.device(f"cuda:{get_rank()}") + + +def load_module_weights( + path, module_name=None, ignore_modules=None, mapping=None, map_location=None +) -> Tuple[dict, int, int]: + if module_name is not None and ignore_modules is not None: + raise ValueError("module_name and ignore_modules cannot be both set") + if map_location is None: + map_location = get_device() + + ckpt = torch.load(path, map_location=map_location) + state_dict = ckpt["state_dict"] + + if mapping is not None: + state_dict_to_load = {} + for k, v in state_dict.items(): + if any([k.startswith(m["to"]) for m in mapping]): + pass + else: + state_dict_to_load[k] = v + for k, v in state_dict.items(): + for m in mapping: + if k.startswith(m["from"]): + k_dest = k.replace(m["from"], m["to"]) + lrm.info(f"Mapping {k} => {k_dest}") + state_dict_to_load[k_dest] = v.clone() + state_dict = state_dict_to_load + + state_dict_to_load = state_dict + + if ignore_modules is not None: + state_dict_to_load = {} + for k, v in state_dict.items(): + ignore = any( + [k.startswith(ignore_module + ".") for ignore_module in ignore_modules] + ) + if ignore: + continue + state_dict_to_load[k] = v + + if module_name is not None: + state_dict_to_load = {} + for k, v in state_dict.items(): + m = re.match(rf"^{module_name}\.(.*)$", k) + if m is None: + continue + state_dict_to_load[m.group(1)] = v + + return state_dict_to_load, ckpt["epoch"], ckpt["global_step"] + + +def C(value: Any, epoch: int, global_step: int) -> float: + if isinstance(value, int) or isinstance(value, float): + pass + else: + value = config_to_primitive(value) + if not isinstance(value, list): + raise TypeError("Scalar specification only supports list, got", type(value)) + if len(value) == 3: + value = [0] + value + assert len(value) == 4 + start_step, start_value, end_value, end_step = value + if isinstance(end_step, int): + current_step = global_step + value = start_value + (end_value - start_value) * max( + min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0 + ) + elif isinstance(end_step, float): + current_step = epoch + value = start_value + (end_value - start_value) * max( + min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0 + ) + return value + + +def cleanup(): + gc.collect() + torch.cuda.empty_cache() + try: + import tinycudann as tcnn + + tcnn.free_temporary_memory() + except: + pass + + +def finish_with_cleanup(func: Callable): + def wrapper(*args, **kwargs): + out = func(*args, **kwargs) + cleanup() + return out + + return wrapper + + +def _distributed_available(): + return torch.distributed.is_available() and torch.distributed.is_initialized() + + +def barrier(): + if not _distributed_available(): + return + else: + torch.distributed.barrier() + + +def broadcast(tensor, src=0): + if not _distributed_available(): + return tensor + else: + torch.distributed.broadcast(tensor, src=src) + return tensor + + +def enable_gradient(model, enabled: bool = True) -> None: + for param in model.parameters(): + param.requires_grad_(enabled) + + +class TimeRecorder: + _instance = None + + def __init__(self): + self.items = {} + self.accumulations = defaultdict(list) + self.time_scale = 1000.0 # ms + self.time_unit = "ms" + self.enabled = False + + def __new__(cls): + # singleton + if cls._instance is None: + cls._instance = super(TimeRecorder, cls).__new__(cls) + return cls._instance + + def enable(self, enabled: bool) -> None: + self.enabled = enabled + + def start(self, name: str) -> None: + if not self.enabled: + return + torch.cuda.synchronize() + self.items[name] = time.time() + + def end(self, name: str, accumulate: bool = False) -> float: + if not self.enabled or name not in self.items: + return + torch.cuda.synchronize() + start_time = self.items.pop(name) + delta = time.time() - start_time + if accumulate: + self.accumulations[name].append(delta) + t = delta * self.time_scale + lrm.info(f"{name}: {t:.2f}{self.time_unit}") + + def get_accumulation(self, name: str, average: bool = False) -> float: + if not self.enabled or name not in self.accumulations: + return + acc = self.accumulations.pop(name) + total = sum(acc) + if average: + t = total / len(acc) * self.time_scale + else: + t = total * self.time_scale + lrm.info(f"{name} for {len(acc)} times: {t:.2f}{self.time_unit}") + + +### global time recorder +time_recorder = TimeRecorder() + + +@contextmanager +def time_recorder_enabled(): + enabled = time_recorder.enabled + time_recorder.enable(enabled=True) + try: + yield + finally: + time_recorder.enable(enabled=enabled) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/ops.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/ops.py new file mode 100644 index 0000000..43acc59 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/ops.py @@ -0,0 +1,435 @@ +from collections import defaultdict +import math + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Function +from torch.cuda.amp import custom_bwd, custom_fwd + +from CharacterGen.Stage_3D import lrm +from .typing import * + + +def dot(x, y): + return torch.sum(x * y, -1, keepdim=True) + + +def reflect(x, n): + return 2 * dot(x, n) * n - x + + +ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]] + + +def scale_tensor( + dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale +): + if inp_scale is None: + inp_scale = (0, 1) + if tgt_scale is None: + tgt_scale = (0, 1) + if isinstance(tgt_scale, Tensor): + assert dat.shape[-1] == tgt_scale.shape[-1] + dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0]) + dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0] + return dat + + +class _TruncExp(Function): # pylint: disable=abstract-method + # Implementation from torch-ngp: + # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py + @staticmethod + @custom_fwd(cast_inputs=torch.float32) + def forward(ctx, x): # pylint: disable=arguments-differ + ctx.save_for_backward(x) + return torch.exp(x) + + @staticmethod + @custom_bwd + def backward(ctx, g): # pylint: disable=arguments-differ + x = ctx.saved_tensors[0] + return g * torch.exp(torch.clamp(x, max=15)) + + +trunc_exp = _TruncExp.apply + + +def get_activation(name) -> Callable: + if name is None: + return lambda x: x + name = name.lower() + if name == "none": + return lambda x: x + elif name == "lin2srgb": + return lambda x: torch.where( + x > 0.0031308, + torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055, + 12.92 * x, + ).clamp(0.0, 1.0) + elif name == "exp": + return lambda x: torch.exp(x) + elif name == "shifted_exp": + return lambda x: torch.exp(x - 1.0) + elif name == "trunc_exp": + return trunc_exp + elif name == "shifted_trunc_exp": + return lambda x: trunc_exp(x - 1.0) + elif name == "sigmoid": + return lambda x: torch.sigmoid(x) + elif name == "tanh": + return lambda x: torch.tanh(x) + elif name == "shifted_softplus": + return lambda x: F.softplus(x - 1.0) + elif name == "scale_-11_01": + return lambda x: x * 0.5 + 0.5 + elif name == "negative": + return lambda x: -x + else: + try: + return getattr(F, name) + except AttributeError: + raise ValueError(f"Unknown activation function: {name}") + + +def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any: + if chunk_size <= 0: + return func(*args, **kwargs) + B = None + for arg in list(args) + list(kwargs.values()): + if isinstance(arg, torch.Tensor): + B = arg.shape[0] + break + assert ( + B is not None + ), "No tensor found in args or kwargs, cannot determine batch size." + out = defaultdict(list) + out_type = None + # max(1, B) to support B == 0 + for i in range(0, max(1, B), chunk_size): + out_chunk = func( + *[ + arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg + for arg in args + ], + **{ + k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg + for k, arg in kwargs.items() + }, + ) + if out_chunk is None: + continue + out_type = type(out_chunk) + if isinstance(out_chunk, torch.Tensor): + out_chunk = {0: out_chunk} + elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list): + chunk_length = len(out_chunk) + out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)} + elif isinstance(out_chunk, dict): + pass + else: + print( + f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}." + ) + exit(1) + for k, v in out_chunk.items(): + v = v if torch.is_grad_enabled() else v.detach() + out[k].append(v) + + if out_type is None: + return None + + out_merged: Dict[Any, Optional[torch.Tensor]] = {} + for k, v in out.items(): + if all([vv is None for vv in v]): + # allow None in return value + out_merged[k] = None + elif all([isinstance(vv, torch.Tensor) for vv in v]): + out_merged[k] = torch.cat(v, dim=0) + else: + raise TypeError( + f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}" + ) + + if out_type is torch.Tensor: + return out_merged[0] + elif out_type in [tuple, list]: + return out_type([out_merged[i] for i in range(chunk_length)]) + elif out_type is dict: + return out_merged + + +def get_ray_directions( + H: int, + W: int, + focal: Union[float, Tuple[float, float]], + principal: Optional[Tuple[float, float]] = None, + use_pixel_centers: bool = True, + normalize: bool = True, +) -> Float[Tensor, "H W 3"]: + """ + Get ray directions for all pixels in camera coordinate. + Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/ + ray-tracing-generating-camera-rays/standard-coordinate-systems + + Inputs: + H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers + Outputs: + directions: (H, W, 3), the direction of the rays in camera coordinate + """ + pixel_center = 0.5 if use_pixel_centers else 0 + + if isinstance(focal, float): + fx, fy = focal, focal + cx, cy = W / 2, H / 2 + else: + fx, fy = focal + assert principal is not None + cx, cy = principal + + i, j = torch.meshgrid( + torch.arange(W, dtype=torch.float32) + pixel_center, + torch.arange(H, dtype=torch.float32) + pixel_center, + indexing="xy", + ) + + directions: Float[Tensor, "H W 3"] = torch.stack( + [(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1 + ) + + if normalize: + directions = F.normalize(directions, dim=-1) + + return directions + + +def get_rays( + directions: Float[Tensor, "... 3"], + c2w: Float[Tensor, "... 4 4"], + keepdim=False, + noise_scale=0.0, + normalize=False, +) -> Tuple[Float[Tensor, "... 3"], Float[Tensor, "... 3"]]: + # Rotate ray directions from camera coordinate to the world coordinate + assert directions.shape[-1] == 3 + + if directions.ndim == 2: # (N_rays, 3) + if c2w.ndim == 2: # (4, 4) + c2w = c2w[None, :, :] + assert c2w.ndim == 3 # (N_rays, 4, 4) or (1, 4, 4) + rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1) # (N_rays, 3) + rays_o = c2w[:, :3, 3].expand(rays_d.shape) + elif directions.ndim == 3: # (H, W, 3) + assert c2w.ndim in [2, 3] + if c2w.ndim == 2: # (4, 4) + rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum( + -1 + ) # (H, W, 3) + rays_o = c2w[None, None, :3, 3].expand(rays_d.shape) + elif c2w.ndim == 3: # (B, 4, 4) + rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum( + -1 + ) # (B, H, W, 3) + rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape) + elif directions.ndim == 4: # (B, H, W, 3) + assert c2w.ndim == 3 # (B, 4, 4) + rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum( + -1 + ) # (B, H, W, 3) + rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape) + + # add camera noise to avoid grid-like artifect + # https://github.com/ashawkey/stable-dreamfusion/blob/49c3d4fa01d68a4f027755acf94e1ff6020458cc/nerf/utils.py#L373 + if noise_scale > 0: + rays_o = rays_o + torch.randn(3, device=rays_o.device) * noise_scale + rays_d = rays_d + torch.randn(3, device=rays_d.device) * noise_scale + + if normalize: + rays_d = F.normalize(rays_d, dim=-1) + if not keepdim: + rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3) + + return rays_o, rays_d + + +def get_projection_matrix( + fovy: Union[float, Float[Tensor, "B"]], aspect_wh: float, near: float, far: float +) -> Float[Tensor, "*B 4 4"]: + if isinstance(fovy, float): + proj_mtx = torch.zeros(4, 4, dtype=torch.float32) + proj_mtx[0, 0] = 1.0 / (math.tan(fovy / 2.0) * aspect_wh) + proj_mtx[1, 1] = -1.0 / math.tan( + fovy / 2.0 + ) # add a negative sign here as the y axis is flipped in nvdiffrast output + proj_mtx[2, 2] = -(far + near) / (far - near) + proj_mtx[2, 3] = -2.0 * far * near / (far - near) + proj_mtx[3, 2] = -1.0 + else: + batch_size = fovy.shape[0] + proj_mtx = torch.zeros(batch_size, 4, 4, dtype=torch.float32) + proj_mtx[:, 0, 0] = 1.0 / (torch.tan(fovy / 2.0) * aspect_wh) + proj_mtx[:, 1, 1] = -1.0 / torch.tan( + fovy / 2.0 + ) # add a negative sign here as the y axis is flipped in nvdiffrast output + proj_mtx[:, 2, 2] = -(far + near) / (far - near) + proj_mtx[:, 2, 3] = -2.0 * far * near / (far - near) + proj_mtx[:, 3, 2] = -1.0 + return proj_mtx + + +def get_mvp_matrix( + c2w: Float[Tensor, "*B 4 4"], proj_mtx: Float[Tensor, "*B 4 4"] +) -> Float[Tensor, "*B 4 4"]: + # calculate w2c from c2w: R' = Rt, t' = -Rt * t + # mathematically equivalent to (c2w)^-1 + if c2w.ndim == 2: + assert proj_mtx.ndim == 2 + w2c: Float[Tensor, "4 4"] = torch.zeros(4, 4).to(c2w) + w2c[:3, :3] = c2w[:3, :3].permute(1, 0) + w2c[:3, 3:] = -c2w[:3, :3].permute(1, 0) @ c2w[:3, 3:] + w2c[3, 3] = 1.0 + else: + w2c: Float[Tensor, "B 4 4"] = torch.zeros(c2w.shape[0], 4, 4).to(c2w) + w2c[:, :3, :3] = c2w[:, :3, :3].permute(0, 2, 1) + w2c[:, :3, 3:] = -c2w[:, :3, :3].permute(0, 2, 1) @ c2w[:, :3, 3:] + w2c[:, 3, 3] = 1.0 + # calculate mvp matrix by proj_mtx @ w2c (mv_mtx) + mvp_mtx = proj_mtx @ w2c + return mvp_mtx + + +def get_intrinsic_from_fov(fov, H, W, bs=-1): + focal_length = 0.5 * H / np.tan(0.5 * fov) + intrinsic = np.identity(3, dtype=np.float32) + intrinsic[0, 0] = focal_length + intrinsic[1, 1] = focal_length + intrinsic[0, 2] = W / 2.0 + intrinsic[1, 2] = H / 2.0 + + if bs > 0: + intrinsic = intrinsic[None].repeat(bs, axis=0) + + return torch.from_numpy(intrinsic) + + +def binary_cross_entropy(input, target): + """ + F.binary_cross_entropy is not numerically stable in mixed-precision training. + """ + return -(target * torch.log(input) + (1 - target) * torch.log(1 - input)).mean() + + +def tet_sdf_diff( + vert_sdf: Float[Tensor, "Nv 1"], tet_edges: Integer[Tensor, "Ne 2"] +) -> Float[Tensor, ""]: + sdf_f1x6x2 = vert_sdf[:, 0][tet_edges.reshape(-1)].reshape(-1, 2) + mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1]) + sdf_f1x6x2 = sdf_f1x6x2[mask] + sdf_diff = F.binary_cross_entropy_with_logits( + sdf_f1x6x2[..., 0], (sdf_f1x6x2[..., 1] > 0).float() + ) + F.binary_cross_entropy_with_logits( + sdf_f1x6x2[..., 1], (sdf_f1x6x2[..., 0] > 0).float() + ) + return sdf_diff + + +def validate_empty_rays(ray_indices, t_start, t_end): + if ray_indices.nelement() == 0: + lrm.warn("Empty rays_indices!") + ray_indices = torch.LongTensor([0]).to(ray_indices) + t_start = torch.Tensor([0]).to(ray_indices) + t_end = torch.Tensor([0]).to(ray_indices) + return ray_indices, t_start, t_end + + +def rays_intersect_bbox( + rays_o: Float[Tensor, "N 3"], + rays_d: Float[Tensor, "N 3"], + radius: Float, + near: Float = 0.0, + valid_thresh: Float = 0.01, + background: bool = False, +): + input_shape = rays_o.shape[:-1] + rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3) + rays_d_valid = torch.where( + rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d + ) + if type(radius) in [int, float]: + radius = torch.FloatTensor( + [[-radius, radius], [-radius, radius], [-radius, radius]] + ).to(rays_o.device) + radius = ( + 1.0 - 1.0e-3 + ) * radius # tighten the radius to make sure the intersection point lies in the bounding box + interx0 = (radius[..., 1] - rays_o) / rays_d_valid + interx1 = (radius[..., 0] - rays_o) / rays_d_valid + t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near) + t_far = torch.maximum(interx0, interx1).amin(dim=-1) + + # check wheter a ray intersects the bbox or not + rays_valid = t_far - t_near > valid_thresh + + t_near_valid, t_far_valid = t_near[rays_valid], t_far[rays_valid] + global_near = t_near_valid.min().item() + global_far = t_far_valid.max().item() + + t_near[torch.where(~rays_valid)] = 0.0 + t_far[torch.where(~rays_valid)] = 0.0 + + t_near = t_near.view(*input_shape, 1) + t_far = t_far.view(*input_shape, 1) + rays_valid = rays_valid.view(*input_shape) + + return t_near, t_far, rays_valid + + +def get_plucker_rays(rays_o: Float[Tensor, "*N 3"], rays_d: Float[Tensor, "*N 3"]) -> Float[Tensor, "*N 6"]: + rays_o = F.normalize(rays_o, dim=-1) + rays_d = F.normalize(rays_d, dim=-1) + return torch.cat([ + rays_o.cross(rays_d), + rays_d + ], dim=-1) + + +def c2w_to_polar(c2w: Float[Tensor, "4 4"]) -> Tuple[float, float, float]: + cam_pos = c2w[:3, 3] + x, y, z = cam_pos.tolist() + distance = cam_pos.norm().item() + elevation = math.asin(z / distance) + if abs(x) < 1.0e-5 and abs(y) < 1.0e-5: + azimuth = 0 + else: + azimuth = math.atan2(y, x) + if azimuth < 0: + azimuth += 2 * math.pi + + return elevation, azimuth, distance + + +def polar_to_c2w(elevation: float, azimuth: float, distance: float) -> Float[Tensor, "4 4"]: + """ + Compute L = p - C. + Normalize L. + Compute s = L x u. (cross product) + Normalize s. + Compute u' = s x L. + rotation = [s, u, -l] + """ + z = distance * math.sin(elevation) + x = distance * math.cos(elevation) * math.cos(azimuth) + y = distance * math.cos(elevation) * math.sin(azimuth) + l = -torch.as_tensor([x, y, z]).float() + l = F.normalize(l, dim=0) + u = torch.as_tensor([0.0, 0.0, 1.0]).float() + s = l.cross(u) + s = F.normalize(s, dim=0) + u = s.cross(l) + rot = torch.stack([s, u, -l], dim=0).T + c2w = torch.zeros((4, 4), dtype=torch.float32) + c2w[:3, :3] = rot + c2w[:3, 3] = torch.as_tensor([x, y, z]) + c2w[3, 3] = 1 + return c2w diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/rasterize.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/rasterize.py new file mode 100644 index 0000000..dff8eab --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/rasterize.py @@ -0,0 +1,81 @@ +import nvdiffrast.torch as dr +import torch + +from .typing import * + + +class NVDiffRasterizerContext: + def __init__(self, context_type: str, device: torch.device) -> None: + self.device = device + self.ctx = self.initialize_context(context_type, device) + + def initialize_context( + self, context_type: str, device: torch.device + ) -> Union[dr.RasterizeGLContext, dr.RasterizeCudaContext]: + context_type = "cuda" # This will throw: RuntimeError: Error building extension 'nvdiffrast_plugin', 'cl.exe' is not recognized as an internal or external command + if context_type == "gl": + return dr.RasterizeGLContext(device=device) + elif context_type == "cuda": + return dr.RasterizeCudaContext(device=device) + else: + raise ValueError(f"Unknown rasterizer context type: {context_type}") + + def vertex_transform( + self, verts: Float[Tensor, "Nv 3"], mvp_mtx: Float[Tensor, "B 4 4"] + ) -> Float[Tensor, "B Nv 4"]: + with torch.cuda.amp.autocast(enabled=False): + verts_homo = torch.cat( + [verts, torch.ones([verts.shape[0], 1]).to(verts)], dim=-1 + ) + verts_clip = torch.matmul(verts_homo, mvp_mtx.permute(0, 2, 1)) + return verts_clip + + def rasterize( + self, + pos: Float[Tensor, "B Nv 4"], + tri: Integer[Tensor, "Nf 3"], + resolution: Union[int, Tuple[int, int]], + ): + # rasterize in instance mode (single topology) + return dr.rasterize(self.ctx, pos.float(), tri.int(), resolution, grad_db=True) + + def rasterize_one( + self, + pos: Float[Tensor, "Nv 4"], + tri: Integer[Tensor, "Nf 3"], + resolution: Union[int, Tuple[int, int]], + ): + # rasterize one single mesh under a single viewpoint + rast, rast_db = self.rasterize(pos[None, ...], tri, resolution) + return rast[0], rast_db[0] + + def antialias( + self, + color: Float[Tensor, "B H W C"], + rast: Float[Tensor, "B H W 4"], + pos: Float[Tensor, "B Nv 4"], + tri: Integer[Tensor, "Nf 3"], + ) -> Float[Tensor, "B H W C"]: + return dr.antialias(color.float(), rast, pos.float(), tri.int()) + + def interpolate( + self, + attr: Float[Tensor, "B Nv C"], + rast: Float[Tensor, "B H W 4"], + tri: Integer[Tensor, "Nf 3"], + rast_db=None, + diff_attrs=None, + ) -> Float[Tensor, "B H W C"]: + return dr.interpolate( + attr.float(), rast, tri.int(), rast_db=rast_db, diff_attrs=diff_attrs + ) + + def interpolate_one( + self, + attr: Float[Tensor, "Nv C"], + rast: Float[Tensor, "B H W 4"], + tri: Integer[Tensor, "Nf 3"], + rast_db=None, + diff_attrs=None, + ) -> Float[Tensor, "B H W C"]: + return self.interpolate(attr[None, ...], rast, tri, rast_db, diff_attrs) diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/saving.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/saving.py new file mode 100644 index 0000000..400c2c4 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/saving.py @@ -0,0 +1,725 @@ +import json +import os +import re +import shutil + +import cv2 +import imageio +import matplotlib.pyplot as plt +import numpy as np +import torch +from matplotlib import cm +from matplotlib.colors import LinearSegmentedColormap +from PIL import Image, ImageDraw +from pytorch_lightning.loggers import WandbLogger + +from CharacterGen.Stage_3D import lrm +from ..models.mesh import Mesh +from ..utils.typing import * + + +class SaverMixin: + import wandb + _save_dir: Optional[str] = None + _wandb_logger: Optional[WandbLogger] = None + + def set_save_dir(self, save_dir: str): + self._save_dir = save_dir + + def get_save_dir(self): + if self._save_dir is None: + raise ValueError("Save dir is not set") + return self._save_dir + + def convert_data(self, data): + if data is None: + return None + elif isinstance(data, np.ndarray): + return data + elif isinstance(data, torch.Tensor): + if data.dtype in [torch.float16, torch.bfloat16]: + data = data.float() + return data.detach().cpu().numpy() + elif isinstance(data, list): + return [self.convert_data(d) for d in data] + elif isinstance(data, dict): + return {k: self.convert_data(v) for k, v in data.items()} + else: + raise TypeError( + "Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting", + type(data), + ) + + def get_save_path(self, filename): + save_path = os.path.join(self.get_save_dir(), filename) + os.makedirs(os.path.dirname(save_path), exist_ok=True) + return save_path + + DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)} + DEFAULT_UV_KWARGS = { + "data_format": "HWC", + "data_range": (0, 1), + "cmap": "checkerboard", + } + DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"} + DEFAULT_GRID_KWARGS = {"align": "max"} + + def get_rgb_image_(self, img, data_format, data_range, rgba=False): + img = self.convert_data(img) + assert data_format in ["CHW", "HWC"] + if data_format == "CHW": + img = img.transpose(1, 2, 0) + if img.dtype != np.uint8: + img = img.clip(min=data_range[0], max=data_range[1]) + img = ( + (img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0 + ).astype(np.uint8) + nc = 4 if rgba else 3 + imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)] + imgs = [ + img_ + if img_.shape[-1] == nc + else np.concatenate( + [ + img_, + np.zeros( + (img_.shape[0], img_.shape[1], nc - img_.shape[2]), + dtype=img_.dtype, + ), + ], + axis=-1, + ) + for img_ in imgs + ] + img = np.concatenate(imgs, axis=1) + if rgba: + img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA) + else: + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + return img + + def _save_rgb_image( + self, + filename, + img, + data_format, + data_range, + name: Optional[str] = None, + step: Optional[int] = None, + ): + img = self.get_rgb_image_(img, data_format, data_range) + cv2.imwrite(filename, img) + if name and self._wandb_logger: + self._wandb_logger.log_image( + key=name, images=[self.get_save_path(filename)], step=step + ) + + def save_rgb_image( + self, + filename, + img, + data_format=DEFAULT_RGB_KWARGS["data_format"], + data_range=DEFAULT_RGB_KWARGS["data_range"], + name: Optional[str] = None, + step: Optional[int] = None, + ) -> str: + save_path = self.get_save_path(filename) + self._save_rgb_image(save_path, img, data_format, data_range, name, step) + return save_path + + def get_uv_image_(self, img, data_format, data_range, cmap): + img = self.convert_data(img) + assert data_format in ["CHW", "HWC"] + if data_format == "CHW": + img = img.transpose(1, 2, 0) + img = img.clip(min=data_range[0], max=data_range[1]) + img = (img - data_range[0]) / (data_range[1] - data_range[0]) + assert cmap in ["checkerboard", "color"] + if cmap == "checkerboard": + n_grid = 64 + mask = (img * n_grid).astype(int) + mask = (mask[..., 0] + mask[..., 1]) % 2 == 0 + img = np.ones((img.shape[0], img.shape[1], 3), dtype=np.uint8) * 255 + img[mask] = np.array([255, 0, 255], dtype=np.uint8) + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + elif cmap == "color": + img_ = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8) + img_[..., 0] = (img[..., 0] * 255).astype(np.uint8) + img_[..., 1] = (img[..., 1] * 255).astype(np.uint8) + img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR) + img = img_ + return img + + def save_uv_image( + self, + filename, + img, + data_format=DEFAULT_UV_KWARGS["data_format"], + data_range=DEFAULT_UV_KWARGS["data_range"], + cmap=DEFAULT_UV_KWARGS["cmap"], + ) -> str: + save_path = self.get_save_path(filename) + img = self.get_uv_image_(img, data_format, data_range, cmap) + cv2.imwrite(save_path, img) + return save_path + + def get_grayscale_image_(self, img, data_range, cmap): + img = self.convert_data(img) + img = np.nan_to_num(img) + if data_range is None: + img = (img - img.min()) / (img.max() - img.min()) + else: + img = img.clip(data_range[0], data_range[1]) + img = (img - data_range[0]) / (data_range[1] - data_range[0]) + assert cmap in [None, "jet", "magma", "spectral"] + if cmap == None: + img = (img * 255.0).astype(np.uint8) + img = np.repeat(img[..., None], 3, axis=2) + elif cmap == "jet": + img = (img * 255.0).astype(np.uint8) + img = cv2.applyColorMap(img, cv2.COLORMAP_JET) + elif cmap == "magma": + img = 1.0 - img + base = cm.get_cmap("magma") + num_bins = 256 + colormap = LinearSegmentedColormap.from_list( + f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins + )(np.linspace(0, 1, num_bins))[:, :3] + a = np.floor(img * 255.0) + b = (a + 1).clip(max=255.0) + f = img * 255.0 - a + a = a.astype(np.uint16).clip(0, 255) + b = b.astype(np.uint16).clip(0, 255) + img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None] + img = (img * 255.0).astype(np.uint8) + elif cmap == "spectral": + colormap = plt.get_cmap("Spectral") + + def blend_rgba(image): + image = image[..., :3] * image[..., -1:] + ( + 1.0 - image[..., -1:] + ) # blend A to RGB + return image + + img = colormap(img) + img = blend_rgba(img) + img = (img * 255).astype(np.uint8) + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + return img + + def _save_grayscale_image( + self, + filename, + img, + data_range, + cmap, + name: Optional[str] = None, + step: Optional[int] = None, + ): + img = self.get_grayscale_image_(img, data_range, cmap) + cv2.imwrite(filename, img) + if name and self._wandb_logger: + self._wandb_logger.log_image( + key=name, images=[self.get_save_path(filename)], step=step + ) + + def save_grayscale_image( + self, + filename, + img, + data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"], + cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"], + name: Optional[str] = None, + step: Optional[int] = None, + ) -> str: + save_path = self.get_save_path(filename) + self._save_grayscale_image(save_path, img, data_range, cmap, name, step) + return save_path + + def get_image_grid_(self, imgs, align): + if isinstance(imgs[0], list): + return np.concatenate( + [self.get_image_grid_(row, align) for row in imgs], axis=0 + ) + cols = [] + for col in imgs: + assert col["type"] in ["rgb", "uv", "grayscale"] + if col["type"] == "rgb": + rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy() + rgb_kwargs.update(col["kwargs"]) + cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs)) + elif col["type"] == "uv": + uv_kwargs = self.DEFAULT_UV_KWARGS.copy() + uv_kwargs.update(col["kwargs"]) + cols.append(self.get_uv_image_(col["img"], **uv_kwargs)) + elif col["type"] == "grayscale": + grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy() + grayscale_kwargs.update(col["kwargs"]) + cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs)) + + if align == "max": + h = max([col.shape[0] for col in cols]) + elif align == "min": + h = min([col.shape[0] for col in cols]) + elif isinstance(align, int): + h = align + else: + raise ValueError( + f"Unsupported image grid align: {align}, should be min, max, or int" + ) + + for i in range(len(cols)): + if cols[i].shape[0] != h: + w = int(cols[i].shape[1] * h / cols[i].shape[0]) + cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_CUBIC) + return np.concatenate(cols, axis=1) + + def save_image_grid( + self, + filename, + imgs, + align=DEFAULT_GRID_KWARGS["align"], + name: Optional[str] = None, + step: Optional[int] = None, + texts: Optional[List[float]] = None, + ): + save_path = self.get_save_path(filename) + img = self.get_image_grid_(imgs, align=align) + + if texts is not None: + img = Image.fromarray(img) + draw = ImageDraw.Draw(img) + black, white = (0, 0, 0), (255, 255, 255) + for i, text in enumerate(texts): + draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white) + draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white) + draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white) + draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white) + draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black) + img = np.asarray(img) + + cv2.imwrite(save_path, img) + if name and self._wandb_logger: + self._wandb_logger.log_image(key=name, images=[save_path], step=step) + return save_path + + def save_image(self, filename, img) -> str: + save_path = self.get_save_path(filename) + img = self.convert_data(img) + assert img.dtype == np.uint8 or img.dtype == np.uint16 + if img.ndim == 3 and img.shape[-1] == 3: + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + elif img.ndim == 3 and img.shape[-1] == 4: + img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA) + cv2.imwrite(save_path, img) + return save_path + + def save_cubemap(self, filename, img, data_range=(0, 1), rgba=False) -> str: + save_path = self.get_save_path(filename) + img = self.convert_data(img) + assert img.ndim == 4 and img.shape[0] == 6 and img.shape[1] == img.shape[2] + + imgs_full = [] + for start in range(0, img.shape[-1], 3): + img_ = img[..., start : start + 3] + img_ = np.stack( + [ + self.get_rgb_image_(img_[i], "HWC", data_range, rgba=rgba) + for i in range(img_.shape[0]) + ], + axis=0, + ) + size = img_.shape[1] + placeholder = np.zeros((size, size, 3), dtype=np.float32) + img_full = np.concatenate( + [ + np.concatenate( + [placeholder, img_[2], placeholder, placeholder], axis=1 + ), + np.concatenate([img_[1], img_[4], img_[0], img_[5]], axis=1), + np.concatenate( + [placeholder, img_[3], placeholder, placeholder], axis=1 + ), + ], + axis=0, + ) + imgs_full.append(img_full) + + imgs_full = np.concatenate(imgs_full, axis=1) + cv2.imwrite(save_path, imgs_full) + return save_path + + def save_data(self, filename, data) -> str: + data = self.convert_data(data) + if isinstance(data, dict): + if not filename.endswith(".npz"): + filename += ".npz" + save_path = self.get_save_path(filename) + np.savez(save_path, **data) + else: + if not filename.endswith(".npy"): + filename += ".npy" + save_path = self.get_save_path(filename) + np.save(save_path, data) + return save_path + + def save_state_dict(self, filename, data) -> str: + save_path = self.get_save_path(filename) + torch.save(data, save_path) + return save_path + + def save_img_sequence( + self, + filename, + img_dir, + matcher, + save_format="mp4", + fps=30, + name: Optional[str] = None, + step: Optional[int] = None, + ) -> str: + assert save_format in ["gif", "mp4"] + if not filename.endswith(save_format): + filename += f".{save_format}" + save_path = self.get_save_path(filename) + matcher = re.compile(matcher) + img_dir = os.path.join(self.get_save_dir(), img_dir) + imgs = [] + for f in os.listdir(img_dir): + if matcher.search(f): + imgs.append(f) + imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0])) + imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs] + + if save_format == "gif": + imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs] + imageio.mimsave(save_path, imgs, fps=fps, palettesize=256) + elif save_format == "mp4": + imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs] + imageio.mimsave(save_path, imgs, fps=fps) + if name and self._wandb_logger: + lrm.warn("Wandb logger does not support video logging yet!") + return save_path + + def save_img_sequences( + self, + seq_dir, + matcher, + save_format="mp4", + fps=30, + delete=True, + name: Optional[str] = None, + step: Optional[int] = None, + ): + seq_dir_ = os.path.join(self.get_save_dir(), seq_dir) + for f in os.listdir(seq_dir_): + img_dir_ = os.path.join(seq_dir_, f) + if not os.path.isdir(img_dir_): + continue + try: + self.save_img_sequence( + os.path.join(seq_dir, f), + os.path.join(seq_dir, f), + matcher, + save_format=save_format, + fps=fps, + name=f"{name}_{f}", + step=step, + ) + if delete: + shutil.rmtree(img_dir_) + except: + lrm.warn(f"Video saving for directory {seq_dir_} failed!") + + def save_mesh(self, filename, v_pos, t_pos_idx, v_tex=None, t_tex_idx=None) -> str: + import trimesh + + save_path = self.get_save_path(filename) + v_pos = self.convert_data(v_pos) + t_pos_idx = self.convert_data(t_pos_idx) + mesh = trimesh.Trimesh(vertices=v_pos, faces=t_pos_idx) + mesh.export(save_path) + return save_path + + def save_obj( + self, + filename: str, + mesh: Mesh, + save_mat: bool = False, + save_normal: bool = False, + save_uv: bool = False, + save_vertex_color: bool = True, + map_Kd: Optional[Float[Tensor, "H W 3"]] = None, + map_Ks: Optional[Float[Tensor, "H W 3"]] = None, + map_Bump: Optional[Float[Tensor, "H W 3"]] = None, + map_Pm: Optional[Float[Tensor, "H W 1"]] = None, + map_Pr: Optional[Float[Tensor, "H W 1"]] = None, + map_format: str = "jpg", + ) -> List[str]: + save_paths: List[str] = [] + if not filename.endswith(".obj"): + filename += ".obj" + v_pos, t_pos_idx = self.convert_data(mesh.v_pos), self.convert_data( + mesh.t_pos_idx + ) + v_nrm, v_tex, t_tex_idx, v_rgb = None, None, None, None + if save_normal: + v_nrm = self.convert_data(mesh.v_nrm) + if save_uv: + v_tex, t_tex_idx = self.convert_data(mesh.v_tex), self.convert_data( + mesh.t_tex_idx + ) + if save_vertex_color: + v_rgb = self.convert_data(mesh.v_rgb) + matname, mtllib = None, None + if save_mat: + matname = "default" + mtl_filename = filename.replace(".obj", ".mtl") + mtllib = os.path.basename(mtl_filename) + mtl_save_paths = self._save_mtl( + mtl_filename, + matname, + map_Kd=self.convert_data(map_Kd), + map_Ks=self.convert_data(map_Ks), + map_Bump=self.convert_data(map_Bump), + map_Pm=self.convert_data(map_Pm), + map_Pr=self.convert_data(map_Pr), + map_format=map_format, + ) + save_paths += mtl_save_paths + obj_save_path = self._save_obj( + filename, + v_pos, + t_pos_idx, + v_nrm=v_nrm, + v_tex=v_tex, + t_tex_idx=t_tex_idx, + v_rgb=v_rgb, + matname=matname, + mtllib=mtllib, + ) + save_paths.append(obj_save_path) + return save_paths + + def _save_obj( + self, + filename, + v_pos, + t_pos_idx, + v_nrm=None, + v_tex=None, + t_tex_idx=None, + v_rgb=None, + matname=None, + mtllib=None, + ) -> str: + obj_str = "" + if matname is not None: + obj_str += f"mtllib {mtllib}\n" + obj_str += f"g object\n" + obj_str += f"usemtl {matname}\n" + for i in range(len(v_pos)): + obj_str += f"v {v_pos[i][0]} {v_pos[i][1]} {v_pos[i][2]}" + if v_rgb is not None: + obj_str += f" {v_rgb[i][0]} {v_rgb[i][1]} {v_rgb[i][2]}" + obj_str += "\n" + if v_nrm is not None: + for v in v_nrm: + obj_str += f"vn {v[0]} {v[1]} {v[2]}\n" + if v_tex is not None: + for v in v_tex: + obj_str += f"vt {v[0]} {1.0 - v[1]}\n" + + for i in range(len(t_pos_idx)): + obj_str += "f" + for j in range(3): + obj_str += f" {t_pos_idx[i][j] + 1}/" + if v_tex is not None: + obj_str += f"{t_tex_idx[i][j] + 1}" + obj_str += "/" + if v_nrm is not None: + obj_str += f"{t_pos_idx[i][j] + 1}" + obj_str += "\n" + + save_path = self.get_save_path(filename) + with open(save_path, "w") as f: + f.write(obj_str) + return save_path + + def _save_mtl( + self, + filename, + matname, + Ka=(0.0, 0.0, 0.0), + Kd=(1.0, 1.0, 1.0), + Ks=(0.0, 0.0, 0.0), + map_Kd=None, + map_Ks=None, + map_Bump=None, + map_Pm=None, + map_Pr=None, + map_format="jpg", + step: Optional[int] = None, + ) -> List[str]: + mtl_save_path = self.get_save_path(filename) + save_paths = [mtl_save_path] + mtl_str = f"newmtl {matname}\n" + mtl_str += f"Ka {Ka[0]} {Ka[1]} {Ka[2]}\n" + if map_Kd is not None: + map_Kd_save_path = os.path.join( + os.path.dirname(mtl_save_path), f"texture_kd.{map_format}" + ) + mtl_str += f"map_Kd texture_kd.{map_format}\n" + self._save_rgb_image( + map_Kd_save_path, + map_Kd, + data_format="HWC", + data_range=(0, 1), + name=f"{matname}_Kd", + step=step, + ) + save_paths.append(map_Kd_save_path) + else: + mtl_str += f"Kd {Kd[0]} {Kd[1]} {Kd[2]}\n" + if map_Ks is not None: + map_Ks_save_path = os.path.join( + os.path.dirname(mtl_save_path), f"texture_ks.{map_format}" + ) + mtl_str += f"map_Ks texture_ks.{map_format}\n" + self._save_rgb_image( + map_Ks_save_path, + map_Ks, + data_format="HWC", + data_range=(0, 1), + name=f"{matname}_Ks", + step=step, + ) + save_paths.append(map_Ks_save_path) + else: + mtl_str += f"Ks {Ks[0]} {Ks[1]} {Ks[2]}\n" + if map_Bump is not None: + map_Bump_save_path = os.path.join( + os.path.dirname(mtl_save_path), f"texture_nrm.{map_format}" + ) + mtl_str += f"map_Bump texture_nrm.{map_format}\n" + self._save_rgb_image( + map_Bump_save_path, + map_Bump, + data_format="HWC", + data_range=(0, 1), + name=f"{matname}_Bump", + step=step, + ) + save_paths.append(map_Bump_save_path) + if map_Pm is not None: + map_Pm_save_path = os.path.join( + os.path.dirname(mtl_save_path), f"texture_metallic.{map_format}" + ) + mtl_str += f"map_Pm texture_metallic.{map_format}\n" + self._save_grayscale_image( + map_Pm_save_path, + map_Pm, + data_range=(0, 1), + cmap=None, + name=f"{matname}_refl", + step=step, + ) + save_paths.append(map_Pm_save_path) + if map_Pr is not None: + map_Pr_save_path = os.path.join( + os.path.dirname(mtl_save_path), f"texture_roughness.{map_format}" + ) + mtl_str += f"map_Pr texture_roughness.{map_format}\n" + self._save_grayscale_image( + map_Pr_save_path, + map_Pr, + data_range=(0, 1), + cmap=None, + name=f"{matname}_Ns", + step=step, + ) + save_paths.append(map_Pr_save_path) + with open(self.get_save_path(filename), "w") as f: + f.write(mtl_str) + return save_paths + + def save_glb( + self, + filename: str, + mesh: Mesh, + save_mat: bool = False, + save_normal: bool = True, + save_uv: bool = True, + save_vertex_color: bool = True, + map_Kd: Optional[Float[Tensor, "H W 3"]] = None, + map_Ks: Optional[Float[Tensor, "H W 3"]] = None, + map_Bump: Optional[Float[Tensor, "H W 3"]] = None, + map_Pm: Optional[Float[Tensor, "H W 1"]] = None, + map_Pr: Optional[Float[Tensor, "H W 1"]] = None, + map_format: str = "jpg", + ) -> List[str]: + save_paths: List[str] = [] + if not filename.endswith(".glb"): + filename += ".glb" + v_pos, t_pos_idx = self.convert_data(mesh.v_pos), self.convert_data( + mesh.t_pos_idx + ) + v_nrm, v_tex, t_tex_idx, v_rgb = None, None, None, None + if save_normal: + v_nrm = self.convert_data(mesh.v_nrm) + if save_uv: + v_tex, t_tex_idx = self.convert_data(mesh.v_tex), self.convert_data( + mesh.t_tex_idx + ) + if save_vertex_color: + v_rgb = self.convert_data(mesh.v_rgb) + + obj_save_path = self._save_glb( + filename, + v_pos, + t_pos_idx, + v_nrm=v_nrm, + v_tex=v_tex, + t_tex_idx=t_tex_idx, + v_rgb=v_rgb, + ) + save_paths.append(obj_save_path) + return save_paths + + def _save_glb( + self, + filename, + v_pos, + t_pos_idx, + v_nrm=None, + v_tex=None, + t_tex_idx=None, + v_rgb=None, + matname=None, + mtllib=None, + ) -> str: + import trimesh + + mesh = trimesh.Trimesh( + vertices=v_pos, faces=t_pos_idx, vertex_normals=v_nrm, vertex_colors=v_rgb + ) + # not tested + if v_tex is not None: + mesh.visual = trimesh.visual.TextureVisuals(uv=v_tex) + + save_path = self.get_save_path(filename) + mesh.export(save_path) + return save_path + + def save_file(self, filename, src_path, delete=False) -> str: + save_path = self.get_save_path(filename) + shutil.copyfile(src_path, save_path) + if delete: + os.remove(src_path) + return save_path + + def save_json(self, filename, payload) -> str: + save_path = self.get_save_path(filename) + with open(save_path, "w") as f: + f.write(json.dumps(payload)) + return save_path diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/sdf.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/sdf.py new file mode 100644 index 0000000..93f3e02 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/sdf.py @@ -0,0 +1,13 @@ +import pySDF as SDF +import numpy as np +from CharacterGen.Stage_3D.lrm.models.isosurface import MarchingTetrahedraHelper + +def get_tetra_for_mesh(mesh_path, resolution=128): + isosurface_helper = MarchingTetrahedraHelper(resolution, f"load/{resolution}_tets.npz") + isosurface_helper.points_range = (-1, 1) + mesh = trimesh.load(mesh_path) + dmtet = np.load(f"") + sdf = SDF(mesh.vertices, mesh.faces) + sdf_gt = sdf(isosurface_helper.grid_vertices.numpy()) + return sdf_gt + diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/typing.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/typing.py new file mode 100644 index 0000000..dee9f96 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/utils/typing.py @@ -0,0 +1,40 @@ +""" +This module contains type annotations for the project, using +1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects +2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors + +Two types of typing checking can be used: +1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode) +2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking) +""" + +# Basic types +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + Literal, + NamedTuple, + NewType, + Optional, + Sized, + Tuple, + Type, + TypeVar, + Union, +) + +# Tensor dtype +# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md +from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt + +# Config type +from omegaconf import DictConfig + +# PyTorch Tensor type +from torch import Tensor + +# Runtime type checking decorator +from typeguard import typechecked as typechecker diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/1.png b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/1.png new file mode 100644 index 0000000..4bbee21 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/1.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/2.png b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/2.png new file mode 100644 index 0000000..9ad67e2 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/2.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/3.png b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/3.png new file mode 100644 index 0000000..5a403ea Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/3.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/4.png b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/4.png new file mode 100644 index 0000000..6a359b5 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/examples/1/4.png differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/meta.json b/Gen_3D_Modules/CharacterGen/Stage_3D/material/meta.json new file mode 100644 index 0000000..9ca2e42 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/material/meta.json @@ -0,0 +1,36 @@ +{ + "locations": [ + { + "transform_matrix": [ + [ 1.0, 0.0, 0.0, 0.0 ], + [ 0.0, 0.0, 1.0, 1.5 ], + [ 0.0, 1.0, 0.0, 0.0 ], + [ 0.0, 0.0, 0.0, 1.0 ] + ] + }, + { + "transform_matrix": [ + [ -1.0, 0.0, 0.0, 0.0 ], + [ 0.0, 0.0, -1.0, -1.5 ], + [ 0.0, 1.0, 0.0, 0.0 ], + [ 0.0, 0.0, 0.0, 1.0 ] + ] + }, + { + "transform_matrix": [ + [ 0.0, 0.0, 1.0, 1.5 ], + [ -1.0, 0.0, 0.0, 0.0 ], + [ 0.0, 1.0, 0.0, 0.0 ], + [ 0.0, 0.0, 0.0, 1.0 ] + ] + }, + { + "transform_matrix": [ + [ 0.0, 0.0, -1.0, -1.5 ], + [ 1.0, 0.0, 0.0, 0.0 ], + [ 0.0, 1.0, 0.0, 0.0 ], + [ 0.0, 0.0, 0.0, 1.0 ] + ] + } + ] +} \ No newline at end of file diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/128_tets.npz b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/128_tets.npz new file mode 100644 index 0000000..2ea2f1d Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/128_tets.npz differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/256_tets.npz b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/256_tets.npz new file mode 100644 index 0000000..fb10dee Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/256_tets.npz differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/32_tets.npz b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/32_tets.npz new file mode 100644 index 0000000..54f7ea5 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/32_tets.npz differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/64_tets.npz b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/64_tets.npz new file mode 100644 index 0000000..ec37c89 Binary files /dev/null and b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/64_tets.npz differ diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/generate_tets.py b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/generate_tets.py new file mode 100644 index 0000000..424d852 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/material/tets/generate_tets.py @@ -0,0 +1,58 @@ +# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import os + +import numpy as np + +""" +This code segment shows how to use Quartet: https://github.com/crawforddoran/quartet, +to generate a tet grid +1) Download, compile and run Quartet as described in the link above. Example usage `quartet meshes/cube.obj 0.5 cube_5.tet` +2) Run the function below to generate a file `cube_32_tet.tet` +""" + + +def generate_tetrahedron_grid_file(res=32, root=".."): + frac = 1.0 / res + command = f"cd {root}; ./quartet meshes/cube.obj {frac} meshes/cube_{res}_tet.tet -s meshes/cube_boundary_{res}.obj" + os.system(command) + + +""" +This code segment shows how to convert from a quartet .tet file to compressed npz file +""" + + +def convert_from_quartet_to_npz(quartetfile="cube_32_tet.tet", npzfile="32_tets"): + file1 = open(quartetfile, "r") + header = file1.readline() + numvertices = int(header.split(" ")[1]) + numtets = int(header.split(" ")[2]) + print(numvertices, numtets) + + # load vertices + vertices = np.loadtxt(quartetfile, skiprows=1, max_rows=numvertices) + print(vertices.shape) + + # load indices + indices = np.loadtxt( + quartetfile, dtype=int, skiprows=1 + numvertices, max_rows=numtets + ) + print(indices.shape) + + np.savez_compressed(npzfile, vertices=vertices, indices=indices) + + +root = "/home/gyc/quartet" +for res in [300, 350, 400]: + generate_tetrahedron_grid_file(res, root) + convert_from_quartet_to_npz( + os.path.join(root, f"meshes/cube_{res}_tet.tet"), npzfile=f"{res}_tets" + ) diff --git a/Gen_3D_Modules/CharacterGen/character_inference.py b/Gen_3D_Modules/CharacterGen/character_inference.py new file mode 100644 index 0000000..4ed3525 --- /dev/null +++ b/Gen_3D_Modules/CharacterGen/character_inference.py @@ -0,0 +1,236 @@ +import sys +import os +from os.path import dirname +from typing import Dict, List +import numpy as np + +import torch +import torch.utils.checkpoint + +from diffusers import AutoencoderKL, DDIMScheduler +from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection +from torchvision import transforms + +from CharacterGen.Stage_2D.tuneavideo.models.unet_mv2d_condition import UNetMV2DConditionModel +from CharacterGen.Stage_2D.tuneavideo.models.unet_mv2d_ref import UNetMV2DRefModel +from CharacterGen.Stage_2D.tuneavideo.models.PoseGuider import PoseGuider +from CharacterGen.Stage_2D.tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline + +from einops import rearrange +from PIL import Image +import json + +import trimesh +from CharacterGen.Stage_3D import lrm +from datetime import datetime +from pygltflib import GLTF2 +import pymeshlab + +CHARACTER_GEN_ROOT_ABS_PATH = dirname(__file__) +CHARACTER_GEN_2D_DATA_ABS_PATH = os.path.join(CHARACTER_GEN_ROOT_ABS_PATH, "Stage_2D/material") +CHARACTER_GEN_3D_ABS_PATH = os.path.join(CHARACTER_GEN_ROOT_ABS_PATH, "Stage_3D") +CHARACTER_GEN_3D_DATA_ABS_PATH = os.path.join(CHARACTER_GEN_3D_ABS_PATH, "material") + +sys.path.append(CHARACTER_GEN_3D_ABS_PATH) + +def get_bg_color(bg_color): + if bg_color == 'white': + bg_color = np.array([1., 1., 1.], dtype=np.float32) + elif bg_color == 'black': + bg_color = np.array([0., 0., 0.], dtype=np.float32) + elif bg_color == 'gray': + bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32) + elif bg_color == 'random': + bg_color = np.random.rand(3) + elif isinstance(bg_color, float): + bg_color = np.array([bg_color] * 3, dtype=np.float32) + else: + raise NotImplementedError + return bg_color + +def process_image(image, totensor, w=512, h=768): + if not image.mode == "RGBA": + image = image.convert("RGBA") + + # Find non-transparent pixels + non_transparent = np.nonzero(np.array(image)[..., 3]) + min_x, max_x = non_transparent[1].min(), non_transparent[1].max() + min_y, max_y = non_transparent[0].min(), non_transparent[0].max() + image = image.crop((min_x, min_y, max_x, max_y)) + + # paste to center + max_dim = max(image.width, image.height) + max_height = max_dim + max_width = int(max_dim / 3 * 2) + new_image = Image.new("RGBA", (max_width, max_height)) + left = (max_width - image.width) // 2 + top = (max_height - image.height) // 2 + new_image.paste(image, (left, top)) + + image = new_image.resize((w, h), resample=Image.Resampling.BICUBIC) + image = np.array(image) + image = image.astype(np.float32) / 255. + assert image.shape[-1] == 4 # RGBA + alpha = image[..., 3:4] + bg_color = get_bg_color("gray") + image = image[..., :3] * alpha + bg_color * (1 - alpha) + # save image + new_image = Image.fromarray((image * 255).astype(np.uint8)) + return totensor(image) + +class Inference2D_API: + + def __init__(self, + pretrained_model_path: str, + checkpoint_root_path: str, + image_encoder_path: str, + ckpt_dir: str, + validation: Dict, + local_crossattn: bool = True, + unet_from_pretrained_kwargs=None, + unet_condition_type=None, + use_pose_guider=False, + use_shifted_noise=False, + use_noise=True, + device="cuda", + weight_dtype=torch.float16 + ): + self.validation = validation + self.use_noise = use_noise + self.use_shifted_noise = use_shifted_noise + self.unet_condition_type = unet_condition_type + self.device = device + self.weight_dtype = weight_dtype + image_encoder_path = os.path.join(checkpoint_root_path, image_encoder_path) + ckpt_dir = os.path.join(checkpoint_root_path, ckpt_dir) + + self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer") + text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder") + image_encoder = CLIPVisionModelWithProjection.from_pretrained(image_encoder_path) + feature_extractor = CLIPImageProcessor() + vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae") + unet = UNetMV2DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) + ref_unet = UNetMV2DRefModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", local_crossattn=local_crossattn, **unet_from_pretrained_kwargs) + if use_pose_guider: + pose_guider = PoseGuider(noise_latent_channels=4).to(device) + else: + pose_guider = None + + unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model.bin"), map_location="cpu") + if use_pose_guider: + pose_guider_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_1.bin"), map_location="cpu") + ref_unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_2.bin"), map_location="cpu") + pose_guider.load_state_dict(pose_guider_params) + else: + ref_unet_params = torch.load(os.path.join(ckpt_dir, "pytorch_model_1.bin"), map_location="cpu") + unet.load_state_dict(unet_params) + ref_unet.load_state_dict(ref_unet_params) + + text_encoder.to(device, dtype=weight_dtype) + image_encoder.to(device, dtype=weight_dtype) + vae.to(device, dtype=weight_dtype) + ref_unet.to(device, dtype=weight_dtype) + unet.to(device, dtype=weight_dtype) + + vae.requires_grad_(False) + unet.requires_grad_(False) + ref_unet.requires_grad_(False) + + noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler") + self.validation_pipeline = TuneAVideoPipeline( + vae=vae, text_encoder=text_encoder, tokenizer=self.tokenizer, unet=unet, ref_unet=ref_unet,feature_extractor=feature_extractor,image_encoder=image_encoder, + scheduler=noise_scheduler + ) + self.validation_pipeline.enable_vae_slicing() + self.validation_pipeline.set_progress_bar_config(disable=True) + + @torch.no_grad() + def inference( + self, input_image, val_width, val_height, prompt="high quality, best quality", prompt_neg=None, + guidance_scale=7.5, num_inference_steps=40, seed=100, use_shifted_noise=False, crop=False + ): + + generator = torch.Generator(device=self.device).manual_seed(seed) + totensor = transforms.ToTensor() + + metas = json.load(open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, "pose.json"), "r")) + cameras = [] + pose_images = [] + for lm in metas: + cameras.append(torch.tensor(np.array(lm[0]).reshape(4, 4).transpose(1,0)[:3, :4]).reshape(-1)) + if not crop: + pose_images.append(totensor(np.asarray(Image.open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, lm[1])).resize( + (val_height, val_width), resample=Image.Resampling.BICUBIC)).astype(np.float32) / 255.)) + else: + pose_image = Image.open(os.path.join(CHARACTER_GEN_2D_DATA_ABS_PATH, lm[1])) + crop_area = (128, 0, 640, 768) + pose_images.append(totensor(np.array(pose_image.crop(crop_area)).astype(np.float32)) / 255.) + camera_matrixs = torch.stack(cameras).unsqueeze(0).to(self.device) + pose_imgs_in = torch.stack(pose_images).to(self.device) + prompt_ids = self.tokenizer( + prompt, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" + ).input_ids[0] + + # (B*Nv, 3, H, W) B==1 + imgs_in = process_image(input_image, totensor, val_width, val_height) + imgs_in = rearrange(imgs_in.unsqueeze(0).unsqueeze(0), "B Nv C H W -> (B Nv) C H W") + + with torch.autocast(self.device, dtype=self.weight_dtype): + imgs_in = imgs_in.to(self.device) + # B*Nv images + out = self.validation_pipeline(prompt=prompt, negative_prompt=prompt_neg, image=imgs_in.to(self.weight_dtype), generator=generator, + num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, + camera_matrixs=camera_matrixs.to(self.weight_dtype), prompt_ids=prompt_ids, + height=val_height, width=val_width, unet_condition_type=self.unet_condition_type, + pose_guider=None, pose_image=pose_imgs_in, use_noise=self.use_noise, + use_shifted_noise=use_shifted_noise, **self.validation).videos + out = rearrange(out, "B C f H W -> (B f) H W C", f=self.validation.video_length) + + torch.cuda.empty_cache() + return out + +class Inference3D_API: + + def __init__(self, checkpoint_root_path, cfg, device="cuda"): + self.cfg = cfg + self.device = device + self.cfg.system.weights = os.path.join(checkpoint_root_path, self.cfg.system.weights) + self.cfg.system.image_tokenizer.pretrained_model_name_or_path = os.path.join(checkpoint_root_path, self.cfg.system.image_tokenizer.pretrained_model_name_or_path) + self.cfg.system.renderer.tet_dir = os.path.join(CHARACTER_GEN_3D_DATA_ABS_PATH, self.cfg.system.renderer.tet_dir) + self.system = lrm.find(self.cfg.system_cls)(self.cfg.system).to(self.device) + self.system.eval() + + def inference(self, mv_imgs: List[Image.Image]): + meta = json.load(open(os.path.join(CHARACTER_GEN_3D_DATA_ABS_PATH, "meta.json"))) + c2w_cond = [np.array(loc["transform_matrix"]) for loc in meta["locations"]] + c2w_cond = torch.from_numpy(np.stack(c2w_cond, axis=0)).float()[None].to(self.device) + + rgb_cond = [] + new_images = [] + for pil_img in mv_imgs: + image = np.array(pil_img) + image = Image.fromarray(image) + if image.width != image.height: + max_dim = max(image.width, image.height) + new_image = Image.new("RGBA", (max_dim, max_dim)) + left = (max_dim - image.width) // 2 + top = (max_dim - image.height) // 2 + new_image.paste(image, (left, top)) + image = new_image + + image = image.convert('RGB') + rgb = image.resize((self.cfg.data.cond_width, self.cfg.data.cond_height), Image.Resampling.BILINEAR) + rgb = np.array(rgb).astype(np.float32) / 255.0 + new_images.append(image) + rgb_cond.append(rgb) + assert len(rgb_cond) == 4, "Please provide 4 images" + + rgb_cond = torch.from_numpy(np.stack(rgb_cond, axis=0)).float()[None].to(self.device) + + scene_code = self.system({"rgb_cond": rgb_cond, "c2w_cond": c2w_cond})[0] + mesh = self.system.get_geometry(scene_code) + + # Flip mesh horizontally + #mesh.v_pos[:, 0] = -mesh.v_pos[:, 0] + + return mesh.v_pos.detach(), mesh.t_pos_idx.detach() diff --git a/README.md b/README.md index dd9b1c5..a33e52f 100644 --- a/README.md +++ b/README.md @@ -25,26 +25,28 @@ This is an extensive node suite that enables ComfyUI to process 3D inputs (Mesh - **Note:** at this moment, you'll still need to install [Visual Studio Build Tools for windows](_Pre_Builds/README.md#build-for-windows) and [install `gcc g++` for Linux](_Pre_Builds/README.md#build-for-linux) in order for `InstantNGP & Convert 3DGS to Mesh with NeRF and Marching_Cubes` nodes to work, since those two nodes used JIT torch cpp extension that builds in runtime, but I plan to replace those nodes soon ## Features: -- For use cases please check out [Example Workflows](./_Example_Workflows/). [**Last update: 04/July/2024**] +- For use cases please check out [Example Workflows](./_Example_Workflows/). [**Last update: 09/July/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_) +- **CharacterGen**: [zjp-shadow/CharacterGen](https://github.com/zjp-shadow/CharacterGen) + - Single front view of a character with arbitrary pose + - Can [combine with Unique3D workflow](./_Example_Workflows/CharacterGen/CharacterGen_to_Unique3D.json) for better result + - **Unique3D**: [AiuniAI/Unique3D](https://github.com/AiuniAI/Unique3D) - Four stages pipeline: 1. Single image to 4 multi-view images with resulution: 256X256 2. Consistent Multi-view images Upscale to 512X512, super resolution to 2048X2048 3. Multi-view images to Normal maps with resulution: 512X512, super resolution to 2048X2048 4. Multi-view images & Normal maps to 3D mesh with texture - - To use the [pure Unique3D workflow](./_Example_Workflows/Unique3D/Unique3D_All_Stages.json), Download Models: - - [img2mvimg](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/img2mvimg) and put it into [./checkpoints/Wuvin/Unique3D/image2mvimage](./checkpoints/Wuvin/Unique3D/image2mvimage) - - [image2normal](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/image2normal) and put it into [./checkpoints/Wuvin/Unique3D/image2normal](./checkpoints/Wuvin/Unique3D/image2normal) + - To use the [All stage Unique3D workflow](./_Example_Workflows/Unique3D/Unique3D_All_Stages.json), Download Models: - [fine-tuned controlnet-tile](https://huggingface.co/spaces/Wuvin/Unique3D/tree/main/ckpt/controlnet-tile) and put it into `Your ComfyUI root directory/ComfyUI/models/controlnet` - [ip-adapter_sd15](https://huggingface.co/h94/IP-Adapter/blob/main/models/ip-adapter_sd15.safetensors) and put it into `Your ComfyUI root directory/ComfyUI/models/ipadapter` - [RealESRGAN_x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth) and put it into `Your ComfyUI root directory/ComfyUI/models/upscale_models` -- **Era3D Diffusion Model**: [pengHTYX/Era3D](https://github.com/pengHTYX/Era3D) +- **Era3D MVDiffusion 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* @@ -155,7 +157,7 @@ This is an extensive node suite that enables ComfyUI to process 3D inputs (Mesh - Custom clients IP address ## Roadmap: -- [ ] Integrate [CharacterGen](https://github.com/zjp-shadow/CharacterGen) +- [X] Integrate [CharacterGen](https://github.com/zjp-shadow/CharacterGen) - [ ] Improve 3DGS/Nerf to Mesh conversion algorithms: - Find better methods to converts 3DGS or Points Cloud to Mesh (Normal maps reconstruction maybe?) diff --git a/_Example_Workflows/CRM MVDiffusion_to_InstantMesh.json b/_Example_Workflows/CRM MVDiffusion_to_InstantMesh.json index b5e7b25..1f4fb29 100644 --- a/_Example_Workflows/CRM MVDiffusion_to_InstantMesh.json +++ b/_Example_Workflows/CRM MVDiffusion_to_InstantMesh.json @@ -1,40 +1,7 @@ { - "last_node_id": 30, - "last_link_id": 46, + "last_node_id": 32, + "last_link_id": 48, "nodes": [ - { - "id": 2, - "type": "[Comfy3D] Load CRM MVDiffusion Model", - "pos": [ - -102, - 131 - ], - "size": { - "0": 489.21466064453125, - "1": 95.50737762451172 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "crm_mvdiffusion_sampler", - "type": "CRM_MVDIFFUSION_SAMPLER", - "links": [ - 1 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "[Comfy3D] Load CRM MVDiffusion Model" - }, - "widgets_values": [ - "pixel-diffusion.pth", - "configs/crm_configs/sd_v2_base_ipmv_zero_SNR.yaml" - ] - }, { "id": 24, "type": "[Comfy3D] Save 3D Mesh", @@ -47,7 +14,7 @@ "1": 58 }, "flags": {}, - "order": 11, + "order": 12, "mode": 0, "inputs": [ { @@ -86,7 +53,7 @@ "1": 130 }, "flags": {}, - "order": 10, + "order": 11, "mode": 0, "inputs": [ { @@ -116,69 +83,6 @@ false ] }, - { - "id": 3, - "type": "[Comfy3D] CRM Images MVDiffusion Model", - "pos": [ - 600, - 300 - ], - "size": { - "0": 506.4000244140625, - "1": 262 - }, - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [ - { - "name": "crm_mvdiffusion_sampler", - "type": "CRM_MVDIFFUSION_SAMPLER", - "link": 1 - }, - { - "name": "reference_image", - "type": "IMAGE", - "link": 2 - }, - { - "name": "reference_mask", - "type": "MASK", - "link": 28 - } - ], - "outputs": [ - { - "name": "multiview_images", - "type": "IMAGE", - "links": [ - 39 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "orbit_camposes", - "type": "ORBIT_CAMPOSES", - "links": [ - 41 - ], - "shape": 3, - "slot_index": 1 - } - ], - "properties": { - "Node name for S&R": "[Comfy3D] CRM Images MVDiffusion Model" - }, - "widgets_values": [ - "3D assets", - "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.", - 1234, - "fixed", - 5.5, - 50 - ] - }, { "id": 25, "type": "[Comfy3D] Preview 3DMesh", @@ -191,7 +95,7 @@ 500 ], "flags": {}, - "order": 12, + "order": 13, "mode": 0, "inputs": [ { @@ -211,45 +115,6 @@ null ] }, - { - "id": 27, - "type": "GetImagesFromBatchIndexed", - "pos": [ - 1141, - 214 - ], - "size": { - "0": 255.99082946777344, - "1": 76.0000228881836 - }, - "flags": {}, - "order": 5, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 39 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 44 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "GetImagesFromBatchIndexed" - }, - "widgets_values": [ - "5, 3, 2, 0" - ] - }, { "id": 22, "type": "[Comfy3D] Load InstantMesh Reconstruction Model", @@ -262,7 +127,7 @@ "1": 58 }, "flags": {}, - "order": 1, + "order": 0, "mode": 0, "outputs": [ { @@ -294,7 +159,7 @@ "1": 76 }, "flags": {}, - "order": 6, + "order": 7, "mode": 0, "inputs": [ { @@ -333,7 +198,7 @@ "1": 250 }, "flags": {}, - "order": 7, + "order": 8, "mode": 0, "inputs": [ { @@ -381,7 +246,7 @@ "1": 122 }, "flags": {}, - "order": 9, + "order": 10, "mode": 0, "inputs": [ { @@ -467,7 +332,7 @@ "1": 545.7767944335938 }, "flags": {}, - "order": 2, + "order": 1, "mode": 0, "outputs": [ { @@ -497,19 +362,180 @@ "image" ] }, + { + "id": 3, + "type": "[Comfy3D] CRM Images MVDiffusion Model", + "pos": [ + 600, + 300 + ], + "size": { + "0": 506.4000244140625, + "1": 262 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "crm_mvdiffusion_sampler", + "type": "CRM_MVDIFFUSION_SAMPLER", + "link": 48 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 2 + }, + { + "name": "reference_mask", + "type": "MASK", + "link": 28 + } + ], + "outputs": [ + { + "name": "multiview_images", + "type": "IMAGE", + "links": [ + 39, + 47 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [ + 41 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] CRM Images MVDiffusion Model" + }, + "widgets_values": [ + "3D assets", + "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.", + 1234, + "fixed", + 5.5, + 50 + ] + }, + { + "id": 27, + "type": "GetImagesFromBatchIndexed", + "pos": [ + 1152, + 214 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 39 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 44 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "5, 3, 2, 0" + ] + }, + { + "id": 32, + "type": "[Comfy3D] Load CRM MVDiffusion Model", + "pos": [ + -157, + 163 + ], + "size": [ + 435.4573393577026, + 83.6087870826683 + ], + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "crm_mvdiffusion_sampler", + "type": "CRM_MVDIFFUSION_SAMPLER", + "links": [ + 48 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load CRM MVDiffusion Model" + }, + "widgets_values": [ + "pixel-diffusion.pth", + "sd_v2_base_ipmv_zero_SNR.yaml" + ] + }, + { + "id": 31, + "type": "PreviewImage", + "pos": [ + 1103, + -110 + ], + "size": [ + 685.57771038656, + 142.33505712129315 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 47 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, { "id": 15, "type": "PreviewImage", "pos": [ - 1805, - 494 + 1820, + 35 + ], + "size": [ + 675.7496017470316, + 189.98764924222667 ], - "size": { - "0": 770.9938354492188, - "1": 246 - }, "flags": {}, - "order": 8, + "order": 9, "mode": 0, "inputs": [ { @@ -524,14 +550,6 @@ } ], "links": [ - [ - 1, - 2, - 0, - 3, - 0, - "CRM_MVDIFFUSION_SAMPLER" - ], [ 2, 1, @@ -635,16 +653,32 @@ 23, 1, "IMAGE" + ], + [ + 47, + 3, + 0, + 31, + 0, + "IMAGE" + ], + [ + 48, + 32, + 0, + 3, + 0, + "CRM_MVDIFFUSION_SAMPLER" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.5644739300537773, + "scale": 0.42409761837248683, "offset": [ - 148.6926373068756, - 262.952479355259 + -1895.6373310621698, + 224.7040148153894 ] } }, diff --git a/_Example_Workflows/CharacterGen/CharacterGen_InstantMesh.json b/_Example_Workflows/CharacterGen/CharacterGen_InstantMesh.json new file mode 100644 index 0000000..2a44a6d --- /dev/null +++ b/_Example_Workflows/CharacterGen/CharacterGen_InstantMesh.json @@ -0,0 +1,942 @@ +{ + "last_node_id": 43, + "last_link_id": 56, + "nodes": [ + { + "id": 18, + "type": "[Comfy3D] Switch Mesh Axis", + "pos": [ + 1680, + 520 + ], + "size": { + "0": 315, + "1": 154 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 21 + } + ], + "outputs": [ + { + "name": "switched_mesh", + "type": "MESH", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Switch Mesh Axis" + }, + "widgets_values": [ + "+y", + "+z", + "+x", + false, + 1 + ] + }, + { + "id": 14, + "type": "[Comfy3D] Load InstantMesh Reconstruction Model", + "pos": [ + 930, + 300 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "lrm_model", + "type": "LRM_MODEL", + "links": [ + 18 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load InstantMesh Reconstruction Model" + }, + "widgets_values": [ + "instant_mesh_large.ckpt" + ] + }, + { + "id": 32, + "type": "IMG_padder", + "pos": [ + 210, + 440 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 43 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 44 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IMG_padder" + }, + "widgets_values": [ + 0.25, + "left" + ] + }, + { + "id": 34, + "type": "[Comfy3D] Load CharacterGen MVDiffusion Model", + "pos": [ + -1900, + 360 + ], + "size": { + "0": 294, + "1": 26 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "links": [ + 46 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load CharacterGen MVDiffusion Model" + } + }, + { + "id": 36, + "type": "InvertMask", + "pos": [ + -1820, + 530 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 45 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 48 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 39, + "type": "Image Rembg (Remove Background)", + "pos": [ + -750, + 510 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 50 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 51, + 53 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Image Rembg (Remove Background)" + }, + "widgets_values": [ + false, + "isnet-anime", + false, + true, + false, + 240, + 10, + 10, + "white" + ] + }, + { + "id": 40, + "type": "ImageToMask", + "pos": [ + -400, + 500 + ], + "size": { + "0": 210, + "1": 58 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 51 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 42 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageToMask" + }, + "widgets_values": [ + "red" + ] + }, + { + "id": 31, + "type": "[Comfy3D] Image Add Pure Color Background", + "pos": [ + -160, + 440 + ], + "size": { + "0": 315, + "1": 126 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 41 + }, + { + "name": "masks", + "type": "MASK", + "link": 42 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 43 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Image Add Pure Color Background" + }, + "widgets_values": [ + 255, + 255, + 255 + ] + }, + { + "id": 33, + "type": "IMG_padder", + "pos": [ + 570, + 440 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 44 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 40 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IMG_padder" + }, + "widgets_values": [ + 0.2, + "right" + ] + }, + { + "id": 30, + "type": "ImageScale", + "pos": [ + 940, + 440 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 40 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 54 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "lanczos", + 1024, + 1024, + "disabled" + ] + }, + { + "id": 38, + "type": "GetImagesFromBatchIndexed", + "pos": [ + -1060, + 440 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 49 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 41, + 50, + 52 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "3, 1, 2, 0" + ] + }, + { + "id": 37, + "type": "[Comfy3D] CharacterGen MVDiffusion Model", + "pos": [ + -1570, + 440 + ], + "size": { + "0": 456, + "1": 334 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "link": 46 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 47 + }, + { + "name": "reference_mask", + "type": "MASK", + "link": 48 + } + ], + "outputs": [ + { + "name": "multiviews", + "type": "IMAGE", + "links": [ + 49 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [ + 55 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] CharacterGen MVDiffusion Model" + }, + "widgets_values": [ + 512, + 768, + 2333, + "fixed", + 5, + 40, + "high quality, best quality", + "", + 1.5 + ] + }, + { + "id": 41, + "type": "PreviewImage", + "pos": [ + -1273, + 109 + ], + "size": [ + 718.229248046875, + 246 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 52 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 42, + "type": "PreviewImage", + "pos": [ + -457, + 104 + ], + "size": [ + 608.1014404296875, + 246 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 53 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 15, + "type": "[Comfy3D] InstantMesh Reconstruction Model", + "pos": [ + 1330, + 520 + ], + "size": { + "0": 315, + "1": 122 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "lrm_model", + "type": "LRM_MODEL", + "link": 18 + }, + { + "name": "multiview_images", + "type": "IMAGE", + "link": 54 + }, + { + "name": "orbit_camera_poses", + "type": "ORBIT_CAMPOSES", + "link": 56 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 21 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] InstantMesh Reconstruction Model" + }, + "widgets_values": [ + 40, + 1024 + ] + }, + { + "id": 43, + "type": "[Comfy3D] Get Camposes From List Indexed", + "pos": [ + -988, + 821 + ], + "size": { + "0": 531.5999755859375, + "1": 200 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "original_orbit_camera_poses", + "type": "ORBIT_CAMPOSES", + "link": 55 + } + ], + "outputs": [ + { + "name": "ORBIT_CAMPOSES", + "type": "ORBIT_CAMPOSES", + "links": [ + 56 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Get Camposes From List Indexed" + }, + "widgets_values": [ + "3, 1, 2, 0" + ] + }, + { + "id": 35, + "type": "LoadImage", + "pos": [ + -2230, + 460 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 47 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 45 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "CharacterGen_Test_2.png", + "image" + ] + }, + { + "id": 16, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2021, + 520 + ], + "size": [ + 288.29027143395615, + 58 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 19 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 20 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "InstantMesh_test\\Mesh_Girl.glb" + ] + }, + { + "id": 17, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 2344, + 519 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 20, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + } + ], + "links": [ + [ + 18, + 14, + 0, + 15, + 0, + "LRM_MODEL" + ], + [ + 19, + 18, + 0, + 16, + 0, + "MESH" + ], + [ + 20, + 16, + 0, + 17, + 0, + "STRING" + ], + [ + 21, + 15, + 0, + 18, + 0, + "MESH" + ], + [ + 40, + 33, + 0, + 30, + 0, + "IMAGE" + ], + [ + 41, + 38, + 0, + 31, + 0, + "IMAGE" + ], + [ + 42, + 40, + 0, + 31, + 1, + "MASK" + ], + [ + 43, + 31, + 0, + 32, + 0, + "IMAGE" + ], + [ + 44, + 32, + 0, + 33, + 0, + "IMAGE" + ], + [ + 45, + 35, + 1, + 36, + 0, + "MASK" + ], + [ + 46, + 34, + 0, + 37, + 0, + "CHARACTER_MV_GEN_PIPE" + ], + [ + 47, + 35, + 0, + 37, + 1, + "IMAGE" + ], + [ + 48, + 36, + 0, + 37, + 2, + "MASK" + ], + [ + 49, + 37, + 0, + 38, + 0, + "IMAGE" + ], + [ + 50, + 38, + 0, + 39, + 0, + "IMAGE" + ], + [ + 51, + 39, + 0, + 40, + 0, + "IMAGE" + ], + [ + 52, + 38, + 0, + 41, + 0, + "IMAGE" + ], + [ + 53, + 39, + 0, + 42, + 0, + "IMAGE" + ], + [ + 54, + 30, + 0, + 15, + 1, + "IMAGE" + ], + [ + 55, + 37, + 1, + 43, + 0, + "ORBIT_CAMPOSES" + ], + [ + 56, + 43, + 0, + 15, + 2, + "ORBIT_CAMPOSES" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -1287.2902103988, + -400.040713409522 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/CharacterGen/CharacterGen_Vanilla.json b/_Example_Workflows/CharacterGen/CharacterGen_Vanilla.json new file mode 100644 index 0000000..098e046 --- /dev/null +++ b/_Example_Workflows/CharacterGen/CharacterGen_Vanilla.json @@ -0,0 +1,680 @@ +{ + "last_node_id": 32, + "last_link_id": 45, + "nodes": [ + { + "id": 6, + "type": "InvertMask", + "pos": [ + -130, + 510 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 13 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 16 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 1, + "type": "[Comfy3D] Load CharacterGen MVDiffusion Model", + "pos": [ + -210, + 340 + ], + "size": { + "0": 294, + "1": 26 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "links": [ + 14 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load CharacterGen MVDiffusion Model" + } + }, + { + "id": 24, + "type": "[Comfy3D] Load CharacterGen Reconstruction Model", + "pos": [ + 923, + 351 + ], + "size": { + "0": 319.20001220703125, + "1": 28.166826248168945 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "character_lrm_pipe", + "type": "CHARACTER_LRM_PIPE", + "links": [ + 34 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load CharacterGen Reconstruction Model" + } + }, + { + "id": 25, + "type": "[Comfy3D] CharacterGen Reconstruction Model", + "pos": [ + 1621, + 400 + ], + "size": { + "0": 292.20001220703125, + "1": 66 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "character_lrm_pipe", + "type": "CHARACTER_LRM_PIPE", + "link": 34 + }, + { + "name": "multiview_images", + "type": "IMAGE", + "link": 38 + }, + { + "name": "multiview_masks", + "type": "MASK", + "link": 37 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 43 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] CharacterGen Reconstruction Model" + } + }, + { + "id": 32, + "type": "[Comfy3D] Switch Mesh Axis", + "pos": [ + 1956, + 400 + ], + "size": { + "0": 315, + "1": 154 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 43 + } + ], + "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": [ + "+x", + "-z", + "+y", + false, + 1 + ] + }, + { + "id": 28, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2320, + 400 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 44 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 41 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "CharacterGen\\Mesh\\Mesh_Girl.glb" + ] + }, + { + "id": 13, + "type": "[Comfy3D] CharacterGen MVDiffusion Model", + "pos": [ + 120, + 420 + ], + "size": { + "0": 456, + "1": 334 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "link": 14 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 15 + }, + { + "name": "reference_mask", + "type": "MASK", + "link": 16 + } + ], + "outputs": [ + { + "name": "multiviews", + "type": "IMAGE", + "links": [ + 26 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] CharacterGen MVDiffusion Model" + }, + "widgets_values": [ + 512, + 768, + 2333, + "fixed", + 5, + 40, + "high quality, best quality", + "", + 1.5 + ] + }, + { + "id": 21, + "type": "Image Rembg (Remove Background)", + "pos": [ + 940, + 483 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 24 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 36, + 42 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Image Rembg (Remove Background)" + }, + "widgets_values": [ + false, + "isnet-anime", + false, + true, + false, + 240, + 10, + 10, + "white" + ] + }, + { + "id": 19, + "type": "GetImagesFromBatchIndexed", + "pos": [ + 619, + 420 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 26 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 24, + 38, + 45 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "2, 3, 0, 1" + ] + }, + { + "id": 26, + "type": "ImageToMask", + "pos": [ + 1290, + 486 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 36 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 37 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageToMask" + }, + "widgets_values": [ + "red" + ] + }, + { + "id": 4, + "type": "SaveImage", + "pos": [ + 930, + 820 + ], + "size": { + "0": 526.6419677734375, + "1": 270 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 45 + } + ], + "properties": {}, + "widgets_values": [ + "CharacterGen\\Multi_Views\\image" + ] + }, + { + "id": 30, + "type": "SaveImage", + "pos": [ + 1507, + 825 + ], + "size": { + "0": 526.6419677734375, + "1": 270 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 42 + } + ], + "properties": {}, + "widgets_values": [ + "CharacterGen\\Multi_View_Masks\\image" + ] + }, + { + "id": 29, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 2663, + 398 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 41, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 3, + "type": "LoadImage", + "pos": [ + -540, + 440 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 15 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 13 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Test_ChunLi 768X512.png", + "image" + ] + } + ], + "links": [ + [ + 13, + 3, + 1, + 6, + 0, + "MASK" + ], + [ + 14, + 1, + 0, + 13, + 0, + "CHARACTER_MV_GEN_PIPE" + ], + [ + 15, + 3, + 0, + 13, + 1, + "IMAGE" + ], + [ + 16, + 6, + 0, + 13, + 2, + "MASK" + ], + [ + 24, + 19, + 0, + 21, + 0, + "IMAGE" + ], + [ + 26, + 13, + 0, + 19, + 0, + "IMAGE" + ], + [ + 34, + 24, + 0, + 25, + 0, + "CHARACTER_LRM_PIPE" + ], + [ + 36, + 21, + 0, + 26, + 0, + "IMAGE" + ], + [ + 37, + 26, + 0, + 25, + 2, + "MASK" + ], + [ + 38, + 19, + 0, + 25, + 1, + "IMAGE" + ], + [ + 41, + 28, + 0, + 29, + 0, + "STRING" + ], + [ + 42, + 21, + 0, + 30, + 0, + "IMAGE" + ], + [ + 43, + 25, + 0, + 32, + 0, + "MESH" + ], + [ + 44, + 32, + 0, + 28, + 0, + "MESH" + ], + [ + 45, + 19, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.7513148009015781, + "offset": [ + 621.9534614969469, + -192.59110955183658 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/CharacterGen/CharacterGen_to_Unique3D.json b/_Example_Workflows/CharacterGen/CharacterGen_to_Unique3D.json new file mode 100644 index 0000000..bbe4ce8 --- /dev/null +++ b/_Example_Workflows/CharacterGen/CharacterGen_to_Unique3D.json @@ -0,0 +1,2501 @@ +{ + "last_node_id": 362, + "last_link_id": 718, + "nodes": [ + { + "id": 289, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 1400, + 2290 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 45, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 597, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 300, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 2080, + 2290 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 49, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 612, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 286, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 710, + 2300 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 42, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 594, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 314, + "type": "[Comfy3D] Load Unique3D Custom UNet", + "pos": [ + -133.973961145846, + 652.576119046973 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "pipe", + "type": "DIFFUSERS_PIPE", + "link": 625 + } + ], + "outputs": [ + { + "name": "pipe", + "type": "DIFFUSERS_PIPE", + "links": [ + 624 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load Unique3D Custom UNet" + }, + "widgets_values": [ + "image2normal" + ] + }, + { + "id": 320, + "type": "ImageUpscaleWithModel", + "pos": [ + 1156.0260388541537, + 722.5761190469728 + ], + "size": { + "0": 241.79998779296875, + "1": 46 + }, + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 687 + }, + { + "name": "image", + "type": "IMAGE", + "link": 693 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 632, + 634 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageUpscaleWithModel" + } + }, + { + "id": 312, + "type": "[Comfy3D] Load Diffusers Pipeline", + "pos": [ + -603.9739611458458, + 650.576119046973 + ], + "size": { + "0": 430.03009033203125, + "1": 130 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "pipe", + "type": "DIFFUSERS_PIPE", + "links": [ + 625 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load Diffusers Pipeline" + }, + "widgets_values": [ + "Unique3DImageCustomPipeline", + "Wuvin/Unique3D/image2normal", + "", + false + ] + }, + { + "id": 274, + "type": "[Comfy3D] Fast Normal Maps To Mesh", + "pos": [ + 850, + 1820 + ], + "size": { + "0": 284.5090026855469, + "1": 46 + }, + "flags": {}, + "order": 37, + "mode": 0, + "inputs": [ + { + "name": "front_side_back_normal_maps", + "type": "IMAGE", + "link": 648 + }, + { + "name": "front_side_back_normal_masks", + "type": "MASK", + "link": 676 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 593, + 598 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Fast Normal Maps To Mesh" + } + }, + { + "id": 323, + "type": "Reroute", + "pos": [ + -750, + 1556 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 671 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 639 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 298, + "type": "[Comfy3D] Fast Clean Mesh", + "pos": [ + 1810, + 1670 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 44, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 610 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 601, + 611 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Fast Clean Mesh" + }, + "widgets_values": [ + true, + 1, + true, + 0.25 + ] + }, + { + "id": 321, + "type": "[Comfy3D] Rotate Normal Maps Horizontally", + "pos": [ + 1756.0260388541537, + 722.5761190469728 + ], + "size": { + "0": 315, + "1": 78 + }, + "flags": {}, + "order": 28, + "mode": 0, + "inputs": [ + { + "name": "normal_maps", + "type": "IMAGE", + "link": 632 + }, + { + "name": "normal_masks", + "type": "MASK", + "link": 633 + } + ], + "outputs": [ + { + "name": "normal_maps", + "type": "IMAGE", + "links": [ + 626, + 645 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Rotate Normal Maps Horizontally" + }, + "widgets_values": [ + true + ] + }, + { + "id": 295, + "type": "Reroute", + "pos": [ + 1230, + 2030 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 38, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 677 + } + ], + "outputs": [ + { + "name": "", + "type": "MASK", + "links": [ + 600, + 653 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 330, + "type": "Reroute", + "pos": [ + 2230, + 2030 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 41, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 653 + } + ], + "outputs": [ + { + "name": "", + "type": "MASK", + "links": [ + 654 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 297, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 3330, + 1530 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 50, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 609, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 322, + "type": "[Comfy3D] Get Masks From Normal Maps", + "pos": [ + 1426.0260388541537, + 792.5761190469728 + ], + "size": { + "0": 304.79998779296875, + "1": 26 + }, + "flags": {}, + "order": 27, + "mode": 0, + "inputs": [ + { + "name": "normal_maps", + "type": "IMAGE", + "link": 634 + } + ], + "outputs": [ + { + "name": "normal_masks", + "type": "MASK", + "links": [ + 633, + 663 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Get Masks From Normal Maps" + } + }, + { + "id": 273, + "type": "PreviewImage", + "pos": [ + 250, + 2050 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 35, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 580 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 277, + "type": "PreviewImage", + "pos": [ + -90, + 2060 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 33, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 584 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 281, + "type": "ImageToMask", + "pos": [ + 440, + 1890 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 36, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 678 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 676, + 677 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageToMask" + }, + "widgets_values": [ + "red" + ] + }, + { + "id": 269, + "type": "Morphology", + "pos": [ + -100, + 1890 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 34, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 680 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 580, + 678 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Morphology" + }, + "widgets_values": [ + "erode", + 3 + ] + }, + { + "id": 282, + "type": "MaskToImage", + "pos": [ + -370, + 1890 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 32, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 664 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 584, + 680 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MaskToImage" + } + }, + { + "id": 318, + "type": "ImageScale", + "pos": [ + 352, + 817 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 644 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 695 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "lanczos", + 512, + 512, + "disabled" + ] + }, + { + "id": 313, + "type": "[Comfy3D] Set Diffusers Pipeline Scheduler", + "pos": [ + 226.02603885415397, + 652.576119046973 + ], + "size": { + "0": 412.3726501464844, + "1": 58 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "pipe", + "type": "DIFFUSERS_PIPE", + "link": 624 + } + ], + "outputs": [ + { + "name": "pipe", + "type": "DIFFUSERS_PIPE", + "links": [ + 696 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Set Diffusers Pipeline Scheduler" + }, + "widgets_values": [ + "EulerAncestralDiscreteScheduler" + ] + }, + { + "id": 350, + "type": "IMG_padder", + "pos": [ + -329.80954879404095, + -427.3873149879231 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 700 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 701 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IMG_padder" + }, + "widgets_values": [ + 0.25, + "left" + ] + }, + { + "id": 351, + "type": "IMG_padder", + "pos": [ + 30.190451205959107, + -427.3873149879231 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 701 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 702 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IMG_padder" + }, + "widgets_values": [ + 0.2, + "right" + ] + }, + { + "id": 348, + "type": "[Comfy3D] Image Add Pure Color Background", + "pos": [ + -699.8095487940408, + -427.3873149879231 + ], + "size": { + "0": 315, + "1": 126 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 714 + }, + { + "name": "masks", + "type": "MASK", + "link": 713 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 700 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Image Add Pure Color Background" + }, + "widgets_values": [ + 255, + 255, + 255 + ] + }, + { + "id": 257, + "type": "ImageUpscaleWithModel", + "pos": [ + 780.1904512059592, + -447.3873149879231 + ], + "size": { + "0": 226.8000030517578, + "1": 46 + }, + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 547 + }, + { + "name": "image", + "type": "IMAGE", + "link": 715 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 686 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageUpscaleWithModel" + } + }, + { + "id": 326, + "type": "Reroute", + "pos": [ + -729, + 818 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 670 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 644 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 291, + "type": "[Comfy3D] ExplicitTarget Color Projection", + "pos": [ + 2480, + 1530 + ], + "size": { + "0": 418.1999816894531, + "1": 142 + }, + "flags": {}, + "order": 46, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 601 + }, + { + "name": "reference_images", + "type": "IMAGE", + "link": 639 + }, + { + "name": "reference_masks", + "type": "MASK", + "link": 654 + }, + { + "name": "reference_orbit_camera_poses", + "type": "ORBIT_CAMPOSES", + "link": null + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 608 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] ExplicitTarget Color Projection" + }, + "widgets_values": [ + 1024, + true + ] + }, + { + "id": 337, + "type": "Reroute", + "pos": [ + -461, + 1731 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 29, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 663 + } + ], + "outputs": [ + { + "name": "", + "type": "MASK", + "links": [ + 664 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 327, + "type": "Reroute", + "pos": [ + 839, + 1486 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 31, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 645 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 648, + 649 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 345, + "type": "[Comfy3D] Unique3D MVDiffusion Model", + "pos": [ + 712, + 764 + ], + "size": { + "0": 380.4000244140625, + "1": 222 + }, + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "unique3d_pipe", + "type": "DIFFUSERS_PIPE", + "link": 696 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 695 + } + ], + "outputs": [ + { + "name": "multiviews", + "type": "IMAGE", + "links": [ + 693 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": null, + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Unique3D MVDiffusion Model" + }, + "widgets_values": [ + 1145, + "fixed", + 1.5, + 30, + 512, + 4, + false + ] + }, + { + "id": 285, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 710, + 2170 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 39, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 593 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 594 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "CharacterGen_Unique3D\\Mesh_Gen_0\\mesh.obj" + ] + }, + { + "id": 288, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 1380, + 2160 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 43, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 596 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 597 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "CharacterGen_Unique3D\\Mesh_Gen_1\\mesh.obj" + ] + }, + { + "id": 299, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2060, + 2170 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 47, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 611 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 612 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "CharacterGen_Unique3D\\Mesh_Gen_2\\mesh.obj" + ] + }, + { + "id": 195, + "type": "SaveImage", + "pos": [ + 1450.1904512059589, + -447.3873149879231 + ], + "size": { + "0": 310.5008850097656, + "1": 341.5386962890625 + }, + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 703 + } + ], + "title": "Save Image [Sys] [Imgs_Out]", + "properties": { + "ttNbgOverride": { + "color": "#223", + "bgcolor": "#335", + "groupcolor": "#88A" + } + }, + "widgets_values": [ + "CharacterGen_Unique3D\\Upscales\\rgb_highres" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 315, + "type": "SaveImage", + "pos": [ + 2106.0260388541546, + 722.5761190469728 + ], + "size": { + "0": 340.4141540527344, + "1": 350.8065490722656 + }, + "flags": {}, + "order": 30, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 626 + } + ], + "properties": {}, + "widgets_values": [ + "CharacterGen_Unique3D/Normals/normal" + ] + }, + { + "id": 290, + "type": "[Comfy3D] ExplicitTarget Mesh Optimization", + "pos": [ + 1330, + 1670 + ], + "size": { + "0": 380.4000244140625, + "1": 238 + }, + "flags": {}, + "order": 40, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 598 + }, + { + "name": "normal_maps", + "type": "IMAGE", + "link": 649 + }, + { + "name": "normal_masks", + "type": "MASK", + "link": 600 + }, + { + "name": "normal_orbit_camera_poses", + "type": "ORBIT_CAMPOSES", + "link": null + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 596, + 610 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] ExplicitTarget Mesh Optimization" + }, + "widgets_values": [ + 200, + 512, + 0.1, + 100, + 5, + 20 + ] + }, + { + "id": 194, + "type": "UpscaleModelLoader", + "pos": [ + 310.19045120595905, + -207.3873149879232 + ], + "size": { + "0": 394.51678466796875, + "1": 58 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 547, + 687 + ], + "shape": 3, + "slot_index": 0 + } + ], + "title": "Load Upscale Model [Var]", + "properties": { + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "realesrganX4plus_v1.pt" + ] + }, + { + "id": 296, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2960, + 1530 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 48, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 608 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 609 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "CharacterGen_Unique3D\\Mesh_Gen_3\\mesh.obj" + ] + }, + { + "id": 353, + "type": "[Comfy3D] Load CharacterGen MVDiffusion Model", + "pos": [ + -2439.809548794041, + -507.38731498792276 + ], + "size": { + "0": 294, + "1": 26 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "links": [ + 707 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load CharacterGen MVDiffusion Model" + } + }, + { + "id": 355, + "type": "InvertMask", + "pos": [ + -2359.809548794041, + -337.3873149879231 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 706 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 709 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 356, + "type": "[Comfy3D] CharacterGen MVDiffusion Model", + "pos": [ + -2109.809548794041, + -427.3873149879231 + ], + "size": { + "0": 456, + "1": 334 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "character_mv_gen_pipe", + "type": "CHARACTER_MV_GEN_PIPE", + "link": 707 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 708 + }, + { + "name": "reference_mask", + "type": "MASK", + "link": 709 + } + ], + "outputs": [ + { + "name": "multiviews", + "type": "IMAGE", + "links": [ + 710 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] CharacterGen MVDiffusion Model" + }, + "widgets_values": [ + 512, + 768, + 2333, + "fixed", + 5, + 40, + "high quality, best quality", + "", + 1.5 + ] + }, + { + "id": 354, + "type": "LoadImage", + "pos": [ + -2769.809548794041, + -407.38731498792305 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 708 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 706 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "CharacterGen_Test_2.png", + "image" + ] + }, + { + "id": 340, + "type": "ImageScale", + "pos": [ + 400.19045120595905, + -427.3873149879231 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 702 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 715 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "lanczos", + 1024, + 1024, + "disabled" + ] + }, + { + "id": 343, + "type": "ImageScale", + "pos": [ + 1060.1904512059589, + -447.3873149879231 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 686 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 689, + 703 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "lanczos", + 2048, + 2048, + "disabled" + ] + }, + { + "id": 325, + "type": "Reroute", + "pos": [ + -955, + 292 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 716 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 670, + 671 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 324, + "type": "Reroute", + "pos": [ + 1325, + 290 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 689 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 716 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 357, + "type": "GetImagesFromBatchIndexed", + "pos": [ + -1599.8095487940411, + -427.3873149879231 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 710 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 711, + 714, + 717 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "3, 1, 2, 0" + ] + }, + { + "id": 358, + "type": "Image Rembg (Remove Background)", + "pos": [ + -1289.8095487940411, + -357.3873149879231 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 711 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 712, + 718 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Image Rembg (Remove Background)" + }, + "widgets_values": [ + false, + "isnet-anime", + false, + true, + false, + 240, + 10, + 10, + "white" + ] + }, + { + "id": 360, + "type": "ImageToMask", + "pos": [ + -945, + -364 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 712 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 713 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageToMask" + }, + "widgets_values": [ + "red" + ] + }, + { + "id": 361, + "type": "PreviewImage", + "pos": [ + -1370, + -62 + ], + "size": [ + 718.2292537113422, + 195.954588702313 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 717 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 362, + "type": "PreviewImage", + "pos": [ + -607, + -57 + ], + "size": [ + 608.1014177002792, + 187.80788955609944 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 718 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + } + ], + "links": [ + [ + 547, + 194, + 0, + 257, + 0, + "UPSCALE_MODEL" + ], + [ + 580, + 269, + 0, + 273, + 0, + "IMAGE" + ], + [ + 584, + 282, + 0, + 277, + 0, + "IMAGE" + ], + [ + 593, + 274, + 0, + 285, + 0, + "MESH" + ], + [ + 594, + 285, + 0, + 286, + 0, + "STRING" + ], + [ + 596, + 290, + 0, + 288, + 0, + "MESH" + ], + [ + 597, + 288, + 0, + 289, + 0, + "STRING" + ], + [ + 598, + 274, + 0, + 290, + 0, + "MESH" + ], + [ + 600, + 295, + 0, + 290, + 2, + "MASK" + ], + [ + 601, + 298, + 0, + 291, + 0, + "MESH" + ], + [ + 608, + 291, + 0, + 296, + 0, + "MESH" + ], + [ + 609, + 296, + 0, + 297, + 0, + "STRING" + ], + [ + 610, + 290, + 0, + 298, + 0, + "MESH" + ], + [ + 611, + 298, + 0, + 299, + 0, + "MESH" + ], + [ + 612, + 299, + 0, + 300, + 0, + "STRING" + ], + [ + 624, + 314, + 0, + 313, + 0, + "DIFFUSERS_PIPE" + ], + [ + 625, + 312, + 0, + 314, + 0, + "DIFFUSERS_PIPE" + ], + [ + 626, + 321, + 0, + 315, + 0, + "IMAGE" + ], + [ + 632, + 320, + 0, + 321, + 0, + "IMAGE" + ], + [ + 633, + 322, + 0, + 321, + 1, + "MASK" + ], + [ + 634, + 320, + 0, + 322, + 0, + "IMAGE" + ], + [ + 639, + 323, + 0, + 291, + 1, + "IMAGE" + ], + [ + 644, + 326, + 0, + 318, + 0, + "IMAGE" + ], + [ + 645, + 321, + 0, + 327, + 0, + "*" + ], + [ + 648, + 327, + 0, + 274, + 0, + "IMAGE" + ], + [ + 649, + 327, + 0, + 290, + 1, + "IMAGE" + ], + [ + 653, + 295, + 0, + 330, + 0, + "*" + ], + [ + 654, + 330, + 0, + 291, + 2, + "MASK" + ], + [ + 663, + 322, + 0, + 337, + 0, + "*" + ], + [ + 664, + 337, + 0, + 282, + 0, + "MASK" + ], + [ + 670, + 325, + 0, + 326, + 0, + "*" + ], + [ + 671, + 325, + 0, + 323, + 0, + "*" + ], + [ + 676, + 281, + 0, + 274, + 1, + "MASK" + ], + [ + 677, + 281, + 0, + 295, + 0, + "*" + ], + [ + 678, + 269, + 0, + 281, + 0, + "IMAGE" + ], + [ + 680, + 282, + 0, + 269, + 0, + "IMAGE" + ], + [ + 686, + 257, + 0, + 343, + 0, + "IMAGE" + ], + [ + 687, + 194, + 0, + 320, + 0, + "UPSCALE_MODEL" + ], + [ + 689, + 343, + 0, + 324, + 0, + "*" + ], + [ + 693, + 345, + 0, + 320, + 1, + "IMAGE" + ], + [ + 695, + 318, + 0, + 345, + 1, + "IMAGE" + ], + [ + 696, + 313, + 0, + 345, + 0, + "DIFFUSERS_PIPE" + ], + [ + 700, + 348, + 0, + 350, + 0, + "IMAGE" + ], + [ + 701, + 350, + 0, + 351, + 0, + "IMAGE" + ], + [ + 702, + 351, + 0, + 340, + 0, + "IMAGE" + ], + [ + 703, + 343, + 0, + 195, + 0, + "IMAGE" + ], + [ + 706, + 354, + 1, + 355, + 0, + "MASK" + ], + [ + 707, + 353, + 0, + 356, + 0, + "CHARACTER_MV_GEN_PIPE" + ], + [ + 708, + 354, + 0, + 356, + 1, + "IMAGE" + ], + [ + 709, + 355, + 0, + 356, + 2, + "MASK" + ], + [ + 710, + 356, + 0, + 357, + 0, + "IMAGE" + ], + [ + 711, + 357, + 0, + 358, + 0, + "IMAGE" + ], + [ + 712, + 358, + 0, + 360, + 0, + "IMAGE" + ], + [ + 713, + 360, + 0, + 348, + 1, + "MASK" + ], + [ + 714, + 357, + 0, + 348, + 0, + "IMAGE" + ], + [ + 715, + 340, + 0, + 257, + 1, + "IMAGE" + ], + [ + 716, + 324, + 0, + 325, + 0, + "*" + ], + [ + 717, + 357, + 0, + 361, + 0, + "IMAGE" + ], + [ + 718, + 358, + 0, + 362, + 0, + "IMAGE" + ] + ], + "groups": [ + { + "title": "Stage_4: Mesh & RGB Texture Reconstruction from Multi-View Images & Normal Maps", + "bounding": [ + -897, + 1302, + 5999, + 1562 + ], + "color": "#3f789e", + "font_size": 48 + }, + { + "title": "Stage_3: Multi-View Normal Maps Prediction from Multi-View Images", + "bounding": [ + -814, + 389, + 3277, + 755 + ], + "color": "#3f789e", + "font_size": 48 + }, + { + "title": "Stage_0, Multi-Views Generation and Super-resolution", + "bounding": [ + -2770, + -674, + 4630, + 889 + ], + "color": "#3f789e", + "font_size": 48 + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.5054470284993127, + "offset": [ + 2906.4865159897017, + 762.6902948157394 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/Era3D_to_InstantMesh.json b/_Example_Workflows/Era3D_to_InstantMesh.json index f451081..280a182 100644 --- a/_Example_Workflows/Era3D_to_InstantMesh.json +++ b/_Example_Workflows/Era3D_to_InstantMesh.json @@ -128,7 +128,7 @@ }, { "id": 22, - "type": "[Comfy3D] Era3D Diffusion Model", + "type": "[Comfy3D] Era3D MVDiffusion Model", "pos": [ 1118, 232 @@ -188,7 +188,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Era3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Era3D MVDiffusion Model" }, "widgets_values": [ 420, diff --git a/_Example_Workflows/Unique3D/0_MV_RGB.json b/_Example_Workflows/Unique3D/0_MV_RGB.json index 3c0a870..996c9a8 100644 --- a/_Example_Workflows/Unique3D/0_MV_RGB.json +++ b/_Example_Workflows/Unique3D/0_MV_RGB.json @@ -82,7 +82,7 @@ }, { "id": 23, - "type": "[Comfy3D] Unique3D Diffusion Model", + "type": "[Comfy3D] Unique3D MVDiffusion Model", "pos": [ 1374, 159 @@ -124,7 +124,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Unique3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Unique3D MVDiffusion Model" }, "widgets_values": [ 1145, diff --git a/_Example_Workflows/Unique3D/2_MV_Normal.json b/_Example_Workflows/Unique3D/2_MV_Normal.json index 2b4f74e..1686751 100644 --- a/_Example_Workflows/Unique3D/2_MV_Normal.json +++ b/_Example_Workflows/Unique3D/2_MV_Normal.json @@ -220,7 +220,7 @@ }, { "id": 21, - "type": "[Comfy3D] Unique3D Diffusion Model", + "type": "[Comfy3D] Unique3D MVDiffusion Model", "pos": [ 1397, 254 @@ -262,7 +262,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Unique3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Unique3D MVDiffusion Model" }, "widgets_values": [ 1145, diff --git a/_Example_Workflows/Unique3D/Unique3D_All_Stages.json b/_Example_Workflows/Unique3D/Unique3D_All_Stages.json index 1bfee42..fcb33db 100644 --- a/_Example_Workflows/Unique3D/Unique3D_All_Stages.json +++ b/_Example_Workflows/Unique3D/Unique3D_All_Stages.json @@ -1516,7 +1516,7 @@ }, { "id": 317, - "type": "[Comfy3D] Unique3D Diffusion Model", + "type": "[Comfy3D] Unique3D MVDiffusion Model", "pos": [ 726.0260388541542, 742.5761190469728 @@ -1558,7 +1558,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Unique3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Unique3D MVDiffusion Model" }, "widgets_values": [ 1145, @@ -1607,7 +1607,7 @@ }, { "id": 307, - "type": "[Comfy3D] Unique3D Diffusion Model", + "type": "[Comfy3D] Unique3D MVDiffusion Model", "pos": [ -67, -1160 @@ -1650,7 +1650,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Unique3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Unique3D MVDiffusion Model" }, "widgets_values": [ 1145, diff --git a/_Example_Workflows/Unique3D/Unique3D_with_Era3D_All.json b/_Example_Workflows/Unique3D/Unique3D_with_Era3D_All.json index 768df31..be6ec12 100644 --- a/_Example_Workflows/Unique3D/Unique3D_with_Era3D_All.json +++ b/_Example_Workflows/Unique3D/Unique3D_with_Era3D_All.json @@ -240,7 +240,7 @@ }, { "id": 336, - "type": "[Comfy3D] Era3D Diffusion Model", + "type": "[Comfy3D] Era3D MVDiffusion Model", "pos": [ -643.2975797406153, -610.8013992433897 @@ -297,7 +297,7 @@ } ], "properties": { - "Node name for S&R": "[Comfy3D] Era3D Diffusion Model" + "Node name for S&R": "[Comfy3D] Era3D MVDiffusion Model" }, "widgets_values": [ 420, diff --git a/_Example_Workflows/_Example_Inputs_Files/CharacterGen_Test_2.png b/_Example_Workflows/_Example_Inputs_Files/CharacterGen_Test_2.png new file mode 100644 index 0000000..c42d98c Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/CharacterGen_Test_2.png differ diff --git a/_Example_Workflows/_Example_Inputs_Files/Test_ChunLi 768X512.png b/_Example_Workflows/_Example_Inputs_Files/Test_ChunLi 768X512.png new file mode 100644 index 0000000..dab8cd7 Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/Test_ChunLi 768X512.png differ diff --git a/checkpoints/Diffusers/Wuvin/Unique3D/controlnet-tile/config.json b/checkpoints/Diffusers/Wuvin/Unique3D/controlnet-tile/config.json deleted file mode 100644 index 754ad3a..0000000 --- a/checkpoints/Diffusers/Wuvin/Unique3D/controlnet-tile/config.json +++ /dev/null @@ -1,52 +0,0 @@ -{ - "_class_name": "ControlNetModel", - "_diffusers_version": "0.27.2", - "_name_or_path": "lllyasviel/control_v11f1e_sd15_tile", - "act_fn": "silu", - "addition_embed_type": null, - "addition_embed_type_num_heads": 64, - "addition_time_embed_dim": null, - "attention_head_dim": 8, - "block_out_channels": [ - 320, - 640, - 1280, - 1280 - ], - "class_embed_type": null, - "conditioning_channels": 3, - "conditioning_embedding_out_channels": [ - 16, - 32, - 96, - 256 - ], - "controlnet_conditioning_channel_order": "rgb", - "cross_attention_dim": 768, - "down_block_types": [ - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", - "DownBlock2D" - ], - "downsample_padding": 1, - "encoder_hid_dim": null, - "encoder_hid_dim_type": null, - "flip_sin_to_cos": true, - "freq_shift": 0, - "global_pool_conditions": false, - "in_channels": 4, - "layers_per_block": 2, - "mid_block_scale_factor": 1, - "mid_block_type": "UNetMidBlock2DCrossAttn", - "norm_eps": 1e-05, - "norm_num_groups": 32, - "num_attention_heads": null, - "num_class_embeds": null, - "only_cross_attention": false, - "projection_class_embeddings_input_dim": null, - "resnet_time_scale_shift": "default", - "transformer_layers_per_block": 1, - "upcast_attention": false, - "use_linear_projection": false -} diff --git a/gen_3d_modules/Unique3D/mesh_reconstruction/recon.py b/gen_3d_modules/Unique3D/mesh_reconstruction/recon.py index 062f42b..7d2ad25 100644 --- a/gen_3d_modules/Unique3D/mesh_reconstruction/recon.py +++ b/gen_3d_modules/Unique3D/mesh_reconstruction/recon.py @@ -12,7 +12,6 @@ from .render import NormalsRenderer, Pytorch3DNormalsRenderer from Unique3D.scripts.utils import to_py3d_mesh, init_target def reconstruct_stage1(pils: List[Image.Image], steps=100, vertices=None, faces=None, start_edge_len=0.1, end_edge_len=0.02, decay=0.995, return_mesh=False, loss_expansion_weight=0.1, gain=0.05): - vertices, faces = vertices.to("cuda"), faces.to("cuda") assert len(pils) == 4 mv,proj = make_star_cameras_orthographic(4, 1) renderer = NormalsRenderer(mv,proj,list(pils[0].size)) @@ -24,7 +23,7 @@ def reconstruct_stage1(pils: List[Image.Image], steps=100, vertices=None, faces= target_images = target_images[[0, 3, 2, 1]] # 2. init from coarse mesh - opt = MeshOptimizer(vertices,faces, local_edgelen=False, gain=gain, edge_len_lims=(end_edge_len, start_edge_len)) + opt = MeshOptimizer(vertices, faces, local_edgelen=False, gain=gain, edge_len_lims=(end_edge_len, start_edge_len)) vertices = opt.vertices diff --git a/gen_3d_modules/Unique3D/mesh_reconstruction/remesh.py b/gen_3d_modules/Unique3D/mesh_reconstruction/remesh.py index 2e2faa8..4979bc9 100644 --- a/gen_3d_modules/Unique3D/mesh_reconstruction/remesh.py +++ b/gen_3d_modules/Unique3D/mesh_reconstruction/remesh.py @@ -95,7 +95,6 @@ def calc_vertex_normals( faces:torch.Tensor, #F,3 long, first face may be all zero face_normals:torch.Tensor=None, #F,3, not normalized )->torch.Tensor: #F,3 - F = faces.shape[0] if face_normals is None: diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py index 3dcc08c..85c1895 100644 --- a/mesh_processer/mesh.py +++ b/mesh_processer/mesh.py @@ -678,7 +678,7 @@ class Mesh: ) # append texture info - if self.vt is not None: + if self.vt is not None and self.albedo is not None: vt_np = self.vt.detach().cpu().numpy().astype(np.float32) vt_np_blob = vt_np.tobytes() diff --git a/nodes.py b/nodes.py index d4a9f74..aedb9ba 100644 --- a/nodes.py +++ b/nodes.py @@ -29,6 +29,7 @@ from diffusers import ( KDPM2AncestralDiscreteScheduler, KDPM2DiscreteScheduler, ) +from huggingface_hub import snapshot_download from plyfile import PlyData from PIL import Image @@ -77,6 +78,8 @@ from Unique3D.scripts.utils import from_py3d_mesh, to_py3d_mesh, to_pyml_mesh, s from Unique3D.scripts.project_mesh import multiview_color_projection, get_cameras_list from Unique3D.mesh_reconstruction.recon import reconstruct_stage1 from Unique3D.mesh_reconstruction.refine import run_mesh_refine +from CharacterGen.character_inference import Inference2D_API, Inference3D_API +from CharacterGen.Stage_3D.lrm.utils.config import load_config as load_config_cg3d from .shared_utils.image_utils import ( prepare_torch_img, torch_imgs_to_pils, troch_image_dilate, @@ -588,6 +591,7 @@ class Switch_Mesh_Axis: "axis_y_to": (["+y", "-y", "+z", "-z", "+x", "-x"],), "axis_z_to": (["+z", "-z", "+x", "-x", "+y", "-y"],), "flip_normal": ("BOOLEAN", {"default": False},), + "scale": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100, "step": 0.01}), }, } @@ -600,12 +604,12 @@ class Switch_Mesh_Axis: FUNCTION = "switch_axis_and_scale" CATEGORY = "Comfy3D/Preprocessor" - def switch_axis_and_scale(self, mesh, axis_x_to, axis_y_to, axis_z_to, flip_normal): + def switch_axis_and_scale(self, mesh, axis_x_to, axis_y_to, axis_z_to, flip_normal, scale): switched_mesh = None if axis_x_to[1] != axis_y_to[1] and axis_x_to[1] != axis_z_to[1] and axis_y_to[1] != axis_z_to[1]: - target_axis, target_scale, coordinate_invert_count = get_target_axis_and_scale([axis_x_to, axis_y_to, axis_z_to]) + target_axis, target_scale, coordinate_invert_count = get_target_axis_and_scale([axis_x_to, axis_y_to, axis_z_to], scale) switched_mesh = switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal) else: cstr(f"[{self.__class__.__name__}] axis_x_to: {axis_x_to}, axis_y_to: {axis_y_to}, axis_z_to: {axis_z_to} have to be on separated axis").error.print() @@ -1292,12 +1296,12 @@ class Load_Triplane_Gaussian_Transformers: checkpoints_dir = "TriplaneGaussian" default_ckpt_name = "model_lvis_rel.ckpt" default_repo_id = "VAST-AI/TriplaneGaussian" - tgs_config_path = "TriplaneGaussian_config.yaml" + config_path = "TriplaneGaussian_config.yaml" @classmethod def INPUT_TYPES(cls): cls.checkpoints_dir_abs = os.path.join(CKPT_ROOT_PATH, cls.checkpoints_dir) - cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.tgs_config_path) + cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS) if cls.default_ckpt_name not in all_models_names: all_models_names += [cls.default_ckpt_name] @@ -1334,11 +1338,11 @@ class Load_Triplane_Gaussian_Transformers: class Triplane_Gaussian_Transformers: - tgs_config_path = "TriplaneGaussian_config.yaml" + config_path = "TriplaneGaussian_config.yaml" @classmethod def INPUT_TYPES(cls): - cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.tgs_config_path) + cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) return { "required": { "reference_image": ("IMAGE", ), @@ -1384,7 +1388,7 @@ class Load_Diffusers_Pipeline: return { "required": { "diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),), - "model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}), + "repo_id": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}), "custom_pipeline": ("STRING", {"default": "", "multiline": False}), "download_from_remote": ("BOOLEAN", {"default": True}), }, @@ -1399,14 +1403,12 @@ class Load_Diffusers_Pipeline: FUNCTION = "load_diffusers_pipe" CATEGORY = "Comfy3D/Import|Export" - def load_diffusers_pipe(self, diffusers_pipeline_name, model_name, custom_pipeline, download_from_remote): + def load_diffusers_pipe(self, diffusers_pipeline_name, repo_id, custom_pipeline, download_from_remote): - # resume pretrained checkpoint - ckpt_path = os.path.join(CKPT_DIFFUSERS_PATH, model_name) - + # resume download pretrained checkpoint + ckpt_path = os.path.join(CKPT_DIFFUSERS_PATH, repo_id) if download_from_remote: - from huggingface_hub import snapshot_download - snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model", ignore_patterns=["*.json", "*.py"]) + snapshot_download(repo_id=repo_id, local_dir=ckpt_path, repo_type="model", ignore_patterns=["*.json", "*.py"]) diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name] @@ -1487,12 +1489,12 @@ class Set_Diffusers_Pipeline_State_Dict: class Wonder3D_MVDiffusion_Model: - wonder3d_config_path = "Wonder3D_config.yaml" + config_path = "Wonder3D_config.yaml" fix_cam_pose_dir = "Wonder3D/data/fixed_poses/nine_views" @classmethod def INPUT_TYPES(cls): - cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.wonder3d_config_path) + cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) cls.fix_cam_pose_dir_abs = os.path.join(MODULE_ROOT_PATH, cls.fix_cam_pose_dir) return { "required": { @@ -1801,7 +1803,7 @@ class Load_TripoSR_Model: checkpoints_dir = "TripoSR" default_ckpt_name = "model.ckpt" default_repo_id = "stabilityai/TripoSR" - tsr_config_path = "TripoSR_config.yaml" + config_path = "TripoSR_config.yaml" @classmethod def INPUT_TYPES(cls): @@ -1810,7 +1812,7 @@ class Load_TripoSR_Model: if cls.default_ckpt_name not in all_models_names: all_models_names += [cls.default_ckpt_name] - cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.tsr_config_path) + cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) return { "required": { "model_name": (all_models_names, ), @@ -1899,7 +1901,7 @@ class Load_CRM_MVDiffusion_Model: default_ckpt_name = ["pixel-diffusion.pth", "ccm-diffusion.pth"] default_conf_name = ["sd_v2_base_ipmv_zero_SNR.yaml", "sd_v2_base_ipmv_chin8_zero_snr.yaml"] default_repo_id = "Zhengyi/CRM" - crm_config_path = "CRM_configs" + config_path = "CRM_configs" @classmethod def INPUT_TYPES(cls): @@ -1909,7 +1911,7 @@ class Load_CRM_MVDiffusion_Model: if ckpt_name not in all_models_names: all_models_names += [ckpt_name] - cls.config_root_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.crm_config_path) + cls.config_root_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) return { "required": { "model_name": (all_models_names, ), @@ -2323,12 +2325,12 @@ class InstantMesh_Reconstruction_Model: mesh.auto_normal() return (mesh,) -class Era3D_Diffusion_Model: +class Era3D_MVDiffusion_Model: - era3d_config_path = "Era3D_config.yaml" + config_path = "Era3D_config.yaml" @classmethod def INPUT_TYPES(cls): - cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.era3d_config_path) + cls.config_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) return { "required": { "era3d_pipe": ("DIFFUSERS_PIPE",), @@ -2574,6 +2576,7 @@ class FlexiCubes_MVS: class Load_Unique3D_Custom_UNet: checkpoints_dir = "Wuvin/Unique3D" + default_repo_id = "MrForExample/Unique3D" config_root_dir = "Unique3D_configs" @classmethod @@ -2600,9 +2603,12 @@ class Load_Unique3D_Custom_UNet: from Unique3D.custum_3d_diffusion.trainings.config_classes import ExprimentConfig from Unique3D.custum_3d_diffusion.custum_modules.unifield_processor import AttnConfig, ConfigurableUNet2DConditionModel from Unique3D.custum_3d_diffusion.trainings.utils import load_config + # Download models and configs + ckpt_path = os.path.join(CKPT_DIFFUSERS_PATH, self.checkpoints_dir) + snapshot_download(repo_id=self.default_repo_id, local_dir=ckpt_path, repo_type="model", ignore_patterns=["*.json", "*.py"]) cfg_path = os.path.join(self.config_path_abs, config_name + ".yaml") - checkpoint_dir_path = os.path.join(CKPT_DIFFUSERS_PATH, self.checkpoints_dir, config_name) + checkpoint_dir_path = os.path.join(ckpt_path, config_name) checkpoint_path = os.path.join(checkpoint_dir_path, "unet_state_dict.pth") cfg: ExprimentConfig = load_config(ExprimentConfig, cfg_path) @@ -2620,7 +2626,7 @@ class Load_Unique3D_Custom_UNet: pipe.unet = configurable_unet.unet return (pipe, ) -class Unique3D_Diffusion_Model: +class Unique3D_MVDiffusion_Model: @classmethod def INPUT_TYPES(cls): @@ -2730,10 +2736,10 @@ class ExplicitTarget_Mesh_Optimization: "mesh": ("MESH",), "normal_maps": ("IMAGE",), "normal_masks": ("MASK",), - "reconstruction_steps": ("INT", {"default": 200, "min": 1, "max": 0xffffffffffffffff}), + "reconstruction_steps": ("INT", {"default": 200, "min": 0, "max": 0xffffffffffffffff}), "coarse_reconstruct_resolution": ("INT", {"default": 512, "min": 128, "max": 8192}), "loss_expansion_weight": ("FLOAT", {"default": 0.1, "min": 0.01, "step": 0.01}), - "refinement_steps": ("INT", {"default": 100, "min": 1, "max": 0xffffffffffffffff}), + "refinement_steps": ("INT", {"default": 100, "min": 0, "max": 0xffffffffffffffff}), "target_warmup_update_num": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}), "target_update_interval": ("INT", {"default": 20, "min": 1, "max": 0xffffffffffffffff}), }, @@ -2769,9 +2775,12 @@ class ExplicitTarget_Mesh_Optimization: pil_normal_list = torch_imgs_to_pils(normal_maps, normal_masks) normal_stg1 = [img.resize((coarse_reconstruct_resolution, coarse_reconstruct_resolution)) for img in pil_normal_list] with torch.inference_mode(False): - vertices, faces = reconstruct_stage1(normal_stg1, steps=reconstruction_steps, vertices=mesh.v, faces=mesh.f, loss_expansion_weight=loss_expansion_weight) - - vertices, faces = run_mesh_refine(vertices, faces, pil_normal_list, steps=refinement_steps, update_normal_interval=target_update_interval, update_warmup=target_warmup_update_num, ) + vertices, faces = mesh.v.detach().clone().to(DEVICE), mesh.f.detach().clone().to(DEVICE).type(torch.int64) + if reconstruction_steps > 0: + vertices, faces = reconstruct_stage1(normal_stg1, steps=reconstruction_steps, vertices=vertices, faces=faces, loss_expansion_weight=loss_expansion_weight) + + if refinement_steps > 0: + vertices, faces = run_mesh_refine(vertices, faces, pil_normal_list, steps=refinement_steps, update_normal_interval=target_update_interval, update_warmup=target_warmup_update_num, ) mesh = Mesh(v=vertices, f=faces, device=DEVICE) return (mesh,) @@ -2833,4 +2842,165 @@ class ExplicitTarget_Color_Projection: mesh = Mesh(v=vertices, f=faces, vc=vertex_colors, device=DEVICE) mesh.auto_normal() + return (mesh,) + +class Load_CharacterGen_MVDiffusion_Model: + checkpoints_dir = "CharacterGen" + default_repo_id = "zjpshadow/CharacterGen" + config_path = "CharacterGen_configs/Stage_2D_infer.yaml" + + @classmethod + def INPUT_TYPES(cls): + cls.checkpoints_dir_abs = os.path.join(CKPT_ROOT_PATH, cls.checkpoints_dir) + cls.config_root_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) + return { + "required": { + }, + } + + RETURN_TYPES = ( + "CHARACTER_MV_GEN_PIPE", + ) + RETURN_NAMES = ( + "character_mv_gen_pipe", + ) + FUNCTION = "load_model" + CATEGORY = "Comfy3D/Import|Export" + + def load_model(self): + # Download checkpoints + snapshot_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, repo_type="model", ignore_patterns=["*.json", "*.py"]) + # Load pre-trained models + character_mv_gen_pipe = Inference2D_API(checkpoint_root_path=self.checkpoints_dir_abs, **OmegaConf.load(self.config_root_path_abs)) + return (character_mv_gen_pipe,) + +class CharacterGen_MVDiffusion_Model: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "character_mv_gen_pipe": ("CHARACTER_MV_GEN_PIPE",), + "reference_image": ("IMAGE", ), + "reference_mask": ("MASK",), + "target_image_width": ("INT", {"default": 512, "min": 128, "max": 8192}), + "target_image_height": ("INT", {"default": 768, "min": 128, "max": 8192}), + "seed": ("INT", {"default": 2333, "min": 0, "max": 0xffffffffffffffff}), + "guidance_scale": ("FLOAT", {"default": 5.0, "min": 0.0, "step": 0.01}), + "num_inference_steps": ("INT", {"default": 40, "min": 1}), + "prompt": ("STRING", { + "default": "high quality, best quality", + "multiline": True + }), + "prompt_neg": ("STRING", { + "default": "", + "multiline": True + }), + "radius": ("FLOAT", {"default": 1.5, "min": 0.1, "step": 0.01}) + }, + } + + RETURN_TYPES = ( + "IMAGE", + "ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z] + ) + RETURN_NAMES = ( + "multiviews", + "orbit_camposes", + ) + FUNCTION = "run_model" + CATEGORY = "Comfy3D/Algorithm" + + @torch.no_grad() + def run_model( + self, + character_mv_gen_pipe, + reference_image, + reference_mask, + target_image_width, + target_image_height, + seed, + guidance_scale, + num_inference_steps, + prompt, + prompt_neg, + radius + ): + single_image = torch_imgs_to_pils(reference_image, reference_mask)[0] + + multiview_images = character_mv_gen_pipe.inference( + single_image, target_image_width, target_image_height, prompt=prompt, prompt_neg=prompt_neg, + guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, seed=seed + ) + + azimuths = [-90, 90, 180, 0] + elevations = [0.0] * 4 + radius = [radius] * 4 + center = [0.0] * 4 + + orbit_camposes = [azimuths, elevations, radius, center, center, center] + + return (multiview_images, orbit_camposes) + +class Load_CharacterGen_Reconstruction_Model: + checkpoints_dir = "CharacterGen" + default_repo_id = "zjpshadow/CharacterGen" + config_path = "CharacterGen_configs/Stage_3D_infer.yaml" + + @classmethod + def INPUT_TYPES(cls): + cls.checkpoints_dir_abs = os.path.join(CKPT_ROOT_PATH, cls.checkpoints_dir) + cls.config_root_path_abs = os.path.join(CONFIG_ROOT_PATH, cls.config_path) + return { + "required": { + }, + } + + RETURN_TYPES = ( + "CHARACTER_LRM_PIPE", + ) + RETURN_NAMES = ( + "character_lrm_pipe", + ) + FUNCTION = "load_model" + CATEGORY = "Comfy3D/Import|Export" + + def load_model(self): + # Download checkpoints + snapshot_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, repo_type="model", ignore_patterns=["*.json", "*.py"]) + # Load pre-trained models + character_lrm_pipe = Inference3D_API(checkpoint_root_path=self.checkpoints_dir_abs, cfg=load_config_cg3d(self.config_root_path_abs)) + return (character_lrm_pipe,) + +class CharacterGen_Reconstruction_Model: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "character_lrm_pipe": ("CHARACTER_LRM_PIPE", ), + "multiview_images": ("IMAGE",), + "multiview_masks": ("MASK",), + } + } + + RETURN_TYPES = ( + "MESH", + ) + RETURN_NAMES = ( + "mesh", + ) + + FUNCTION = "run_LRM" + CATEGORY = "Comfy3D/Algorithm" + + @torch.no_grad() + def run_LRM(self, character_lrm_pipe, multiview_images, multiview_masks): + pil_mv_image_list = torch_imgs_to_pils(multiview_images, multiview_masks, alpha_min=0.2) + + vertices, faces = character_lrm_pipe.inference(pil_mv_image_list) + + mesh = Mesh(v=vertices, f=faces.to(torch.int64), device=DEVICE) + mesh.auto_normal() + mesh.auto_uv() + return (mesh,) \ No newline at end of file diff --git a/shared_utils/common_utils.py b/shared_utils/common_utils.py index 92ca14b..7d69013 100644 --- a/shared_utils/common_utils.py +++ b/shared_utils/common_utils.py @@ -52,13 +52,13 @@ def get_list_filenames(directory, extension_filter=None): return [] # Download pre-trained model if it not exist locally -def resume_or_download_model_from_hf(checkpoints_dir_abs, default_repo_id, model_name, class_name=""): +def resume_or_download_model_from_hf(checkpoints_dir_abs, repo_id, model_name, class_name="", repo_type="model"): ckpt_path = os.path.join(checkpoints_dir_abs, model_name) if not os.path.isfile(ckpt_path): - cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {default_repo_id} instead").warning.print() + cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {repo_id} instead").warning.print() from huggingface_hub import hf_hub_download - hf_hub_download(repo_id=default_repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type="model") + hf_hub_download(repo_id=repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type=repo_type) return ckpt_path diff --git a/shared_utils/image_utils.py b/shared_utils/image_utils.py index e0325ab..dbf7207 100644 --- a/shared_utils/image_utils.py +++ b/shared_utils/image_utils.py @@ -13,7 +13,7 @@ def prepare_torch_img(img, size_H, size_W, device="cuda", keep_shape=False): img_new = img_new.permute(0, 2, 3, 1) return img_new -def torch_imgs_to_pils(images, masks=None): +def torch_imgs_to_pils(images, masks=None, alpha_min=0.1): """ images (torch): [N, H, W, C] or [H, W, C] masks (torch): [N, H, W] or [H, W] @@ -24,8 +24,11 @@ def torch_imgs_to_pils(images, masks=None): if masks is not None: if len(masks.shape) == 2: masks = masks.unsqueeze(0) - masks = masks.unsqueeze(3) + inv_mask_index = masks < alpha_min + images[inv_mask_index] = 0. + + masks = masks.unsqueeze(3) images = torch.cat((images, masks), dim=3) mode="RGBA" else: