Merge pull request #234 from MrForExample/dev

Add Craftman

.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

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

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)

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)

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

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

Gen_3D_Modules/craftsman/data/__init__.py

@@ -0,0 +1,3 @@
+from . import (
+    objaverse
+)

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)

Gen_3D_Modules/craftsman/models/__init__.py

@@ -0,0 +1,5 @@
+from . import (
+    autoencoders,
+    conditional_encoders,
+    denoisers,
+)

Gen_3D_Modules/craftsman/models/autoencoders/__init__.py

@@ -0,0 +1,3 @@
+from . import (
+    michelangelo_autoencoder,
+)

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
+    

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

Gen_3D_Modules/craftsman/models/conditional_encoders/__init__.py

@@ -0,0 +1,3 @@
+from . import (
+    clip_encoder,
+)

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)

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,
+        )

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

Gen_3D_Modules/craftsman/models/denoisers/__init__.py

@@ -0,0 +1,3 @@
+from . import (
+    simple_denoiser,
+)

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

Gen_3D_Modules/craftsman/models/geometry/__init__.py

@@ -0,0 +1,3 @@
+from . import (
+    base
+)

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,
+            ),
+        )

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

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

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

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)))

Gen_3D_Modules/craftsman/systems/__init__.py

@@ -0,0 +1,4 @@
+from . import (
+    shape_autoencoder,
+    shape_diffusion,
+)

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

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

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

Gen_3D_Modules/craftsman/utils/__init__.py

@@ -0,0 +1 @@
+from . import base

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

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 ...")

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

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

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

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

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

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

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

Gen_3D_Modules/craftsman/utils/visualizers/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

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
+

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"""
+    <html>
+      <body>
+        {content}
+      </body>
+    </html>
+    """
+
+    return html_frame
+
+
+def to_single_row_table(caption: str, content: str):
+
+    table_html = f"""
+    <table border = "1">
+        <caption>{caption}</caption>
+        <tr>
+            <td>{content}</td>
+        </tr>
+    </table>
+    """
+
+    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"""
+    <img src="data:image/png;base64,{image_base64}" alt="Embedded Image">
+    """
+
+    return image_tag

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 = "<html>\n<body>\n" + s + "\n</body>\n</html>"
+
+        # 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)

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)

_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
+}

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:
@@ -3006,6 +3015,106 @@ class CharacterGen_Reconstruction_Model:
         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,)
 
 ORBITPOSE_PRESET = ["Custom", "CRM(6)", "Zero123Plus(6)", "Wonder3D(6)", "Era3D(6)", "MVDream(4)", "Unique3D(4)", "CharacterGen(4)"]
 
@@ -3040,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):
         
@@ -3092,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):
         
@@ -3119,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):
         

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 []