update

ckpts/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

ckpts/README.md

@@ -0,0 +1,117 @@
+---
+library_name: craftsman-v1-5
+license: creativeml-openrail-m
+license_name: creativeml-openrail-m
+license_link: https://raw.githubusercontent.com/CompVis/stable-diffusion/refs/heads/main/LICENSE
+pipeline_tag: image-to-3d
+language:
+- en
+- zh
+---
+
+## **CraftsMan-v1-5**
+
+<p align="center">
+  <img src="./assets/teaser.png"  height=200>
+</p>
+
+### <div align="center">CraftsMan: High-fidelity Mesh Generation <br> with 3D Native Generation and Interactive Geometry Refiner<div> 
+#####  <p align="center"> [Weiyu Li<sup>*1,2</sup>](https://wyysf-98.github.io/), Jiarui Liu<sup>*1,2</sup>, Hongyu Yan<sup>*1,2</sup>, [Rui Chen<sup>1,2</sup>](https://aruichen.github.io/), [Yixun Liang<sup>2,3</sup>](https://yixunliang.github.io/), [Xuelin Chen<sup>4</sup>](https://xuelin-chen.github.io/), [Ping Tan<sup>1,2</sup>](https://ece.hkust.edu.hk/pingtan), [Xiaoxiao Long<sup>1,2</sup>](https://www.xxlong.site/)</p>
+#####  <p align="center"> <sup>1</sup>HKUST, <sup>2</sup>LightIllusions, <sup>3</sup>HKUST(GZ), <sup>4</sup>Tencent AI Lab</p>
+<div align="center">
+  <a href="https://craftsman3d.github.io/"><img src="https://img.shields.io/static/v1?label=Project%20Page&message=Github&color=blue&logo=github-pages"></a> &ensp;
+  <a href="http://algodemo.bj.lightions.top:24926"><img src="https://www.gradio.app/_app/immutable/assets/gradio.CHB5adID.svg" height="25"/></a> &ensp;
+  <a href="https://arxiv.org/pdf/2405.14979"><img src="https://img.shields.io/static/v1?label=Paper&message=Arxiv&color=red&logo=arxiv"></a> &ensp;
+</div>
+
+# Usage
+
+To use the model, please refer to the [official repository](https://github.com/wyysf-98/CraftsMan) for installation and usage instructions.
+
+```
+
+from craftsman import CraftsManPipeline
+import torch
+
+pipeline = CraftsManPipeline.from_pretrained("./ckpts/craftsman-v1-5", device="cuda:0", torch_dtype=torch.float32) # load from local ckpt
+mesh = pipeline("https://pub-f9073a756ec645d692ce3d171c2e1232.r2.dev/data/werewolf.png").meshes[0]
+mesh.export("werewolf.obj")
+
+```
+
+## 🔥🔥🔥 News!!
+
+* Nov 16, 2024: 💬 We release the CraftsMan-v1-5
+
+
+## 📑 Open-source Plan
+
+- [x] Inference 
+- [x] Checkpoints
+- [x] Training
+- [ ] ComfyUI
+
+## 🎉 **CraftMan-v1-5 Architecture**
+
+<p align="center">
+  <img src="./assets/arch.png"  height=400>
+</p>
+
+
+## Get Started
+
+#### Begin by cloning the repository:
+
+```shell
+git clone https://github.com/wyysf-98/CraftsMan
+cd CraftsMan
+```
+
+#### Installation Guide for Linux
+
+We provide an env_install.sh script file for setting up environment. 
+
+```
+# step 1, create conda env
+conda create -n CraftsMan python=3.10
+conda activate CraftsMan
+
+
+# step 2. install torch realated package
+conda install -c pytorch pytorch=2.3.0 torchvision=0.18.0 cudatoolkit=11.8
+
+# step 3. install other packages
+pip install -r docker/requirements.txt
+```
+<details>
+
+
+#### Using Gradio
+
+We have prepared a gradio demo for you to try out the model. You can run the following command to start the demo.
+
+```shell
+# std 
+python3 gradio.py
+```
+
+Then the demo can be accessed through the output link.
+
+
+## Citation
+
+If you found this repository helpful, please cite our report:
+```bibtex
+@misc{li2024craftsman,
+title         = {CraftsMan: High-fidelity Mesh Generation with 3D Native Generation and Interactive Geometry Refiner}, 
+author        = {Weiyu Li and Jiarui Liu and Rui Chen and Yixun Liang and Xuelin Chen and Ping Tan and Xiaoxiao Long},
+year          = {2024},
+archivePrefix = {arXiv preprint arXiv:2405.14979},
+primaryClass  = {cs.CG}
+}
+```
+
+
+# License
+
+[creativeml-openrail-m](https://raw.githubusercontent.com/CompVis/stable-diffusion/refs/heads/main/LICENSE)

ckpts/craftsman-v1-5

@@ -0,0 +1 @@
+Subproject commit 9a5e9189c2dfab20cf838885dd6acaf99b41844e

craftsman/data/base.py

@@ -53,7 +53,7 @@ class BaseDataModuleConfig:
     # for occupancy and sdf data
     n_samples: int = 4096                # number of points in input point cloud
     upsample_ratio: int = 1              # upsample ratio for input point cloud
-    sampling_strategy: str = "random"    # sampling strategy for input point cloud
+    sampling_strategy: Optional[str] = None    # sampling strategy for input point cloud
     scale: float = 1.0                   # scale of the input point cloud and target supervision
     load_supervision: bool = True        # whether to load supervision
     supervision_type: str = "occupancy"  # occupancy, sdf, tsdf
@@ -70,6 +70,8 @@ class BaseDataModuleConfig:
     idx: Optional[List[int]] = None      # index of the image to load
     n_views: int = 1                     # number of views
     marign_pix_dis: int = 30             # margin of the bounding box
+    batch_size: int = 32
+    num_workers: int = 8
 
 
 class BaseDataset(Dataset):
@@ -78,7 +80,7 @@ class BaseDataset(Dataset):
         self.cfg: BaseDataModuleConfig = cfg
         self.split = split
 
-        self.uids = json.load(open(f'{cfg.root_dir}/{split}.json'))
+        self.uids = json.load(open(f'{cfg.local_dir}/{split}.json'))
         print(f"Loaded {len(self.uids)} {split} uids")
     
     def __len__(self):
@@ -94,10 +96,7 @@ class BaseDataset(Dataset):
             surface = np.concatenate([surface, normal], axis=1)
         elif self.cfg.geo_data_type == "sdf":
             # for sdf data with our own format
-            if re.match(r"\.\.", self.uids[index]):
-                data = np.load(f'{self.cfg.geo_data_path}/{self.uids[index]}.npz')
-            else:
-                data = np.load(f'{self.uids[index]}.npz')
+            data = np.load(f'{self.cfg.geo_data_path}/{self.uids[index]}.npz')
             # for input point cloud
             surface = data["surface"]
         else:
@@ -112,6 +111,8 @@ class BaseDataset(Dataset):
             import fpsample
             kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(surface[:, :3], self.cfg.n_samples, h=5)
             surface = surface[kdline_fps_samples_idx]
+        elif self.cfg.sampling_strategy is None:
+            pass
         else:
             raise NotImplementedError(f"sampling strategy {self.cfg.sampling_strategy} not implemented")
         # rescale data
@@ -189,9 +190,9 @@ class BaseDataset(Dataset):
             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"/{'{:04d}'.format(sel_idx)}_rgb.png"
+                img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{'{:04d}'.format(sel_idx)}_rgb.jpeg"
             elif self.cfg.image_type == "normal":
-                img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{'{:04d}'.format(sel_idx)}_normal.png"
+                img_path = f'{self.cfg.image_data_path}/' + "/".join(self.uids[index].split('/')[-2:]) + f"/{'{:04d}'.format(sel_idx)}_normal.jpeg"
             ret["image"], ret["mask"] = _load_single_image(img_path, background_color, self.cfg.marign_pix_dis)
 
         else:

craftsman/models/conditional_encoders/cond_encoder.py

@@ -46,7 +46,6 @@ class CondEmbedder(BaseEmbedder):
         enable_gradient_checkpointing: bool = False
         embeds_fusion_mode: int = 1 # 0: sum | 1: concat
         linear_proj_init: str = "constant"
-        text_model_type: str = "clip"
         text_max_length: int = 77
         image_size_clip: int = 224
         image_size_dino: int = 224
@@ -277,29 +276,9 @@ class CondEmbedder(BaseEmbedder):
         else:
             return vision_outputs.last_hidden_state
 
-    def post_process_embeds(self, text_embeds, visual_embeds):
-        clip_embeds, dino_embeds = visual_embeds.chunk(2, dim=2)
-        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 clip_embeds is not None:
-                clip_embeds = clip_embeds / clip_embeds.norm(dim=-1, keepdim=True)
-            if dino_embeds is not None:
-                dino_embeds = dino_embeds / dino_embeds.norm(dim=-1, keepdim=True)
-
-        assert text_embeds is not None or dino_embeds is not None or clip_embeds is not None
-            
-        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 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:
         clip_embeds = self.encode_image_clip(images, cameras)
         dino_embeds = self.encode_image_dino(images, cameras)
         dino_embeds = self.linear_proj(dino_embeds)
         visual_embeds = torch.cat([clip_embeds, dino_embeds], dim=1)
-        return visual_embeds
+        return visual_embeds

craftsman/models/denoisers/pixart_denoiser.py

@@ -25,15 +25,11 @@ class PixArtDinoDenoiser(BaseModule):
         context_dim: int = 1024
         n_views: int = 1
         context_ln: bool = True
-        skip_ln: bool = False
         init_scale: float = 0.25
         use_checkpoint: bool = False
         drop_path: float = 0.
-        variance_type: str = ""
-        img_pos_embed: bool = False
         clip_weight: float = 1.0
         dino_weight: float = 1.0
-        dit_block: str = ""
 
     cfg: Config
 
@@ -63,9 +59,8 @@ class PixArtDinoDenoiser(BaseModule):
 
         init_scale = self.cfg.init_scale * math.sqrt(1.0 / self.cfg.width)
         drop_path = [x.item() for x in torch.linspace(0, self.cfg.drop_path, self.cfg.layers)]
-        ditblock =  getattr(importlib.import_module("craftsman.models.denoisers.utils"), self.cfg.dit_block)
         self.blocks = nn.ModuleList([
-            ditblock(
+            DiTBlock(
                     width=self.cfg.width, 
                     heads=self.cfg.heads, 
                     init_scale=init_scale, 
@@ -82,11 +77,7 @@ class PixArtDinoDenoiser(BaseModule):
                     )
         
          # final layer
-        if self.cfg.variance_type.upper() in ["LEARNED", "LEARNED_RANGE"]:
-            self.output_channels = self.cfg.output_channels * 2
-        else:
-            self.output_channels = self.cfg.output_channels
-        self.final_layer = T2IFinalLayer(self.cfg.width, self.output_channels)
+        self.final_layer = T2IFinalLayer(self.cfg.width, self.cfg.output_channels)
 
         self.identity_initialize()
 
@@ -99,17 +90,6 @@ class PixArtDinoDenoiser(BaseModule):
                     self.denoiser_ckpt[k.replace('denoiser_model.', '')] = v
             self.load_state_dict(self.denoiser_ckpt, strict=False)
 
-    def forward_with_dpmsolver(self, model_input, timestep, context):
-        """
-        dpm solver donnot need variance prediction
-        """
-        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
-        model_out = self.forward(model_input, timestep, context)
-        if self.cfg.variance_type.upper() in ["LEARNED", "LEARNED_RANGE"]:
-            return model_out.chunk(2, dim=-1)[0]
-        else:
-            return model_out
-
     def identity_initialize(self):
         for block in self.blocks:
             nn.init.constant_(block.attn.c_proj.weight, 0)

craftsman/models/denoisers/utils.py

@@ -10,126 +10,6 @@ from timm.models.layers import DropPath
 from craftsman.models.transformers.utils import MLP
 from craftsman.models.transformers.attention import MultiheadAttention, MultiheadCrossAttention
 
-class PatchEmbed(nn.Module):
-    """ 2D Image to Patch Embedding
-    """
-    def __init__(
-            self,
-            patch_size=16,
-            in_chans=3,
-            embed_dim=768,
-            norm_layer=None,
-            flatten=True,
-            bias=True,
-    ):
-        super().__init__()
-        patch_size = to_2tuple(patch_size)
-        self.patch_size = patch_size
-        self.flatten = flatten
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-
-    def forward(self, x):
-        x = self.proj(x)
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        x = self.norm(x)
-        return x
-
-class DiTBlock(nn.Module):
-    """
-    A PixArt block with adaptive layer norm (adaLN-single) conditioning.
-    """
-
-    def __init__(self, width, heads, init_scale=1.0, qkv_bias=True, use_flash=True, drop_path=0.0):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(width, elementwise_affine=True, eps=1e-6)
-        self.attn = MultiheadAttention(
-            n_ctx=None,
-            width=width,
-            heads=heads,
-            init_scale=init_scale,
-            qkv_bias=qkv_bias,
-            use_flash=use_flash
-        )
-        self.cross_attn = MultiheadCrossAttention(
-            n_data=None,
-            width=width,
-            heads=heads,
-            data_width=None,
-            init_scale=init_scale,
-            qkv_bias=qkv_bias,
-            use_flash=use_flash,
-        )
-
-        self.norm2 = nn.LayerNorm(width, elementwise_affine=True, eps=1e-6)
-
-        self.mlp = MLP(width=width, init_scale=init_scale)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.scale_shift_table = nn.Parameter(torch.randn(6, width) / width ** 0.5)
-
-    def forward(self, x, visual_cond, t, **kwargs):
-        B, N, C = x.shape
-
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1)
-        x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa)).reshape(B, N, C))
-        x = x + self.cross_attn(x, visual_cond)
-        x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))
-
-        return x
-
-class DiTBlock_text(nn.Module):
-    """
-    A PixArt block with adaptive layer norm (adaLN-single) conditioning.
-    """
-
-    def __init__(self, width, heads, init_scale=1.0, qkv_bias=True, use_flash=True, drop_path=0.0):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(width, elementwise_affine=True, eps=1e-6)
-        self.attn = MultiheadAttention(
-            n_ctx=None,
-            width=width,
-            heads=heads,
-            init_scale=init_scale,
-            qkv_bias=qkv_bias,
-            use_flash=use_flash
-        )
-        self.cross_attn = MultiheadCrossAttention(
-            n_data=None,
-            width=width,
-            heads=heads,
-            data_width=None,
-            init_scale=init_scale,
-            qkv_bias=qkv_bias,
-            use_flash=use_flash,
-        )
-
-        self.cross_attn_extra = MultiheadCrossAttention(
-            n_data=None,
-            width=width,
-            heads=heads,
-            data_width=None,
-            init_scale=init_scale,
-            qkv_bias=qkv_bias,
-            use_flash=use_flash,
-        )
-        self.norm2 = nn.LayerNorm(width, elementwise_affine=True, eps=1e-6)
-
-        self.mlp = MLP(width=width, init_scale=init_scale)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
-        self.scale_shift_table = nn.Parameter(torch.randn(6, width) / width ** 0.5)
-
-    def forward(self, x, visual_cond, text_cond, t, **kwargs):
-        B, N, C = x.shape
-
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1)
-        x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa)).reshape(B, N, C))
-        x = x + self.cross_attn(x, visual_cond)
-        x = x + self.cross_attn_extra(x, text_cond)
-        x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))
-
-        return x
-
 class DiTBlock(nn.Module):
     """
     A DiT block with adaptive layer norm (adaLN-single) conditioning.
@@ -174,11 +54,6 @@ class DiTBlock(nn.Module):
 def t2i_modulate(x, shift, scale):
     return x * (1 + scale) + shift
 
-# def t2i_modulate(x, shift, scale):
-#     a = torch.ones_like(scale)
-#     a[..., 768:] = 0
-#     return x * (a + scale) + shift
-
 def auto_grad_checkpoint(module, *args, **kwargs):
     if getattr(module, 'grad_checkpointing', False):
         if not isinstance(module, Iterable):
@@ -268,63 +143,4 @@ class T2IFinalLayer(nn.Module):
         shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
         x = t2i_modulate(self.norm_final(x), shift, scale)
         x = self.linear(x)
-        return x
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """
-    embed_dim: output dimension for each position
-    pos: a list of positions to be encoded: size (M,)
-    out: (M, D)
-    """
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.
-    omega = 1. / 10000 ** omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out)  # (M, D/2)
-    emb_cos = np.cos(out)  # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    return emb
-
-def _ntuple(n):
-    def parse(x):
-        if isinstance(x, Iterable) and not isinstance(x, str):
-            return x
-        return tuple(repeat(x, n))
-    return parse
-
-to_1tuple = _ntuple(1)
-to_2tuple = _ntuple(2)
-
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0, base_size=16):
-    """
-    grid_size: int of the grid height and width
-    return:
-    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_size = to_2tuple(grid_size)
-    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0]/base_size) / pe_interpolation
-    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1]/base_size) / pe_interpolation
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
-
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
+        return x

craftsman/models/transformers/attention.py

@@ -9,126 +9,6 @@ from craftsman.utils.checkpoint import checkpoint
 from .utils import init_linear, MLP
 from timm.models.vision_transformer import Attention
 
-def scaled_dot_product_gqa(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    dropout: float = 0.0,
-    scale: Optional[float] = None,
-    mask: Optional[Tensor] = None,
-    is_causal: Optional[bool] = None,
-    need_weights: bool = False,
-    average_attn_weights: bool = False,
-    force_grouped: bool = False,
-):
-    """Scaled dot product attention with support for grouped queries.
-
-    Einstein notation:
-    - b: batch size
-    - n / s: sequence length
-    - h: number of heads
-    - g: number of groups
-    - d: dimension of query/key/value
-
-    Args:
-        query: Query tensor of shape (b, n, h, d)
-        key: Key tensor of shape (b, s, h, d)
-        value: Value tensor of shape (b, s, h, d)
-        dropout: Dropout probability (default: 0.0)
-        scale: Scale factor for query (default: d_query ** 0.5)
-        mask: Mask tensor of shape (b, n, s) or (b, s). If 'ndim == 2', the mask is
-            applied to all 'n' rows of the attention matrix. (default: None)
-        force_grouped: If True, apply grouped-query attention even if the number of
-            heads is equal for query, key, and value. (default: False)
-
-    Returns:
-        2-tuple of:
-        - Attention output with shape (b, n, h, d)
-        - (Optional) Attention weights with shape (b, h, n, s). Only returned if
-          'need_weights' is True.
-    """
-    if (mask is not None) and (is_causal is not None):
-        raise ValueError(
-            "Only one of 'mask' and 'is_causal' should be provided, but got both."
-        )
-    elif not query.ndim == key.ndim == value.ndim == 4:
-        raise ValueError(
-            f"Expected query, key, and value to be 4-dimensional, but got shapes "
-            f"{query.shape}, {key.shape}, and {value.shape}."
-        )
-
-    # Move sequence length dimension to axis 2.
-    # This makes the attention operations below *much* faster.
-    query = rearrange(query, "b n h d -> b h n d")
-    key = rearrange(key, "b s h d -> b h s d")
-    value = rearrange(value, "b s h d -> b h s d")
-
-    bq, hq, nq, dq = query.shape
-    bk, hk, nk, dk = key.shape
-    bv, hv, nv, dv = value.shape
-    if not (bq == bk == bv and dq == dk == dv):
-        raise ValueError(
-            "Expected query, key, and value to have the same batch size (dim=0) and "
-            f"embedding dimension (dim=3), but got query: {query.shape}, "
-            f"key: {key.shape}, and value: {value.shape}."
-        )
-    elif (hk != hv) or (nk != nv):
-        raise ValueError(
-            "Expected key and value to have the same size in dimensions 1 and 2, but "
-            f"got key: {key.shape} and value: {value.shape}."
-        )
-    elif hq % hk != 0:
-        raise ValueError(
-            "Expected query heads to be a multiple of key/value heads, but got "
-            f"query: {query.shape} and key/value: {key.shape}."
-        )
-
-    if scale is None:
-        scale = query.size(-1) ** 0.5
-    query = query / scale
-
-    num_head_groups = hq // hk
-    query = rearrange(query, "b (h g) n d -> b g h n d", g=num_head_groups)
-    similarity = einsum(query, key, "b g h n d, b h s d -> b g h n s")
-
-    if is_causal:
-        # Mask out the upper triangular portion of the attention matrix. This prevents
-        # the model from attending to tokens in the future.
-        mask = torch.ones((bq, nq, nk), device=query.device, dtype=torch.bool).tril_()
-
-    if mask is not None:
-        # Expand mask to match the shape of the attention matrix.
-        # If mask is 2D, assume that it is applied to the key/value sequence dimension.
-        # Else if mask is 3D, assume that it is applied to the query/key/value sequence
-        # dimension for all attention heads.
-        #
-        if mask.ndim == 2:
-            mask = rearrange(mask, "b s -> b () () () s")
-        elif mask.ndim == 3:
-            mask = rearrange(mask, "b n s -> b () () n s")
-        # Mask similarity values by setting them to negative infinity.  This guarantees
-        # that they will not contribute to the softmax computation below.
-        similarity.masked_fill_(~mask, torch.finfo(similarity.dtype).min)
-
-    attention = F.softmax(similarity, dim=-1)
-    if dropout > 0.0:
-        attention = F.dropout(attention, p=dropout)
-
-    # Apply attention matrix to the value Tensor.
-    out = einsum(attention, value, "b g h n s, b h s d -> b g h n d")
-    # Move head dimension back to axis 2
-    out = rearrange(out, "b g h n d -> b n (h g) d")
-
-    attn_weights: Optional[Tensor] = None
-    if need_weights:
-        # Move the sequence dimensions back to positions 1, 2.  Move the head dimension
-        # to position 3.  This more closely matches the return shape of the attention
-        # output: (b, n, h, d).
-        attn_weights = rearrange(attention, "b g h n s -> b n s (h g)")
-        if average_attn_weights:
-            attn_weights = attn_weights.mean(dim=1)
-
-    return out, attn_weights
 
 class MultiheadAttention(nn.Module):
     def __init__(
@@ -327,4 +207,4 @@ class ResidualCrossAttentionBlock(nn.Module):
     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
+        return x

craftsman/pipeline.py

@@ -158,6 +158,7 @@ class CraftsManPipeline():
         background_color: List[int] = [255, 255, 255],
         foreground_ratio: float = 0.95,
         mc_depth: int = 8,
+        only_max_component: bool = False,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -198,6 +199,9 @@ class CraftsManPipeline():
             mc_depth (`int`, *optional*, defaults to 8):
                 The resolution of the Marching Cubes algorithm. The resolution is the number of cubes in the x, y, and z.
                 8 means 2^8 = 256 cubes in each dimension. The higher the resolution, the more detailed the mesh will be.
+            only_max_component (`bool`, *optional*, defaults to `False`):
+                Whether to only keep the largest connected component of the mesh. This is useful when the mesh has
+                multiple components and only the largest one is needed.
         Examples:
 
         Returns:
@@ -258,6 +262,15 @@ class CraftsManPipeline():
                 if output_type == "trimesh":
                     import trimesh
                     cur_mesh = trimesh.Trimesh(vertices=mesh_v_f[0][0], faces=mesh_v_f[0][1])
+                    if only_max_component:
+                        components = cur_mesh.split(only_watertight=False)
+                        bbox = []
+                        for c in components:
+                            bbmin = c.vertices.min(0)
+                            bbmax = c.vertices.max(0)
+                            bbox.append((bbmax - bbmin).max())
+                        max_component = np.argmax(bbox)
+                        cur_mesh = components[max_component]
                     mesh.append(cur_mesh)
                 elif output_type == "np":
                     mesh.append(mesh_v_f[0])

craftsman/systems/pixart_diffusion.py

@@ -251,9 +251,9 @@ class PixArtDiffusionSystem(BaseSystem):
         return {
             "loss_diffusion": loss,
             "latents": latents,
-            "x_t": x_t,
+            "x_t": noisy_z,
             "noise": noise,
-            "noise_pred": pred_noise,
+            "noise_pred": noise_pred,
             "timesteps": timesteps,
             }
 
@@ -373,4 +373,4 @@ class PixArtDiffusionSystem(BaseSystem):
         return outputs
 
     def on_validation_epoch_end(self):
-        pass
+        pass