From 1927e3021ccdaacdc528b5e6b76bb64ce2a4f6a7 Mon Sep 17 00:00:00 2001 From: "Mr. For Example" <62230687+MrForExample@users.noreply.github.com> Date: Tue, 23 Jul 2024 15:08:09 -0600 Subject: [PATCH 1/3] Add Craftman --- .gitignore | 4 +- .../Put Craftsman Model here.txt | 0 Configs/Craftsman_config.yaml | 143 + .../CharacterGen/Stage_3D/lrm/models/mesh.py | 2 +- .../models/geometry/render/neural_render.py | 2 +- .../Unique3D/mesh_reconstruction/remesh.py | 6 +- Gen_3D_Modules/craftsman/__init__.py | 52 + Gen_3D_Modules/craftsman/data/__init__.py | 3 + Gen_3D_Modules/craftsman/data/objaverse.py | 311 +++ Gen_3D_Modules/craftsman/models/__init__.py | 5 + .../craftsman/models/autoencoders/__init__.py | 3 + .../autoencoders/michelangelo_autoencoder.py | 403 +++ .../craftsman/models/autoencoders/utils.py | 302 +++ .../models/conditional_encoders/__init__.py | 3 + .../models/conditional_encoders/base.py | 138 + .../clip/modeling_clip.py | 1419 ++++++++++ .../clip/modeling_conditional_clip.py | 385 +++ .../conditional_encoders/clip_encoder.py | 172 ++ .../craftsman/models/denoisers/__init__.py | 3 + .../models/denoisers/simple_denoiser.py | 191 ++ .../craftsman/models/geometry/__init__.py | 3 + .../craftsman/models/geometry/base.py | 194 ++ .../craftsman/models/geometry/utils.py | 559 ++++ .../models/transformers/attention.py | 207 ++ .../models/transformers/perceiver_1d.py | 48 + .../craftsman/models/transformers/utils.py | 21 + Gen_3D_Modules/craftsman/systems/__init__.py | 4 + Gen_3D_Modules/craftsman/systems/base.py | 209 ++ .../craftsman/systems/shape_autoencoder.py | 137 + .../craftsman/systems/shape_diffusion.py | 396 +++ Gen_3D_Modules/craftsman/utils/__init__.py | 1 + Gen_3D_Modules/craftsman/utils/base.py | 118 + Gen_3D_Modules/craftsman/utils/callbacks.py | 156 ++ Gen_3D_Modules/craftsman/utils/checkpoint.py | 68 + Gen_3D_Modules/craftsman/utils/config.py | 128 + Gen_3D_Modules/craftsman/utils/misc.py | 168 ++ Gen_3D_Modules/craftsman/utils/ops.py | 167 ++ Gen_3D_Modules/craftsman/utils/saving.py | 653 +++++ Gen_3D_Modules/craftsman/utils/scheduler.py | 104 + Gen_3D_Modules/craftsman/utils/typing.py | 41 + .../craftsman/utils/visualizers/__init__.py | 1 + .../craftsman/utils/visualizers/color_util.py | 43 + .../craftsman/utils/visualizers/html_util.py | 49 + .../utils/visualizers/pythreejs_viewer.py | 537 ++++ .../DiffRastMesh/diff_mesh_renderer.py | 2 +- .../CRM MVDiffusion_to_InstantMesh.json | 0 .../CRM/CRM MVDiffusion_to_Unique3D.json | 2309 +++++++++++++++++ ...RM_Img_to_Multi-views_to_CCMs_to_Mesh.json | 0 nodes.py | 119 +- shared_utils/common_utils.py | 24 +- 50 files changed, 9996 insertions(+), 17 deletions(-) create mode 100644 Checkpoints/Craftsman/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/Put Craftsman Model here.txt create mode 100644 Configs/Craftsman_config.yaml create mode 100644 Gen_3D_Modules/craftsman/__init__.py create mode 100644 Gen_3D_Modules/craftsman/data/__init__.py create mode 100644 Gen_3D_Modules/craftsman/data/objaverse.py create mode 100644 Gen_3D_Modules/craftsman/models/__init__.py create mode 100644 Gen_3D_Modules/craftsman/models/autoencoders/__init__.py create mode 100644 Gen_3D_Modules/craftsman/models/autoencoders/michelangelo_autoencoder.py create mode 100644 Gen_3D_Modules/craftsman/models/autoencoders/utils.py create mode 100644 Gen_3D_Modules/craftsman/models/conditional_encoders/__init__.py create mode 100644 Gen_3D_Modules/craftsman/models/conditional_encoders/base.py create mode 100644 Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_clip.py create mode 100644 Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_conditional_clip.py create mode 100644 Gen_3D_Modules/craftsman/models/conditional_encoders/clip_encoder.py create mode 100644 Gen_3D_Modules/craftsman/models/denoisers/__init__.py create mode 100644 Gen_3D_Modules/craftsman/models/denoisers/simple_denoiser.py create mode 100644 Gen_3D_Modules/craftsman/models/geometry/__init__.py create mode 100644 Gen_3D_Modules/craftsman/models/geometry/base.py create mode 100644 Gen_3D_Modules/craftsman/models/geometry/utils.py create mode 100644 Gen_3D_Modules/craftsman/models/transformers/attention.py create mode 100644 Gen_3D_Modules/craftsman/models/transformers/perceiver_1d.py create mode 100644 Gen_3D_Modules/craftsman/models/transformers/utils.py create mode 100644 Gen_3D_Modules/craftsman/systems/__init__.py create mode 100644 Gen_3D_Modules/craftsman/systems/base.py create mode 100644 Gen_3D_Modules/craftsman/systems/shape_autoencoder.py create mode 100644 Gen_3D_Modules/craftsman/systems/shape_diffusion.py create mode 100644 Gen_3D_Modules/craftsman/utils/__init__.py create mode 100644 Gen_3D_Modules/craftsman/utils/base.py create mode 100644 Gen_3D_Modules/craftsman/utils/callbacks.py create mode 100644 Gen_3D_Modules/craftsman/utils/checkpoint.py create mode 100644 Gen_3D_Modules/craftsman/utils/config.py create mode 100644 Gen_3D_Modules/craftsman/utils/misc.py create mode 100644 Gen_3D_Modules/craftsman/utils/ops.py create mode 100644 Gen_3D_Modules/craftsman/utils/saving.py create mode 100644 Gen_3D_Modules/craftsman/utils/scheduler.py create mode 100644 Gen_3D_Modules/craftsman/utils/typing.py create mode 100644 Gen_3D_Modules/craftsman/utils/visualizers/__init__.py create mode 100644 Gen_3D_Modules/craftsman/utils/visualizers/color_util.py create mode 100644 Gen_3D_Modules/craftsman/utils/visualizers/html_util.py create mode 100644 Gen_3D_Modules/craftsman/utils/visualizers/pythreejs_viewer.py rename _Example_Workflows/{ => CRM}/CRM MVDiffusion_to_InstantMesh.json (100%) create mode 100644 _Example_Workflows/CRM/CRM MVDiffusion_to_Unique3D.json rename _Example_Workflows/{ => CRM}/CRM_Img_to_Multi-views_to_CCMs_to_Mesh.json (100%) diff --git a/.gitignore b/.gitignore index c10fa8f..df54d65 100644 --- a/.gitignore +++ b/.gitignore @@ -161,13 +161,11 @@ cython_debug/ # Ignore test outputs outputs/ +output/ # Ignore folders only used for build dependencies _Pre_Builds/_Build_Dependencies _Pre_Builds/_Build_Wheels _Pre_Builds/_Libs -models/ -output/ - .idea diff --git a/Checkpoints/Craftsman/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/Put Craftsman Model here.txt b/Checkpoints/Craftsman/image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/Put Craftsman Model here.txt new file mode 100644 index 0000000..e69de29 diff --git a/Configs/Craftsman_config.yaml b/Configs/Craftsman_config.yaml new file mode 100644 index 0000000..f34e177 --- /dev/null +++ b/Configs/Craftsman_config.yaml @@ -0,0 +1,143 @@ +name: michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k +description: '' +tag: michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue +seed: 0 +use_timestamp: true +timestamp: '' +exp_root_dir: outputs +exp_dir: outputs/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k +trial_name: michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue +trial_dir: outputs/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k/michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue +n_gpus: 8 +resume: ./ckpts/3DNativeGeneration/michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k.ckpt +data_type: objaverse-datamodule +data: + root_dir: data/objaverse_clean/cap3d_high_quality_170k_images + data_type: occupancy + n_samples: 4096 + noise_sigma: 0.0 + load_supervision: false + supervision_type: occupancy + n_supervision: 10000 + load_image: true + image_data_path: data/objaverse_clean/raw_data/images/cap3d_high_quality_170k + image_type: mvrgb + idx: + - 0 + - 4 + - 8 + - 12 + - 16 + n_views: 4 + load_caption: false + rotate_points: false + batch_size: 32 + num_workers: 16 +system_type: shape-diffusion-system +system: + val_samples_json: val_data/mv_images/val_samples_rgb_mvimage.json + z_scale_factor: 1.0 + guidance_scale: 7.5 + num_inference_steps: 50 + eta: 0.0 + shape_model_type: michelangelo-aligned-autoencoder + shape_model: + num_latents: 256 + embed_dim: 64 + point_feats: 3 + out_dim: 1 + num_freqs: 8 + include_pi: false + heads: 12 + width: 768 + num_encoder_layers: 8 + num_decoder_layers: 16 + use_ln_post: true + init_scale: 0.25 + qkv_bias: false + use_flash: true + use_checkpoint: true + condition_model_type: clip-embedder + condition_model: + pretrained_model_name_or_path: openai/clip-vit-large-patch14 + encode_camera: true + camera_embeds_dim: 32 + n_views: 4 + empty_embeds_ratio: 0.1 + normalize_embeds: false + zero_uncond_embeds: true + denoiser_model_type: simple-denoiser + denoiser_model: + input_channels: 64 + output_channels: 64 + n_ctx: 256 + width: 768 + layers: 6 + heads: 12 + context_dim: 1024 + init_scale: 1.0 + skip_ln: true + use_checkpoint: true + noise_scheduler_type: diffusers.schedulers.DDPMScheduler + noise_scheduler: + num_train_timesteps: 1000 + beta_start: 0.00085 + beta_end: 0.012 + beta_schedule: scaled_linear + variance_type: fixed_small + clip_sample: false + denoise_scheduler_type: diffusers.schedulers.DDIMScheduler + denoise_scheduler: + num_train_timesteps: 1000 + beta_start: 0.00085 + beta_end: 0.012 + beta_schedule: scaled_linear + clip_sample: false + set_alpha_to_one: false + steps_offset: 1 + loggers: + wandb: + enable: false + project: JiangXin + name: text-to-shape-diffusion+michelangelo-image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-170k+michelangelo-aligned-autoencoder+n4096+noise0.0+pfeat3+zeroemb0.0+normembFalse+lr5e-05+qkvbiasFalse+nfreq8+ln_postTrue + loss: + loss_type: mse + lambda_diffusion: 1.0 + optimizer: + name: AdamW + args: + lr: 5.0e-05 + betas: + - 0.9 + - 0.99 + eps: 1.0e-06 + scheduler: + name: SequentialLR + interval: step + schedulers: + - name: LinearLR + interval: step + args: + start_factor: 1.0e-06 + end_factor: 1.0 + total_iters: 5000 + - name: CosineAnnealingLR + interval: step + args: + T_max: 5000 + eta_min: 0.0 + milestones: + - 5000 +trainer: + num_nodes: 2 + max_epochs: 100000 + log_every_n_steps: 5 + num_sanity_val_steps: 1 + check_val_every_n_epoch: 3 + enable_progress_bar: true + precision: 16-mixed + strategy: ddp_find_unused_parameters_true +checkpoint: + save_last: true + save_top_k: -1 + every_n_train_steps: 5000 diff --git a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py index 771e436..0446735 100644 --- a/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py +++ b/Gen_3D_Modules/CharacterGen/Stage_3D/lrm/models/mesh.py @@ -194,7 +194,7 @@ class Mesh: v1 = self.v_pos[i1, :] v2 = self.v_pos[i2, :] - face_normals = torch.cross(v1 - v0, v2 - v0) + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) # Splat face normals to vertices v_nrm = torch.zeros_like(self.v_pos) diff --git a/Gen_3D_Modules/InstantMesh/models/geometry/render/neural_render.py b/Gen_3D_Modules/InstantMesh/models/geometry/render/neural_render.py index 4734644..ff7cf4b 100644 --- a/Gen_3D_Modules/InstantMesh/models/geometry/render/neural_render.py +++ b/Gen_3D_Modules/InstantMesh/models/geometry/render/neural_render.py @@ -48,7 +48,7 @@ def compute_vertex_normal(v_pos, t_pos_idx): v1 = v_pos[i1, :] v2 = v_pos[i2, :] - face_normals = torch.cross(v1 - v0, v2 - v0) + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) # Splat face normals to vertices v_nrm = torch.zeros_like(v_pos) diff --git a/Gen_3D_Modules/Unique3D/mesh_reconstruction/remesh.py b/Gen_3D_Modules/Unique3D/mesh_reconstruction/remesh.py index 4979bc9..ccd14ab 100644 --- a/Gen_3D_Modules/Unique3D/mesh_reconstruction/remesh.py +++ b/Gen_3D_Modules/Unique3D/mesh_reconstruction/remesh.py @@ -350,9 +350,9 @@ def flip_edges( e1 = v[:,1] cl = v[:,2] cr = v[:,3] - n = torch.cross(e1,cl) + torch.cross(cr,e1) #sum of old normal vectors - flip.logical_and_(torch.sum(n*torch.cross(cr,cl),dim=-1)>0) #first new face - flip.logical_and_(torch.sum(n*torch.cross(cl-e1,cr-e1),dim=-1)>0) #second new face + n = torch.cross(e1,cl,dim=-1) + torch.cross(cr,e1,dim=-1) #sum of old normal vectors + flip.logical_and_(torch.sum(n*torch.cross(cr,cl,dim=-1),dim=-1)>0) #first new face + flip.logical_and_(torch.sum(n*torch.cross(cl-e1,cr-e1,dim=-1),dim=-1)>0) #second new face flip_edges_neighbors = edges_neighbors[flip] #E",4 flip_edge_to_face = edge_to_face[candidates,:,0][flip] #E",2 diff --git a/Gen_3D_Modules/craftsman/__init__.py b/Gen_3D_Modules/craftsman/__init__.py new file mode 100644 index 0000000..aaae191 --- /dev/null +++ b/Gen_3D_Modules/craftsman/__init__.py @@ -0,0 +1,52 @@ +import importlib + +__modules__ = {} + + +def register(name): + def decorator(cls): + if name in __modules__: + raise ValueError( + f"Module {name} already exists! Names of extensions conflict!" + ) + else: + __modules__[name] = cls + return cls + + return decorator + + +def find(name): + if name in __modules__: + return __modules__[name] + else: + try: + module_string = ".".join(name.split(".")[:-1]) + cls_name = name.split(".")[-1] + module = importlib.import_module(module_string, package=None) + return getattr(module, cls_name) + except Exception as e: + raise ValueError(f"Module {name} not found!") + + +### 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) + + +from . import data, models, systems diff --git a/Gen_3D_Modules/craftsman/data/__init__.py b/Gen_3D_Modules/craftsman/data/__init__.py new file mode 100644 index 0000000..6548321 --- /dev/null +++ b/Gen_3D_Modules/craftsman/data/__init__.py @@ -0,0 +1,3 @@ +from . import ( + objaverse +) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/data/objaverse.py b/Gen_3D_Modules/craftsman/data/objaverse.py new file mode 100644 index 0000000..7501f3a --- /dev/null +++ b/Gen_3D_Modules/craftsman/data/objaverse.py @@ -0,0 +1,311 @@ +import math +import os +import json +from dataclasses import dataclass, field + +import random +import imageio +import numpy as np +import pytorch_lightning as pl +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader, Dataset +from torchvision import transforms +from PIL import Image +from transformers import CLIPImageProcessor, CLIPTokenizer + +from craftsman import register +from craftsman.utils.base import Updateable +from craftsman.utils.config import parse_structured +from craftsman.utils.typing import * + +def rot2eul(R): + beta = -np.arcsin(R[2,0]) + alpha = np.arctan2(R[2,1]/np.cos(beta),R[2,2]/np.cos(beta)) + gamma = np.arctan2(R[1,0]/np.cos(beta),R[0,0]/np.cos(beta)) + return np.array((alpha, beta, gamma)) + +def eul2rot(theta) : + R = np.array([[np.cos(theta[1])*np.cos(theta[2]), np.sin(theta[0])*np.sin(theta[1])*np.cos(theta[2]) - np.sin(theta[2])*np.cos(theta[0]), np.sin(theta[1])*np.cos(theta[0])*np.cos(theta[2]) + np.sin(theta[0])*np.sin(theta[2])], + [np.sin(theta[2])*np.cos(theta[1]), np.sin(theta[0])*np.sin(theta[1])*np.sin(theta[2]) + np.cos(theta[0])*np.cos(theta[2]), np.sin(theta[1])*np.sin(theta[2])*np.cos(theta[0]) - np.sin(theta[0])*np.cos(theta[2])], + [-np.sin(theta[1]), np.sin(theta[0])*np.cos(theta[1]), np.cos(theta[0])*np.cos(theta[1])]]) + return R + +@dataclass +class ObjaverseDataModuleConfig: + root_dir: str = None + data_type: str = "occupancy" # occupancy or sdf + n_samples: int = 4096 # number of points in input point cloud + scale: float = 1.0 # scale of the input point cloud and target supervision + noise_sigma: float = 0.0 # noise level of the input point cloud + + load_supervision: bool = True # whether to load supervision + supervision_type: str = "occupancy" # occupancy, sdf, tsdf, tsdf_w_surface + n_supervision: int = 10000 # number of points in supervision + + load_image: bool = False # whether to load images + image_data_path: str = "" # path to the image data + image_type: str = "rgb" # rgb, normal + background_color: Tuple[float, float, float] = field( + default_factory=lambda: (1.0, 1.0, 1.0) + ) + idx: Optional[List[int]] = None # index of the image to load + n_views: int = 1 # number of views + rotate_points: bool = False # whether to rotate the input point cloud and the supervision + + load_caption: bool = False # whether to load captions + caption_type: str = "text" # text, clip_embeds + tokenizer_pretrained_model_name_or_path: str = "" + + batch_size: int = 32 + num_workers: int = 0 + +class ObjaverseDataset(Dataset): + def __init__(self, cfg: Any, split: str) -> None: + super().__init__() + self.cfg: ObjaverseDataModuleConfig = cfg + self.split = split + + self.uids = json.load(open(f'{cfg.root_dir}/{split}.json')) + print(f"Loaded {len(self.uids)} {split} uids") + + if self.cfg.load_caption: + self.tokenizer = CLIPTokenizer.from_pretrained(self.cfg.tokenizer_pretrained_model_name_or_path) + + self.background_color = torch.as_tensor(self.cfg.background_color) + self.distance = 1.0 + self.camera_embedding = torch.as_tensor([ + [[1, 0, 0, 0], + [0, 0, -1, -self.distance], + [0, 1, 0, 0], + [0, 0, 0, 1]], # front to back + + [[0, 0, 1, self.distance], + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 0, 1]], # right to left + + [[-1, 0, 0, 0], + [0, 0, 1, self.distance], + [0, 1, 0, 0], + [0, 0, 0, 1]], # back to front + + [[0, 0, -1, -self.distance], + [-1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 0, 1]], # left to right + ], dtype=torch.float32) + if self.cfg.n_views != 1: + assert self.cfg.n_views == self.camera_embedding.shape[0] + + def __len__(self): + return len(self.uids) + + def _load_shape(self, index: int) -> Dict[str, Any]: + if self.cfg.data_type == "occupancy": + # for input point cloud + pointcloud = np.load(f'{self.cfg.root_dir}/{self.uids[index]}/pointcloud.npz') + surface = np.asarray(pointcloud['points']) * 2 # range from -1 to 1 + normal = np.asarray(pointcloud['normals']) + surface = np.concatenate([surface, normal], axis=1) + elif self.cfg.data_type == "sdf": + data = np.load(f'{self.cfg.root_dir}/{self.uids[index]}.npz') + # for input point cloud + surface = data["surface"] + else: + raise NotImplementedError(f"Data type {self.cfg.data_type} not implemented") + + # random sampling + rng = np.random.default_rng() + ind = rng.choice(surface.shape[0], self.cfg.n_samples, replace=False) + surface = surface[ind] + # rescale data + surface[:, :3] = surface[:, :3] * self.cfg.scale # target scale + # add noise to input point cloud + surface[:, :3] += (np.random.rand(surface.shape[0], 3) * 2 - 1) * self.cfg.noise_sigma + ret = { + "uid": self.uids[index].split('/')[-1], + "surface": surface.astype(np.float32), + } + + return ret + + def _load_shape_supervision(self, index: int) -> Dict[str, Any]: + # for supervision + ret = {} + if self.cfg.data_type == "occupancy": + points = np.load(f'{self.cfg.root_dir}/{self.uids[index]}/points.npz') + rand_points = np.asarray(points['points']) * 2 # range from -1.1 to 1.1 + occupancies = np.asarray(points['occupancies']) + occupancies = np.unpackbits(occupancies) + elif self.cfg.data_type == "sdf": + data = np.load(f'{self.cfg.root_dir}/{self.uids[index]}.npz') + rand_points = data['rand_points'] + sdfs = data['sdfs'] + else: + raise NotImplementedError(f"Data type {self.cfg.data_type} not implemented") + + # random sampling + rng = np.random.default_rng() + ind = rng.choice(rand_points.shape[0], self.cfg.n_supervision, replace=False) + rand_points = rand_points[ind] + rand_points = rand_points * self.cfg.scale + ret["rand_points"] = rand_points.astype(np.float32) + + if self.cfg.data_type == "occupancy": + assert self.cfg.supervision_type == "occupancy", "Only occupancy supervision is supported for occupancy data" + occupancies = occupancies[ind] + ret["occupancies"] = occupancies.astype(np.float32) + elif self.cfg.data_type == "sdf": + if self.cfg.supervision_type == "sdf": + ret["sdf"] = sdfs[ind].flatten().astype(np.float32) + elif self.cfg.supervision_type == "occupancy": + ret["occupancies"] = np.where(sdfs[ind].flatten() < 1e-3, 0, 1).astype(np.float32) + else: + raise NotImplementedError(f"Supervision type {self.cfg.supervision_type} not implemented") + + return ret + + def _load_image(self, index: int) -> Dict[str, Any]: + def _load_single_image(img_path): + img = torch.from_numpy( + np.asarray( + Image.fromarray(imageio.v2.imread(img_path)) + .convert("RGBA") + ) + / 255.0 + ).float() + mask: Float[Tensor, "H W 1"] = img[:, :, -1:] + image: Float[Tensor, "H W 3"] = img[:, :, :3] * mask + self.background_color[ + None, None, : + ] * (1 - mask) + return image + + ret = {} + if self.cfg.image_type == "rgb" or self.cfg.image_type == "normal": + assert self.cfg.n_views == 1, "Only single view is supported for single image" + sel_idx = random.choice(self.cfg.idx) + ret["sel_image_idx"] = sel_idx + if self.cfg.image_type == "rgb": + img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{sel_idx}.png" + elif self.cfg.image_type == "normal": + img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{sel_idx}_normal.png" + ret["image"] = _load_single_image(img_path) + ret["c2w"] = self.camera_embedding[sel_idx % 4] + elif self.cfg.image_type == "mvrgb" or self.cfg.image_type == "mvnormal": + sel_idx = random.choice(self.cfg.idx) + ret["sel_image_idx"] = sel_idx + mvimages = [] + for i in range(self.cfg.n_views): + if self.cfg.image_type == "mvrgb": + img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{sel_idx+i}.png" + elif self.cfg.image_type == "mvnormal": + img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{sel_idx+i}_normal.png" + mvimages.append(_load_single_image(img_path)) + ret["mvimages"] = torch.stack(mvimages) + ret["c2ws"] = self.camera_embedding + else: + raise NotImplementedError(f"Image type {self.cfg.image_type} not implemented") + + return ret + + def _load_caption(self, index: int, drop_text_embed: bool = False) -> Dict[str, Any]: + ret = {} + if self.cfg.caption_type == "text": + caption = eval(json.load(open(f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f'/annotation.json'))) + texts = [v for k, v in caption.items()] + sel_idx = random.randint(0, len(texts) - 1) + ret["sel_caption_idx"] = sel_idx + ret['text_input_ids'] = self.tokenizer( + texts[sel_idx] if not drop_text_embed else "", + max_length=self.tokenizer.model_max_length, + padding="max_length", + truncation=True, + return_tensors="pt" + ).input_ids.detach() + else: + raise NotImplementedError(f"Caption type {self.cfg.caption_type} not implemented") + + return ret + + def get_data(self, index): + # load shape + ret = self._load_shape(index) + + # load supervision for shape + if self.cfg.load_supervision: + ret.update(self._load_shape_supervision(index)) + + # load image + if self.cfg.load_image: + ret.update(self._load_image(index)) + + # load the rotation of the object and rotate the camera + rots = np.load(f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f'/rots.npy')[ret['sel_image_idx']].astype(np.float32) + rots = torch.tensor(rots[:3, :3], dtype=torch.float32) + if "c2ws" in ret.keys(): + ret["c2ws"][:, :3, :3] = torch.matmul(rots, ret["c2ws"][:, :3, :3]) + ret["c2ws"][:, :3, 3] = torch.matmul(rots, ret["c2ws"][:, :3, 3].unsqueeze(-1)).squeeze(-1) + elif "c2w" in ret.keys(): + ret["c2w"][:3, :3] = torch.matmul(rots, ret["c2w"][:3, :3]) + ret["c2w"][:3, 3] = torch.matmul(rots, ret["c2w"][:3, 3].unsqueeze(-1)).squeeze(-1) + + # load caption + if self.cfg.load_caption: + ret.update(self._load_caption(index)) + + return ret + + def __getitem__(self, index): + try: + return self.get_data(index) + except Exception as e: + print(f"Error in {self.uids[index]}: {e}") + return self.__getitem__(np.random.randint(len(self))) + + + def collate(self, batch): + batch = torch.utils.data.default_collate(batch) + return batch + + +@register("objaverse-datamodule") +class ObjaverseDataModule(pl.LightningDataModule): + cfg: ObjaverseDataModuleConfig + + def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None: + super().__init__() + self.cfg = parse_structured(ObjaverseDataModuleConfig, cfg) + + def setup(self, stage=None) -> None: + if stage in [None, "fit"]: + self.train_dataset = ObjaverseDataset(self.cfg, "train") + if stage in [None, "fit", "validate"]: + self.val_dataset = ObjaverseDataset(self.cfg, "val") + if stage in [None, "test", "predict"]: + self.test_dataset = ObjaverseDataset(self.cfg, "test") + + def prepare_data(self): + pass + + def general_loader(self, dataset, batch_size, collate_fn=None, num_workers=0) -> DataLoader: + return DataLoader( + dataset, batch_size=batch_size, collate_fn=collate_fn, num_workers=num_workers + ) + + def train_dataloader(self) -> DataLoader: + return self.general_loader( + self.train_dataset, + batch_size=self.cfg.batch_size, + collate_fn=self.train_dataset.collate, + num_workers=self.cfg.num_workers + ) + + def val_dataloader(self) -> DataLoader: + return self.general_loader(self.val_dataset, batch_size=1) + + def test_dataloader(self) -> DataLoader: + return self.general_loader(self.test_dataset, batch_size=1) + + def predict_dataloader(self) -> DataLoader: + return self.general_loader(self.test_dataset, batch_size=1) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/__init__.py b/Gen_3D_Modules/craftsman/models/__init__.py new file mode 100644 index 0000000..73b20fb --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/__init__.py @@ -0,0 +1,5 @@ +from . import ( + autoencoders, + conditional_encoders, + denoisers, +) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/autoencoders/__init__.py b/Gen_3D_Modules/craftsman/models/autoencoders/__init__.py new file mode 100644 index 0000000..623726e --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/autoencoders/__init__.py @@ -0,0 +1,3 @@ +from . import ( + michelangelo_autoencoder, +) diff --git a/Gen_3D_Modules/craftsman/models/autoencoders/michelangelo_autoencoder.py b/Gen_3D_Modules/craftsman/models/autoencoders/michelangelo_autoencoder.py new file mode 100644 index 0000000..a805eeb --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/autoencoders/michelangelo_autoencoder.py @@ -0,0 +1,403 @@ +from dataclasses import dataclass +import math + +import torch +import torch.nn as nn +from einops import repeat, rearrange +from transformers import CLIPModel + +import craftsman +from craftsman.models.transformers.perceiver_1d import Perceiver +from craftsman.models.transformers.attention import ResidualCrossAttentionBlock +from craftsman.utils.checkpoint import checkpoint +from craftsman.utils.base import BaseModule +from craftsman.utils.typing import * + +from .utils import AutoEncoder, FourierEmbedder, get_embedder + +class PerceiverCrossAttentionEncoder(nn.Module): + def __init__(self, + use_downsample: bool, + num_latents: int, + embedder: FourierEmbedder, + point_feats: int, + embed_point_feats: bool, + width: int, + heads: int, + layers: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + use_ln_post: bool = False, + use_flash: bool = False, + use_checkpoint: bool = False): + + super().__init__() + + self.use_checkpoint = use_checkpoint + self.num_latents = num_latents + self.use_downsample = use_downsample + self.embed_point_feats = embed_point_feats + + if not self.use_downsample: + self.query = nn.Parameter(torch.randn((num_latents, width)) * 0.02) + + self.embedder = embedder + if self.embed_point_feats: + self.input_proj = nn.Linear(self.embedder.out_dim * 2, width) + else: + self.input_proj = nn.Linear(self.embedder.out_dim + point_feats, width) + + self.cross_attn = ResidualCrossAttentionBlock( + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash, + ) + + self.self_attn = Perceiver( + n_ctx=num_latents, + width=width, + layers=layers, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash, + use_checkpoint=False + ) + + if use_ln_post: + self.ln_post = nn.LayerNorm(width) + else: + self.ln_post = None + + def _forward(self, pc, feats): + """ + + Args: + pc (torch.FloatTensor): [B, N, 3] + feats (torch.FloatTensor or None): [B, N, C] + + Returns: + + """ + + bs, N, D = pc.shape + + data = self.embedder(pc) + if feats is not None: + if self.embed_point_feats: + feats = self.embedder(feats) + data = torch.cat([data, feats], dim=-1) + data = self.input_proj(data) + + if self.use_downsample: + ###### fps + from torch_cluster import fps + flattened = pc.view(bs*N, D) + + batch = torch.arange(bs).to(pc.device) + batch = torch.repeat_interleave(batch, N) + + pos = flattened + + ratio = 1.0 * self.num_latents / N + + idx = fps(pos, batch, ratio=ratio) + + query = data.view(bs*N, -1)[idx].view(bs, -1, data.shape[-1]) + else: + query = self.query + query = repeat(query, "m c -> b m c", b=bs) + + latents = self.cross_attn(query, data) + latents = self.self_attn(latents) + + if self.ln_post is not None: + latents = self.ln_post(latents) + + return latents + + def forward(self, pc: torch.FloatTensor, feats: Optional[torch.FloatTensor] = None): + """ + + Args: + pc (torch.FloatTensor): [B, N, 3] + feats (torch.FloatTensor or None): [B, N, C] + + Returns: + dict + """ + + return checkpoint(self._forward, (pc, feats), self.parameters(), self.use_checkpoint) + + +class PerceiverCrossAttentionDecoder(nn.Module): + + def __init__(self, + num_latents: int, + out_dim: int, + embedder: FourierEmbedder, + width: int, + heads: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + use_flash: bool = False, + use_checkpoint: bool = False): + + super().__init__() + + self.use_checkpoint = use_checkpoint + self.embedder = embedder + + self.query_proj = nn.Linear(self.embedder.out_dim, width) + + self.cross_attn_decoder = ResidualCrossAttentionBlock( + n_data=num_latents, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash + ) + + self.ln_post = nn.LayerNorm(width) + self.output_proj = nn.Linear(width, out_dim) + + def _forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor): + queries = self.query_proj(self.embedder(queries)) + x = self.cross_attn_decoder(queries, latents) + x = self.ln_post(x) + x = self.output_proj(x) + return x + + def forward(self, queries: torch.FloatTensor, latents: torch.FloatTensor): + return checkpoint(self._forward, (queries, latents), self.parameters(), self.use_checkpoint) + + +@craftsman.register("michelangelo-autoencoder") +class MichelangeloAutoencoder(AutoEncoder): + r""" + A VAE model for encoding shapes into latents and decoding latent representations into shapes. + """ + + @dataclass + class Config(BaseModule.Config): + pretrained_model_name_or_path: str = "" + use_downsample: bool = False + num_latents: int = 256 + point_feats: int = 0 + embed_point_feats: bool = False + out_dim: int = 1 + embed_dim: int = 64 + embed_type: str = "fourier" + num_freqs: int = 8 + include_pi: bool = True + width: int = 768 + heads: int = 12 + num_encoder_layers: int = 8 + num_decoder_layers: int = 16 + init_scale: float = 0.25 + qkv_bias: bool = True + use_ln_post: bool = False + use_flash: bool = False + use_checkpoint: bool = True + + cfg: Config + + def configure(self) -> None: + super().configure() + + self.embedder = get_embedder(embed_type=self.cfg.embed_type, num_freqs=self.cfg.num_freqs, include_pi=self.cfg.include_pi) + + # encoder + self.cfg.init_scale = self.cfg.init_scale * math.sqrt(1.0 / self.cfg.width) + self.encoder = PerceiverCrossAttentionEncoder( + use_downsample=self.cfg.use_downsample, + embedder=self.embedder, + num_latents=self.cfg.num_latents, + point_feats=self.cfg.point_feats, + embed_point_feats=self.cfg.embed_point_feats, + width=self.cfg.width, + heads=self.cfg.heads, + layers=self.cfg.num_encoder_layers, + init_scale=self.cfg.init_scale, + qkv_bias=self.cfg.qkv_bias, + use_ln_post=self.cfg.use_ln_post, + use_flash=self.cfg.use_flash, + use_checkpoint=self.cfg.use_checkpoint + ) + + if self.cfg.embed_dim > 0: + # VAE embed + self.pre_kl = nn.Linear(self.cfg.width, self.cfg.embed_dim * 2) + self.post_kl = nn.Linear(self.cfg.embed_dim, self.cfg.width) + self.latent_shape = (self.cfg.num_latents, self.cfg.embed_dim) + else: + self.latent_shape = (self.cfg.num_latents, self.cfg.width) + + self.transformer = Perceiver( + n_ctx=self.cfg.num_latents, + width=self.cfg.width, + layers=self.cfg.num_decoder_layers, + heads=self.cfg.heads, + init_scale=self.cfg.init_scale, + qkv_bias=self.cfg.qkv_bias, + use_flash=self.cfg.use_flash, + use_checkpoint=self.cfg.use_checkpoint + ) + + # decoder + self.decoder = PerceiverCrossAttentionDecoder( + embedder=self.embedder, + out_dim=self.cfg.out_dim, + num_latents=self.cfg.num_latents, + width=self.cfg.width, + heads=self.cfg.heads, + init_scale=self.cfg.init_scale, + qkv_bias=self.cfg.qkv_bias, + use_flash=self.cfg.use_flash, + use_checkpoint=self.cfg.use_checkpoint + ) + + if self.cfg.pretrained_model_name_or_path != "": + print(f"Loading pretrained model from {self.cfg.pretrained_model_name_or_path}") + pretrained_ckpt = torch.load(self.cfg.pretrained_model_name_or_path, map_location="cpu") + if 'state_dict' in pretrained_ckpt: + _pretrained_ckpt = {} + for k, v in pretrained_ckpt['state_dict'].items(): + if k.startswith('shape_model.'): + _pretrained_ckpt[k.replace('shape_model.', '')] = v + pretrained_ckpt = _pretrained_ckpt + self.load_state_dict(pretrained_ckpt, strict=True) + + + def encode(self, + surface: torch.FloatTensor, + sample_posterior: bool = True): + """ + Args: + surface (torch.FloatTensor): [B, N, 3+C] + sample_posterior (bool): + + Returns: + shape_latents (torch.FloatTensor): [B, num_latents, width] + kl_embed (torch.FloatTensor): [B, num_latents, embed_dim] + posterior (DiagonalGaussianDistribution or None): + """ + assert surface.shape[-1] == 3 + self.cfg.point_feats, f"\ + Expected {3 + self.cfg.point_feats} channels, got {surface.shape[-1]}" + + pc, feats = surface[..., :3], surface[..., 3:] # B, n_samples, 3 + shape_latents = self.encoder(pc, feats) # B, num_latents, width + kl_embed, posterior = self.encode_kl_embed(shape_latents, sample_posterior) # B, num_latents, embed_dim + + return shape_latents, kl_embed, posterior + + + def decode(self, + latents: torch.FloatTensor): + """ + Args: + latents (torch.FloatTensor): [B, embed_dim] + + Returns: + latents (torch.FloatTensor): [B, embed_dim] + """ + latents = self.post_kl(latents) # [B, num_latents, embed_dim] -> [B, num_latents, width] + + return self.transformer(latents) + + + def query(self, + queries: torch.FloatTensor, + latents: torch.FloatTensor): + """ + Args: + queries (torch.FloatTensor): [B, N, 3] + latents (torch.FloatTensor): [B, embed_dim] + + Returns: + logits (torch.FloatTensor): [B, N], occupancy logits + """ + + logits = self.decoder(queries, latents).squeeze(-1) + + return logits + + + +@craftsman.register("michelangelo-aligned-autoencoder") +class MichelangeloAlignedAutoencoder(MichelangeloAutoencoder): + r""" + A VAE model for encoding shapes into latents and decoding latent representations into shapes. + """ + @dataclass + class Config(MichelangeloAutoencoder.Config): + clip_model_version: Optional[str] = None + + cfg: Config + + def configure(self) -> None: + if self.cfg.clip_model_version is not None: + self.clip_model: CLIPModel = CLIPModel.from_pretrained(self.cfg.clip_model_version) + self.projection = nn.Parameter(torch.empty(self.cfg.width, self.clip_model.projection_dim)) + self.logit_scale = torch.exp(self.clip_model.logit_scale.data) + nn.init.normal_(self.projection, std=self.clip_model.projection_dim ** -0.5) + else: + self.projection = nn.Parameter(torch.empty(self.cfg.width, 768)) + nn.init.normal_(self.projection, std=768 ** -0.5) + + self.cfg.num_latents = self.cfg.num_latents + 1 + + super().configure() + + def encode(self, + surface: torch.FloatTensor, + sample_posterior: bool = True): + """ + Args: + surface (torch.FloatTensor): [B, N, 3+C] + sample_posterior (bool): + + Returns: + latents (torch.FloatTensor) + posterior (DiagonalGaussianDistribution or None): + """ + assert surface.shape[-1] == 3 + self.cfg.point_feats, f"\ + Expected {3 + self.cfg.point_feats} channels, got {surface.shape[-1]}" + + pc, feats = surface[..., :3], surface[..., 3:] # B, n_samples, 3 + shape_latents = self.encoder(pc, feats) # B, num_latents, width + shape_embeds = shape_latents[:, 0] # B, width + shape_latents = shape_latents[:, 1:] # B, num_latents-1, width + kl_embed, posterior = self.encode_kl_embed(shape_latents, sample_posterior) # B, num_latents, embed_dim + + shape_embeds = shape_embeds @ self.projection + return shape_embeds, kl_embed, posterior + + def forward(self, + surface: torch.FloatTensor, + queries: torch.FloatTensor, + sample_posterior: bool = True): + """ + Args: + surface (torch.FloatTensor): [B, N, 3+C] + queries (torch.FloatTensor): [B, P, 3] + sample_posterior (bool): + + Returns: + shape_embeds (torch.FloatTensor): [B, width] + latents (torch.FloatTensor): [B, num_latents, embed_dim] + posterior (DiagonalGaussianDistribution or None). + logits (torch.FloatTensor): [B, P] + """ + + shape_embeds, kl_embed, posterior = self.encode(surface, sample_posterior=sample_posterior) + + latents = self.decode(kl_embed) # [B, num_latents - 1, width] + + logits = self.query(queries, latents) # [B,] + + return shape_embeds, latents, posterior, logits + diff --git a/Gen_3D_Modules/craftsman/models/autoencoders/utils.py b/Gen_3D_Modules/craftsman/models/autoencoders/utils.py new file mode 100644 index 0000000..7c0b3a4 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/autoencoders/utils.py @@ -0,0 +1,302 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn +from torch import distributed as tdist +from torch.nn import functional as F +import math +import mcubes +import numpy as np +from einops import repeat, rearrange +from skimage import measure + +from craftsman.utils.base import BaseModule +from craftsman.utils.typing import * +from craftsman.utils.misc import get_world_size +from craftsman.utils.ops import generate_dense_grid_points + +VALID_EMBED_TYPES = ["identity", "fourier", "hashgrid", "sphere_harmonic", "triplane_fourier"] + +class FourierEmbedder(nn.Module): + def __init__(self, + num_freqs: int = 6, + logspace: bool = True, + input_dim: int = 3, + include_input: bool = True, + include_pi: bool = True) -> None: + super().__init__() + + if logspace: + frequencies = 2.0 ** torch.arange( + num_freqs, + dtype=torch.float32 + ) + else: + frequencies = torch.linspace( + 1.0, + 2.0 ** (num_freqs - 1), + num_freqs, + dtype=torch.float32 + ) + + if include_pi: + frequencies *= torch.pi + + self.register_buffer("frequencies", frequencies, persistent=False) + self.include_input = include_input + self.num_freqs = num_freqs + + self.out_dim = self.get_dims(input_dim) + + def get_dims(self, input_dim): + temp = 1 if self.include_input or self.num_freqs == 0 else 0 + out_dim = input_dim * (self.num_freqs * 2 + temp) + + return out_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.num_freqs > 0: + embed = (x[..., None].contiguous() * self.frequencies).view(*x.shape[:-1], -1) + if self.include_input: + return torch.cat((x, embed.sin(), embed.cos()), dim=-1) + else: + return torch.cat((embed.sin(), embed.cos()), dim=-1) + else: + return x + + +class LearnedFourierEmbedder(nn.Module): + def __init__(self, input_dim, dim): + super().__init__() + assert (dim % 2) == 0 + half_dim = dim // 2 + per_channel_dim = half_dim // input_dim + self.weights = nn.Parameter(torch.randn(per_channel_dim)) + + self.out_dim = self.get_dims(input_dim) + + def forward(self, x): + # [b, t, c, 1] * [1, d] = [b, t, c, d] -> [b, t, c * d] + freqs = (x[..., None] * self.weights[None] * 2 * np.pi).view(*x.shape[:-1], -1) + fouriered = torch.cat((x, freqs.sin(), freqs.cos()), dim=-1) + return fouriered + + def get_dims(self, input_dim): + return input_dim * (self.weights.shape[0] * 2 + 1) + +class Sine(nn.Module): + def __init__(self, w0 = 1.): + super().__init__() + self.w0 = w0 + def forward(self, x): + return torch.sin(self.w0 * x) + +class Siren(nn.Module): + def __init__( + self, + in_dim, + out_dim, + w0 = 1., + c = 6., + is_first = False, + use_bias = True, + activation = None, + dropout = 0. + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.is_first = is_first + + weight = torch.zeros(out_dim, in_dim) + bias = torch.zeros(out_dim) if use_bias else None + self.init_(weight, bias, c = c, w0 = w0) + + self.weight = nn.Parameter(weight) + self.bias = nn.Parameter(bias) if use_bias else None + self.activation = Sine(w0) if activation is None else activation + self.dropout = nn.Dropout(dropout) + + def init_(self, weight, bias, c, w0): + dim = self.in_dim + + w_std = (1 / dim) if self.is_first else (math.sqrt(c / dim) / w0) + weight.uniform_(-w_std, w_std) + + if bias is not None: + bias.uniform_(-w_std, w_std) + + def forward(self, x): + out = F.linear(x, self.weight, self.bias) + out = self.activation(out) + out = self.dropout(out) + return out + +def get_embedder(embed_type="fourier", num_freqs=-1, input_dim=3, include_pi=True): + if embed_type == "identity" or (embed_type == "fourier" and num_freqs == -1): + return nn.Identity(), input_dim + + elif embed_type == "fourier": + embedder_obj = FourierEmbedder(num_freqs=num_freqs, include_pi=include_pi) + + elif embed_type == "learned_fourier": + embedder_obj = LearnedFourierEmbedder(in_channels=input_dim, dim=num_freqs) + + elif embed_type == "siren": + embedder_obj = Siren(in_dim=input_dim, out_dim=num_freqs * input_dim * 2 + input_dim) + + elif embed_type == "hashgrid": + raise NotImplementedError + + elif embed_type == "sphere_harmonic": + raise NotImplementedError + + else: + raise ValueError(f"{embed_type} is not valid. Currently only supprts {VALID_EMBED_TYPES}") + return embedder_obj + + +###################### AutoEncoder +class AutoEncoder(BaseModule): + @dataclass + class Config(BaseModule.Config): + pretrained_model_name_or_path: str = "" + num_latents: int = 256 + embed_dim: int = 64 + width: int = 768 + + cfg: Config + + def configure(self) -> None: + super().configure() + + def encode(self, x: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + raise NotImplementedError + + def decode(self, z: torch.FloatTensor) -> torch.FloatTensor: + raise NotImplementedError + + def encode_kl_embed(self, latents: torch.FloatTensor, sample_posterior: bool = True): + posterior = None + if self.cfg.embed_dim > 0: + moments = self.pre_kl(latents) + posterior = DiagonalGaussianDistribution(moments, feat_dim=-1) + if sample_posterior: + kl_embed = posterior.sample() + else: + kl_embed = posterior.mode() + else: + kl_embed = latents + return kl_embed, posterior + + def forward(self, + surface: torch.FloatTensor, + queries: torch.FloatTensor, + sample_posterior: bool = True): + shape_latents, kl_embed, posterior = self.encode(surface, sample_posterior=sample_posterior) + + latents = self.decode(kl_embed) # [B, num_latents, width] + + logits = self.query(queries, latents) # [B,] + + return shape_latents, latents, posterior, logits + + def query(self, queries: torch.FloatTensor, latents: torch.FloatTensor) -> torch.FloatTensor: + raise NotImplementedError + + @torch.no_grad() + def extract_geometry(self, + latents: torch.FloatTensor, + bounds: Union[Tuple[float], List[float], float] = (-1.05, -1.05, -1.05, 1.05, 1.05, 1.05), + grids_resolution: int = 256, + num_chunks: int = 10000, + ): + + if isinstance(bounds, float): + bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds] + + bbox_min = np.array(bounds[0:3]) + bbox_max = np.array(bounds[3:6]) + bbox_size = bbox_max - bbox_min + + xyz_samples, grid_size, length = generate_dense_grid_points( + bbox_min=bbox_min, + bbox_max=bbox_max, + grids_resolution=grids_resolution, + indexing="ij" + ) + xyz_samples = torch.FloatTensor(xyz_samples) + batch_size = latents.shape[0] + + batch_logits = [] + for start in range(0, xyz_samples.shape[0], num_chunks): + queries = xyz_samples[start: start + num_chunks, :].to(latents) + batch_queries = repeat(queries, "p c -> b p c", b=batch_size) + + logits = self.query(batch_queries, latents) + batch_logits.append(logits.cpu()) + + grid_logits = torch.cat(batch_logits, dim=1).view((batch_size, grid_size[0], grid_size[1], grid_size[2])).float().numpy() + + mesh_v_f = [] + has_surface = np.zeros((batch_size,), dtype=np.bool_) + for i in range(batch_size): + try: + vertices, faces, normals, _ = measure.marching_cubes(grid_logits[i], 0, method="lewiner") + # vertices, faces = mcubes.marching_cubes(grid_logits[i], 0) + vertices = vertices / grid_size * bbox_size + bbox_min + faces = faces[:, [2, 1, 0]] + mesh_v_f.append((vertices.astype(np.float32), np.ascontiguousarray(faces))) + has_surface[i] = True + except: + mesh_v_f.append((None, None)) + has_surface[i] = False + + return mesh_v_f, has_surface + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters: Union[torch.Tensor, List[torch.Tensor]], deterministic=False, feat_dim=1): + self.feat_dim = feat_dim + self.parameters = parameters + + if isinstance(parameters, list): + self.mean = parameters[0] + self.logvar = parameters[1] + else: + self.mean, self.logvar = torch.chunk(parameters, 2, dim=feat_dim) + + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean) + + def sample(self): + x = self.mean + self.std * torch.randn_like(self.mean) + return x + + def kl(self, other=None, dims=(1, 2)): + if self.deterministic: + return torch.Tensor([0.]) + else: + if other is None: + return 0.5 * torch.mean(torch.pow(self.mean, 2) + + self.var - 1.0 - self.logvar, + dim=dims) + else: + return 0.5 * torch.mean( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=dims) + + def nll(self, sample, dims=(1, 2)): + if self.deterministic: + return torch.Tensor([0.]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims) + + def mode(self): + return self.mean diff --git a/Gen_3D_Modules/craftsman/models/conditional_encoders/__init__.py b/Gen_3D_Modules/craftsman/models/conditional_encoders/__init__.py new file mode 100644 index 0000000..be7f5c6 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/conditional_encoders/__init__.py @@ -0,0 +1,3 @@ +from . import ( + clip_encoder, +) diff --git a/Gen_3D_Modules/craftsman/models/conditional_encoders/base.py b/Gen_3D_Modules/craftsman/models/conditional_encoders/base.py new file mode 100644 index 0000000..1ff6ba6 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/conditional_encoders/base.py @@ -0,0 +1,138 @@ +import random +import torch +import torch.nn as nn +import numpy as np +from PIL import Image +from dataclasses import dataclass +from torchvision.transforms import Normalize +from torchvision.transforms import InterpolationMode +from torchvision.transforms.transforms import _interpolation_modes_from_int + +from transformers import CLIPModel, CLIPTokenizer, CLIPImageProcessor +from transformers.utils import ModelOutput +from typing import Iterable, Optional, Union, List + +from craftsman.utils.base import BaseModule +from craftsman.utils.typing import * + +ImageType = Union[np.ndarray, torch.Tensor, Image.Image] + + +class BaseEmbedder(BaseModule): + @dataclass + class Config(BaseModule.Config): + pretrained_model_name_or_path: Optional[str] = None # the pretrained model name or path + + encode_camera: bool = False # whether to encode camera + camera_embeds_type: str = "sincos" # the type of camera embeds + camera_embeds_dim: Optional[int] = None # the dimension of camera embeds + n_views: int = 1 # the number of views + + empty_embeds_ratio: float = 0.1 # the ratio of empty embeds + zero_uncond_embeds: bool = True + + normalize_embeds: bool = False # whether to normalize the embeds + + cfg: Config + + def configure(self) -> None: + super().configure() + + if self.cfg.encode_camera: + self.distance = 1.0 + self.register_buffer( + "cameras", + torch.as_tensor([ + [[1, 0, 0, 0], + [0, 0, -1, -self.distance], + [0, 1, 0, 0], + [0, 0, 0, 1]], # front to back + + [[0, 0, 1, self.distance], + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 0, 1]], # right to left + + [[-1, 0, 0, 0], + [0, 0, 1, self.distance], + [0, 1, 0, 0], + [0, 0, 0, 1]], # back to front + + [[0, 0, -1, -self.distance], + [-1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 0, 1]], # left to right + ], dtype=torch.float32), + ) + + def encode_image(self, images: Iterable[Optional[ImageType]], camera_embeds: Optional[torch.Tensor] = None, **kwargs) -> torch.FloatTensor: + pass + + def encode_text(self, texts: List[str], **kwargs) -> torch.FloatTensor: + pass + + def encode_camera(self, c2ws: torch.Tensor): + if self.cfg.camera_embeds_type == "sincos": + assert c2ws.shape[-1] == 4 and c2ws.shape[-2] == 4, f"Invalid c2ws shape: {c2ws.shape}" + c2ws = c2ws.view(-1, 16) + return torch.cat([torch.sin(c2ws), torch.cos(c2ws)], dim=-1) + else: + raise NotImplementedError(f"Unknown camera_embeds_type: {self.cfg.camera_embeds_type}") + + def post_process_embeds(self, text_embeds, visual_embeds): + bs = text_embeds.shape[0] if text_embeds is not None else visual_embeds.shape[0] + + if self.cfg.normalize_embeds: + # post-process the text/visual embeds + if text_embeds is not None: + text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True) + if visual_embeds is not None: + visual_embeds = visual_embeds / visual_embeds.norm(dim=-1, keepdim=True) + + assert text_embeds is not None or visual_embeds is not None + + # return text_embeds, visual_embeds + if text_embeds is not None and visual_embeds is not None: + return torch.cat([text_embeds, visual_embeds], dim=1) + elif text_embeds is not None: + return text_embeds + else: + return visual_embeds + + def forward(self, batch): + bs = batch["surface"].shape[0] + + text_embeds, visual_embeds = None, None + + if random.random() < self.cfg.empty_embeds_ratio: + if "text_input_ids" in batch or "text_embeds" in batch: + if self.empty_text_embeds is None: + if not self.cfg.zero_uncond_embeds: + self.empty_text_embeds = self.encode_text([""]).detach() # [1, 77, 768] + text_embeds = self.empty_text_embeds.repeat(bs, 1, 1) + if "image" in batch or "image_embeds" in batch: + visual_embeds = self.empty_image_embeds.repeat(bs, 1, 1) + elif "mvimages" in batch or "mvimage_embeds" in batch: + visual_embeds = self.empty_image_embeds.unsqueeze(1).repeat(bs, 1, 1, 1) + else: + # for text inputs + if "text_input_ids" in batch: + text_embeds = self.encode_text(batch["text_input_ids"]) + + # for visual inputs + if "image" in batch: + if self.cfg.encode_camera: + visual_embeds = self.encode_image(batch["image"], cameras=batch["c2w"]) + else: + visual_embeds = self.encode_image(batch["image"]) + elif "mvimages" in batch: + n_views = batch["mvimages"].shape[1] + if self.cfg.encode_camera: + visual_embeds = self.encode_image( + batch["mvimages"].view(-1, *batch["mvimages"].shape[-3:]), \ + cameras=batch["c2ws"]).view(bs, n_views, *self.empty_image_embeds.shape[-2:]) + else: + visual_embeds = self.encode_image( + batch["mvimages"].view(-1, *batch["mvimages"].shape[-3:])).view(bs, n_views, *self.empty_image_embeds.shape[-2:]) + + return self.post_process_embeds(text_embeds, visual_embeds) diff --git a/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_clip.py b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_clip.py new file mode 100644 index 0000000..cefb5c9 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_clip.py @@ -0,0 +1,1419 @@ +# coding=utf-8 +# Copyright 2021 The OpenAI Team Authors and 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. +""" PyTorch CLIP model.""" + + +from dataclasses import dataclass +from typing import Any, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask +from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + ModelOutput, + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from transformers.models.clip.configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig + + +logger = logging.get_logger(__name__) + +# General docstring +_CONFIG_FOR_DOC = "CLIPConfig" +_CHECKPOINT_FOR_DOC = "openai/clip-vit-base-patch32" + +# Image classification docstring +_IMAGE_CLASS_CHECKPOINT = "openai/clip-vit-base-patch32" +_IMAGE_CLASS_EXPECTED_OUTPUT = "LABEL_0" + + +# contrastive loss function, adapted from +# https://sachinruk.github.io/blog/2021-03-07-clip.html +def contrastive_loss(logits: torch.Tensor) -> torch.Tensor: + return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device)) + + +def clip_loss(similarity: torch.Tensor) -> torch.Tensor: + caption_loss = contrastive_loss(similarity) + image_loss = contrastive_loss(similarity.t()) + return (caption_loss + image_loss) / 2.0 + + +@dataclass +class CLIPVisionModelOutput(ModelOutput): + """ + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + + Args: + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + image_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class CLIPTextModelOutput(ModelOutput): + """ + Base class for text model's outputs that also contains a pooling of the last hidden states. + + Args: + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + text_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class CLIPOutput(ModelOutput): + """ + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image:(`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text:(`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`CLIPTextModel`]. + image_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`): + The image embeddings obtained by applying the projection layer to the pooled output of [`CLIPVisionModel`]. + text_model_output(`BaseModelOutputWithPooling`): + The output of the [`CLIPTextModel`]. + vision_model_output(`BaseModelOutputWithPooling`): + The output of the [`CLIPVisionModel`]. + """ + + loss: Optional[torch.FloatTensor] = None + logits_per_image: torch.FloatTensor = None + logits_per_text: torch.FloatTensor = None + text_embeds: torch.FloatTensor = None + image_embeds: torch.FloatTensor = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPooling = None + + def to_tuple(self) -> Tuple[Any]: + return tuple( + self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple() + for k in self.keys() + ) + + +class CLIPVisionEmbeddings(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: + batch_size = pixel_values.shape[0] + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +class CLIPAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Input shape: Batch x Time x Channel""" + + bsz, tgt_len, embed_dim = hidden_states.size() + + # get query proj + query_states = self.q_proj(hidden_states) * self.scale + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_states = value_states.view(*proj_shape) + + src_len = key_states.size(1) + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) + + if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" + f" {attn_weights.size()}" + ) + + # apply the causal_attention_mask first + if causal_attention_mask is not None: + if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" + f" {causal_attention_mask.size()}" + ) + attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, tgt_len, src_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + if output_attentions: + # this operation is a bit akward, but it's required to + # make sure that attn_weights keeps its gradient. + # In order to do so, attn_weights have to reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) + else: + attn_weights_reshaped = None + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + + attn_output = torch.bmm(attn_probs, value_states) + + if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, tgt_len, embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights_reshaped + + +class CLIPMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class CLIPEncoderLayer(nn.Module): + def __init__(self, config: CLIPConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = CLIPAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + causal_attention_mask: torch.Tensor, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + `(config.encoder_attention_heads,)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class CLIPPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = CLIPConfig + base_model_prefix = "clip" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor + if isinstance(module, CLIPTextEmbeddings): + module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02) + module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02) + elif isinstance(module, CLIPVisionEmbeddings): + factor = self.config.initializer_factor + nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + elif isinstance(module, CLIPAttention): + factor = self.config.initializer_factor + in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor + out_proj_std = (module.embed_dim**-0.5) * factor + nn.init.normal_(module.q_proj.weight, std=in_proj_std) + nn.init.normal_(module.k_proj.weight, std=in_proj_std) + nn.init.normal_(module.v_proj.weight, std=in_proj_std) + nn.init.normal_(module.out_proj.weight, std=out_proj_std) + elif isinstance(module, CLIPMLP): + factor = self.config.initializer_factor + in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor + fc_std = (2 * module.config.hidden_size) ** -0.5 * factor + nn.init.normal_(module.fc1.weight, std=fc_std) + nn.init.normal_(module.fc2.weight, std=in_proj_std) + elif isinstance(module, CLIPModel): + nn.init.normal_( + module.text_projection.weight, + std=module.text_embed_dim**-0.5 * self.config.initializer_factor, + ) + nn.init.normal_( + module.visual_projection.weight, + std=module.vision_embed_dim**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPVisionModelWithProjection): + nn.init.normal_( + module.visual_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPTextModelWithProjection): + nn.init.normal_( + module.text_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPForImageClassification): + nn.init.normal_( + module.classifier.weight, + std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor, + ) + + if isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + + +CLIP_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also 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 ([`CLIPConfig`]): 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. +""" + +CLIP_TEXT_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + 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. +""" + +CLIP_VISION_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. + 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. +""" + +CLIP_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + 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. +""" + + +class CLIPEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__(self, config: CLIPConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Causal mask for the text model. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + 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. + """ + 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 + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + encoder_layer.__call__, + hidden_states, + attention_mask, + causal_attention_mask, + output_attentions, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + causal_attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + +class CLIPTextTransformer(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + self.embeddings = CLIPTextEmbeddings(config) + self.encoder = CLIPEncoder(config) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + 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 input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + causal_attention_mask = _create_4d_causal_attention_mask( + input_shape, hidden_states.dtype, device=hidden_states.device + ) + # expand attention_mask + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id) + .int() + .argmax(dim=-1), + ] + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@add_start_docstrings( + """The text model from CLIP without any head or projection on top.""", + CLIP_START_DOCSTRING, +) +class CLIPTextModel(CLIPPreTrainedModel): + config_class = CLIPTextConfig + + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + self.text_model = CLIPTextTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPTextModel + + >>> model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class CLIPVisionTransformer(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = CLIPEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPVisionConfig) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + 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") + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + pooled_output = last_hidden_state[:, 0, :] + pooled_output = self.post_layernorm(pooled_output) + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@add_start_docstrings( + """The vision model from CLIP without any head or projection on top.""", + CLIP_START_DOCSTRING, +) +class CLIPVisionModel(CLIPPreTrainedModel): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + self.vision_model = CLIPVisionTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPVisionConfig) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPVisionModel + + >>> model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled CLS states + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + return self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +@add_start_docstrings(CLIP_START_DOCSTRING) +class CLIPModel(CLIPPreTrainedModel): + config_class = CLIPConfig + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + if not isinstance(config.text_config, CLIPTextConfig): + raise ValueError( + "config.text_config is expected to be of type CLIPTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, CLIPVisionConfig): + raise ValueError( + "config.vision_config is expected to be of type CLIPVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + self.vision_embed_dim = vision_config.hidden_size + + self.text_model = CLIPTextTransformer(text_config) + self.vision_model = CLIPVisionTransformer(vision_config) + + self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False) + self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False) + self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value)) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + def get_text_features( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by + applying the projection layer to the pooled output of [`CLIPTextModel`]. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + >>> text_features = model.get_text_features(**inputs) + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = text_outputs[1] + text_features = self.text_projection(pooled_output) + + return text_features + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + def get_image_features( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by + applying the projection layer to the pooled output of [`CLIPVisionModel`]. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> image_features = model.get_image_features(**inputs) + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = vision_outputs[1] # pooled_output + image_features = self.visual_projection(pooled_output) + + return image_features + + @add_start_docstrings_to_model_forward(CLIP_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CLIPOutput, config_class=CLIPConfig) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor( + ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True + ... ) + + >>> outputs = model(**inputs) + >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score + >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + image_embeds = vision_outputs[1] + image_embeds = self.visual_projection(image_embeds) + + text_embeds = text_outputs[1] + text_embeds = self.text_projection(text_embeds) + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = clip_loss(logits_per_text) + + if not return_dict: + output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs) + return ((loss,) + output) if loss is not None else output + + return CLIPOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +@add_start_docstrings( + """ + CLIP Text Model with a projection layer on top (a linear layer on top of the pooled output). + """, + CLIP_START_DOCSTRING, +) +class CLIPTextModelWithProjection(CLIPPreTrainedModel): + config_class = CLIPTextConfig + + _no_split_modules = ["CLIPTextEmbeddings", "CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + + self.text_model = CLIPTextTransformer(config) + + self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CLIPTextModelOutput, config_class=CLIPTextConfig) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPTextModelOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPTextModelWithProjection + + >>> model = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> text_embeds = outputs.text_embeds + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = text_outputs[1] + + text_embeds = self.text_projection(pooled_output) + + if not return_dict: + outputs = (text_embeds, text_outputs[0]) + text_outputs[2:] + return tuple(output for output in outputs if output is not None) + + return CLIPTextModelOutput( + text_embeds=text_embeds, + last_hidden_state=text_outputs.last_hidden_state, + hidden_states=text_outputs.hidden_states, + attentions=text_outputs.attentions, + ) + + +@add_start_docstrings( + """ + CLIP Vision Model with a projection layer on top (a linear layer on top of the pooled output). + """, + CLIP_START_DOCSTRING, +) +class CLIPVisionModelWithProjection(CLIPPreTrainedModel): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + + self.vision_model = CLIPVisionTransformer(config) + + self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CLIPVisionModelOutput, config_class=CLIPVisionConfig) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPVisionModelOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPVisionModelWithProjection + + >>> model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> image_embeds = outputs.image_embeds + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = vision_outputs[1] # pooled_output + + image_embeds = self.visual_projection(pooled_output) + + if not return_dict: + outputs = (image_embeds, vision_outputs[0]) + vision_outputs[2:] + return tuple(output for output in outputs if output is not None) + + return CLIPVisionModelOutput( + image_embeds=image_embeds, + last_hidden_state=vision_outputs.last_hidden_state, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + ) + + +@add_start_docstrings( + """ + CLIP vision encoder with an image classification head on top (a linear layer on top of the pooled final hidden states of + the patch tokens) e.g. for ImageNet. + """, + CLIP_START_DOCSTRING, +) +class CLIPForImageClassification(CLIPPreTrainedModel): + main_input_name = "pixel_values" + + def __init__(self, config: CLIPConfig) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.vision_model = CLIPVisionTransformer(config.vision_config) + + # Classifier head + self.classifier = ( + nn.Linear(config.vision_config.hidden_size, 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(CLIP_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_IMAGE_CLASS_CHECKPOINT, + output_type=ImageClassifierOutput, + config_class=_CONFIG_FOR_DOC, + expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT, + ) + def forward( + self, + pixel_values: 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). + """ + 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 + + outputs = self.vision_model( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + # average pool the patch tokens + sequence_output = torch.mean(sequence_output[:, 1:, :], dim=1) + # apply classifier + logits = self.classifier(sequence_output) + + 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, + ) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_conditional_clip.py b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_conditional_clip.py new file mode 100644 index 0000000..1f026e4 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip/modeling_conditional_clip.py @@ -0,0 +1,385 @@ +# 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. +# Reference: +# * transformers/models/dinov2/modeling_dinov2.py +# * https://github.com/facebookresearch/DiT/blob/main/models.py#L101 +# * https://github.com/3DTopia/OpenLRM/tree/main/openlrm/models/encoders/dinov2 +""" PyTorch CLIP model.""" + +from typing import Dict, List, Optional, Set, Tuple, Union + +import torch +import torch.nn as nn + +from .modeling_clip import ( + CLIPConfig, + CLIPTextConfig, + CLIPVisionConfig, + CLIPEncoderLayer, + CLIPTextTransformer, + CLIPVisionTransformer, + CLIPModel, + CLIPVisionEmbeddings, + CLIPVisionModel, + CLIPOutput, + BaseModelOutput, + BaseModelOutputWithPooling +) + + +class ModLN(nn.Module): + def __init__(self, inner_dim: int, mod_dim: int = 32): + super().__init__() + self.mlp = nn.Sequential( + nn.SiLU(), + nn.Linear(mod_dim, inner_dim * 2), + ) + + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.zeros_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x:torch.Tensor, condition:torch.Tensor): + ''' + x: [N, M, C_in], M: num of tokens + condition: [N, C_mod] + ''' + shift, scale = self.mlp(condition).unsqueeze(1).chunk(2, dim=-1) + return x * (1 + scale) + shift + + +class ConditionalCLIPVisionConfig(CLIPVisionConfig): + def __init__(self, modulation_dim: int = 32, *args, **kwargs): + super().__init__(*args, **kwargs) + self.modulation_dim = modulation_dim + + +class ConditionalCLIPEncoderLayer(CLIPEncoderLayer): + """This corresponds to the Block class in the original implementation.""" + + def __init__(self, config: ConditionalCLIPVisionConfig) -> None: + super().__init__(config) + self.mod_norm1 = ModLN(config.hidden_size, config.modulation_dim) + self.mod_norm2 = ModLN(config.hidden_size, config.modulation_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + causal_attention_mask: torch.Tensor, + condition: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + residual = hidden_states + + hidden_states = self.mod_norm1(self.layer_norm1(hidden_states), condition) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.mod_norm2(self.layer_norm2(hidden_states), condition) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class ConditionalCLIPEncoder(nn.Module): + def __init__(self, config: CLIPConfig) -> None: + super().__init__() + self.config = config + self.layers = nn.ModuleList([ConditionalCLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + condition: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[tuple, BaseModelOutput]: + 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 + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + encoder_layer.__call__, + hidden_states, + attention_mask, + causal_attention_mask, + condition=condition, + output_attentions=output_attentions, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + causal_attention_mask, + condition=condition, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + +class ConditionalCLIPVisionTransformer(CLIPVisionTransformer): + def __init__(self, config: ConditionalCLIPVisionConfig): + super().__init__(config) + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = ConditionalCLIPEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + condition: 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") + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + condition=condition, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + pooled_output = last_hidden_state[:, 0, :] + pooled_output = self.post_layernorm(pooled_output) + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +class ConditionalCLIPVisionModel(CLIPVisionModel): + config_class = ConditionalCLIPVisionConfig + + def __init__(self, config: ConditionalCLIPVisionConfig): + super().__init__(config) + self.vision_model = ConditionalCLIPVisionTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + condition: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + return self.vision_model( + pixel_values=pixel_values, + condition=condition, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class ConditionalCLIPModel(CLIPModel): + config_class = CLIPConfig + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + if not isinstance(config.text_config, CLIPTextConfig): + raise ValueError( + "config.text_config is expected to be of type CLIPTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, CLIPVisionConfig): + raise ValueError( + "config.vision_config is expected to be of type CLIPVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + self.vision_embed_dim = vision_config.hidden_size + + self.text_model = CLIPTextTransformer(text_config) + self.vision_model = ConditionalCLIPVisionTransformer(vision_config) + + self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False) + self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False) + self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value)) + + # Initialize weights and apply final processing + self.post_init() + + def get_image_features( + self, + pixel_values: Optional[torch.FloatTensor] = None, + condition: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> torch.FloatTensor: + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + condition=condition, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = vision_outputs[1] # pooled_output + image_features = self.visual_projection(pooled_output) + + return image_features + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + condition: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPOutput]: + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + 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 + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + condition=condition, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + image_embeds = vision_outputs[1] + image_embeds = self.visual_projection(image_embeds) + + text_embeds = text_outputs[1] + text_embeds = self.text_projection(text_embeds) + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = clip_loss(logits_per_text) + + if not return_dict: + output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs) + return ((loss,) + output) if loss is not None else output + + return CLIPOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) diff --git a/Gen_3D_Modules/craftsman/models/conditional_encoders/clip_encoder.py b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip_encoder.py new file mode 100644 index 0000000..326f3b6 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/conditional_encoders/clip_encoder.py @@ -0,0 +1,172 @@ +import random +import torch +from torch import nn +import numpy as np +from PIL import Image +from einops import rearrange +from dataclasses import dataclass +from torchvision.transforms import Normalize +from torchvision.transforms import InterpolationMode +from torchvision.transforms.transforms import _interpolation_modes_from_int +from torchvision import transforms + +from transformers import CLIPTokenizer, CLIPImageProcessor +from transformers.utils import ModelOutput +from typing import Iterable, Optional, Union, List + +import craftsman +from craftsman.utils.typing import * +from .clip.modeling_clip import CLIPModel +from .clip.modeling_conditional_clip import ConditionalCLIPModel +from .base import BaseEmbedder, ImageType + +@dataclass +class CLIPEmbedOutput(ModelOutput): + last_hidden_state: torch.FloatTensor = None + pooler_output: torch.FloatTensor = None + embeds: torch.FloatTensor = None + +@craftsman.register("clip-embedder") +class CLIPEmbedder(BaseEmbedder): + + @dataclass + class Config(BaseEmbedder.Config): + freeze_modulation: bool = False + config_path: str = '' + + cfg: Config + + def configure(self) -> None: + super().configure() + + # Load the CLIP model and processor + if not self.cfg.encode_camera: + self.model: CLIPModel = CLIPModel.from_pretrained(self.cfg.pretrained_model_name_or_path) + else: + if self.cfg.pretrained_model_name_or_path == '': + assert self.cfg.config_path is not None, "The config path should be provided" + conditional_clip_config = ConditionalCLIPModel.config_class.from_json_file(self.cfg.config_path) + conditional_clip_config.vision_config.modulation_dim = self.cfg.camera_embeds_dim + self.model: CLIPModel = ConditionalCLIPModel(conditional_clip_config) + else: + conditional_clip_config = ConditionalCLIPModel.config_class.from_pretrained( + self.cfg.pretrained_model_name_or_path, + ) + conditional_clip_config.vision_config.modulation_dim = self.cfg.camera_embeds_dim + self.model: CLIPModel = ConditionalCLIPModel.from_pretrained( + self.cfg.pretrained_model_name_or_path, + vision_config=conditional_clip_config.vision_config + ) + + self.tokenizer = None + self.image_preprocess = CLIPImageProcessor() + self.transform = transforms.Compose( + [ + transforms.Resize(224, transforms.InterpolationMode.BICUBIC, antialias=True), + transforms.CenterCrop(224), # crop a (224, 224) square + transforms.Normalize( + mean=[0.48145466, 0.4578275, 0.40821073], + std=[0.26862954, 0.26130258, 0.27577711], + ), + ] + ) + + self.logit_scale = self.model.logit_scale.exp() + + if self.cfg.zero_uncond_embeds: + self.empty_text_embeds = torch.zeros((1, 77, 768)).detach() + self.empty_image_embeds = torch.zeros((self.cfg.n_views, 257, 1024)).detach() + else: + try: + self.empty_text_embeds = self.encode_text([""]).detach() # [1, 77, 768] + except: + self.empty_text_embeds = None + if self.cfg.encode_camera: + self.empty_image_embeds = self.encode_image(torch.zeros(self.cfg.n_views, 224, 224, 3), self.cameras[:self.cfg.n_views]).detach() + else: + self.empty_image_embeds = self.encode_image(torch.zeros(self.cfg.n_views, 224, 224, 3)).detach() + + # Freeze the model parameters + self.model.eval() + for k, p in self.model.named_parameters(): + ks = k.split('.') + if 'mod_norm1' in ks or 'mod_norm2' in ks and not self.cfg.freeze_modulation: + p.requires_grad_(True) + else: + p.requires_grad_(False) + + def encode_image(self, images: Iterable[Optional[ImageType]], cameras: Optional[torch.Tensor] = None, force_none_camera_embeds: bool = False, return_dict: bool = False, **kwargs) -> torch.FloatTensor: + camera_embeds = None + if isinstance(images, (np.ndarray, torch.Tensor)): # for training process + assert images.min() >= 0.0 and images.max() <= 1.0, "The pixel values should be in the range of [0, 1]" + do_rescale = False + if self.cfg.encode_camera: + assert cameras is not None, "The cameras should be provided" + camera_embeds = self.encode_camera(cameras) + pixel_values = self.transform(images.permute(0, 3, 1, 2)) + else: # for inference process + do_rescale = True + if self.cfg.encode_camera: + if cameras is None: + bs = len(images) // self.cfg.n_views + cameras = self.cameras[:self.cfg.n_views].repeat(bs, 1, 1).to(self.model.device) + camera_embeds = self.encode_camera(cameras) + pixel_values = self.image_preprocess.preprocess(images, return_tensors='pt', do_rescale=do_rescale).pixel_values + + if force_none_camera_embeds: + camera_embeds = None + + packed = False + if pixel_values.ndim == 4: + packed = True + pixel_values = pixel_values.unsqueeze(1) + if camera_embeds is not None: + camera_embeds = camera_embeds.unsqueeze(1) + + if self.cfg.encode_camera and camera_embeds is not None: + vision_outputs = self.model.vision_model( + pixel_values=rearrange(pixel_values.to(self.model.device), "B N C H W -> (B N) C H W"), + condition=rearrange(camera_embeds, "B N C -> (B N) C") + ) + else: + vision_outputs = self.model.vision_model( + pixel_values=rearrange(pixel_values.to(self.model.device), "B N C H W -> (B N) C H W"), + ) + + if return_dict: + pooler_output = vision_outputs[1] # pooled_output + image_features = self.model.visual_projection(pooler_output) + + return CLIPEmbedOutput( + last_hidden_state=vision_outputs.last_hidden_state, + pooler_output=pooler_output, + embeds=image_features + ) + else: + return vision_outputs.last_hidden_state + + @torch.no_grad() + def encode_text(self, text_inputs: torch.Tensor, return_dict: bool = False) -> torch.FloatTensor: + if self.tokenizer is None: + self.tokenizer = CLIPTokenizer.from_pretrained(self.cfg.pretrained_model_name_or_path) + + if isinstance(text_inputs, list): + text_inputs = self.tokenizer( + text_inputs, + max_length=self.tokenizer.model_max_length, + padding="max_length", + return_tensors="pt" + ).input_ids + text_outputs = self.model.text_model(input_ids=text_inputs.to(self.model.device)) + + pooler_output = text_outputs[1] # pooled_output + text_features = self.model.text_projection(pooler_output) + + if return_dict: + return CLIPEmbedOutput( + last_hidden_state=text_outputs.last_hidden_state, + pooler_output=pooler_output, + embeds=text_features + ) + else: + return text_outputs.last_hidden_state \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/denoisers/__init__.py b/Gen_3D_Modules/craftsman/models/denoisers/__init__.py new file mode 100644 index 0000000..b878f03 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/denoisers/__init__.py @@ -0,0 +1,3 @@ +from . import ( + simple_denoiser, +) diff --git a/Gen_3D_Modules/craftsman/models/denoisers/simple_denoiser.py b/Gen_3D_Modules/craftsman/models/denoisers/simple_denoiser.py new file mode 100644 index 0000000..2b674a3 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/denoisers/simple_denoiser.py @@ -0,0 +1,191 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn +from typing import Optional +from diffusers.models.embeddings import Timesteps +import math + +import craftsman +from craftsman.models.transformers.attention import ResidualAttentionBlock +from craftsman.models.transformers.utils import init_linear, MLP +from craftsman.utils.base import BaseModule + + +class UNetDiffusionTransformer(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + layers: int, + heads: int, + init_scale: float = 0.25, + qkv_bias: bool = False, + skip_ln: bool = False, + use_checkpoint: bool = False + ): + super().__init__() + + self.n_ctx = n_ctx + self.width = width + self.layers = layers + + self.encoder = nn.ModuleList() + for _ in range(layers): + resblock = ResidualAttentionBlock( + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_checkpoint=use_checkpoint + ) + self.encoder.append(resblock) + + self.middle_block = ResidualAttentionBlock( + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_checkpoint=use_checkpoint + ) + + self.decoder = nn.ModuleList() + for _ in range(layers): + resblock = ResidualAttentionBlock( + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_checkpoint=use_checkpoint + ) + linear = nn.Linear(width * 2, width) + init_linear(linear, init_scale) + + layer_norm = nn.LayerNorm(width) if skip_ln else None + + self.decoder.append(nn.ModuleList([resblock, linear, layer_norm])) + + def forward(self, x: torch.Tensor): + + enc_outputs = [] + for block in self.encoder: + x = block(x) + enc_outputs.append(x) + + x = self.middle_block(x) + + for i, (resblock, linear, layer_norm) in enumerate(self.decoder): + x = torch.cat([enc_outputs.pop(), x], dim=-1) + x = linear(x) + + if layer_norm is not None: + x = layer_norm(x) + + x = resblock(x) + + return x + + +@craftsman.register("simple-denoiser") +class SimpleDenoiser(BaseModule): + + @dataclass + class Config(BaseModule.Config): + pretrained_model_name_or_path: Optional[str] = None + input_channels: int = 32 + output_channels: int = 32 + n_ctx: int = 512 + width: int = 768 + layers: int = 6 + heads: int = 12 + context_dim: int = 1024 + context_ln: bool = True + skip_ln: bool = False + init_scale: float = 0.25 + flip_sin_to_cos: bool = False + use_checkpoint: bool = False + + cfg: Config + + def configure(self) -> None: + super().configure() + + init_scale = self.cfg.init_scale * math.sqrt(1.0 / self.cfg.width) + + self.backbone = UNetDiffusionTransformer( + n_ctx=self.cfg.n_ctx, + width=self.cfg.width, + layers=self.cfg.layers, + heads=self.cfg.heads, + skip_ln=self.cfg.skip_ln, + init_scale=init_scale, + use_checkpoint=self.cfg.use_checkpoint + ) + self.ln_post = nn.LayerNorm(self.cfg.width) + self.input_proj = nn.Linear(self.cfg.input_channels, self.cfg.width) + self.output_proj = nn.Linear(self.cfg.width, self.cfg.output_channels) + + # timestep embedding + self.time_embed = Timesteps(self.cfg.width, flip_sin_to_cos=self.cfg.flip_sin_to_cos, downscale_freq_shift=0) + self.time_proj = MLP(width=self.cfg.width, init_scale=init_scale) + + if self.cfg.context_ln: + self.context_embed = nn.Sequential( + nn.LayerNorm(self.cfg.context_dim), + nn.Linear(self.cfg.context_dim, self.cfg.width), + ) + else: + self.context_embed = nn.Linear(self.cfg.context_dim, self.cfg.width) + + if self.cfg.pretrained_model_name_or_path: + pretrained_ckpt = torch.load(self.cfg.pretrained_model_name_or_path, map_location="cpu") + _pretrained_ckpt = {} + for k, v in pretrained_ckpt.items(): + if k.startswith('denoiser_model.'): + _pretrained_ckpt[k.replace('denoiser_model.', '')] = v + pretrained_ckpt = _pretrained_ckpt + if 'state_dict' in pretrained_ckpt: + _pretrained_ckpt = {} + for k, v in pretrained_ckpt['state_dict'].items(): + if k.startswith('denoiser_model.'): + _pretrained_ckpt[k.replace('denoiser_model.', '')] = v + pretrained_ckpt = _pretrained_ckpt + self.load_state_dict(pretrained_ckpt, strict=True) + + def forward(self, + model_input: torch.FloatTensor, + timestep: torch.LongTensor, + context: torch.FloatTensor): + + r""" + Args: + model_input (torch.FloatTensor): [bs, n_data, c] + timestep (torch.LongTensor): [bs,] + context (torch.FloatTensor): [bs, context_tokens, c] + + Returns: + sample (torch.FloatTensor): [bs, n_data, c] + + """ + + _, n_data, _ = model_input.shape + + # 1. time + t_emb = self.time_proj(self.time_embed(timestep)).unsqueeze(dim=1) + + # 2. conditions projector + context = self.context_embed(context) + + # 3. denoiser + x = self.input_proj(model_input) + x = torch.cat([t_emb, context, x], dim=1) + x = self.backbone(x) + x = self.ln_post(x) + x = x[:, -n_data:] # B, n_data, width + sample = self.output_proj(x) # B, n_data, embed_dim + + return sample \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/geometry/__init__.py b/Gen_3D_Modules/craftsman/models/geometry/__init__.py new file mode 100644 index 0000000..c8cc981 --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/geometry/__init__.py @@ -0,0 +1,3 @@ +from . import ( + base +) diff --git a/Gen_3D_Modules/craftsman/models/geometry/base.py b/Gen_3D_Modules/craftsman/models/geometry/base.py new file mode 100644 index 0000000..aabea8a --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/geometry/base.py @@ -0,0 +1,194 @@ +from dataclasses import dataclass, field + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +import craftsman +from .utils import ( + Mesh, + IsosurfaceHelper, + MarchingCubeCPUHelper, + MarchingTetrahedraHelper, +) + +from craftsman.utils.base import BaseModule +from craftsman.utils.ops import chunk_batch, scale_tensor +from craftsman.utils.typing import * + +class BaseGeometry(BaseModule): + @dataclass + class Config(BaseModule.Config): + pass + + cfg: Config + + @staticmethod + def create_from( + other: "BaseGeometry", cfg: Optional[Union[dict, DictConfig]] = None, **kwargs + ) -> "BaseGeometry": + raise TypeError( + f"Cannot create {BaseGeometry.__name__} from {other.__class__.__name__}" + ) + + def export(self, *args, **kwargs) -> Dict[str, Any]: + return {} + + +class BaseImplicitGeometry(BaseGeometry): + @dataclass + class Config(BaseGeometry.Config): + radius: float = 1.0 + isosurface: bool = True + isosurface_method: str = "mt" + isosurface_resolution: int = 128 + isosurface_threshold: Union[float, str] = 0.0 + isosurface_chunk: int = 0 + isosurface_coarse_to_fine: bool = True + isosurface_deformable_grid: bool = False + isosurface_remove_outliers: bool = True + isosurface_outlier_n_faces_threshold: Union[int, float] = 0.01 + + cfg: Config + + def configure(self) -> None: + 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, + ), + ) + self.isosurface_helper: Optional[IsosurfaceHelper] = None + self.unbounded: bool = False + + def _initilize_isosurface_helper(self): + if self.cfg.isosurface and self.isosurface_helper is None: + if self.cfg.isosurface_method == "mc-cpu": + self.isosurface_helper = MarchingCubeCPUHelper( + self.cfg.isosurface_resolution + ).to(self.device) + elif self.cfg.isosurface_method == "mt": + self.isosurface_helper = MarchingTetrahedraHelper( + self.cfg.isosurface_resolution, + f"load/tets/{self.cfg.isosurface_resolution}_tets.npz", + ).to(self.device) + else: + raise AttributeError( + "Unknown isosurface method {self.cfg.isosurface_method}" + ) + + def forward( + self, points: Float[Tensor, "*N Di"], output_normal: bool = False + ) -> Dict[str, Float[Tensor, "..."]]: + raise NotImplementedError + + def forward_field( + self, points: Float[Tensor, "*N Di"] + ) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]: + # return the value of the implicit field, could be density / signed distance + # also return a deformation field if the grid vertices can be optimized + raise NotImplementedError + + def forward_level( + self, field: Float[Tensor, "*N 1"], threshold: float + ) -> Float[Tensor, "*N 1"]: + # return the value of the implicit field, where the zero level set represents the surface + raise NotImplementedError + + def _isosurface(self, bbox: Float[Tensor, "2 3"], fine_stage: bool = False) -> Mesh: + def batch_func(x): + # scale to bbox as the input vertices are in [0, 1] + field, deformation = self.forward_field( + scale_tensor( + x.to(bbox.device), self.isosurface_helper.points_range, bbox + ), + ) + field = field.to( + x.device + ) # move to the same device as the input (could be CPU) + if deformation is not None: + deformation = deformation.to(x.device) + return field, deformation + + assert self.isosurface_helper is not None + + field, deformation = chunk_batch( + batch_func, + self.cfg.isosurface_chunk, + self.isosurface_helper.grid_vertices, + ) + + threshold: float + + if isinstance(self.cfg.isosurface_threshold, float): + threshold = self.cfg.isosurface_threshold + elif self.cfg.isosurface_threshold == "auto": + eps = 1.0e-5 + threshold = field[field > eps].mean().item() + craftsman.info( + f"Automatically determined isosurface threshold: {threshold}" + ) + else: + raise TypeError( + f"Unknown isosurface_threshold {self.cfg.isosurface_threshold}" + ) + + level = self.forward_level(field, threshold) + mesh: Mesh = self.isosurface_helper(level, deformation=deformation) + mesh.v_pos = scale_tensor( + mesh.v_pos, self.isosurface_helper.points_range, bbox + ) # scale to bbox as the grid vertices are in [0, 1] + mesh.add_extra("bbox", bbox) + + if self.cfg.isosurface_remove_outliers: + # remove outliers components with small number of faces + # only enabled when the mesh is not differentiable + mesh = mesh.remove_outlier(self.cfg.isosurface_outlier_n_faces_threshold) + + return mesh + + def isosurface(self) -> Mesh: + if not self.cfg.isosurface: + raise NotImplementedError( + "Isosurface is not enabled in the current configuration" + ) + self._initilize_isosurface_helper() + if self.cfg.isosurface_coarse_to_fine: + craftsman.debug("First run isosurface to get a tight bounding box ...") + with torch.no_grad(): + mesh_coarse = self._isosurface(self.bbox) + vmin, vmax = mesh_coarse.v_pos.amin(dim=0), mesh_coarse.v_pos.amax(dim=0) + vmin_ = (vmin - (vmax - vmin) * 0.1).max(self.bbox[0]) + vmax_ = (vmax + (vmax - vmin) * 0.1).min(self.bbox[1]) + craftsman.debug("Run isosurface again with the tight bounding box ...") + mesh = self._isosurface(torch.stack([vmin_, vmax_], dim=0), fine_stage=True) + else: + mesh = self._isosurface(self.bbox) + return mesh + + +class BaseExplicitGeometry(BaseGeometry): + @dataclass + class Config(BaseGeometry.Config): + radius: float = 1.0 + + cfg: Config + + def configure(self) -> None: + 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, + ), + ) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/geometry/utils.py b/Gen_3D_Modules/craftsman/models/geometry/utils.py new file mode 100644 index 0000000..2df9e5e --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/geometry/utils.py @@ -0,0 +1,559 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +import craftsman +from craftsman.utils.typing import * + + +def dot(x, y): + return torch.sum(x * y, -1, keepdim=True) + + +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): + self.extras[k] = v + + def remove_outlier(self, outlier_n_faces_threshold: Union[int, float]): + if self.requires_grad: + craftsman.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 + craftsman.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 + craftsman.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 + craftsman.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 + craftsman.debug( + f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}" + ) + return clean_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, dim=-1) + + # 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( + self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {} + ): + craftsman.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) + + setattr(co, 'max_cost', 2.0) + setattr(po, 'resolution', 4096) + + atlas.generate(co, po, verbose=True) + 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( + self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {} + ): + self._v_tex, self._t_tex_idx = self._unwrap_uv( + xatlas_chart_options, xatlas_pack_options + ) + + def set_vertex_color(self, v_rgb): + assert v_rgb.shape[0] == self.v_pos.shape[0] + self._v_rgb = v_rgb + + 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 + +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: + craftsman.warn( + f"{self.__class__.__name__} does not support deformation. Ignoring." + ) + level = -level.view(self.resolution, self.resolution, self.resolution) + v_pos, t_pos_idx = self.mc_func( + level.detach().cpu().numpy(), 0.0 + ) # transform to numpy + 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) + + +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 + + 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/craftsman/models/transformers/attention.py b/Gen_3D_Modules/craftsman/models/transformers/attention.py new file mode 100644 index 0000000..7d4119b --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/transformers/attention.py @@ -0,0 +1,207 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +from craftsman.utils.typing import * +from craftsman.utils.checkpoint import checkpoint + +from .utils import init_linear, MLP + +class MultiheadAttention(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + heads: int, + init_scale: float, + qkv_bias: bool, + use_flash: bool = False + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.heads = heads + self.c_qkv = nn.Linear(width, width * 3, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + self.attention = QKVMultiheadAttention(heads=heads, n_ctx=n_ctx, use_flash=use_flash) + init_linear(self.c_qkv, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x): + x = self.c_qkv(x) + x = checkpoint(self.attention, (x,), (), True) + x = self.c_proj(x) + return x + + +class QKVMultiheadAttention(nn.Module): + def __init__(self, *, heads: int, n_ctx: int, use_flash: bool = False): + super().__init__() + self.heads = heads + self.n_ctx = n_ctx + self.use_flash = use_flash + + def forward(self, qkv): + bs, n_ctx, width = qkv.shape + attn_ch = width // self.heads // 3 + scale = 1 / math.sqrt(math.sqrt(attn_ch)) + qkv = qkv.view(bs, n_ctx, self.heads, -1) + q, k, v = torch.split(qkv, attn_ch, dim=-1) + + if self.use_flash: + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + out = F.scaled_dot_product_attention(q, k, v).permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) + else: + weight = torch.einsum( + "bthc,bshc->bhts", q * scale, k * scale + ) # More stable with f16 than dividing afterwards + wdtype = weight.dtype + weight = torch.softmax(weight.float(), dim=-1).type(wdtype) + out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1) + + return out + +class ResidualAttentionBlock(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + heads: int, + init_scale: float = 1.0, + qkv_bias: bool = True, + use_flash: bool = False, + use_checkpoint: bool = False + ): + super().__init__() + + self.use_checkpoint = use_checkpoint + + self.attn = MultiheadAttention( + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash + ) + self.ln_1 = nn.LayerNorm(width) + self.mlp = MLP(width=width, init_scale=init_scale) + self.ln_2 = nn.LayerNorm(width) + + def _forward(self, x: torch.Tensor): + x = x + self.attn(self.ln_1(x)) + x = x + self.mlp(self.ln_2(x)) + return x + + def forward(self, x: torch.Tensor): + return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint) + + +class MultiheadCrossAttention(nn.Module): + def __init__( + self, + *, + width: int, + heads: int, + init_scale: float, + qkv_bias: bool = True, + use_flash: bool = False, + n_data: Optional[int] = None, + data_width: Optional[int] = None, + ): + super().__init__() + self.n_data = n_data + self.width = width + self.heads = heads + self.data_width = width if data_width is None else data_width + self.c_q = nn.Linear(width, width, bias=qkv_bias) + self.c_kv = nn.Linear(self.data_width, width * 2, bias=qkv_bias) + self.c_proj = nn.Linear(width, width) + self.attention = QKVMultiheadCrossAttention( + heads=heads, n_data=n_data, use_flash=use_flash + ) + init_linear(self.c_q, init_scale) + init_linear(self.c_kv, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x, data): + x = self.c_q(x) + data = self.c_kv(data) + x = checkpoint(self.attention, (x, data), (), True) + x = self.c_proj(x) + return x + + +class QKVMultiheadCrossAttention(nn.Module): + def __init__(self, *, heads: int, use_flash: bool = False, n_data: Optional[int] = None): + + super().__init__() + self.heads = heads + self.n_data = n_data + self.use_flash = use_flash + + def forward(self, q, kv): + _, n_ctx, _ = q.shape + bs, n_data, width = kv.shape + attn_ch = width // self.heads // 2 + scale = 1 / math.sqrt(math.sqrt(attn_ch)) + q = q.view(bs, n_ctx, self.heads, -1) + kv = kv.view(bs, n_data, self.heads, -1) + k, v = torch.split(kv, attn_ch, dim=-1) + + if self.use_flash: + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + v = v.permute(0, 2, 1, 3) + out = F.scaled_dot_product_attention(q, k, v).permute(0, 2, 1, 3).reshape(bs, n_ctx, -1) + else: + weight = torch.einsum( + "bthc,bshc->bhts", q * scale, k * scale + ) # More stable with f16 than dividing afterwards + wdtype = weight.dtype + weight = torch.softmax(weight.float(), dim=-1).type(wdtype) + out = torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1) + + return out + + +class ResidualCrossAttentionBlock(nn.Module): + def __init__( + self, + *, + n_data: Optional[int] = None, + width: int, + heads: int, + data_width: Optional[int] = None, + init_scale: float = 0.25, + qkv_bias: bool = True, + use_flash: bool = False + ): + super().__init__() + + if data_width is None: + data_width = width + + self.attn = MultiheadCrossAttention( + n_data=n_data, + width=width, + heads=heads, + data_width=data_width, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash, + ) + self.ln_1 = nn.LayerNorm(width) + self.ln_2 = nn.LayerNorm(data_width) + self.mlp = MLP(width=width, init_scale=init_scale) + self.ln_3 = nn.LayerNorm(width) + + def forward(self, x: torch.Tensor, data: torch.Tensor): + x = x + self.attn(self.ln_1(x), self.ln_2(data)) + x = x + self.mlp(self.ln_3(x)) + return x \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/transformers/perceiver_1d.py b/Gen_3D_Modules/craftsman/models/transformers/perceiver_1d.py new file mode 100644 index 0000000..6705a2b --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/transformers/perceiver_1d.py @@ -0,0 +1,48 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +from craftsman.utils.typing import * +from craftsman.utils.checkpoint import checkpoint + +from .utils import init_linear +from .attention import ResidualAttentionBlock + + +class Perceiver(nn.Module): + def __init__( + self, + *, + n_ctx: int, + width: int, + layers: int, + heads: int, + init_scale: float = 0.25, + qkv_bias: bool = True, + use_flash: bool = False, + use_checkpoint: bool = False + ): + super().__init__() + self.n_ctx = n_ctx + self.width = width + self.layers = layers + self.resblocks = nn.ModuleList( + [ + ResidualAttentionBlock( + n_ctx=n_ctx, + width=width, + heads=heads, + init_scale=init_scale, + qkv_bias=qkv_bias, + use_flash=use_flash, + use_checkpoint=use_checkpoint + ) + for _ in range(layers) + ] + ) + + def forward(self, x: torch.Tensor): + for block in self.resblocks: + x = block(x) + return x \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/models/transformers/utils.py b/Gen_3D_Modules/craftsman/models/transformers/utils.py new file mode 100644 index 0000000..1ed8bbf --- /dev/null +++ b/Gen_3D_Modules/craftsman/models/transformers/utils.py @@ -0,0 +1,21 @@ +import torch.nn as nn + +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) + +class MLP(nn.Module): + def __init__(self, *, + width: int, + init_scale: float): + super().__init__() + self.width = width + self.c_fc = nn.Linear(width, width * 4) + self.c_proj = nn.Linear(width * 4, width) + self.gelu = nn.GELU() + init_linear(self.c_fc, init_scale) + init_linear(self.c_proj, init_scale) + + def forward(self, x): + return self.c_proj(self.gelu(self.c_fc(x))) diff --git a/Gen_3D_Modules/craftsman/systems/__init__.py b/Gen_3D_Modules/craftsman/systems/__init__.py new file mode 100644 index 0000000..c2b32ff --- /dev/null +++ b/Gen_3D_Modules/craftsman/systems/__init__.py @@ -0,0 +1,4 @@ +from . import ( + shape_autoencoder, + shape_diffusion, +) \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/systems/base.py b/Gen_3D_Modules/craftsman/systems/base.py new file mode 100644 index 0000000..c7d07f4 --- /dev/null +++ b/Gen_3D_Modules/craftsman/systems/base.py @@ -0,0 +1,209 @@ +import os +from dataclasses import dataclass, field + +import pytorch_lightning as pl +import torch.nn.functional as F + +from craftsman.utils.base import ( + Updateable, + update_end_if_possible, + update_if_possible, +) +from craftsman.utils.scheduler import parse_optimizer, parse_scheduler +from craftsman.utils.config import parse_structured +from craftsman.utils.misc import C, cleanup, get_device, load_module_weights +#from craftsman.utils.saving import SaverMixin +from craftsman.utils.typing import * + + +#class BaseSystem(pl.LightningModule, Updateable, SaverMixin): +class BaseSystem(pl.LightningModule, Updateable): + @dataclass + class Config: + loggers: dict = field(default_factory=dict) + loss: dict = field(default_factory=dict) + optimizer: dict = field(default_factory=dict) + scheduler: Optional[dict] = None + weights: Optional[str] = None + weights_ignore_modules: Optional[List[str]] = None + cleanup_after_validation_step: bool = False + cleanup_after_test_step: bool = False + + pretrained_model_path: Optional[str] = None + strict_load: bool = True + cfg: Config + + def __init__(self, cfg, resumed=False) -> None: + super().__init__() + self.cfg = parse_structured(self.Config, cfg) + self._save_dir: Optional[str] = None + self._resumed: bool = resumed + self._resumed_eval: bool = False + self._resumed_eval_status: dict = {"global_step": 0, "current_epoch": 0} + #if "loggers" in cfg: + # self.create_loggers(cfg.loggers) + + self.configure() + if self.cfg.weights is not None: + self.load_weights(self.cfg.weights, self.cfg.weights_ignore_modules) + self.post_configure() + + def load_weights(self, weights: str, ignore_modules: Optional[List[str]] = None): + state_dict, epoch, global_step = load_module_weights( + weights, ignore_modules=ignore_modules, 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 training_step(self, batch, batch_idx): + raise NotImplementedError + + def validation_step(self, batch, batch_idx): + raise NotImplementedError + + 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_validation_epoch_end(self): + raise NotImplementedError + + def test_step(self, batch, batch_idx): + raise NotImplementedError + + 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_test_epoch_end(self): + pass + + def predict_step(self, batch, batch_idx): + raise NotImplementedError + + 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 on_predict_epoch_end(self): + pass + + def preprocess_data(self, batch, stage): + pass + + """ + Implementing on_after_batch_transfer of DataModule does the same. + But on_after_batch_transfer does not support DP. + """ + + def on_train_batch_start(self, batch, batch_idx, unused=0): + self.preprocess_data(batch, "train") + self.dataset = self.trainer.train_dataloader.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_validation_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "validation") + self.dataset = self.trainer.val_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_test_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "test") + self.dataset = self.trainer.test_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def on_predict_batch_start(self, batch, batch_idx, dataloader_idx=0): + self.preprocess_data(batch, "predict") + self.dataset = self.trainer.predict_dataloaders.dataset + update_if_possible(self.dataset, self.true_current_epoch, self.true_global_step) + self.do_update_step(self.true_current_epoch, self.true_global_step) + + def 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) + """ + pass diff --git a/Gen_3D_Modules/craftsman/systems/shape_autoencoder.py b/Gen_3D_Modules/craftsman/systems/shape_autoencoder.py new file mode 100644 index 0000000..383a1ec --- /dev/null +++ b/Gen_3D_Modules/craftsman/systems/shape_autoencoder.py @@ -0,0 +1,137 @@ +from dataclasses import dataclass, field +import numpy as np +import torch +from skimage import measure +from einops import repeat, rearrange + +import craftsman +from craftsman.systems.base import BaseSystem +from craftsman.utils.ops import generate_dense_grid_points +from craftsman.utils.typing import * +from craftsman.utils.misc import get_rank + + +@craftsman.register("shape-autoencoder-system") +class ShapeAutoEncoderSystem(BaseSystem): + @dataclass + class Config(BaseSystem.Config): + shape_model_type: str = None + shape_model: dict = field(default_factory=dict) + + sample_posterior: bool = True + + cfg: Config + + def configure(self): + super().configure() + + self.shape_model = craftsman.find(self.cfg.shape_model_type)(self.cfg.shape_model) + + def forward(self, batch: Dict[str, Any]) -> Dict[str, Any]: + if "xyz" in batch: + if "sdf" in batch: + bs = batch["sdf"].shape[0] + rand_points = torch.cat([batch["xyz"].view(bs, -1, 3), batch["patch_xyz"].view(bs, -1, 3)], dim=1) + target = torch.cat([batch["sdf"].view(bs, -1, 1), batch["patch_sdf"].view(bs, -1, 1)], dim=1).squeeze(-1) + criteria = torch.nn.MSELoss() + elif "occupancy" in batch: + bs = batch["occupancy"].shape[0] + rand_points = torch.cat([batch["xyz"].view(bs, -1, 3), batch["patch_xyz"].view(bs, -1, 3)], dim=1) + target = torch.cat([batch["occupancy"].view(bs, -1, 1), batch["patch_occupancy"].view(bs, -1, 1)], dim=1).squeeze(-1) + criteria = torch.nn.BCEWithLogitsLoss() + else: + raise NotImplementedError + else: + rand_points = batch["rand_points"] + if "sdf" in batch: + target = batch["sdf"] + criteria = torch.nn.MSELoss() + elif "occupancies" in batch: + target = batch["occupancies"] + criteria = torch.nn.BCEWithLogitsLoss() + else: + raise NotImplementedError + + # forward pass + _, latents, posterior, logits = self.shape_model( + batch["surface"][..., :3 + self.cfg.shape_model.point_feats], + rand_points, + sample_posterior=self.cfg.sample_posterior + ) + + if self.cfg.sample_posterior: + loss_kl = posterior.kl() + loss_kl = torch.sum(loss_kl) / loss_kl.shape[0] + + return { + "loss_logits": criteria(logits, target).mean(), + "loss_kl": loss_kl, + "logits": logits, + "target": target, + "latents": latents, + } + else: + return { + "loss_logits": criteria(logits, target).mean(), + "latents": latents, + "logits": logits, + } + + def training_step(self, batch, batch_idx): + """ + Description: + + Args: + batch: + batch_idx: + Returns: + loss: + """ + out = self(batch) + + loss = 0. + for name, value in out.items(): + if name.startswith("loss_"): + self.log(f"train/{name}", value) + loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")]) + + for name, value in self.cfg.loss.items(): + self.log(f"train_params/{name}", self.C(value)) + + return {"loss": loss} + + @torch.no_grad() + def validation_step(self, batch, batch_idx): + self.eval() + out = self(batch) + # self.save_state_dict("latest-weights", self.state_dict()) + + mesh_v_f, has_surface = self.shape_model.extract_geometry(out["latents"]) + self.save_mesh( + f"it{self.true_global_step}/{batch['uid'][0]}.obj", + mesh_v_f[0][0], mesh_v_f[0][1] + ) + + threshold = 0 + outputs = out["logits"] + labels = out["target"] + pred = torch.zeros_like(outputs) + pred[outputs>=threshold] = 1 + + accuracy = (pred==labels).float().sum(dim=1) / labels.shape[1] + accuracy = accuracy.mean() + intersection = (pred * labels).sum(dim=1) + union = (pred + labels).gt(0).sum(dim=1) + iou = intersection * 1.0 / union + 1e-5 + iou = iou.mean() + + self.log("val/accuracy", accuracy) + self.log("val/iou", iou) + + torch.cuda.empty_cache() + + return {"val/loss": out["loss_logits"], "val/accuracy": accuracy, "val/iou": iou} + + + def on_validation_epoch_end(self): + pass \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/systems/shape_diffusion.py b/Gen_3D_Modules/craftsman/systems/shape_diffusion.py new file mode 100644 index 0000000..c327156 --- /dev/null +++ b/Gen_3D_Modules/craftsman/systems/shape_diffusion.py @@ -0,0 +1,396 @@ +from dataclasses import dataclass, field + +import numpy as np +import json +import copy +import torch +import torch.nn.functional as F +from skimage import measure +from einops import repeat +from tqdm import tqdm +from PIL import Image + +from diffusers import ( + DDPMScheduler, + DDIMScheduler, + UniPCMultistepScheduler, + KarrasVeScheduler, + DPMSolverMultistepScheduler +) + +import craftsman +from craftsman.systems.base import BaseSystem +from craftsman.utils.ops import generate_dense_grid_points +from craftsman.utils.misc import get_rank +from craftsman.utils.typing import * + +import comfy.utils + +def compute_snr(noise_scheduler, timesteps): + """ + Computes SNR as per + https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 + """ + alphas_cumprod = noise_scheduler.alphas_cumprod + sqrt_alphas_cumprod = alphas_cumprod**0.5 + sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 + + # Expand the tensors. + # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026 + sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float() + while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape): + sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None] + alpha = sqrt_alphas_cumprod.expand(timesteps.shape) + + sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float() + while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape): + sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None] + sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape) + + # Compute SNR. + snr = (alpha / sigma) ** 2 + return snr + +def ddim_sample(ddim_scheduler: DDIMScheduler, + diffusion_model: torch.nn.Module, + shape: Union[List[int], Tuple[int]], + cond: torch.FloatTensor, + steps: int, + eta: float = 0.0, + guidance_scale: float = 3.0, + do_classifier_free_guidance: bool = True, + generator: Optional[torch.Generator] = None, + device: torch.device = "cuda:0", + disable_prog: bool = True): + + assert steps > 0, f"{steps} must > 0." + + # init latents + bsz = cond.shape[0] + if do_classifier_free_guidance: + bsz = bsz // 2 + + latents = torch.randn( + (bsz, *shape), + generator=generator, + device=cond.device, + dtype=cond.dtype, + ) + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * ddim_scheduler.init_noise_sigma + # set timesteps + ddim_scheduler.set_timesteps(steps) + timesteps = ddim_scheduler.timesteps.to(device) + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, and between [0, 1] + extra_step_kwargs = { + # "eta": eta, + "generator": generator + } + + # reverse + comfy_pbar = comfy.utils.ProgressBar(steps) + for i, t in enumerate(tqdm(timesteps, disable=disable_prog, desc="DDIM Sampling:", leave=False)): + # expand the latents if we are doing classifier free guidance + latent_model_input = ( + torch.cat([latents] * 2) + if do_classifier_free_guidance + else latents + ) + # predict the noise residual + timestep_tensor = torch.tensor([t], dtype=torch.long, device=device) + timestep_tensor = timestep_tensor.expand(latent_model_input.shape[0]) + noise_pred = diffusion_model.forward(latent_model_input, timestep_tensor, cond) + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = ddim_scheduler.step( + noise_pred, t, latents, **extra_step_kwargs + ).prev_sample + + comfy_pbar.update_absolute(i + 1) + + yield latents, t + + +@craftsman.register("shape-diffusion-system") +class ShapeDiffusionSystem(BaseSystem): + @dataclass + class Config(BaseSystem.Config): + val_samples_json: str = None + z_scale_factor: float = 1.0 + guidance_scale: float = 7.5 + num_inference_steps: int = 50 + eta: float = 0.0 + snr_gamma: float = 5.0 + + # shape vae model + shape_model_type: str = None + shape_model: dict = field(default_factory=dict) + + # condition model + condition_model_type: str = None + condition_model: dict = field(default_factory=dict) + + # diffusion model + denoiser_model_type: str = None + denoiser_model: dict = field(default_factory=dict) + + # noise scheduler + noise_scheduler_type: str = None + noise_scheduler: dict = field(default_factory=dict) + + # denoise scheduler + denoise_scheduler_type: str = None + denoise_scheduler: dict = field(default_factory=dict) + + cfg: Config + + def configure(self): + super().configure() + + self.shape_model = craftsman.find(self.cfg.shape_model_type)(self.cfg.shape_model) + self.shape_model.eval() + self.shape_model.requires_grad_(False) + + self.condition = craftsman.find(self.cfg.condition_model_type)(self.cfg.condition_model) + + self.denoiser_model = craftsman.find(self.cfg.denoiser_model_type)(self.cfg.denoiser_model) + + self.noise_scheduler = craftsman.find(self.cfg.noise_scheduler_type)(**self.cfg.noise_scheduler) + self.denoise_scheduler = craftsman.find(self.cfg.denoise_scheduler_type)(**self.cfg.denoise_scheduler) + + self.z_scale_factor = self.cfg.z_scale_factor + + def forward(self, batch: Dict[str, Any]) -> Dict[str, Any]: + # encode shape latents + shape_embeds, kl_embed, posterior = self.shape_model.encode( + batch["surface"][..., :3 + self.cfg.shape_model.point_feats], + sample_posterior=True + ) + latents = kl_embed * self.z_scale_factor + + cond_latents = self.condition(batch) + cond_latents = cond_latents.to(latents).view(latents.shape[0], -1, cond_latents.shape[-1]) + + # Sample noise that we"ll add to the latents + # [batch_size, n_token, latent_dim] + noise = torch.randn_like(latents).to(latents) + bs = latents.shape[0] + # Sample a random timestep for each motion + timesteps = torch.randint( + 0, + self.noise_scheduler.config.num_train_timesteps, + (bs,), + device=latents.device, + ) + # import pdb; pdb.set_trace() + + timesteps = timesteps.long() + # Add noise to the latents according to the noise magnitude at each timestep + # x_t + noisy_z = self.noise_scheduler.add_noise(latents, noise, timesteps) + + # diffusion model forward + noise_pred = self.denoiser_model(noisy_z, timesteps, cond_latents) + + # compute loss + if self.noise_scheduler.config.prediction_type == "epsilon": + target = noise + elif self.noise_scheduler.config.prediction_type == "v_prediction": + target = self.noise_scheduler.get_velocity(latents, noise, timesteps) + else: + raise ValueError(f"Prediction Type: {self.noise_scheduler.prediction_type} not supported.") + if self.cfg.snr_gamma == 0: + if self.cfg.loss.loss_type == "l1": + loss = F.l1_loss(noise_pred, target, reduction="mean") + elif self.cfg.loss.loss_type in ["mse", "l2"]: + loss = F.mse_loss(noise_pred, target, reduction="mean") + else: + raise ValueError(f"Loss Type: {self.cfg.loss.loss_type} not supported.") + else: + # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556. + # Since we predict the noise instead of x_0, the original formulation is slightly changed. + # This is discussed in Section 4.2 of the same paper. + snr = compute_snr(self.noise_scheduler, timesteps) + mse_loss_weights = torch.stack([snr, self.cfg.snr_gamma * torch.ones_like(timesteps)], dim=1).min( + dim=1 + )[0] + if self.noise_scheduler.config.prediction_type == "epsilon": + mse_loss_weights = mse_loss_weights / snr + elif noise_scheduler.config.prediction_type == "v_prediction": + mse_loss_weights = mse_loss_weights / (snr + 1) + + if self.cfg.loss.loss_type == "l1": + loss = F.l1_loss(noise_pred, target, reduction="none") + elif self.cfg.loss.loss_type in ["mse", "l2"]: + loss = F.mse_loss(noise_pred, target, reduction="none") + else: + raise ValueError(f"Loss Type: {self.cfg.loss.loss_type} not supported.") + loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights + loss = loss.mean() + + return { + "loss_diffusion": loss, + "latents": latents, + "x_0": noisy_z, + "noise": noise, + "noise_pred": noise_pred, + "timesteps": timesteps, + } + + def training_step(self, batch, batch_idx): + out = self(batch) + + loss = 0. + for name, value in out.items(): + if name.startswith("loss_"): + self.log(f"train/{name}", value) + loss += value * self.C(self.cfg.loss[name.replace("loss_", "lambda_")]) + + for name, value in self.cfg.loss.items(): + if name.startswith("lambda_"): + self.log(f"train_params/{name}", self.C(value)) + + return {"loss": loss} + + @torch.no_grad() + def validation_step(self, batch, batch_idx): + self.eval() + + if get_rank() == 0: + sample_inputs = json.loads(open(self.cfg.val_samples_json).read()) # condition + sample_inputs_ = copy.deepcopy(sample_inputs) + sample_outputs = self.sample(sample_inputs) # list + for i, sample_output in enumerate(sample_outputs): + mesh_v_f, has_surface = self.shape_model.extract_geometry(sample_output, octree_depth=7) + for j in range(len(mesh_v_f)): + if "text" in sample_inputs_ and "image" in sample_inputs_: + name = sample_inputs_["image"][j].split("/")[-1].replace(".png", "") + elif "text" in sample_inputs_ and "mvimage" in sample_inputs_: + name = sample_inputs_["mvimages"][j][0].split("/")[-2].replace(".png", "") + elif "text" in sample_inputs_: + name = sample_inputs_["text"][j].replace(" ", "_") + elif "image" in sample_inputs_: + name = sample_inputs_["image"][j].split("/")[-1].replace(".png", "") + elif "mvimages" in sample_inputs_: + name = sample_inputs_["mvimages"][j][0].split("/")[-2].replace(".png", "") + self.save_mesh( + f"it{self.true_global_step}/{name}_{i}.obj", + mesh_v_f[j][0], mesh_v_f[j][1] + ) + + out = self(batch) + if self.global_step == 0: + latents = self.shape_model.decode(out["latents"]) + mesh_v_f, has_surface = self.shape_model.extract_geometry(latents) + self.save_mesh( + f"it{self.true_global_step}/{batch['uid'][0]}_{batch['sel_idx'][0] if 'sel_idx' in batch.keys() else 0}.obj", + mesh_v_f[0][0], mesh_v_f[0][1] + ) + # exit() + torch.cuda.empty_cache() + + return {"val/loss": out["loss_diffusion"]} + + @torch.no_grad() + def sample(self, + sample_inputs: Dict[str, Union[torch.FloatTensor, List[str]]], + sample_times: int = 1, + steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + eta: float = 0.0, + return_intermediates: bool = False, + camera_embeds: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + **kwargs): + + if steps is None: + steps = self.cfg.num_inference_steps + if guidance_scale is None: + guidance_scale = self.cfg.guidance_scale + do_classifier_free_guidance = guidance_scale > 0 + + # conditional encode + if "image" in sample_inputs: + sample_inputs["image"] = [Image.open(img) for img in sample_inputs["image"]] + cond = self.condition.encode_image(sample_inputs["image"]) + if do_classifier_free_guidance: + un_cond = self.condition.empty_image_embeds.repeat(len(sample_inputs["image"]), 1, 1).to(cond) + cond = torch.cat([un_cond, cond], dim=0) + elif "mvimages" in sample_inputs: # by default 4 views + bs = len(sample_inputs["mvimages"]) + cond = [] + for image in sample_inputs["mvimages"]: + if isinstance(image, list) and isinstance(image[0], str): + sample_inputs["image"] = [Image.open(img) for img in image] # List[PIL] + else: + sample_inputs["image"] = image + cond += [self.condition.encode_image(sample_inputs["image"])] + cond = torch.stack(cond, dim=0)# tensor shape 为[len(sample_inputs["mvimages"], 4*(num_latents+1), context_dim] + if do_classifier_free_guidance: + un_cond = self.condition.empty_image_embeds.unsqueeze(0).repeat(len(sample_inputs["mvimages"]), cond.shape[1] // self.condition.cfg.n_views, 1, 1).to(cond) # shape 为[len(sample_inputs["mvimages"], 4*(num_latents+1), context_dim] + cond = torch.cat([un_cond, cond], dim=0).view(bs * 2, -1, cond[0].shape[-1]) + else: + raise NotImplementedError("Only text, image or mvimages condition is supported.") + + outputs = [] + latents = None + + if seed != None: + generator = torch.Generator(device="cuda").manual_seed(seed) + else: + generator = None + + if not return_intermediates: + for _ in range(sample_times): + sample_loop = ddim_sample( + self.denoise_scheduler, + self.denoiser_model.eval(), + shape=self.shape_model.latent_shape, + cond=cond, + steps=steps, + guidance_scale=guidance_scale, + do_classifier_free_guidance=do_classifier_free_guidance, + device=self.device, + eta=eta, + disable_prog=False, + generator= generator + ) + for sample, t in sample_loop: + latents = sample + outputs.append(self.shape_model.decode(latents / self.z_scale_factor, **kwargs)) + else: + sample_loop = ddim_sample( + self.denoise_scheduler, + self.denoiser_model.eval(), + shape=self.shape_model.latent_shape, + cond=cond, + steps=steps, + guidance_scale=guidance_scale, + do_classifier_free_guidance=do_classifier_free_guidance, + device=self.device, + eta=eta, + disable_prog=False, + generator= generator + ) + + iter_size = steps // sample_times + i = 0 + for sample, t in sample_loop: + latents = sample + if i % iter_size == 0 or i == steps - 1: + outputs.append(self.shape_model.decode(latents / self.z_scale_factor, **kwargs)) + i += 1 + + return outputs + + + def on_validation_epoch_end(self): + pass diff --git a/Gen_3D_Modules/craftsman/utils/__init__.py b/Gen_3D_Modules/craftsman/utils/__init__.py new file mode 100644 index 0000000..0e44449 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/__init__.py @@ -0,0 +1 @@ +from . import base diff --git a/Gen_3D_Modules/craftsman/utils/base.py b/Gen_3D_Modules/craftsman/utils/base.py new file mode 100644 index 0000000..2fa680e --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/base.py @@ -0,0 +1,118 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from craftsman.utils.config import parse_structured +from craftsman.utils.misc import get_device, load_module_weights +from craftsman.utils.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.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 + # dummy tensor to indicate model state + self._dummy: Float[Tensor, "..."] + self.register_buffer("_dummy", torch.zeros(0).float(), persistent=False) + + def configure(self, *args, **kwargs) -> None: + pass diff --git a/Gen_3D_Modules/craftsman/utils/callbacks.py b/Gen_3D_Modules/craftsman/utils/callbacks.py new file mode 100644 index 0000000..c0382c5 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/callbacks.py @@ -0,0 +1,156 @@ +import os +import shutil +import subprocess + +import pytorch_lightning + +from craftsman.utils.config import dump_config +from craftsman.utils.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/craftsman/utils/checkpoint.py b/Gen_3D_Modules/craftsman/utils/checkpoint.py new file mode 100644 index 0000000..363e2bf --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/checkpoint.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +""" +Adapted from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/nn.py#L124 +""" + +import torch +from craftsman.utils.typing import * + +def checkpoint( + func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]], + inputs: Sequence[torch.Tensor], + params: Iterable[torch.Tensor], + flag: bool, + use_deepspeed: bool = False +): + """ + Evaluate a function without caching intermediate activations, allowing for + reduced memory at the expense of extra compute in the backward pass. + :param func: the function to evaluate. + :param inputs: the argument sequence to pass to `func`. + :param params: a sequence of parameters `func` depends on but does not + explicitly take as arguments. + :param flag: if False, disable gradient checkpointing. + :param use_deepspeed: if True, use deepspeed + """ + if flag: + if use_deepspeed: + import deepspeed + return deepspeed.checkpointing.checkpoint(func, *inputs) + + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + @torch.cuda.amp.custom_fwd + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + @torch.cuda.amp.custom_bwd + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/utils/config.py b/Gen_3D_Modules/craftsman/utils/config.py new file mode 100644 index 0000000..1f2f510 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/config.py @@ -0,0 +1,128 @@ +import os +from dataclasses import dataclass, field +from datetime import datetime + +from omegaconf import OmegaConf + +import craftsman +from craftsman.utils.typing import * + +# ============ Register OmegaConf Recolvers ============= # +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: str(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("cmaxgt0", lambda s: C_max(s) > 0, replace=True) +OmegaConf.register_new_resolver("not", lambda s: not s, replace=True) +OmegaConf.register_new_resolver( + "cmaxgt0orcmaxgt0", lambda a, b: C_max(a) > 0 or C_max(b) > 0, replace=True +) +# ======================================================= # + + +def C_max(value: Any) -> 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) >= 6: + max_value = value[2] + for i in range(4, len(value), 2): + max_value = max(max_value, value[i]) + value = [value[0], value[1], max_value, value[3]] + if len(value) == 3: + value = [0] + value + assert len(value) == 4 + start_step, start_value, end_value, end_step = value + value = max(start_value, end_value) + return value + + +@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_type: str = "" + data: dict = field(default_factory=dict) + + system_type: 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 __post_init__(self): + if not self.tag and not self.use_timestamp: + raise ValueError("Either tag is specified or use_timestamp is True.") + self.trial_name = self.tag + # if resume from an existing config, self.timestamp should not be None + if self.timestamp is None: + self.timestamp = "" + if self.use_timestamp: + if self.n_gpus > 1: + craftsman.warn( + "Timestamp is disabled when using multiple GPUs, please make sure you have a unique tag." + ) + else: + self.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S") + self.trial_name += self.timestamp + self.exp_dir = os.path.join(self.exp_root_dir, self.name) + self.trial_dir = os.path.join(self.exp_dir, self.trial_name) + os.makedirs(self.trial_dir, exist_ok=True) + + +def load_config(*yamls: str, cli_args: list = [], from_string=False, **kwargs) -> Any: + if from_string: + yaml_confs = [OmegaConf.create(s) for s in yamls] + else: + yaml_confs = [OmegaConf.load(f) for f in yamls] + cli_conf = OmegaConf.from_cli(cli_args) + cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs) + OmegaConf.resolve(cfg) + assert isinstance(cfg, DictConfig) + scfg = parse_structured(ExperimentConfig, cfg) + 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.structured(fields(**cfg)) + return scfg \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/utils/misc.py b/Gen_3D_Modules/craftsman/utils/misc.py new file mode 100644 index 0000000..732835e --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/misc.py @@ -0,0 +1,168 @@ +import gc +import os +import re + +import torch +import torch.distributed as dist +from packaging import version + +from craftsman.utils.config import config_to_primitive +from craftsman.utils.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_world_size(): + world_size_keys = ("WORLD_SIZE", "SLURM_NTASKS", "JSM_NAMESPACE_SIZE") + for key in world_size_keys: + world_size = os.environ.get(key) + if world_size is not None: + return int(world_size) + return 1 + +def get_device(): + return torch.device(f"cuda:{get_rank()}") + + +def load_module_weights( + path, module_name=None, ignore_modules=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"] + 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() + tcnn.free_temporary_memory() + + +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) + + +def all_gather_batch(tensors): + """ + Performs all_gather operation on the provided tensors. + """ + # Queue the gathered tensors + world_size = get_world_size() + # There is no need for reduction in the single-proc case + if world_size == 1: + if isinstance(tensors, list): + return tensors + return tensors + if not isinstance(tensors, list): + is_list = False + tensors = [tensors] + else: + is_list = True + output_tensor = [] + tensor_list = [] + for tensor in tensors: + tensor_all = [torch.ones_like(tensor) for _ in range(world_size)] + dist.all_gather( + tensor_all, + tensor, + async_op=False # performance opt + ) + + tensor_list.append(tensor_all) + + for tensor_all in tensor_list: + output_tensor.append(torch.cat(tensor_all, dim=0)) + if not is_list: + return output_tensor[0] + return output_tensor \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/utils/ops.py b/Gen_3D_Modules/craftsman/utils/ops.py new file mode 100644 index 0000000..2a46206 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/ops.py @@ -0,0 +1,167 @@ +from collections import defaultdict + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from craftsman.utils.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 + + +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 randn_tensor( + shape: Union[Tuple, List], + generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None, + device: Optional["torch.device"] = None, + dtype: Optional["torch.dtype"] = None, + layout: Optional["torch.layout"] = None, +): + """A helper function to create random tensors on the desired `device` with the desired `dtype`. When + passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor + is always created on the CPU. + """ + # device on which tensor is created defaults to device + rand_device = device + batch_size = shape[0] + + layout = layout or torch.strided + device = device or torch.device("cpu") + + if generator is not None: + gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type + if gen_device_type != device.type and gen_device_type == "cpu": + rand_device = "cpu" + if device != "mps": + logger.info( + f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." + f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" + f" slighly speed up this function by passing a generator that was created on the {device} device." + ) + elif gen_device_type != device.type and gen_device_type == "cuda": + raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.") + + # make sure generator list of length 1 is treated like a non-list + if isinstance(generator, list) and len(generator) == 1: + generator = generator[0] + + if isinstance(generator, list): + shape = (1,) + shape[1:] + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout) + 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, layout=layout).to(device) + + return latents + + +def generate_dense_grid_points( + bbox_min: np.ndarray, + bbox_max: np.ndarray, + grids_resolution: int, + indexing: str = "ij" +): + length = bbox_max - bbox_min + num_cells = grids_resolution + x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32) + y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32) + z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32) + [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing) + xyz = np.stack((xs, ys, zs), axis=-1) + xyz = xyz.reshape(-1, 3) + grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1] + + return xyz, grid_size, length \ No newline at end of file diff --git a/Gen_3D_Modules/craftsman/utils/saving.py b/Gen_3D_Modules/craftsman/utils/saving.py new file mode 100644 index 0000000..5c714c5 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/saving.py @@ -0,0 +1,653 @@ +import json +import os +import re +import shutil + +import cv2 +import imageio +import matplotlib.pyplot as plt +import numpy as np +import torch +import trimesh +from matplotlib import cm +from matplotlib.colors import LinearSegmentedColormap +from PIL import Image, ImageDraw +from pytorch_lightning.loggers import WandbLogger + +from craftsman.models.geometry.utils import Mesh +from craftsman.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): + 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 + + def create_loggers(self, cfg_loggers: DictConfig) -> None: + if "wandb" in cfg_loggers.keys() and cfg_loggers.wandb.enable: + self._wandb_logger = WandbLogger( + project=cfg_loggers.wandb.project, name=cfg_loggers.wandb.name + ) + + def get_loggers(self) -> List: + if self._wandb_logger: + return [self._wandb_logger] + else: + return [] + + 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: + wandb.log( + { + name: wandb.Image(self.get_save_path(filename)), + "trainer/global_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: + wandb.log( + { + name: wandb.Image(self.get_save_path(filename)), + "trainer/global_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]) + w = max([col.shape[1] for col in cols]) + elif align == "min": + h = min([col.shape[0] for col in cols]) + w = min([col.shape[1] for col in cols]) + elif isinstance(align, int): + h = align + w = align + elif ( + isinstance(align, tuple) + and isinstance(align[0], int) + and isinstance(align[1], int) + ): + h, w = align + else: + raise ValueError( + f"Unsupported image grid align: {align}, should be min, max, int or (int, int)" + ) + + for i in range(len(cols)): + if cols[i].shape[0] != h or cols[i].shape[1] != w: + cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_LINEAR) + 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: + wandb.log({name: wandb.Image(save_path), "trainer/global_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: + wandb.log( + { + name: wandb.Video(save_path, format="mp4"), + "trainer/global_step": step, + } + ) + return save_path + + def save_mesh(self, filename, v_pos, t_pos_idx, v_tex=None, t_tex_idx=None) -> str: + 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 = False, + 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_file(self, filename, src_path) -> str: + save_path = self.get_save_path(filename) + shutil.copyfile(src_path, save_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/craftsman/utils/scheduler.py b/Gen_3D_Modules/craftsman/utils/scheduler.py new file mode 100644 index 0000000..d0da7c4 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/scheduler.py @@ -0,0 +1,104 @@ +import sys +import warnings +from bisect import bisect_right + +import torch +import torch.nn as nn +from torch.optim import lr_scheduler + +import craftsman + + +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() + ] + craftsman.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 ["Adan"]: + from craftsman.systems import optimizers + + optim = getattr(optimizers, config.name)(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/craftsman/utils/typing.py b/Gen_3D_Modules/craftsman/utils/typing.py new file mode 100644 index 0000000..21b1bb2 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/typing.py @@ -0,0 +1,41 @@ +""" +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, + Sequence, +) + +# 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/craftsman/utils/visualizers/__init__.py b/Gen_3D_Modules/craftsman/utils/visualizers/__init__.py new file mode 100644 index 0000000..40a96af --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/visualizers/__init__.py @@ -0,0 +1 @@ +# -*- coding: utf-8 -*- diff --git a/Gen_3D_Modules/craftsman/utils/visualizers/color_util.py b/Gen_3D_Modules/craftsman/utils/visualizers/color_util.py new file mode 100644 index 0000000..7983243 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/visualizers/color_util.py @@ -0,0 +1,43 @@ +import numpy as np +import matplotlib.pyplot as plt + + +# Helper functions +def get_colors(inp, colormap="viridis", normalize=True, vmin=None, vmax=None): + colormap = plt.cm.get_cmap(colormap) + if normalize: + vmin = np.min(inp) + vmax = np.max(inp) + + norm = plt.Normalize(vmin, vmax) + return colormap(norm(inp))[:, :3] + + +def gen_checkers(n_checkers_x, n_checkers_y, width=256, height=256): + # tex dims need to be power of two. + array = np.ones((width, height, 3), dtype='float32') + + # width in texels of each checker + checker_w = width / n_checkers_x + checker_h = height / n_checkers_y + + for y in range(height): + for x in range(width): + color_key = int(x / checker_w) + int(y / checker_h) + if color_key % 2 == 0: + array[x, y, :] = [1., 0.874, 0.0] + else: + array[x, y, :] = [0., 0., 0.] + return array + + +def gen_circle(width=256, height=256): + xx, yy = np.mgrid[:width, :height] + circle = (xx - width / 2 + 0.5) ** 2 + (yy - height / 2 + 0.5) ** 2 + array = np.ones((width, height, 4), dtype='float32') + array[:, :, 0] = (circle <= width) + array[:, :, 1] = (circle <= width) + array[:, :, 2] = (circle <= width) + array[:, :, 3] = circle <= width + return array + diff --git a/Gen_3D_Modules/craftsman/utils/visualizers/html_util.py b/Gen_3D_Modules/craftsman/utils/visualizers/html_util.py new file mode 100644 index 0000000..f90fe6c --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/visualizers/html_util.py @@ -0,0 +1,49 @@ +# -*- coding: utf-8 -*- +import io +import base64 +import numpy as np +from PIL import Image + + +def to_html_frame(content): + + html_frame = f""" + + + {content} + + + """ + + return html_frame + + +def to_single_row_table(caption: str, content: str): + + table_html = f""" + + + + + +
{caption}
{content}
+ """ + + return table_html + + +def to_image_embed_tag(image: np.ndarray): + + # Convert np.ndarray to bytes + img = Image.fromarray(image) + raw_bytes = io.BytesIO() + img.save(raw_bytes, "PNG") + + # Encode bytes to base64 + image_base64 = base64.b64encode(raw_bytes.getvalue()).decode("utf-8") + + image_tag = f""" + Embedded Image + """ + + return image_tag diff --git a/Gen_3D_Modules/craftsman/utils/visualizers/pythreejs_viewer.py b/Gen_3D_Modules/craftsman/utils/visualizers/pythreejs_viewer.py new file mode 100644 index 0000000..70abbb7 --- /dev/null +++ b/Gen_3D_Modules/craftsman/utils/visualizers/pythreejs_viewer.py @@ -0,0 +1,537 @@ +import numpy as np +from ipywidgets import embed +import pythreejs as p3s +import uuid + +from .color_util import get_colors, gen_circle, gen_checkers + + +EMBED_URL = "https://cdn.jsdelivr.net/npm/@jupyter-widgets/html-manager@1.0.1/dist/embed-amd.js" + + +class PyThreeJSViewer(object): + + def __init__(self, settings, render_mode="WEBSITE"): + self.render_mode = render_mode + self.__update_settings(settings) + self._light = p3s.DirectionalLight(color='white', position=[0, 0, 1], intensity=0.6) + self._light2 = p3s.AmbientLight(intensity=0.5) + self._cam = p3s.PerspectiveCamera(position=[0, 0, 1], lookAt=[0, 0, 0], fov=self.__s["fov"], + aspect=self.__s["width"] / self.__s["height"], children=[self._light]) + self._orbit = p3s.OrbitControls(controlling=self._cam) + self._scene = p3s.Scene(children=[self._cam, self._light2], background=self.__s["background"]) # "#4c4c80" + self._renderer = p3s.Renderer(camera=self._cam, scene=self._scene, controls=[self._orbit], + width=self.__s["width"], height=self.__s["height"], + antialias=self.__s["antialias"]) + + self.__objects = {} + self.__cnt = 0 + + def jupyter_mode(self): + self.render_mode = "JUPYTER" + + def offline(self): + self.render_mode = "OFFLINE" + + def website(self): + self.render_mode = "WEBSITE" + + def __get_shading(self, shading): + shad = {"flat": True, "wireframe": False, "wire_width": 0.03, "wire_color": "black", + "side": 'DoubleSide', "colormap": "viridis", "normalize": [None, None], + "bbox": False, "roughness": 0.5, "metalness": 0.25, "reflectivity": 1.0, + "line_width": 1.0, "line_color": "black", + "point_color": "red", "point_size": 0.01, "point_shape": "circle", + "text_color": "red" + } + for k in shading: + shad[k] = shading[k] + return shad + + def __update_settings(self, settings={}): + sett = {"width": 600, "height": 600, "antialias": True, "scale": 1.5, "background": "#ffffff", + "fov": 30} + for k in settings: + sett[k] = settings[k] + self.__s = sett + + def __add_object(self, obj, parent=None): + if not parent: # Object is added to global scene and objects dict + self.__objects[self.__cnt] = obj + self.__cnt += 1 + self._scene.add(obj["mesh"]) + else: # Object is added to parent object and NOT to objects dict + parent.add(obj["mesh"]) + + self.__update_view() + + if self.render_mode == "JUPYTER": + return self.__cnt - 1 + elif self.render_mode == "WEBSITE": + return self + + def __add_line_geometry(self, lines, shading, obj=None): + lines = lines.astype("float32", copy=False) + mi = np.min(lines, axis=0) + ma = np.max(lines, axis=0) + + geometry = p3s.LineSegmentsGeometry(positions=lines.reshape((-1, 2, 3))) + material = p3s.LineMaterial(linewidth=shading["line_width"], color=shading["line_color"]) + # , vertexColors='VertexColors'), + lines = p3s.LineSegments2(geometry=geometry, material=material) # type='LinePieces') + line_obj = {"geometry": geometry, "mesh": lines, "material": material, + "max": ma, "min": mi, "type": "Lines", "wireframe": None} + + if obj: + return self.__add_object(line_obj, obj), line_obj + else: + return self.__add_object(line_obj) + + def __update_view(self): + if len(self.__objects) == 0: + return + ma = np.zeros((len(self.__objects), 3)) + mi = np.zeros((len(self.__objects), 3)) + for r, obj in enumerate(self.__objects): + ma[r] = self.__objects[obj]["max"] + mi[r] = self.__objects[obj]["min"] + ma = np.max(ma, axis=0) + mi = np.min(mi, axis=0) + diag = np.linalg.norm(ma - mi) + mean = ((ma - mi) / 2 + mi).tolist() + scale = self.__s["scale"] * (diag) + self._orbit.target = mean + self._cam.lookAt(mean) + self._cam.position = [mean[0], mean[1], mean[2] + scale] + self._light.position = [mean[0], mean[1], mean[2] + scale] + + self._orbit.exec_three_obj_method('update') + self._cam.exec_three_obj_method('updateProjectionMatrix') + + def __get_bbox(self, v): + m = np.min(v, axis=0) + M = np.max(v, axis=0) + + # Corners of the bounding box + v_box = np.array([[m[0], m[1], m[2]], [M[0], m[1], m[2]], [M[0], M[1], m[2]], [m[0], M[1], m[2]], + [m[0], m[1], M[2]], [M[0], m[1], M[2]], [M[0], M[1], M[2]], [m[0], M[1], M[2]]]) + + f_box = np.array([[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6], [6, 7], [7, 4], + [0, 4], [1, 5], [2, 6], [7, 3]], dtype=np.uint32) + return v_box, f_box + + def __get_colors(self, v, f, c, sh): + coloring = "VertexColors" + if type(c) == np.ndarray and c.size == 3: # Single color + colors = np.ones_like(v) + colors[:, 0] = c[0] + colors[:, 1] = c[1] + colors[:, 2] = c[2] + # print("Single colors") + elif type(c) == np.ndarray and len(c.shape) == 2 and c.shape[1] == 3: # Color values for + if c.shape[0] == f.shape[0]: # faces + colors = np.hstack([c, c, c]).reshape((-1, 3)) + coloring = "FaceColors" + # print("Face color values") + elif c.shape[0] == v.shape[0]: # vertices + colors = c + # print("Vertex color values") + else: # Wrong size, fallback + print("Invalid color array given! Supported are numpy arrays.", type(c)) + colors = np.ones_like(v) + colors[:, 0] = 1.0 + colors[:, 1] = 0.874 + colors[:, 2] = 0.0 + elif type(c) == np.ndarray and c.size == f.shape[0]: # Function values for faces + normalize = sh["normalize"][0] != None and sh["normalize"][1] != None + cc = get_colors(c, sh["colormap"], normalize=normalize, + vmin=sh["normalize"][0], vmax=sh["normalize"][1]) + # print(cc.shape) + colors = np.hstack([cc, cc, cc]).reshape((-1, 3)) + coloring = "FaceColors" + # print("Face function values") + elif type(c) == np.ndarray and c.size == v.shape[0]: # Function values for vertices + normalize = sh["normalize"][0] != None and sh["normalize"][1] != None + colors = get_colors(c, sh["colormap"], normalize=normalize, + vmin=sh["normalize"][0], vmax=sh["normalize"][1]) + # print("Vertex function values") + + else: + colors = np.ones_like(v) + # colors[:, 0] = 1.0 + # colors[:, 1] = 0.874 + # colors[:, 2] = 0.0 + colors[:, 0] = 1 + colors[:, 1] = 1 + colors[:, 2] = 1 + + # No color + if c is not None: + print("Invalid color array given! Supported are numpy arrays.", type(c)) + + return colors, coloring + + def __get_point_colors(self, v, c, sh): + v_color = True + if c is None: # No color given, use global color + # conv = mpl.colors.ColorConverter() + colors = sh["point_color"] # np.array(conv.to_rgb(sh["point_color"])) + v_color = False + elif isinstance(c, str): # No color given, use global color + # conv = mpl.colors.ColorConverter() + colors = c # np.array(conv.to_rgb(c)) + v_color = False + elif type(c) == np.ndarray and len(c.shape) == 2 and c.shape[0] == v.shape[0] and c.shape[1] == 3: + # Point color + colors = c.astype("float32", copy=False) + + elif isinstance(c, np.ndarray) and len(c.shape) == 2 and c.shape[0] == v.shape[0] and c.shape[1] != 3: + # Function values for vertices, but the colors are features + c_norm = np.linalg.norm(c, ord=2, axis=-1) + normalize = sh["normalize"][0] != None and sh["normalize"][1] != None + colors = get_colors(c_norm, sh["colormap"], normalize=normalize, + vmin=sh["normalize"][0], vmax=sh["normalize"][1]) + colors = colors.astype("float32", copy=False) + + elif type(c) == np.ndarray and c.size == v.shape[0]: # Function color + normalize = sh["normalize"][0] != None and sh["normalize"][1] != None + colors = get_colors(c, sh["colormap"], normalize=normalize, + vmin=sh["normalize"][0], vmax=sh["normalize"][1]) + colors = colors.astype("float32", copy=False) + # print("Vertex function values") + + else: + print("Invalid color array given! Supported are numpy arrays.", type(c)) + colors = sh["point_color"] + v_color = False + + return colors, v_color + + def add_mesh(self, v, f, c=None, uv=None, n=None, shading={}, texture_data=None, **kwargs): + shading.update(kwargs) + sh = self.__get_shading(shading) + mesh_obj = {} + + # it is a tet + if v.shape[1] == 3 and f.shape[1] == 4: + f_tmp = np.ndarray([f.shape[0] * 4, 3], dtype=f.dtype) + for i in range(f.shape[0]): + f_tmp[i * 4 + 0] = np.array([f[i][1], f[i][0], f[i][2]]) + f_tmp[i * 4 + 1] = np.array([f[i][0], f[i][1], f[i][3]]) + f_tmp[i * 4 + 2] = np.array([f[i][1], f[i][2], f[i][3]]) + f_tmp[i * 4 + 3] = np.array([f[i][2], f[i][0], f[i][3]]) + f = f_tmp + + if v.shape[1] == 2: + v = np.append(v, np.zeros([v.shape[0], 1]), 1) + + # Type adjustment vertices + v = v.astype("float32", copy=False) + + # Color setup + colors, coloring = self.__get_colors(v, f, c, sh) + + # Type adjustment faces and colors + c = colors.astype("float32", copy=False) + + # Material and geometry setup + ba_dict = {"color": p3s.BufferAttribute(c)} + if coloring == "FaceColors": + verts = np.zeros((f.shape[0] * 3, 3), dtype="float32") + for ii in range(f.shape[0]): + # print(ii*3, f[ii]) + verts[ii * 3] = v[f[ii, 0]] + verts[ii * 3 + 1] = v[f[ii, 1]] + verts[ii * 3 + 2] = v[f[ii, 2]] + v = verts + else: + f = f.astype("uint32", copy=False).ravel() + ba_dict["index"] = p3s.BufferAttribute(f, normalized=False) + + ba_dict["position"] = p3s.BufferAttribute(v, normalized=False) + + if uv is not None: + uv = (uv - np.min(uv)) / (np.max(uv) - np.min(uv)) + if texture_data is None: + texture_data = gen_checkers(20, 20) + tex = p3s.DataTexture(data=texture_data, format="RGBFormat", type="FloatType") + material = p3s.MeshStandardMaterial(map=tex, reflectivity=sh["reflectivity"], side=sh["side"], + roughness=sh["roughness"], metalness=sh["metalness"], + flatShading=sh["flat"], + polygonOffset=True, polygonOffsetFactor=1, polygonOffsetUnits=5) + ba_dict["uv"] = p3s.BufferAttribute(uv.astype("float32", copy=False)) + else: + material = p3s.MeshStandardMaterial(vertexColors=coloring, reflectivity=sh["reflectivity"], + side=sh["side"], roughness=sh["roughness"], metalness=sh["metalness"], + flatShading=sh["flat"], + polygonOffset=True, polygonOffsetFactor=1, polygonOffsetUnits=5) + + if type(n) != type(None) and coloring == "VertexColors": # TODO: properly handle normals for FaceColors as well + ba_dict["normal"] = p3s.BufferAttribute(n.astype("float32", copy=False), normalized=True) + + geometry = p3s.BufferGeometry(attributes=ba_dict) + + if coloring == "VertexColors" and type(n) == type(None): + geometry.exec_three_obj_method('computeVertexNormals') + elif coloring == "FaceColors" and type(n) == type(None): + geometry.exec_three_obj_method('computeFaceNormals') + + # Mesh setup + mesh = p3s.Mesh(geometry=geometry, material=material) + + # Wireframe setup + mesh_obj["wireframe"] = None + if sh["wireframe"]: + wf_geometry = p3s.WireframeGeometry(mesh.geometry) # WireframeGeometry + wf_material = p3s.LineBasicMaterial(color=sh["wire_color"], linewidth=sh["wire_width"]) + wireframe = p3s.LineSegments(wf_geometry, wf_material) + mesh.add(wireframe) + mesh_obj["wireframe"] = wireframe + + # Bounding box setup + if sh["bbox"]: + v_box, f_box = self.__get_bbox(v) + _, bbox = self.add_edges(v_box, f_box, sh, mesh) + mesh_obj["bbox"] = [bbox, v_box, f_box] + + # Object setup + mesh_obj["max"] = np.max(v, axis=0) + mesh_obj["min"] = np.min(v, axis=0) + mesh_obj["geometry"] = geometry + mesh_obj["mesh"] = mesh + mesh_obj["material"] = material + mesh_obj["type"] = "Mesh" + mesh_obj["shading"] = sh + mesh_obj["coloring"] = coloring + mesh_obj["arrays"] = [v, f, c] # TODO replays with proper storage or remove if not needed + + return self.__add_object(mesh_obj) + + def add_lines(self, beginning, ending, shading={}, obj=None, **kwargs): + shading.update(kwargs) + if len(beginning.shape) == 1: + if len(beginning) == 2: + beginning = np.array([[beginning[0], beginning[1], 0]]) + else: + if beginning.shape[1] == 2: + beginning = np.append( + beginning, np.zeros([beginning.shape[0], 1]), 1) + if len(ending.shape) == 1: + if len(ending) == 2: + ending = np.array([[ending[0], ending[1], 0]]) + else: + if ending.shape[1] == 2: + ending = np.append( + ending, np.zeros([ending.shape[0], 1]), 1) + + sh = self.__get_shading(shading) + lines = np.hstack([beginning, ending]) + lines = lines.reshape((-1, 3)) + return self.__add_line_geometry(lines, sh, obj) + + def add_edges(self, vertices, edges, shading={}, obj=None, **kwargs): + shading.update(kwargs) + if vertices.shape[1] == 2: + vertices = np.append( + vertices, np.zeros([vertices.shape[0], 1]), 1) + sh = self.__get_shading(shading) + lines = np.zeros((edges.size, 3)) + cnt = 0 + for e in edges: + lines[cnt, :] = vertices[e[0]] + lines[cnt + 1, :] = vertices[e[1]] + cnt += 2 + return self.__add_line_geometry(lines, sh, obj) + + def add_points(self, points, c=None, shading={}, obj=None, **kwargs): + shading.update(kwargs) + if len(points.shape) == 1: + if len(points) == 2: + points = np.array([[points[0], points[1], 0]]) + else: + if points.shape[1] == 2: + points = np.append( + points, np.zeros([points.shape[0], 1]), 1) + sh = self.__get_shading(shading) + points = points.astype("float32", copy=False) + mi = np.min(points, axis=0) + ma = np.max(points, axis=0) + + g_attributes = {"position": p3s.BufferAttribute(points, normalized=False)} + m_attributes = {"size": sh["point_size"]} + + if sh["point_shape"] == "circle": # Plot circles + tex = p3s.DataTexture(data=gen_circle(16, 16), format="RGBAFormat", type="FloatType") + m_attributes["map"] = tex + m_attributes["alphaTest"] = 0.5 + m_attributes["transparency"] = True + else: # Plot squares + pass + + colors, v_colors = self.__get_point_colors(points, c, sh) + if v_colors: # Colors per point + m_attributes["vertexColors"] = 'VertexColors' + g_attributes["color"] = p3s.BufferAttribute(colors, normalized=False) + + else: # Colors for all points + m_attributes["color"] = colors + + material = p3s.PointsMaterial(**m_attributes) + geometry = p3s.BufferGeometry(attributes=g_attributes) + points = p3s.Points(geometry=geometry, material=material) + point_obj = {"geometry": geometry, "mesh": points, "material": material, + "max": ma, "min": mi, "type": "Points", "wireframe": None} + + if obj: + return self.__add_object(point_obj, obj), point_obj + else: + return self.__add_object(point_obj) + + def remove_object(self, obj_id): + if obj_id not in self.__objects: + print("Invalid object id. Valid ids are: ", list(self.__objects.keys())) + return + self._scene.remove(self.__objects[obj_id]["mesh"]) + del self.__objects[obj_id] + self.__update_view() + + def reset(self): + for obj_id in list(self.__objects.keys()).copy(): + self._scene.remove(self.__objects[obj_id]["mesh"]) + del self.__objects[obj_id] + self.__update_view() + + def update_object(self, oid=0, vertices=None, colors=None, faces=None): + obj = self.__objects[oid] + if type(vertices) != type(None): + if obj["coloring"] == "FaceColors": + f = obj["arrays"][1] + verts = np.zeros((f.shape[0] * 3, 3), dtype="float32") + for ii in range(f.shape[0]): + # print(ii*3, f[ii]) + verts[ii * 3] = vertices[f[ii, 0]] + verts[ii * 3 + 1] = vertices[f[ii, 1]] + verts[ii * 3 + 2] = vertices[f[ii, 2]] + v = verts + + else: + v = vertices.astype("float32", copy=False) + obj["geometry"].attributes["position"].array = v + # self.wireframe.attributes["position"].array = v # Wireframe updates? + obj["geometry"].attributes["position"].needsUpdate = True + # obj["geometry"].exec_three_obj_method('computeVertexNormals') + if type(colors) != type(None): + colors, coloring = self.__get_colors(obj["arrays"][0], obj["arrays"][1], colors, obj["shading"]) + colors = colors.astype("float32", copy=False) + obj["geometry"].attributes["color"].array = colors + obj["geometry"].attributes["color"].needsUpdate = True + if type(faces) != type(None): + if obj["coloring"] == "FaceColors": + print("Face updates are currently only possible in vertex color mode.") + return + f = faces.astype("uint32", copy=False).ravel() + print(obj["geometry"].attributes) + obj["geometry"].attributes["index"].array = f + # self.wireframe.attributes["position"].array = v # Wireframe updates? + obj["geometry"].attributes["index"].needsUpdate = True + # obj["geometry"].exec_three_obj_method('computeVertexNormals') + # self.mesh.geometry.verticesNeedUpdate = True + # self.mesh.geometry.elementsNeedUpdate = True + # self.update() + if self.render_mode == "WEBSITE": + return self + + # def update(self): + # self.mesh.exec_three_obj_method('update') + # self.orbit.exec_three_obj_method('update') + # self.cam.exec_three_obj_method('updateProjectionMatrix') + # self.scene.exec_three_obj_method('update') + + def add_text(self, text, shading={}, **kwargs): + shading.update(kwargs) + sh = self.__get_shading(shading) + tt = p3s.TextTexture(string=text, color=sh["text_color"]) + sm = p3s.SpriteMaterial(map=tt) + text = p3s.Sprite(material=sm, scaleToTexture=True) + self._scene.add(text) + + # def add_widget(self, widget, callback): + # self.widgets.append(widget) + # widget.observe(callback, names='value') + + # def add_dropdown(self, options, default, desc, cb): + # widget = widgets.Dropdown(options=options, value=default, description=desc) + # self.__widgets.append(widget) + # widget.observe(cb, names="value") + # display(widget) + + # def add_button(self, text, cb): + # button = widgets.Button(description=text) + # self.__widgets.append(button) + # button.on_click(cb) + # display(button) + + def to_html(self, imports=True, html_frame=True): + # Bake positions (fixes centering bug in offline rendering) + if len(self.__objects) == 0: + return + ma = np.zeros((len(self.__objects), 3)) + mi = np.zeros((len(self.__objects), 3)) + for r, obj in enumerate(self.__objects): + ma[r] = self.__objects[obj]["max"] + mi[r] = self.__objects[obj]["min"] + ma = np.max(ma, axis=0) + mi = np.min(mi, axis=0) + diag = np.linalg.norm(ma - mi) + mean = (ma - mi) / 2 + mi + for r, obj in enumerate(self.__objects): + v = self.__objects[obj]["geometry"].attributes["position"].array + v -= mean + v += np.array([0.0, .9, 0.0]) #! to move the obj to the center of window + + scale = self.__s["scale"] * (diag) + self._orbit.target = [0.0, 0.0, 0.0] + self._cam.lookAt([0.0, 0.0, 0.0]) + # self._cam.position = [0.0, 0.0, scale] + self._cam.position = [0.0, 0.5, scale * 1.3] #! show four complete meshes in the window + self._light.position = [0.0, 0.0, scale] + + state = embed.dependency_state(self._renderer) + + # Somehow these entries are missing when the state is exported in python. + # Exporting from the GUI works, so we are inserting the missing entries. + for k in state: + if state[k]["model_name"] == "OrbitControlsModel": + state[k]["state"]["maxAzimuthAngle"] = "inf" + state[k]["state"]["maxDistance"] = "inf" + state[k]["state"]["maxZoom"] = "inf" + state[k]["state"]["minAzimuthAngle"] = "-inf" + + tpl = embed.load_requirejs_template + if not imports: + embed.load_requirejs_template = "" + + s = embed.embed_snippet(self._renderer, state=state, embed_url=EMBED_URL) + # s = embed.embed_snippet(self.__w, state=state) + embed.load_requirejs_template = tpl + + if html_frame: + s = "\n\n" + s + "\n\n" + + # Revert changes + for r, obj in enumerate(self.__objects): + v = self.__objects[obj]["geometry"].attributes["position"].array + v += mean + self.__update_view() + + return s + + def save(self, filename=""): + if filename == "": + uid = str(uuid.uuid4()) + ".html" + else: + filename = filename.replace(".html", "") + uid = filename + '.html' + with open(uid, "w") as f: + f.write(self.to_html()) + print("Plot saved to file %s." % uid) diff --git a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py index d404149..b23dd96 100644 --- a/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py +++ b/MVs_Algorithms/DiffRastMesh/diff_mesh_renderer.py @@ -116,7 +116,7 @@ class DiffRastRenderer(nn.Module): i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long() v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :] - face_normals = torch.cross(v1 - v0, v2 - v0) + face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1) face_normals = safe_normalize(face_normals) vn = torch.zeros_like(v) diff --git a/_Example_Workflows/CRM MVDiffusion_to_InstantMesh.json b/_Example_Workflows/CRM/CRM MVDiffusion_to_InstantMesh.json similarity index 100% rename from _Example_Workflows/CRM MVDiffusion_to_InstantMesh.json rename to _Example_Workflows/CRM/CRM MVDiffusion_to_InstantMesh.json diff --git a/_Example_Workflows/CRM/CRM MVDiffusion_to_Unique3D.json b/_Example_Workflows/CRM/CRM MVDiffusion_to_Unique3D.json new file mode 100644 index 0000000..ac30863 --- /dev/null +++ b/_Example_Workflows/CRM/CRM MVDiffusion_to_Unique3D.json @@ -0,0 +1,2309 @@ +{ + "last_node_id": 375, + "last_link_id": 746, + "nodes": [ + { + "id": 289, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 1400, + 2290 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 41, + "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": 44, + "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": 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": 22, + "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": 323, + "type": "Reroute", + "pos": [ + -750, + 1556 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 671 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 639 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 321, + "type": "[Comfy3D] Rotate Normal Maps Horizontally", + "pos": [ + 1756.0260388541537, + 722.5761190469728 + ], + "size": { + "0": 315, + "1": 78 + }, + "flags": {}, + "order": 24, + "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": 330, + "type": "Reroute", + "pos": [ + 2230, + 2030 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 35, + "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": 46, + "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": 318, + "type": "ImageScale", + "pos": [ + 352, + 817 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 20, + "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": 257, + "type": "ImageUpscaleWithModel", + "pos": [ + 780.1904512059592, + -447.3873149879231 + ], + "size": { + "0": 226.8000030517578, + "1": 46 + }, + "flags": {}, + "order": 13, + "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": 18, + "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": 43, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 739 + }, + { + "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": 345, + "type": "[Comfy3D] Unique3D MVDiffusion Model", + "pos": [ + 712, + 764 + ], + "size": { + "0": 380.4000244140625, + "1": 222 + }, + "flags": {}, + "order": 21, + "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": 340, + "type": "ImageScale", + "pos": [ + 400.19045120595905, + -427.3873149879231 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 729 + } + ], + "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": 14, + "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": 324, + "type": "Reroute", + "pos": [ + 1325, + 290 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 689 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 716 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 194, + "type": "UpscaleModelLoader", + "pos": [ + 310.19045120595905, + -207.3873149879232 + ], + "size": { + "0": 394.51678466796875, + "1": 58 + }, + "flags": {}, + "order": 0, + "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": [ + "RealESRGAN_x4plus.pth" + ] + }, + { + "id": 312, + "type": "[Comfy3D] Load Diffusers Pipeline", + "pos": [ + -603.9739611458458, + 650.576119046973 + ], + "size": { + "0": 430.03009033203125, + "1": 154 + }, + "flags": {}, + "order": 1, + "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", + "MrForExample/Unique3D", + "", + false, + "image2normal" + ] + }, + { + "id": 363, + "type": "LoadImage", + "pos": [ + -2250, + -400 + ], + "size": { + "0": 473.3881530761719, + "1": 545.7767944335938 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 723 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 725 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Mario_Cat.png", + "image" + ] + }, + { + "id": 365, + "type": "InvertMask", + "pos": [ + -1720, + -320 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 725 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 724 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 368, + "type": "[Comfy3D] Load CRM MVDiffusion Model", + "pos": [ + -2220, + -560 + ], + "size": { + "0": 435.45733642578125, + "1": 83.6087875366211 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "crm_mvdiffusion_sampler", + "type": "CRM_MVDIFFUSION_SAMPLER", + "links": [ + 722 + ], + "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": 364, + "type": "[Comfy3D] CRM Images MVDiffusion Model", + "pos": [ + -1450, + -420 + ], + "size": { + "0": 506.4000244140625, + "1": 262 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "crm_mvdiffusion_sampler", + "type": "CRM_MVDIFFUSION_SAMPLER", + "link": 722 + }, + { + "name": "reference_image", + "type": "IMAGE", + "link": 723 + }, + { + "name": "reference_mask", + "type": "MASK", + "link": 724 + } + ], + "outputs": [ + { + "name": "multiview_images", + "type": "IMAGE", + "links": [ + 726, + 727, + 730 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "orbit_camposes", + "type": "ORBIT_CAMPOSES", + "links": [], + "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": 367, + "type": "PreviewImage", + "pos": [ + -809, + -283 + ], + "size": [ + 741.4351806640625, + 246 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 727 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 322, + "type": "[Comfy3D] Get Masks From Normal Maps", + "pos": [ + 1423, + 792 + ], + "size": { + "0": 304.79998779296875, + "1": 26 + }, + "flags": {}, + "order": 23, + "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": 269, + "type": "Morphology", + "pos": [ + 64, + 1891 + ], + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 30, + "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": [ + -201, + 1889 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 28, + "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": 325, + "type": "Reroute", + "pos": [ + -761, + 288 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 716 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 670, + 671 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 366, + "type": "GetImagesFromBatchIndexed", + "pos": [ + -875, + 1699 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 726 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "5, 0, 2, 3" + ] + }, + { + "id": 369, + "type": "GetImagesFromBatchIndexed", + "pos": [ + -325, + -420 + ], + "size": { + "0": 255.99082946777344, + "1": 76.0000228881836 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 730 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 728 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "5, 3, 2, 0" + ] + }, + { + "id": 370, + "type": "Image Rembg (Remove Background)", + "pos": [ + -24, + -420 + ], + "size": { + "0": 315, + "1": 250 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 728 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 729 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Image Rembg (Remove Background)" + }, + "widgets_values": [ + false, + "u2net", + false, + false, + false, + 240, + 10, + 10, + "white" + ] + }, + { + "id": 286, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 745, + 2301 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 40, + "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": 337, + "type": "Reroute", + "pos": [ + -354, + 1732 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 663 + } + ], + "outputs": [ + { + "name": "", + "type": "MASK", + "links": [ + 664 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 298, + "type": "[Comfy3D] Fast Clean Mesh", + "pos": [ + 1806, + 1658 + ], + "size": { + "0": 315, + "1": 130 + }, + "flags": {}, + "order": 39, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 610 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 611, + 739 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Fast Clean Mesh" + }, + "widgets_values": [ + true, + 1, + true, + 0.25 + ] + }, + { + "id": 295, + "type": "Reroute", + "pos": [ + 1197, + 2030 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 33, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 677 + } + ], + "outputs": [ + { + "name": "", + "type": "MASK", + "links": [ + 600, + 653 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 273, + "type": "PreviewImage", + "pos": [ + 442, + 2048 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 31, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 580 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 277, + "type": "PreviewImage", + "pos": [ + 64, + 2046 + ], + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 29, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 584 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 281, + "type": "ImageToMask", + "pos": [ + 447, + 1890 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 32, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 678 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 677, + 745 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageToMask" + }, + "widgets_values": [ + "red" + ] + }, + { + "id": 327, + "type": "Reroute", + "pos": [ + 765, + 1435 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 27, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 645 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 649, + 744 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 375, + "type": "[Comfy3D] Fast Normal Maps To Mesh", + "pos": [ + 828, + 1657 + ], + "size": { + "0": 418.1999816894531, + "1": 46 + }, + "flags": {}, + "order": 34, + "mode": 0, + "inputs": [ + { + "name": "front_side_back_normal_maps", + "type": "IMAGE", + "link": 744 + }, + { + "name": "front_side_back_normal_masks", + "type": "MASK", + "link": 745 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 743, + 746 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Fast Normal Maps To Mesh" + } + }, + { + "id": 290, + "type": "[Comfy3D] ExplicitTarget Mesh Optimization", + "pos": [ + 1329, + 1659 + ], + "size": { + "0": 380.4000244140625, + "1": 238 + }, + "flags": {}, + "order": 36, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 743 + }, + { + "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": 285, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 763, + 2145 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 37, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 746 + } + ], + "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": [ + "Craftsman_Unique3D\\Mesh_Gen_0\\mesh.obj" + ] + }, + { + "id": 288, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 1404, + 2161 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 38, + "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": [ + "CRM_Unique3D\\Mesh_Gen_1\\mesh.obj" + ] + }, + { + "id": 299, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2060, + 2170 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 42, + "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": [ + "CRM_Unique3D\\Mesh_Gen_2\\mesh.obj" + ] + }, + { + "id": 296, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 2960, + 1530 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 45, + "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": [ + "CRM_Unique3D\\Mesh_Gen_3\\mesh.obj" + ] + }, + { + "id": 315, + "type": "SaveImage", + "pos": [ + 2106.0260388541546, + 722.5761190469728 + ], + "size": { + "0": 340.4141540527344, + "1": 350.8065490722656 + }, + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 626 + } + ], + "properties": {}, + "widgets_values": [ + "CRM_Unique3D/Normals/normal" + ] + }, + { + "id": 195, + "type": "SaveImage", + "pos": [ + 1450.1904512059589, + -447.3873149879231 + ], + "size": { + "0": 310.5008850097656, + "1": 341.5386962890625 + }, + "flags": {}, + "order": 16, + "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": [ + "CRM_Unique3D\\Upscales\\rgb_highres" + ], + "color": "#223", + "bgcolor": "#335" + } + ], + "links": [ + [ + 547, + 194, + 0, + 257, + 0, + "UPSCALE_MODEL" + ], + [ + 580, + 269, + 0, + 273, + 0, + "IMAGE" + ], + [ + 584, + 282, + 0, + 277, + 0, + "IMAGE" + ], + [ + 594, + 285, + 0, + 286, + 0, + "STRING" + ], + [ + 596, + 290, + 0, + 288, + 0, + "MESH" + ], + [ + 597, + 288, + 0, + 289, + 0, + "STRING" + ], + [ + 600, + 295, + 0, + 290, + 2, + "MASK" + ], + [ + 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, + "*" + ], + [ + 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, + "*" + ], + [ + 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" + ], + [ + 703, + 343, + 0, + 195, + 0, + "IMAGE" + ], + [ + 715, + 340, + 0, + 257, + 1, + "IMAGE" + ], + [ + 716, + 324, + 0, + 325, + 0, + "*" + ], + [ + 722, + 368, + 0, + 364, + 0, + "CRM_MVDIFFUSION_SAMPLER" + ], + [ + 723, + 363, + 0, + 364, + 1, + "IMAGE" + ], + [ + 724, + 365, + 0, + 364, + 2, + "MASK" + ], + [ + 725, + 363, + 1, + 365, + 0, + "MASK" + ], + [ + 726, + 364, + 0, + 366, + 0, + "IMAGE" + ], + [ + 727, + 364, + 0, + 367, + 0, + "IMAGE" + ], + [ + 728, + 369, + 0, + 370, + 0, + "IMAGE" + ], + [ + 729, + 370, + 0, + 340, + 0, + "IMAGE" + ], + [ + 730, + 364, + 0, + 369, + 0, + "IMAGE" + ], + [ + 739, + 298, + 0, + 291, + 0, + "MESH" + ], + [ + 743, + 375, + 0, + 290, + 0, + "MESH" + ], + [ + 744, + 327, + 0, + 375, + 0, + "IMAGE" + ], + [ + 745, + 281, + 0, + 375, + 1, + "MASK" + ], + [ + 746, + 375, + 0, + 285, + 0, + "MESH" + ] + ], + "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.672749994932603, + "offset": [ + -566.2907157972046, + -991.1602261533006 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/_Example_Workflows/CRM_Img_to_Multi-views_to_CCMs_to_Mesh.json b/_Example_Workflows/CRM/CRM_Img_to_Multi-views_to_CCMs_to_Mesh.json similarity index 100% rename from _Example_Workflows/CRM_Img_to_Multi-views_to_CCMs_to_Mesh.json rename to _Example_Workflows/CRM/CRM_Img_to_Multi-views_to_CCMs_to_Mesh.json diff --git a/nodes.py b/nodes.py index 22c060d..ca9044c 100644 --- a/nodes.py +++ b/nodes.py @@ -53,7 +53,7 @@ from GaussianSplatting.main_3DGS_renderer import GaussianSplattingRenderer from NeRF.Instant_NGP import InstantNGP from TriplaneGaussian.triplane_gaussian_transformers import TGS -from TriplaneGaussian.utils.config import ExperimentConfig, load_config as load_config_tgs +from TriplaneGaussian.utils.config import ExperimentConfig as ExperimentConfigTGS, load_config as load_config_tgs from TriplaneGaussian.data import CustomImageOrbitDataset from TriplaneGaussian.utils.misc import todevice, get_device from LGM.core.options import config_defaults @@ -80,6 +80,9 @@ 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 +import craftsman +from craftsman.systems.base import BaseSystem +from craftsman.utils.config import ExperimentConfig as ExperimentConfigCraftsman, load_config as load_config_craftsman from .shared_utils.image_utils import ( prepare_torch_img, torch_imgs_to_pils, troch_image_dilate, @@ -1325,7 +1328,7 @@ class Load_Triplane_Gaussian_Transformers: device = get_device() - cfg: ExperimentConfig = load_config_tgs(self.config_path_abs) + cfg: ExperimentConfigTGS = load_config_tgs(self.config_path_abs) ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__) @@ -1362,7 +1365,7 @@ class Triplane_Gaussian_Transformers: CATEGORY = "Comfy3D/Algorithm" def run_TGS(self, reference_image, reference_mask, tgs_model, cam_dist): - cfg: ExperimentConfig = load_config_tgs(self.config_path_abs) + cfg: ExperimentConfigTGS = load_config_tgs(self.config_path_abs) cfg.data.cond_camera_distance = cam_dist cfg.data.eval_camera_distance = cam_dist @@ -1874,8 +1877,6 @@ class TripoSR: @torch.no_grad() def run_TSR(self, tsr_model, reference_image, reference_mask, geometry_extract_resolution, marching_cude_threshold): - cstr(f"[{self.__class__.__name__}] Running TripoSR...").msg.print() - mesh = None image = reference_image[0] @@ -2274,6 +2275,8 @@ class Load_InstantMesh_Reconstruction_Model: if is_flexicubes: lrm_model.init_flexicubes_geometry(DEVICE, fovy=30.0) lrm_model = lrm_model.eval() + + cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print() return (lrm_model, ) @@ -2624,6 +2627,8 @@ class Load_Unique3D_Custom_UNet: configurable_unet.unet.to(DEVICE, dtype=WEIGHT_DTYPE) pipe.unet = configurable_unet.unet + + cstr(f"[{self.__class__.__name__}] loaded unet ckpt from {checkpoint_path}").msg.print() return (pipe, ) class Unique3D_MVDiffusion_Model: @@ -2873,6 +2878,8 @@ class Load_CharacterGen_MVDiffusion_Model: snapshot_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, force_download=force_download, 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)) + + cstr(f"[{self.__class__.__name__}] loaded model ckpt from {self.checkpoints_dir_abs}").msg.print() return (character_mv_gen_pipe,) class CharacterGen_MVDiffusion_Model: @@ -2972,6 +2979,8 @@ class Load_CharacterGen_Reconstruction_Model: snapshot_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, force_download=force_download, 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)) + + cstr(f"[{self.__class__.__name__}] loaded model ckpt from {self.checkpoints_dir_abs}").msg.print() return (character_lrm_pipe,) class CharacterGen_Reconstruction_Model: @@ -3005,4 +3014,104 @@ class CharacterGen_Reconstruction_Model: mesh.auto_normal() mesh.auto_uv() + return (mesh,) + +class Load_Craftsman_Shape_Diffusion_Model: + checkpoints_dir = "Craftsman" + default_repo_id = "wyysf/CraftsMan" + default_ckpt_name = "image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/model.ckpt" + config_path = "Craftsman_config.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) + all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS, recursive=True) + if cls.default_ckpt_name not in all_models_names: + all_models_names += [cls.default_ckpt_name] + + return { + "required": { + "model_name": (all_models_names, ), + }, + } + + RETURN_TYPES = ( + "CRAFTSMAN_MODEL", + ) + RETURN_NAMES = ( + "craftsman_model", + ) + FUNCTION = "load_model" + CATEGORY = "Comfy3D/Import|Export" + + def load_model(self, model_name): + ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__) + + cfg: ExperimentConfigCraftsman + cfg = load_config_craftsman(self.config_root_path_abs) + + craftsman_model: BaseSystem = craftsman.find(cfg.system_type)( + cfg.system, + ) + + craftsman_model.load_state_dict(torch.load(ckpt_path, map_location=torch.device('cpu'))['state_dict']) + craftsman_model = craftsman_model.to(DEVICE).eval() + + cstr(f"[{self.__class__.__name__}] loaded model ckpt from {self.checkpoints_dir_abs}").msg.print() + return (craftsman_model,) + +class Craftsman_Shape_Diffusion_Model: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "craftsman_model": ("CRAFTSMAN_MODEL", ), + "multiview_images": ("IMAGE",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "guidance_scale": ("FLOAT", {"default": 5.0, "min": 0.0, "step": 0.01}), + "num_inference_steps": ("INT", {"default": 50, "min": 1}), + "marching_cude_grids_resolution": ("INT", {"default": 256, "min": 1, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ( + "MESH", + ) + RETURN_NAMES = ( + "mesh", + ) + + FUNCTION = "run_model" + CATEGORY = "Comfy3D/Algorithm" + + @torch.no_grad() + def run_model(self, craftsman_model, multiview_images, seed, guidance_scale, num_inference_steps, marching_cude_grids_resolution): + pil_mv_image_list = torch_imgs_to_pils(multiview_images) + + sample_inputs = {"mvimages": [pil_mv_image_list]} # view order: front, right, back, left + + latents = craftsman_model.sample( + sample_inputs, + sample_times=1, + steps=num_inference_steps, + guidance_scale=guidance_scale, + return_intermediates=False, + seed=seed + )[0] + + cstr(f"[{self.__class__.__name__}] Starting to extract mesh...").msg.print() + # decode the latents to mesh + box_v = 1.1 + mesh_outputs, _ = craftsman_model.shape_model.extract_geometry( + latents, + bounds=[-box_v, -box_v, -box_v, box_v, box_v, box_v], + grids_resolution=marching_cude_grids_resolution + ) + vertices, faces = torch.from_numpy(mesh_outputs[0][0]).to(DEVICE), torch.from_numpy(mesh_outputs[0][1]).to(torch.int64).to(DEVICE) + + mesh = Mesh(v=vertices, f=faces, 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 7d69013..a51d217 100644 --- a/shared_utils/common_utils.py +++ b/shared_utils/common_utils.py @@ -44,10 +44,28 @@ def parse_save_filename(save_path, output_directory, supported_extensions, class return None -# Get all files in a folder, if extension_filter is provided then only reture the files with extensions in extension_filter -def get_list_filenames(directory, extension_filter=None): +def get_list_filenames(directory, extension_filter=None, recursive=False): + """ + Recursively finds files with specified extensions in a directory and returns relative paths. + + Args: + directory (str): The directory path to search. + extension_filter (list): List of file extensions (e.g., ['.txt', '.csv']). + + Returns: + list: List of relative file paths matching the specified extensions. + """ if exists(directory): - return [f for f in listdir(directory) if isfile(join(directory, f)) and (extension_filter is None or f.lower().endswith(extension_filter))] + if recursive: + result = [] + for root, _, files in os.walk(directory): + for item in files: + if extension_filter is None or os.path.splitext(item)[1].lower() in extension_filter: + relative_path = os.path.relpath(os.path.join(root, item), directory) + result.append(relative_path) + return result + else: + return [f for f in listdir(directory) if isfile(join(directory, f)) and (extension_filter is None or f.lower().endswith(extension_filter))] else: return [] From f4181c27d23eea6500c93caf00aa0e202d05deaf Mon Sep 17 00:00:00 2001 From: "Mr. For Example" <62230687+MrForExample@users.noreply.github.com> Date: Tue, 23 Jul 2024 15:13:01 -0600 Subject: [PATCH 2/3] Update nodes.py --- nodes.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index 40a1865..c814bb7 100644 --- a/nodes.py +++ b/nodes.py @@ -3149,7 +3149,7 @@ class OrbitPoses_JK: RETURN_NAMES = ("orbit_lists", "orbit_camposes",) FUNCTION = "get_orbit_poses" - CATEGORY = icons.get("JK/3D") + CATEGORY = "Comfy3D/Preprocessor" def get_orbit_poses(self, orbitpose_preset, azimuths, elevations, radius, center): @@ -3201,7 +3201,7 @@ class OrbitLists_to_OrbitPoses_JK: RETURN_NAMES = ("orbit_camposes",) FUNCTION = "convert_orbit_poses" - CATEGORY = icons.get("JK/3D") + CATEGORY = "Comfy3D/Preprocessor" def convert_orbit_poses(self, orbit_lists): @@ -3228,7 +3228,7 @@ class OrbitPoses_to_OrbitLists_JK: RETURN_NAMES = ("orbit_lists",) FUNCTION = "convert_orbit_poses" - CATEGORY = icons.get("JK/3D") + CATEGORY = "Comfy3D/Preprocessor" def convert_orbit_poses(self, orbit_camposes): From f524f3d60c64d7854b2fd55e326c15c1646083b5 Mon Sep 17 00:00:00 2001 From: "Mr. For Example" <62230687+MrForExample@users.noreply.github.com> Date: Tue, 23 Jul 2024 15:14:52 -0600 Subject: [PATCH 3/3] Create CRM MVDiffusion_to_Craftsman.json --- .../CRM/CRM MVDiffusion_to_Craftsman.json | 488 ++++++++++++++++++ 1 file changed, 488 insertions(+) create mode 100644 _Example_Workflows/CRM/CRM MVDiffusion_to_Craftsman.json diff --git a/_Example_Workflows/CRM/CRM MVDiffusion_to_Craftsman.json b/_Example_Workflows/CRM/CRM MVDiffusion_to_Craftsman.json new file mode 100644 index 0000000..9acb07f --- /dev/null +++ b/_Example_Workflows/CRM/CRM MVDiffusion_to_Craftsman.json @@ -0,0 +1,488 @@ +{ + "last_node_id": 34, + "last_link_id": 53, + "nodes": [ + { + "id": 20, + "type": "InvertMask", + "pos": [ + 341, + 400 + ], + "size": { + "0": 210, + "1": 26 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 26 + } + ], + "outputs": [ + { + "name": "MASK", + "type": "MASK", + "links": [ + 28 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "InvertMask" + } + }, + { + "id": 1, + "type": "LoadImage", + "pos": [ + -191, + 320 + ], + "size": { + "0": 473.3881530761719, + "1": 545.7767944335938 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 26 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Mario_Cat.png", + "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": [], + "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": 32, + "type": "[Comfy3D] Load CRM MVDiffusion Model", + "pos": [ + -157, + 163 + ], + "size": { + "0": 435.45733642578125, + "1": 83.6087875366211 + }, + "flags": {}, + "order": 1, + "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": 33, + "type": "[Comfy3D] Craftsman Shape Diffusion Model", + "pos": [ + 1501, + 283 + ], + "size": { + "0": 315, + "1": 174 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "craftsman_model", + "type": "CRAFTSMAN_MODEL", + "link": 51 + }, + { + "name": "multiview_images", + "type": "IMAGE", + "link": 52 + } + ], + "outputs": [ + { + "name": "mesh", + "type": "MESH", + "links": [ + 53 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Craftsman Shape Diffusion Model" + }, + "widgets_values": [ + 0, + "fixed", + 5, + 50, + 256 + ] + }, + { + "id": 34, + "type": "[Comfy3D] Load Craftsman Shape Diffusion Model", + "pos": [ + 807, + 169 + ], + "size": { + "0": 642.2877807617188, + "1": 60.448020935058594 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "craftsman_model", + "type": "CRAFTSMAN_MODEL", + "links": [ + 51 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Load Craftsman Shape Diffusion Model" + }, + "widgets_values": [ + "image-to-shape-diffusion/clip-mvrgb-modln-l256-e64-ne8-nd16-nl6-aligned-vae/model.ckpt" + ] + }, + { + "id": 25, + "type": "[Comfy3D] Preview 3DMesh", + "pos": [ + 2218, + 287 + ], + "size": [ + 600, + 500 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "mesh_file_path", + "type": "STRING", + "link": 35, + "widget": { + "name": "mesh_file_path" + } + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Preview 3DMesh" + }, + "widgets_values": [ + "", + null + ] + }, + { + "id": 27, + "type": "GetImagesFromBatchIndexed", + "pos": [ + 1187, + 300 + ], + "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": [ + 52 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "GetImagesFromBatchIndexed" + }, + "widgets_values": [ + "5, 0, 2, 3" + ] + }, + { + "id": 31, + "type": "PreviewImage", + "pos": [ + 1184, + 515 + ], + "size": [ + 741.4351814481588, + 158.14768246691654 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 47 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 24, + "type": "[Comfy3D] Save 3D Mesh", + "pos": [ + 1861, + 284 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "mesh", + "type": "MESH", + "link": 53 + } + ], + "outputs": [ + { + "name": "save_path", + "type": "STRING", + "links": [ + 35 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[Comfy3D] Save 3D Mesh" + }, + "widgets_values": [ + "Craftsman_test\\Mesh_Cat_CRM.obj" + ] + } + ], + "links": [ + [ + 2, + 1, + 0, + 3, + 1, + "IMAGE" + ], + [ + 26, + 1, + 1, + 20, + 0, + "MASK" + ], + [ + 28, + 20, + 0, + 3, + 2, + "MASK" + ], + [ + 35, + 24, + 0, + 25, + 0, + "STRING" + ], + [ + 39, + 3, + 0, + 27, + 0, + "IMAGE" + ], + [ + 47, + 3, + 0, + 31, + 0, + "IMAGE" + ], + [ + 48, + 32, + 0, + 3, + 0, + "CRM_MVDIFFUSION_SAMPLER" + ], + [ + 51, + 34, + 0, + 33, + 0, + "CRAFTSMAN_MODEL" + ], + [ + 52, + 27, + 0, + 33, + 1, + "IMAGE" + ], + [ + 53, + 33, + 0, + 24, + 0, + "MESH" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + -816.3016446237854, + -1.7073104334998879 + ] + } + }, + "version": 0.4 +} \ No newline at end of file