v1

.gitattributes

@@ -1,37 +1,2 @@
-*.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
 *.png filter=lfs diff=lfs merge=lfs -text
 *.glb filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Junliang Ye
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,14 +1,49 @@
----
-title: ShapeLLM Omni
-emoji: 📈
-colorFrom: purple
-colorTo: green
-sdk: gradio
-sdk_version: 5.31.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: A Native Multimodal LLM for 3DGeneration and Understanding
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<p align="center">
+  <h3 align="center"><strong>ShapeLLM-omni: A Native Multimodal LLM for 3D Generation and Understanding</strong></h3>
+
+<p align="center">
+    <a href="https://jamesyjl.github.io/">Junliang Ye</a><sup>1,2*</sup>,
+    <a href="https://thuwzy.github.io/">Zhengyi Wang</a><sup>1,2*</sup>,
+    <a href="https://zhaorw02.github.io/">Ruowen Zhao</a><sup>1*</sup>,
+    <a href="">Shenghao Xie</a><sup>3</sup>,
+    <a href="https://ml.cs.tsinghua.edu.cn/~jun/index.shtml">Jun Zhu</a><sup>1,2†</sup>
+    <br>
+    <sup>*</sup>Equal Contribution.
+    <br>
+    <sup>†</sup>Corresponding authors.
+    <br>
+    <sup>1</sup>Tsinghua University,
+    <sup>2</sup>ShengShu,
+    <sup>3</sup>Peking University,
+</p>
+
+<div align="center">
+
+<a href='https://arxiv.org/abs/2503.15265'><img src='https://img.shields.io/badge/arXiv-2503.15265-b31b1b.svg'></a> &nbsp;&nbsp;&nbsp;&nbsp;
+ <a href='https://jamesyjl.github.io/ShapeLLM/'><img src='https://img.shields.io/badge/Project-Page-Green'></a> &nbsp;&nbsp;&nbsp;&nbsp;
+ <a><img src='https://img.shields.io/badge/License-MIT-blue'></a> &nbsp;&nbsp;&nbsp;&nbsp;
+<a href="https://huggingface.co/zzzrw/DeepMesh/tree/main"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Weights-HF-orange"></a> &nbsp;&nbsp;&nbsp;&nbsp;
+<a href='https://www.youtube.com/watch?v=6grL7bSbQ2w'><img src='https://img.shields.io/badge/Youtube-Video-b31b1b.svg'>
+
+</div>
+
+## Release
+- [6/03] 🔥🔥We released the pretrained weights for both **ShapeLLM-omni** (7B) and **3DVQVAE**.
+- [6/03] 🔥🔥We released 50k high-quality 3D edited data pairs.
+
+## Important Notes
+- Please refer to our [project_page](https://zhaorw02.github.io/DeepMesh/) for more examples.
+## Todo
+- [ ] Release of training code.
+- [ ] Release of model weights featuring multi-turn dialogue and 3D editing capabilities.
+- [ ] Release of the entire 3D-Alpaca dataset.
+
+## Acknowledgement
+Our code is based on these wonderful repos:
+* **[LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory)**
+* **[TRELLIS](https://github.com/microsoft/TRELLIS)**
+* **[PointLLM](https://github.com/OpenRobotLab/PointLLM)**
+* **[Qwen2.5-VL](https://github.com/QwenLM/Qwen2.5-VL)**
+* [LLaMA-Mesh](https://github.com/nv-tlabs/LLaMA-Mesh)
+
+

app.py

@@ -4,33 +4,84 @@ from threading import Thread
 import gradio as gr
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor,TextIteratorStreamer,AutoTokenizer
 from qwen_vl_utils import process_vision_info
+from trellis.pipelines import TrellisImageTo3DPipeline,TrellisTextTo3DPipeline
+from trellis.utils import render_utils, postprocessing_utils
 import trimesh
 from trimesh.exchange.gltf import export_glb
 import numpy as np
 import tempfile
 import copy
+import plotly.graph_objs as go
+from PIL import Image
+import plotly.express as px
+import random
+import open3d as o3d
+from huggingface_hub import hf_hub_download
 
 def _remove_image_special(text):
     text = text.replace('<ref>', '').replace('</ref>', '')
     return re.sub(r'<box>.*?(</box>|$)', '', text)
 
-
 def is_video_file(filename):
     video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
     return any(filename.lower().endswith(ext) for ext in video_extensions)
 
-def predict(_chatbot, task_history):
+def token_to_mesh(full_response):
+    d1=full_response.split("><mesh")
+    d2=[]
+    for i in range(len(d1)):
+        try:
+            if d1[i][:5]=="<mesh":
+                d2.append(int(d1[i][5:]))
+            else:
+                d2.append(int(d1[i]))
+        except:
+            pass
+    while len(d2)<1024:
+        d2.append(d2[-1])
+    encoding_indices=torch.tensor(d2).unsqueeze(0)
+    return encoding_indices
+
+def save_ply_from_array(verts):    
+    header = [
+        "ply",
+        "format ascii 1.0",
+        f"element vertex {verts.shape[0]}",
+        "property float x",
+        "property float y",
+        "property float z",
+        "end_header"
+    ]
+    tmpf = tempfile.NamedTemporaryFile(suffix=".ply", delete=False)
+    tmpf.write(("\n".join(header) + "\n").encode("utf-8"))
+    np.savetxt(tmpf, verts, fmt="%.6f")
+    tmpf.flush(); tmpf.close()
+    return tmpf.name
+
+def predict(_chatbot,task_history,viewer_voxel,viewer_mesh,task_new,seed,top_k,top_p,temperature):
+    torch.manual_seed(seed)
     chat_query = _chatbot[-1][0]
     query = task_history[-1][0]
+
     if len(chat_query) == 0:
         _chatbot.pop()
         task_history.pop()
-        return _chatbot
+        return _chatbot,task_history,viewer_voxel,viewer_mesh,task_new
     print("User: " + _parse_text(query))
     history_cp = copy.deepcopy(task_history)
     full_response = ""
     messages = []
     content = []
+
+    image_lst = []
+    for q, a in task_new:
+        if isinstance(q, (tuple, list)):
+            if not is_video_file(q[0]):
+                image_lst.append(q[0])
+            else:
+                image_lst.append(q[0])
+
+    task_new.clear()
     for q, a in history_cp:
         if isinstance(q, (tuple, list)):
             if is_video_file(q[0]):
@@ -44,45 +95,92 @@ def predict(_chatbot, task_history):
             content = []
     messages.pop()
     messages = _transform_messages(messages)
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True)
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(text=[text], images=image_inputs,
-                        videos=video_inputs, padding=True, return_tensors='pt')
+    inputs = processor(text=[text], images=image_inputs,videos=video_inputs, padding=True, return_tensors='pt')
     inputs = inputs.to(model.device)
 
-    streamer = TextIteratorStreamer(
-        tokenizer, timeout=2000.0, skip_prompt=True, skip_special_tokens=True)
-
-    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
+    eos_token_id = [tokenizer.eos_token_id,159858]
+    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs = {'max_new_tokens': 2048, 'streamer': streamer,"eos_token_id":eos_token_id,\
+                   "top_k":top_k,"top_p":top_p,"temperature":temperature,"eos_token_id":eos_token_id,**inputs}
 
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
-
-    # 初始化响应文本
     full_response = ""
-    _chatbot[-1] = (_parse_text(chat_query), "")  # 初始化空响应
-
-    # 处理流式输出
+    encoding_indices = None
+    _chatbot[-1] = (_parse_text(chat_query), "")  
     for new_text in streamer:
-        full_response += new_text
-        # 更新最后一条对话的响应部分
-        _chatbot[-1] = (_parse_text(chat_query), _parse_text(full_response))
-        yield _chatbot
+        if new_text:
+            if "<mesh" in new_text:
+                encoding_indices = token_to_mesh(new_text)
+                new_text = new_text.replace("><",",")[1:-1]
+                new_text = new_text.split("mesh-start,")[1].split(",mesh-end")[0]
+                new_text = f"mesh-start\n{new_text}\nmesh-end"
+            full_response += new_text
+            _chatbot[-1] = (_parse_text(chat_query), _parse_text(full_response))
+            yield _chatbot,viewer_voxel,viewer_mesh,task_new
 
-    # 最终处理（如果需要保存完整响应）
     task_history[-1] = (chat_query, full_response)
-    print("Model Output: " + _parse_text(full_response))
-    yield _chatbot
-    """
-    #for new_text in streamer:
-    #    yield new_text
-
-    buffer = []
-    for chunk in streamer:
-        buffer.append(chunk)
-        yield "".join(buffer)
-    """
+    yield _chatbot,viewer_voxel,viewer_mesh,task_new
+
+    if encoding_indices is not None:
+        print("processing mesh...")
+        recon = vqvae.Decode(encoding_indices.to(model.device))
+        z_s           = recon[0].detach().cpu() 
+        z_s           = (z_s>0)*1      
+        indices       = torch.nonzero(z_s[0] == 1)  
+        position_recon= (indices.float() + 0.5) / 64 - 0.5 
+        fig = make_pointcloud_figure(position_recon)
+        yield _chatbot,fig,viewer_mesh,task_new
+
+        position=position_recon
+        coords        = ((position + 0.5) * 64).int().contiguous()
+        ss            = torch.zeros(1, 64, 64, 64, dtype=torch.long)
+        ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+        ss=ss.unsqueeze(0)
+        coords = torch.argwhere(ss>0)[:, [0, 2, 3, 4]].int()
+        coords = coords.to(model.device)
+        try:
+            print("processing mesh...")
+            if len(image_lst) == 0:
+                # text to 3d
+                with torch.no_grad():
+                    prompt  = chat_query
+                    cond    = pipeline_text.get_cond([prompt])
+                    slat    = pipeline_text.sample_slat(cond, coords)
+                    outputs = pipeline_text.decode_slat(slat, ['mesh', 'gaussian'])
+
+                glb = postprocessing_utils.to_glb(
+                    outputs['gaussian'][0],
+                    outputs['mesh'][0],
+                    simplify=0.95,          
+                    texture_size=1024,    
+                    verbose=False  
+                )
+                glb.export(f"temper.glb")
+                print("processing mesh over...")
+                yield _chatbot,fig,"temper.glb"
+            else:
+                # image to 3d
+                with torch.no_grad():
+                    img = pipeline_image.preprocess_image(Image.open(image_lst[-1]))
+                    cond    = pipeline_image.get_cond([img])
+                    slat    = pipeline_image.sample_slat(cond, coords)
+                    outputs = pipeline_image.decode_slat(slat, ['mesh', 'gaussian'])
+                glb = postprocessing_utils.to_glb(
+                    outputs['gaussian'][0],
+                    outputs['mesh'][0],
+                    simplify=0.95,          
+                    texture_size=1024,    
+                    verbose=False  
+                )
+                glb.export(f"temper.glb")
+                print("processing mesh over...")
+                yield _chatbot,fig,"temper.glb",task_new
+        except:
+            print("processing mesh...bug")
+            yield _chatbot,fig,viewer_mesh,task_new
 
 def regenerate(_chatbot, task_history):
     if not task_history:
@@ -131,20 +229,45 @@ def _parse_text(text):
     text = "".join(lines)
     return text
 
-def add_text(history, task_history, text):
+def add_text_prefix(text):
+    text = f"Please generate a 3D asset based on the prompt I provided: {text}"
+    return gr.update(value=text)
+
+def token_to_words(token):
+    mesh             = "<mesh-start>"
+    for j in range(1024):
+        mesh += f"<mesh{token[j]}>"
+    mesh            += "<mesh-end>"
+    return mesh
+
+def add_text(history, task_history, text,task_new):
     task_text = text
     history = history if history is not None else []
     task_history = task_history if task_history is not None else []
     history = history + [(_parse_text(text), None)]
     task_history = task_history + [(task_text, None)]
-    return history, task_history
-
-def add_file(history, task_history, file):
+    task_new     = task_new + [(task_text, None)]
+    return history, task_history,task_new
+
+def add_file(history, task_history, file, task_new, fig, query):
+    if file.name.endswith(('.obj', '.glb')):
+        position_recon = load_vertices(file.name)#(N,3)
+
+        coords        = ((torch.from_numpy(position_recon) + 0.5) * 64).int().contiguous()
+        ss            = torch.zeros(1, 64, 64, 64, dtype=torch.long)
+        ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+        token         = vqvae.Encode(ss.to(dtype=torch.float32).unsqueeze(0).to("cuda"))
+        token         = token[0].cpu().numpy().tolist()
+        words         = token_to_words(token)
+        fig            = make_pointcloud_figure(position_recon,rotate=True)
+        return history, task_history,file.name,task_new,fig,gr.update(
+            value=f"{words}\nGive a quick overview of the object represented by this 3D mesh.")
     history = history if history is not None else []
     task_history = task_history if task_history is not None else []
     history = history + [((file.name,), None)]
     task_history = task_history + [((file.name,), None)]
-    return history, task_history
+    task_new     = task_new + [((file.name,), None)]
+    return history, task_history, file.name, task_new, fig, query
 
 def reset_user_input():
     return gr.update(value="")
@@ -153,6 +276,96 @@ def reset_state(task_history):
     task_history.clear()
     return []
 
+def make_pointcloud_figure(verts,rotate=False):
+    if rotate:
+        verts = verts.copy()
+        verts[:, 0] *= -1.0
+    N      = len(verts)
+    soft_palette = ["#FFEBEE", "#FFF3E0", "#FFFDE7", "#E8F5E9",]
+    palette = px.colors.qualitative.Set3
+    base_colors = [palette[i % len(palette)] for i in range(N)]
+    random.shuffle(base_colors)
+
+    camera = dict(
+        eye=dict(x=0.0, y=2.5, z=0.0),   
+        center=dict(x=0.0, y=0.0, z=0.0),
+        up=dict(x=0.0, y=0.0, z=1.0),    
+        projection=dict(type="orthographic")  
+    )
+
+    scatter = go.Scatter3d(
+        x=verts[:, 0],
+        y=verts[:, 1],
+        z=verts[:, 2],
+        mode='markers',
+        marker=dict(
+            size=2,           
+            color=base_colors,  
+            opacity=1,        
+            line=dict(width=1)  
+        )
+    )
+    layout = go.Layout(
+        width =700,    
+        height=200,
+        scene=dict(
+            xaxis=dict(visible=False),
+            yaxis=dict(visible=False),
+            zaxis=dict(visible=False),
+            camera=camera
+        ),
+        margin=dict(l=0, r=0, b=0, t=0)
+    )
+    fig = go.Figure(data=[scatter], layout=layout)
+    return fig
+
+def rotate_points(points, axis='x', angle_deg=90):
+    angle_rad = np.deg2rad(angle_deg)
+    if axis == 'x':
+        R = trimesh.transformations.rotation_matrix(angle_rad, [1, 0, 0])[:3, :3]
+    elif axis == 'y':
+        R = trimesh.transformations.rotation_matrix(angle_rad, [0, 1, 0])[:3, :3]
+    elif axis == 'z':
+        R = trimesh.transformations.rotation_matrix(angle_rad, [0, 0, 1])[:3, :3]
+    else:
+        raise ValueError("axis must be 'x', 'y', or 'z'")
+    return points @ R.T
+
+def convert_trimesh_to_open3d(trimesh_mesh):
+    o3d_mesh = o3d.geometry.TriangleMesh()
+    o3d_mesh.vertices = o3d.utility.Vector3dVector(
+        np.asarray(trimesh_mesh.vertices, dtype=np.float64)
+    )
+    o3d_mesh.triangles = o3d.utility.Vector3iVector(
+        np.asarray(trimesh_mesh.faces, dtype=np.int32)
+    )
+    return o3d_mesh
+
+def load_vertices(filepath):
+    mesh = trimesh.load(filepath, force='mesh')
+    mesh = convert_trimesh_to_open3d(mesh)
+    vertices = np.asarray(mesh.vertices)
+    min_vals = vertices.min()
+    max_vals = vertices.max()
+    vertices_normalized = (vertices - min_vals) / (max_vals - min_vals)  
+    vertices = vertices_normalized * 1.0 - 0.5  
+    vertices = np.clip(vertices, -0.5 + 1e-6, 0.5 - 1e-6)
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
+    vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
+    assert np.all(vertices >= 0) and np.all(vertices < 64), "Some vertices are out of bounds"
+    vertices = (vertices + 0.5) / 64 - 0.5
+    voxel = rotate_points(vertices, axis='x', angle_deg=90)
+    return voxel
+
+def add_file2(history, task_history, file,task_new):
+    history = history if history is not None else []
+    task_history = task_history if task_history is not None else []
+    history = history + [((file,), None)]
+    task_history = task_history + [((file,), None)]
+    task_new     = task_new + [((file,), None)]
+    return history, task_history,file,task_new
+
 def _transform_messages(original_messages):
     transformed_messages = []
     for message in original_messages:
@@ -173,84 +386,104 @@ def _transform_messages(original_messages):
 
     return transformed_messages
 
-# --------- Configuration & Model Loading ---------
+from trellis.models.sparse_structure_vqvae import VQVAE3D
+device       = torch.device("cuda")
+vqvae        = VQVAE3D(num_embeddings=8192)
+device       = torch.device("cuda")
+vqvae.eval()
+filepath = hf_hub_download(repo_id="yejunliang23/3DVQVAE",filename="3DVQVAE.bin")
+state_dict = torch.load(filepath, map_location="cpu")
+vqvae.load_state_dict(state_dict)
+vqvae=vqvae.to(device)
+
 MODEL_DIR = "yejunliang23/ShapeLLM-7B-omni"
-"Qwen/Qwen2.5-VL-3B-Instruct"
-# Load processor, tokenizer, model for Qwen2.5-VL
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_DIR,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True
-)
-processor = AutoProcessor.from_pretrained(MODEL_DIR)
+model_ckpt_path=MODEL_DIR
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_ckpt_path, torch_dtype="auto", device_map={"": 0})
+processor = AutoProcessor.from_pretrained(model_ckpt_path)
 tokenizer = processor.tokenizer
-#terminators = [tokenizer.eos_token_id]
-
-def chat_qwen_vl(messages: str, history: list, temperature: float = 0.1, max_new_tokens: int = 1024):
-    messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": messages},
-                ],
-            }
-        ]
-    messages = _transform_messages(messages)
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(text=[text], images=image_inputs,
-                        videos=video_inputs, padding=True, return_tensors='pt')
-    inputs = inputs.to(model.device)
-
-    streamer = TextIteratorStreamer(
-        tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-
-    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
-
-    thread = Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
-
-    #for new_text in streamer:
-    #    yield new_text
-
-    buffer = []
-    for chunk in streamer:
-        buffer.append(chunk)
-        yield "".join(buffer)
-
-
-css = """
-h1 { text-align: center; }
-"""
-PLACEHOLDER = (
-    "<div style='padding:30px;text-align:center;display:flex;flex-direction:column;align-items:center;'>"
-    "<h1 style='font-size:28px;opacity:0.55;'>Qwen2.5-VL Local Chat</h1>"
-    "<p style='font-size:18px;opacity:0.65;'>Ask anything or generate images!</p></div>"
-)
-
+from huggingface_hub import hf_hub_download
+
+pipeline_text = TrellisTextTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-text-xlarge")
+pipeline_text.to(device)
+pipeline_image = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline_image.to(device)
+
+_DESCRIPTION = '''
+* Project page of ShapeLLM-Omni: https://jamesyjl.github.io/ShapeLLM/
+* As generation tasks currently lack support for multi-turn dialogue, it's strongly recommended to clear the chat history before starting a new task
+* The model's 3D understanding is limited to shape only, so color and texture should be ignored in 3D captioning tasks
+'''
 with gr.Blocks() as demo:
-    gr.Markdown("""<center><font size=3> ShapeLLM-7B Demo </center>""")
-
-    chatbot = gr.Chatbot(label='ShapeLLM-4o', elem_classes="control-height", height=500)
-    query = gr.Textbox(lines=2, label='Input')
-    task_history = gr.State([])
-
+    gr.Markdown("# ShapeLLM-omni: A Native Multimodal LLM for 3D Generation and Understanding")
+    gr.Markdown(_DESCRIPTION)
     with gr.Row():
-        addfile_btn = gr.UploadButton("📁 Upload (上传文件)", file_types=["image", "video"])
-        submit_btn = gr.Button("🚀 Submit (发送)")
-        regen_btn = gr.Button("🤔️ Regenerate (重试)")
-        empty_bin = gr.Button("🧹 Clear History (清除历史)")
-
-    submit_btn.click(add_text, [chatbot, task_history, query], [chatbot, task_history]).then(
-        predict, [chatbot, task_history], [chatbot], show_progress=True
+        with gr.Column():
+            chatbot = gr.Chatbot(label='ShapeLLM-Omni', elem_classes="control-height", height=500)
+            seed        = gr.Number(value=42, label="seed", precision=0)
+            top_k       = gr.Slider(label="top_k",minimum=1024,maximum=8194,value=1024,step=10)
+            top_p       = gr.Slider(label="top_p",minimum=0.1,maximum=1.0,value=0.1,step=0.05)
+            temperature = gr.Slider(label="temperature",minimum=0.1,maximum=1.0,value=0.1,step=0.05)
+
+            query = gr.Textbox(lines=2, label='Input')
+            image_input = gr.Image(visible=False, type="filepath", label="Image Input")
+            with gr.Column():
+                with gr.Row():
+                    addfile_btn = gr.UploadButton("📁 Upload", file_types=["image", "video",".obj",".glb"])
+                    submit_btn = gr.Button("🚀 Submit")
+                with gr.Row():
+                    regen_btn = gr.Button("🤔️ Regenerate")
+                    empty_bin = gr.Button("🧹 Clear History")
+            task_history = gr.State([])
+            task_new     = gr.State([])
+        with gr.Column():
+            viewer_plot  = gr.Plot(label="Voxel Visual",scale=1.0)
+            viewer_mesh  = gr.Model3D(label="Mesh Visual", height=200,scale=1.0)
+
+            examples_text = gr.Examples(
+                examples=[
+                    ["A drone with four propellers and a central body."],
+                    ["A stone axe with a handle."],
+                    ["the titanic, aerial view."],
+                    ["A 3D model of a small yellow and blue robot with wheels and two pots."],
+                    ["A futuristic vehicle with a sleek design and multiple wheels."],
+                    ["A car with four wheels and a roof."],
+                ],
+                inputs=[query],
+                label="text-to-3d examples",
+                fn=add_text_prefix,
+                outputs=[query],
+                cache_examples=True,
+                )
+
+            examples_text.dataset.click(
+                fn=add_text,
+                inputs=[chatbot, task_history, query,task_new],
+                outputs=[chatbot, task_history,task_new],
+            )
+            examples_image = gr.Examples(
+                label="image-to-3d examples",
+                examples=[os.path.join("examples", i) for i in os.listdir("examples")],
+                inputs=[image_input],
+                examples_per_page = 20,
+            )
+            image_input.change(
+                fn=add_file2,
+                inputs=[chatbot, task_history, image_input,task_new],
+                outputs=[chatbot, task_history,viewer_mesh,task_new],
+                show_progress=True
+            )
+
+    submit_btn.click(add_text, [chatbot, task_history, query,task_new],\
+                               [chatbot, task_history,task_new]).then(
+        predict, [chatbot, task_history,viewer_plot,viewer_mesh,task_new,seed,top_k,top_p,temperature],\
+                 [chatbot,viewer_plot,viewer_mesh,task_new], show_progress=True
     )
     submit_btn.click(reset_user_input, [], [query])
     empty_bin.click(reset_state, [task_history], [chatbot], show_progress=True)
-    regen_btn.click(regenerate, [chatbot, task_history], [chatbot], show_progress=True)
-    addfile_btn.upload(add_file, [chatbot, task_history, addfile_btn], [chatbot, task_history], show_progress=True)
+    regen_btn.click(regenerate,  [chatbot, task_history], [chatbot], show_progress=True)
+    addfile_btn.upload(add_file, [chatbot, task_history, addfile_btn, task_new, viewer_plot, query],\
+                                 [chatbot, task_history, viewer_mesh, task_new, viewer_plot, query],\
+                                  show_progress=True)
     
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()

configs/generation/slat_flow_img_dit_L_64l8p2_fp16.json

@@ -0,0 +1,102 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1024,
+                "cond_channels": 1024,
+                "num_blocks": 24,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 2,
+                "num_io_res_blocks": 2,
+                "io_block_channels": [128],
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "ImageConditionedSLat",
+        "args": {
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768,
+            "image_size": 518,
+            "normalization": {
+                "mean": [
+                    -2.1687545776367188,
+                    -0.004347046371549368,
+                    -0.13352349400520325,
+                    -0.08418072760105133,
+                    -0.5271206498146057,
+                    0.7238689064979553,
+                    -1.1414450407028198,
+                    1.2039363384246826
+                ],
+                "std": [
+                    2.377650737762451,
+                    2.386378288269043,
+                    2.124418020248413,
+                    2.1748552322387695,
+                    2.663944721221924,
+                    2.371192216873169,
+                    2.6217446327209473,
+                    2.684523105621338
+                ]
+            },
+            "pretrained_slat_dec": "microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 4,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "image_cond_model": "dinov2_vitl14_reg"
+        }
+    }
+}

configs/generation/slat_flow_txt_dit_B_64l8p2_fp16.json

@@ -0,0 +1,101 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 768,
+                "cond_channels": 768,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "patch_size": 2,
+                "num_io_res_blocks": 2,
+                "io_block_channels": [128],
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSLat",
+        "args": {
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768,
+            "normalization": {
+                "mean": [
+                    -2.1687545776367188,
+                    -0.004347046371549368,
+                    -0.13352349400520325,
+                    -0.08418072760105133,
+                    -0.5271206498146057,
+                    0.7238689064979553,
+                    -1.1414450407028198,
+                    1.2039363384246826
+                ],
+                "std": [
+                    2.377650737762451,
+                    2.386378288269043,
+                    2.124418020248413,
+                    2.1748552322387695,
+                    2.663944721221924,
+                    2.371192216873169,
+                    2.6217446327209473,
+                    2.684523105621338
+                ]
+            },
+            "pretrained_slat_dec": "microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 16,
+            "batch_split": 4,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/generation/slat_flow_txt_dit_L_64l8p2_fp16.json

@@ -0,0 +1,101 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1024,
+                "cond_channels": 768,
+                "num_blocks": 24,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 2,
+                "num_io_res_blocks": 2,
+                "io_block_channels": [128],
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSLat",
+        "args": {
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768,
+            "normalization": {
+                "mean": [
+                    -2.1687545776367188,
+                    -0.004347046371549368,
+                    -0.13352349400520325,
+                    -0.08418072760105133,
+                    -0.5271206498146057,
+                    0.7238689064979553,
+                    -1.1414450407028198,
+                    1.2039363384246826
+                ],
+                "std": [
+                    2.377650737762451,
+                    2.386378288269043,
+                    2.124418020248413,
+                    2.1748552322387695,
+                    2.663944721221924,
+                    2.371192216873169,
+                    2.6217446327209473,
+                    2.684523105621338
+                ]
+            },
+            "pretrained_slat_dec": "microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 4,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/generation/slat_flow_txt_dit_XL_64l8p2_fp16.json

@@ -0,0 +1,101 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1280,
+                "cond_channels": 768,
+                "num_blocks": 28,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 2,
+                "num_io_res_blocks": 3,
+                "io_block_channels": [256],
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSLat",
+        "args": {
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768,
+            "normalization": {
+                "mean": [
+                    -2.1687545776367188,
+                    -0.004347046371549368,
+                    -0.13352349400520325,
+                    -0.08418072760105133,
+                    -0.5271206498146057,
+                    0.7238689064979553,
+                    -1.1414450407028198,
+                    1.2039363384246826
+                ],
+                "std": [
+                    2.377650737762451,
+                    2.386378288269043,
+                    2.124418020248413,
+                    2.1748552322387695,
+                    2.663944721221924,
+                    2.371192216873169,
+                    2.6217446327209473,
+                    2.684523105621338
+                ]
+            },
+            "pretrained_slat_dec": "microsoft/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 4,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/generation/ss_flow_img_dit_L_16l8_fp16.json

@@ -0,0 +1,70 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1024,
+                "cond_channels": 1024,
+                "num_blocks": 24,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 1,
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "ImageConditionedSparseStructureLatent",
+        "args": {
+            "latent_model": "ss_enc_conv3d_16l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "image_size": 518,
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 1,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "image_cond_model": "dinov2_vitl14_reg"
+        }
+    }
+}

configs/generation/ss_flow_txt_dit_B_16l8_fp16.json

@@ -0,0 +1,69 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 768,
+                "cond_channels": 768,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "patch_size": 1,
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSparseStructureLatent",
+        "args": {
+            "latent_model": "ss_enc_conv3d_16l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 16,
+            "batch_split": 1,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/generation/ss_flow_txt_dit_L_16l8_fp16.json

@@ -0,0 +1,69 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1024,
+                "cond_channels": 768,
+                "num_blocks": 24,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 1,
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSparseStructureLatent",
+        "args": {
+            "latent_model": "ss_enc_conv3d_16l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 1,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/generation/ss_flow_txt_dit_XL_16l8_fp16.json

@@ -0,0 +1,70 @@
+{
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1280,
+                "cond_channels": 768,
+                "num_blocks": 28,
+                "num_heads": 16,
+                "mlp_ratio": 4,
+                "patch_size": 1,
+                "pe_mode": "ape",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "TextConditionedSparseStructureLatent",
+        "args": {
+            "latent_model": "ss_enc_conv3d_16l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "TextConditionedFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 1,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 0.0001,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "text_cond_model": "openai/clip-vit-large-patch14"
+        }
+    }
+}

configs/vae/slat_vae_dec_mesh_swin8_B_64l8_fp16.json

@@ -0,0 +1,73 @@
+{
+    "models": {
+        "decoder": {
+            "name": "ElasticSLatMeshDecoder",
+            "args": {
+                "resolution": 64,
+                "model_channels": 768,
+                "latent_channels": 8,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "attn_mode": "swin",
+                "window_size": 8,
+                "use_fp16": true,
+                "representation_config": {
+                    "use_color": true
+                }
+            }
+        }
+    },
+    "dataset": {
+        "name": "Slat2RenderGeo",
+        "args": {
+            "image_size": 512,
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768
+        }
+    },
+    "trainer": {
+        "name": "SLatVaeMeshDecoderTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 4,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "lambda_ssim": 0.2,
+            "lambda_lpips": 0.2,
+            "lambda_tsdf": 0.01,
+            "lambda_depth": 10.0,
+            "lambda_color": 0.1,
+            "depth_loss_type": "smooth_l1"
+        }
+    }
+}

configs/vae/slat_vae_dec_rf_swin8_B_64l8_fp16.json

@@ -0,0 +1,71 @@
+{
+    "models": {
+        "decoder": {
+            "name": "ElasticSLatRadianceFieldDecoder",
+            "args": {
+                "resolution": 64,
+                "model_channels": 768,
+                "latent_channels": 8,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "attn_mode": "swin",
+                "window_size": 8,
+                "use_fp16": true,
+                "representation_config": {
+                    "rank": 16,
+                    "dim": 8
+                }
+            }
+        }
+    },
+    "dataset": {
+        "name": "SLat2Render",
+        "args": {
+            "image_size": 512,
+            "latent_model": "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16",
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768
+        }
+    },
+    "trainer": {
+        "name": "SLatVaeRadianceFieldDecoderTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 2,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "loss_type": "l1",
+            "lambda_ssim": 0.2,
+            "lambda_lpips": 0.2
+        }
+    }
+}

configs/vae/slat_vae_enc_dec_gs_swin8_B_64l8_fp16.json

@@ -0,0 +1,105 @@
+{
+    "models": {
+        "encoder": {
+            "name": "ElasticSLatEncoder",
+            "args": {
+                "resolution": 64,
+                "in_channels": 1024,
+                "model_channels": 768,
+                "latent_channels": 8,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "attn_mode": "swin",
+                "window_size": 8,
+                "use_fp16": true
+            }
+        },
+        "decoder": {
+            "name": "ElasticSLatGaussianDecoder",
+            "args": {
+                "resolution": 64,
+                "model_channels": 768,
+                "latent_channels": 8,
+                "num_blocks": 12,
+                "num_heads": 12,
+                "mlp_ratio": 4,
+                "attn_mode": "swin",
+                "window_size": 8,
+                "use_fp16": true,
+                "representation_config": {
+                    "lr": {
+                        "_xyz": 1.0,
+                        "_features_dc": 1.0,
+                        "_opacity": 1.0,
+                        "_scaling": 1.0,
+                        "_rotation": 0.1
+                    },
+                    "perturb_offset": true,
+                    "voxel_size": 1.5,
+                    "num_gaussians": 32,
+                    "2d_filter_kernel_size": 0.1,
+                    "3d_filter_kernel_size": 9e-4,
+                    "scaling_bias": 4e-3,
+                    "opacity_bias": 0.1,
+                    "scaling_activation": "softplus"
+                }
+            }
+        }
+    },
+    "dataset": {
+        "name": "SparseFeat2Render",
+        "args": {
+            "image_size": 512,
+            "model": "dinov2_vitl14_reg",
+            "resolution": 64,
+            "min_aesthetic_score": 4.5,
+            "max_num_voxels": 32768
+        }
+    },
+    "trainer": {
+        "name": "SLatVaeGaussianTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 2,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "loss_type": "l1",
+            "lambda_ssim": 0.2,
+            "lambda_lpips": 0.2,
+            "lambda_kl": 1e-06,
+            "regularizations": {
+                "lambda_vol": 10000.0,
+                "lambda_opacity": 0.001
+            }
+        }
+    }
+}

configs/vae/ss_vae_conv3d_16l8_fp16.json

@@ -0,0 +1,65 @@
+{
+    "models": {
+        "encoder": {
+            "name": "SparseStructureEncoder",
+            "args": {
+                "in_channels": 1,
+                "latent_channels": 8,
+                "num_res_blocks": 2,
+                "num_res_blocks_middle": 2,
+                "channels": [32, 128, 512],
+                "use_fp16": true
+            }
+        },
+        "decoder": {
+            "name": "SparseStructureDecoder",
+            "args": {
+                "out_channels": 1,
+                "latent_channels": 8,
+                "num_res_blocks": 2,
+                "num_res_blocks_middle": 2,
+                "channels": [512, 128, 32],
+                "use_fp16": true
+            }
+        }
+    },
+    "dataset": {
+        "name": "SparseStructure",
+        "args": {
+            "resolution": 64,
+            "min_aesthetic_score": 4.5
+        }
+    },
+    "trainer": {
+        "name": "SparseStructureVaeTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 4,
+            "batch_split": 1,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.0
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "fp16_mode": "inflat_all",
+            "fp16_scale_growth": 0.001,
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_log": 500,
+            "i_sample": 10000,
+            "i_save": 10000,
+            "loss_type": "dice",
+            "lambda_kl": 0.001
+        }
+    }
+}

dataset_toolkits/blender_script/render.py

@@ -0,0 +1,528 @@
+import argparse, sys, os, math, re, glob
+from typing import *
+import bpy
+from mathutils import Vector, Matrix
+import numpy as np
+import json
+import glob
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+EXT = {
+    'PNG': 'png',
+    'JPEG': 'jpg',
+    'OPEN_EXR': 'exr',
+    'TIFF': 'tiff',
+    'BMP': 'bmp',
+    'HDR': 'hdr',
+    'TARGA': 'tga'
+}
+
+def init_render(engine='CYCLES', resolution=512, geo_mode=False):
+    bpy.context.scene.render.engine = engine
+    bpy.context.scene.render.resolution_x = resolution
+    bpy.context.scene.render.resolution_y = resolution
+    bpy.context.scene.render.resolution_percentage = 100
+    bpy.context.scene.render.image_settings.file_format = 'PNG'
+    bpy.context.scene.render.image_settings.color_mode = 'RGBA'
+    bpy.context.scene.render.film_transparent = True
+    
+    bpy.context.scene.cycles.device = 'GPU'
+    bpy.context.scene.cycles.samples = 128 if not geo_mode else 1
+    bpy.context.scene.cycles.filter_type = 'BOX'
+    bpy.context.scene.cycles.filter_width = 1
+    bpy.context.scene.cycles.diffuse_bounces = 1
+    bpy.context.scene.cycles.glossy_bounces = 1
+    bpy.context.scene.cycles.transparent_max_bounces = 3 if not geo_mode else 0
+    bpy.context.scene.cycles.transmission_bounces = 3 if not geo_mode else 1
+    bpy.context.scene.cycles.use_denoising = True
+        
+    bpy.context.preferences.addons['cycles'].preferences.get_devices()
+    bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+    
+def init_nodes(save_depth=False, save_normal=False, save_albedo=False, save_mist=False):
+    if not any([save_depth, save_normal, save_albedo, save_mist]):
+        return {}, {}
+    outputs = {}
+    spec_nodes = {}
+    
+    bpy.context.scene.use_nodes = True
+    bpy.context.scene.view_layers['View Layer'].use_pass_z = save_depth
+    bpy.context.scene.view_layers['View Layer'].use_pass_normal = save_normal
+    bpy.context.scene.view_layers['View Layer'].use_pass_diffuse_color = save_albedo
+    bpy.context.scene.view_layers['View Layer'].use_pass_mist = save_mist
+    
+    nodes = bpy.context.scene.node_tree.nodes
+    links = bpy.context.scene.node_tree.links
+    for n in nodes:
+        nodes.remove(n)
+    
+    render_layers = nodes.new('CompositorNodeRLayers')
+    
+    if save_depth:
+        depth_file_output = nodes.new('CompositorNodeOutputFile')
+        depth_file_output.base_path = ''
+        depth_file_output.file_slots[0].use_node_format = True
+        depth_file_output.format.file_format = 'PNG'
+        depth_file_output.format.color_depth = '16'
+        depth_file_output.format.color_mode = 'BW'
+        # Remap to 0-1
+        map = nodes.new(type="CompositorNodeMapRange")
+        map.inputs[1].default_value = 0  # (min value you will be getting)
+        map.inputs[2].default_value = 10 # (max value you will be getting)
+        map.inputs[3].default_value = 0  # (min value you will map to)
+        map.inputs[4].default_value = 1  # (max value you will map to)
+        
+        links.new(render_layers.outputs['Depth'], map.inputs[0])
+        links.new(map.outputs[0], depth_file_output.inputs[0])
+        
+        outputs['depth'] = depth_file_output
+        spec_nodes['depth_map'] = map
+    
+    if save_normal:
+        normal_file_output = nodes.new('CompositorNodeOutputFile')
+        normal_file_output.base_path = ''
+        normal_file_output.file_slots[0].use_node_format = True
+        normal_file_output.format.file_format = 'OPEN_EXR'
+        normal_file_output.format.color_mode = 'RGB'
+        normal_file_output.format.color_depth = '16'
+        
+        links.new(render_layers.outputs['Normal'], normal_file_output.inputs[0])
+        
+        outputs['normal'] = normal_file_output
+    
+    if save_albedo:
+        albedo_file_output = nodes.new('CompositorNodeOutputFile')
+        albedo_file_output.base_path = ''
+        albedo_file_output.file_slots[0].use_node_format = True
+        albedo_file_output.format.file_format = 'PNG'
+        albedo_file_output.format.color_mode = 'RGBA'
+        albedo_file_output.format.color_depth = '8'
+        
+        alpha_albedo = nodes.new('CompositorNodeSetAlpha')
+        
+        links.new(render_layers.outputs['DiffCol'], alpha_albedo.inputs['Image'])
+        links.new(render_layers.outputs['Alpha'], alpha_albedo.inputs['Alpha'])
+        links.new(alpha_albedo.outputs['Image'], albedo_file_output.inputs[0])
+        
+        outputs['albedo'] = albedo_file_output
+        
+    if save_mist:
+        bpy.data.worlds['World'].mist_settings.start = 0
+        bpy.data.worlds['World'].mist_settings.depth = 10
+        
+        mist_file_output = nodes.new('CompositorNodeOutputFile')
+        mist_file_output.base_path = ''
+        mist_file_output.file_slots[0].use_node_format = True
+        mist_file_output.format.file_format = 'PNG'
+        mist_file_output.format.color_mode = 'BW'
+        mist_file_output.format.color_depth = '16'
+        
+        links.new(render_layers.outputs['Mist'], mist_file_output.inputs[0])
+        
+        outputs['mist'] = mist_file_output
+        
+    return outputs, spec_nodes
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # delete all the materials
+    for material in bpy.data.materials:
+        bpy.data.materials.remove(material, do_unlink=True)
+
+    # delete all the textures
+    for texture in bpy.data.textures:
+        bpy.data.textures.remove(texture, do_unlink=True)
+
+    # delete all the images
+    for image in bpy.data.images:
+        bpy.data.images.remove(image, do_unlink=True)
+
+def init_camera():
+    cam = bpy.data.objects.new('Camera', bpy.data.cameras.new('Camera'))
+    bpy.context.collection.objects.link(cam)
+    bpy.context.scene.camera = cam
+    cam.data.sensor_height = cam.data.sensor_width = 32
+    cam_constraint = cam.constraints.new(type='TRACK_TO')
+    cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
+    cam_constraint.up_axis = 'UP_Y'
+    cam_empty = bpy.data.objects.new("Empty", None)
+    cam_empty.location = (0, 0, 0)
+    bpy.context.scene.collection.objects.link(cam_empty)
+    cam_constraint.target = cam_empty
+    return cam
+
+def init_lighting():
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+    
+    # Create key light
+    default_light = bpy.data.objects.new("Default_Light", bpy.data.lights.new("Default_Light", type="POINT"))
+    bpy.context.collection.objects.link(default_light)
+    default_light.data.energy = 1000
+    default_light.location = (4, 1, 6)
+    default_light.rotation_euler = (0, 0, 0)
+    
+    # create top light
+    top_light = bpy.data.objects.new("Top_Light", bpy.data.lights.new("Top_Light", type="AREA"))
+    bpy.context.collection.objects.link(top_light)
+    top_light.data.energy = 10000
+    top_light.location = (0, 0, 10)
+    top_light.scale = (100, 100, 100)
+    
+    # create bottom light
+    bottom_light = bpy.data.objects.new("Bottom_Light", bpy.data.lights.new("Bottom_Light", type="AREA"))
+    bpy.context.collection.objects.link(bottom_light)
+    bottom_light.data.energy = 1000
+    bottom_light.location = (0, 0, -10)
+    bottom_light.rotation_euler = (0, 0, 0)
+    
+    return {
+        "default_light": default_light,
+        "top_light": top_light,
+        "bottom_light": bottom_light
+    }
+
+
+def load_object(object_path: str) -> None:
+    """Loads a model with a supported file extension into the scene.
+
+    Args:
+        object_path (str): Path to the model file.
+
+    Raises:
+        ValueError: If the file extension is not supported.
+
+    Returns:
+        None
+    """
+    file_extension = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS')
+    else:
+        import_function(filepath=object_path)
+        
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    # bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            obj.select_set(True)
+    bpy.ops.object.delete()
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+        
+def split_mesh_normal():
+    bpy.ops.object.select_all(action="DESELECT")
+    objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+    bpy.context.view_layer.objects.active = objs[0]
+    for obj in objs:
+        obj.select_set(True)
+    bpy.ops.object.mode_set(mode="EDIT")
+    bpy.ops.mesh.select_all(action='SELECT')
+    bpy.ops.mesh.split_normals()
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.select_all(action="DESELECT")
+            
+def delete_custom_normals():
+     for this_obj in bpy.data.objects:
+        if this_obj.type == "MESH":
+            bpy.context.view_layer.objects.active = this_obj
+            bpy.ops.mesh.customdata_custom_splitnormals_clear()
+
+def override_material():
+    new_mat = bpy.data.materials.new(name="Override0123456789")
+    new_mat.use_nodes = True
+    new_mat.node_tree.nodes.clear()
+    bsdf = new_mat.node_tree.nodes.new('ShaderNodeBsdfDiffuse')
+    bsdf.inputs[0].default_value = (0.5, 0.5, 0.5, 1)
+    bsdf.inputs[1].default_value = 1
+    output = new_mat.node_tree.nodes.new('ShaderNodeOutputMaterial')
+    new_mat.node_tree.links.new(bsdf.outputs['BSDF'], output.inputs['Surface'])
+    bpy.context.scene.view_layers['View Layer'].material_override = new_mat
+
+def unhide_all_objects() -> None:
+    """Unhides all objects in the scene.
+
+    Returns:
+        None
+    """
+    for obj in bpy.context.scene.objects:
+        obj.hide_set(False)
+        
+def convert_to_meshes() -> None:
+    """Converts all objects in the scene to meshes.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.context.view_layer.objects.active = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"][0]
+    for obj in bpy.context.scene.objects:
+        obj.select_set(True)
+    bpy.ops.object.convert(target="MESH")
+        
+def triangulate_meshes() -> None:
+    """Triangulates all meshes in the scene.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+    bpy.context.view_layer.objects.active = objs[0]
+    for obj in objs:
+        obj.select_set(True)
+    bpy.ops.object.mode_set(mode="EDIT")
+    bpy.ops.mesh.reveal()
+    bpy.ops.mesh.select_all(action="SELECT")
+    bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
+    bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+
+def scene_bbox() -> Tuple[Vector, Vector]:
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            coord = obj.matrix_world @ coord
+            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+    return Vector(bbox_min), Vector(bbox_max)
+
+def normalize_scene() -> Tuple[float, Vector]:
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    bpy.ops.object.select_all(action="DESELECT")
+    
+    return scale, offset
+
+def get_transform_matrix(obj: bpy.types.Object) -> list:
+    pos, rt, _ = obj.matrix_world.decompose()
+    rt = rt.to_matrix()
+    matrix = []
+    for ii in range(3):
+        a = []
+        for jj in range(3):
+            a.append(rt[ii][jj])
+        a.append(pos[ii])
+        matrix.append(a)
+    matrix.append([0, 0, 0, 1])
+    return matrix
+
+def main(arg):
+    os.makedirs(arg.output_folder, exist_ok=True)
+    
+    # Initialize context
+    init_render(engine=arg.engine, resolution=arg.resolution, geo_mode=arg.geo_mode)
+    outputs, spec_nodes = init_nodes(
+        save_depth=arg.save_depth,
+        save_normal=arg.save_normal,
+        save_albedo=arg.save_albedo,
+        save_mist=arg.save_mist
+    )
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+        if arg.split_normal:
+            split_mesh_normal()
+        # delete_custom_normals()
+    print('[INFO] Scene initialized.')
+    
+    # normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Initialize camera and lighting
+    cam = init_camera()
+    init_lighting()
+    print('[INFO] Camera and lighting initialized.')
+
+    # Override material
+    if arg.geo_mode:
+        override_material()
+    
+    # Create a list of views
+    to_export = {
+        "aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        "scale": scale,
+        "offset": [offset.x, offset.y, offset.z],
+        "frames": []
+    }
+    views = json.loads(arg.views)
+    for i, view in enumerate(views):
+        cam.location = (
+            view['radius'] * np.cos(view['yaw']) * np.cos(view['pitch']),
+            view['radius'] * np.sin(view['yaw']) * np.cos(view['pitch']),
+            view['radius'] * np.sin(view['pitch'])
+        )
+        cam.data.lens = 16 / np.tan(view['fov'] / 2)
+        
+        if arg.save_depth:
+            spec_nodes['depth_map'].inputs[1].default_value = view['radius'] - 0.5 * np.sqrt(3)
+            spec_nodes['depth_map'].inputs[2].default_value = view['radius'] + 0.5 * np.sqrt(3)
+        
+        bpy.context.scene.render.filepath = os.path.join(arg.output_folder, f'{i:03d}.png')
+        for name, output in outputs.items():
+            output.file_slots[0].path = os.path.join(arg.output_folder, f'{i:03d}_{name}')
+            
+        # Render the scene
+        bpy.ops.render.render(write_still=True)
+        bpy.context.view_layer.update()
+        for name, output in outputs.items():
+            ext = EXT[output.format.file_format]
+            path = glob.glob(f'{output.file_slots[0].path}*.{ext}')[0]
+            os.rename(path, f'{output.file_slots[0].path}.{ext}')
+            
+        # Save camera parameters
+        metadata = {
+            "file_path": f'{i:03d}.png',
+            "camera_angle_x": view['fov'],
+            "transform_matrix": get_transform_matrix(cam)
+        }
+        if arg.save_depth:
+            metadata['depth'] = {
+                'min': view['radius'] - 0.5 * np.sqrt(3),
+                'max': view['radius'] + 0.5 * np.sqrt(3)
+            }
+        to_export["frames"].append(metadata)
+    
+    # Save the camera parameters
+    with open(os.path.join(arg.output_folder, 'transforms.json'), 'w') as f:
+        json.dump(to_export, f, indent=4)
+        
+    if arg.save_mesh:
+        # triangulate meshes
+        unhide_all_objects()
+        convert_to_meshes()
+        triangulate_meshes()
+        print('[INFO] Meshes triangulated.')
+        
+        # export ply mesh
+        bpy.ops.export_mesh.ply(filepath=os.path.join(arg.output_folder, 'mesh.ply'))
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--views', type=str, help='JSON string of views. Contains a list of {yaw, pitch, radius, fov} object.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--output_folder', type=str, default='/tmp', help='The path the output will be dumped to.')
+    parser.add_argument('--resolution', type=int, default=512, help='Resolution of the images.')
+    parser.add_argument('--engine', type=str, default='CYCLES', help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
+    parser.add_argument('--geo_mode', action='store_true', help='Geometry mode for rendering.')
+    parser.add_argument('--save_depth', action='store_true', help='Save the depth maps.')
+    parser.add_argument('--save_normal', action='store_true', help='Save the normal maps.')
+    parser.add_argument('--save_albedo', action='store_true', help='Save the albedo maps.')
+    parser.add_argument('--save_mist', action='store_true', help='Save the mist distance maps.')
+    parser.add_argument('--split_normal', action='store_true', help='Split the normals of the mesh.')
+    parser.add_argument('--save_mesh', action='store_true', help='Save the mesh as a .ply file.')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    

dataset_toolkits/build_metadata.py

@@ -0,0 +1,270 @@
+import os
+import shutil
+import sys
+import time
+import importlib
+import argparse
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+import utils3d
+
+def get_first_directory(path):  
+    with os.scandir(path) as it:  
+        for entry in it:  
+            if entry.is_dir():  
+                return entry.name  
+    return None
+
+def need_process(key):
+    return key in opt.field or opt.field == ['all']
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--field', type=str, default='all',
+                        help='Fields to process, separated by commas')
+    parser.add_argument('--from_file', action='store_true',
+                        help='Build metadata from file instead of from records of processings.' +
+                             'Useful when some processing fail to generate records but file already exists.')
+    dataset_utils.add_args(parser)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(opt.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(opt.output_dir, 'merged_records'), exist_ok=True)
+
+    opt.field = opt.field.split(',')
+    
+    timestamp = str(int(time.time()))
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        print('Loading previous metadata...')
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        metadata = dataset_utils.get_metadata(**opt)
+    metadata.set_index('sha256', inplace=True)
+    
+    # merge downloaded
+    df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('downloaded_') and f.endswith('.csv')]
+    df_parts = []
+    for f in df_files:
+        try:
+            df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+        except:
+            pass
+    if len(df_parts) > 0:
+        df = pd.concat(df_parts)
+        df.set_index('sha256', inplace=True)
+        if 'local_path' in metadata.columns:
+            metadata.update(df, overwrite=True)
+        else:
+            metadata = metadata.join(df, on='sha256', how='left')
+        for f in df_files:
+            shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+            
+    # detect models
+    image_models = []
+    if os.path.exists(os.path.join(opt.output_dir, 'features')):
+        image_models = os.listdir(os.path.join(opt.output_dir, 'features'))
+    latent_models = []
+    if os.path.exists(os.path.join(opt.output_dir, 'latents')):
+        latent_models = os.listdir(os.path.join(opt.output_dir, 'latents'))
+    ss_latent_models = []
+    if os.path.exists(os.path.join(opt.output_dir, 'ss_latents')):
+        ss_latent_models = os.listdir(os.path.join(opt.output_dir, 'ss_latents'))
+    print(f'Image models: {image_models}')
+    print(f'Latent models: {latent_models}')
+    print(f'Sparse Structure latent models: {ss_latent_models}')
+
+    if 'rendered' not in metadata.columns:
+        metadata['rendered'] = [False] * len(metadata)
+    if 'voxelized' not in metadata.columns:
+        metadata['voxelized'] = [False] * len(metadata)
+    if 'num_voxels' not in metadata.columns:
+        metadata['num_voxels'] = [0] * len(metadata)
+    if 'cond_rendered' not in metadata.columns:
+        metadata['cond_rendered'] = [False] * len(metadata)
+    for model in image_models:
+        if f'feature_{model}' not in metadata.columns:
+            metadata[f'feature_{model}'] = [False] * len(metadata)
+    for model in latent_models:
+        if f'latent_{model}' not in metadata.columns:
+            metadata[f'latent_{model}'] = [False] * len(metadata)
+    for model in ss_latent_models:
+        if f'ss_latent_{model}' not in metadata.columns:
+            metadata[f'ss_latent_{model}'] = [False] * len(metadata)
+    
+    # merge rendered
+    df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('rendered_') and f.endswith('.csv')]
+    df_parts = []
+    for f in df_files:
+        try:
+            df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+        except:
+            pass
+    if len(df_parts) > 0:
+        df = pd.concat(df_parts)
+        df.set_index('sha256', inplace=True)
+        metadata.update(df, overwrite=True)
+        for f in df_files:
+            shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+    
+    # merge voxelized
+    df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('voxelized_') and f.endswith('.csv')]
+    df_parts = []
+    for f in df_files:
+        try:
+            df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+        except:
+            pass
+    if len(df_parts) > 0:
+        df = pd.concat(df_parts)
+        df.set_index('sha256', inplace=True)
+        metadata.update(df, overwrite=True)
+        for f in df_files:
+            shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+    
+    # merge cond_rendered
+    df_files = [f for f in os.listdir(opt.output_dir) if f.startswith('cond_rendered_') and f.endswith('.csv')]
+    df_parts = []
+    for f in df_files:
+        try:
+            df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+        except:
+            pass
+    if len(df_parts) > 0:
+        df = pd.concat(df_parts)
+        df.set_index('sha256', inplace=True)
+        metadata.update(df, overwrite=True)
+        for f in df_files:
+            shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+    
+    # merge features
+    for model in image_models:
+        df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'feature_{model}_') and f.endswith('.csv')]
+        df_parts = []
+        for f in df_files:
+            try:
+                df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+            except:
+                pass
+        if len(df_parts) > 0:
+            df = pd.concat(df_parts)
+            df.set_index('sha256', inplace=True)
+            metadata.update(df, overwrite=True)
+            for f in df_files:
+                shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+                
+    # merge latents
+    for model in latent_models:
+        df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'latent_{model}_') and f.endswith('.csv')]
+        df_parts = []
+        for f in df_files:
+            try:
+                df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+            except:
+                pass
+        if len(df_parts) > 0:
+            df = pd.concat(df_parts)
+            df.set_index('sha256', inplace=True)
+            metadata.update(df, overwrite=True)
+            for f in df_files:
+                shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+                
+    # merge sparse structure latents
+    for model in ss_latent_models:
+        df_files = [f for f in os.listdir(opt.output_dir) if f.startswith(f'ss_latent_{model}_') and f.endswith('.csv')]
+        df_parts = []
+        for f in df_files:
+            try:
+                df_parts.append(pd.read_csv(os.path.join(opt.output_dir, f)))
+            except:
+                pass
+        if len(df_parts) > 0:
+            df = pd.concat(df_parts)
+            df.set_index('sha256', inplace=True)
+            metadata.update(df, overwrite=True)
+            for f in df_files:
+                shutil.move(os.path.join(opt.output_dir, f), os.path.join(opt.output_dir, 'merged_records', f'{timestamp}_{f}'))
+
+    # build metadata from files
+    if opt.from_file:
+        with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+            tqdm(total=len(metadata), desc="Building metadata") as pbar:
+            def worker(sha256):
+                try:
+                    if need_process('rendered') and metadata.loc[sha256, 'rendered'] == False and \
+                        os.path.exists(os.path.join(opt.output_dir, 'renders', sha256, 'transforms.json')):
+                        metadata.loc[sha256, 'rendered'] = True
+                    if need_process('voxelized') and metadata.loc[sha256, 'rendered'] == True and metadata.loc[sha256, 'voxelized'] == False and \
+                        os.path.exists(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply')):
+                        try:
+                            pts = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0]
+                            metadata.loc[sha256, 'voxelized'] = True
+                            metadata.loc[sha256, 'num_voxels'] = len(pts)
+                        except Exception as e:
+                            pass
+                    if need_process('cond_rendered') and metadata.loc[sha256, 'cond_rendered'] == False and \
+                        os.path.exists(os.path.join(opt.output_dir, 'renders_cond', sha256, 'transforms.json')):
+                        metadata.loc[sha256, 'cond_rendered'] = True
+                    for model in image_models:
+                        if need_process(f'feature_{model}') and \
+                            metadata.loc[sha256, f'feature_{model}'] == False and \
+                            metadata.loc[sha256, 'rendered'] == True and \
+                            metadata.loc[sha256, 'voxelized'] == True and \
+                            os.path.exists(os.path.join(opt.output_dir, 'features', model, f'{sha256}.npz')):
+                            metadata.loc[sha256, f'feature_{model}'] = True
+                    for model in latent_models:
+                        if need_process(f'latent_{model}') and \
+                            metadata.loc[sha256, f'latent_{model}'] == False and \
+                            metadata.loc[sha256, 'rendered'] == True and \
+                            metadata.loc[sha256, 'voxelized'] == True and \
+                            os.path.exists(os.path.join(opt.output_dir, 'latents', model, f'{sha256}.npz')):
+                            metadata.loc[sha256, f'latent_{model}'] = True
+                    for model in ss_latent_models:
+                        if need_process(f'ss_latent_{model}') and \
+                            metadata.loc[sha256, f'ss_latent_{model}'] == False and \
+                            metadata.loc[sha256, 'voxelized'] == True and \
+                            os.path.exists(os.path.join(opt.output_dir, 'ss_latents', model, f'{sha256}.npz')):
+                            metadata.loc[sha256, f'ss_latent_{model}'] = True
+                    pbar.update()
+                except Exception as e:
+                    print(f'Error processing {sha256}: {e}')
+                    pbar.update()
+            
+            executor.map(worker, metadata.index)
+            executor.shutdown(wait=True)
+
+    # statistics
+    metadata.to_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    num_downloaded = metadata['local_path'].count() if 'local_path' in metadata.columns else 0
+    with open(os.path.join(opt.output_dir, 'statistics.txt'), 'w') as f:
+        f.write('Statistics:\n')
+        f.write(f'  - Number of assets: {len(metadata)}\n')
+        f.write(f'  - Number of assets downloaded: {num_downloaded}\n')
+        f.write(f'  - Number of assets rendered: {metadata["rendered"].sum()}\n')
+        f.write(f'  - Number of assets voxelized: {metadata["voxelized"].sum()}\n')
+        if len(image_models) != 0:
+            f.write(f'  - Number of assets with image features extracted:\n')
+            for model in image_models:
+                f.write(f'    - {model}: {metadata[f"feature_{model}"].sum()}\n')
+        if len(latent_models) != 0:
+            f.write(f'  - Number of assets with latents extracted:\n')
+            for model in latent_models:
+                f.write(f'    - {model}: {metadata[f"latent_{model}"].sum()}\n')
+        if len(ss_latent_models) != 0:
+            f.write(f'  - Number of assets with sparse structure latents extracted:\n')
+            for model in ss_latent_models:
+                f.write(f'    - {model}: {metadata[f"ss_latent_{model}"].sum()}\n')
+        f.write(f'  - Number of assets with captions: {metadata["captions"].count()}\n')
+        f.write(f'  - Number of assets with image conditions: {metadata["cond_rendered"].sum()}\n')
+        
+    with open(os.path.join(opt.output_dir, 'statistics.txt'), 'r') as f:
+        print(f.read())

dataset_toolkits/datasets/3D-FUTURE.py

@@ -0,0 +1,97 @@
+import os
+import re
+import argparse
+import zipfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/3D-FUTURE.csv")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(output_dir, exist_ok=True)
+
+    if not os.path.exists(os.path.join(output_dir, 'raw', '3D-FUTURE-model.zip')):
+        print("\033[93m")
+        print("3D-FUTURE have to be downloaded manually")
+        print(f"Please download the 3D-FUTURE-model.zip file and place it in the {output_dir}/raw directory")
+        print("Visit https://tianchi.aliyun.com/specials/promotion/alibaba-3d-future for more information")
+        print("\033[0m")
+        raise FileNotFoundError("3D-FUTURE-model.zip not found")
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    with zipfile.ZipFile(os.path.join(output_dir, 'raw', '3D-FUTURE-model.zip')) as zip_ref:
+        all_names = zip_ref.namelist()
+        instances = [instance[:-1] for instance in all_names if re.match(r"^3D-FUTURE-model/[^/]+/$", instance)]
+        instances = list(filter(lambda x: x in metadata.index, instances))
+        
+        with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+            tqdm(total=len(instances), desc="Extracting") as pbar:
+            def worker(instance: str) -> str:
+                try:
+                    instance_files = list(filter(lambda x: x.startswith(f"{instance}/") and not x.endswith("/"), all_names))
+                    zip_ref.extractall(os.path.join(output_dir, 'raw'), members=instance_files)
+                    sha256 = get_file_hash(os.path.join(output_dir, 'raw', f"{instance}/image.jpg"))
+                    pbar.update()
+                    return sha256
+                except Exception as e:
+                    pbar.update()
+                    print(f"Error extracting for {instance}: {e}")
+                    return None
+                
+            sha256s = executor.map(worker, instances)
+            executor.shutdown(wait=True)
+
+    for k, sha256 in zip(instances, sha256s):
+        if sha256 is not None:
+            if sha256 == metadata.loc[k, "sha256"]:
+                downloaded[sha256] = os.path.join("raw", f"{k}/raw_model.obj")
+            else:
+                print(f"Error downloading {k}: sha256s do not match")
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    file = os.path.join(output_dir, local_path)
+                    record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/datasets/ABO.py

@@ -0,0 +1,96 @@
+import os
+import re
+import argparse
+import tarfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ABO.csv")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    if not os.path.exists(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')):
+        try:
+            os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+            os.system(f"wget -O {output_dir}/raw/abo-3dmodels.tar https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-3dmodels.tar")
+        except:
+            print("\033[93m")
+            print("Error downloading ABO dataset. Please check your internet connection and try again.")
+            print("Or, you can manually download the abo-3dmodels.tar file and place it in the {output_dir}/raw directory")
+            print("Visit https://amazon-berkeley-objects.s3.amazonaws.com/index.html for more information")
+            print("\033[0m")
+            raise FileNotFoundError("Error downloading ABO dataset")
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    with tarfile.open(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')) as tar:
+        with ThreadPoolExecutor(max_workers=1) as executor, \
+            tqdm(total=len(metadata), desc="Extracting") as pbar:
+            def worker(instance: str) -> str:
+                try:
+                    tar.extract(f"3dmodels/original/{instance}", path=os.path.join(output_dir, 'raw'))
+                    sha256 = get_file_hash(os.path.join(output_dir, 'raw/3dmodels/original', instance))
+                    pbar.update()
+                    return sha256
+                except Exception as e:
+                    pbar.update()
+                    print(f"Error extracting for {instance}: {e}")
+                    return None
+                
+            sha256s = executor.map(worker, metadata.index)
+            executor.shutdown(wait=True)
+
+    for k, sha256 in zip(metadata.index, sha256s):
+        if sha256 is not None:
+            if sha256 == metadata.loc[k, "sha256"]:
+                downloaded[sha256] = os.path.join('raw/3dmodels/original', k)
+            else:
+                print(f"Error downloading {k}: sha256s do not match")
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    file = os.path.join(output_dir, local_path)
+                    record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/datasets/HSSD.py

@@ -0,0 +1,103 @@
+import os
+import re
+import argparse
+import tarfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+import huggingface_hub
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/HSSD.csv")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    # check login
+    try:
+        huggingface_hub.whoami()
+    except:
+        print("\033[93m")
+        print("Haven't logged in to the Hugging Face Hub.")
+        print("Visit https://huggingface.co/settings/tokens to get a token.")
+        print("\033[0m")
+        huggingface_hub.login()
+        
+    try:
+        huggingface_hub.hf_hub_download(repo_id="hssd/hssd-models", filename="README.md", repo_type="dataset")
+    except:
+        print("\033[93m")
+        print("Error downloading HSSD dataset.")
+        print("Check if you have access to the HSSD dataset.")
+        print("Visit https://huggingface.co/datasets/hssd/hssd-models for more information")
+        print("\033[0m")
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+        tqdm(total=len(metadata), desc="Downloading") as pbar:
+        def worker(instance: str) -> str:
+            try:
+                huggingface_hub.hf_hub_download(repo_id="hssd/hssd-models", filename=instance, repo_type="dataset", local_dir=os.path.join(output_dir, 'raw'))
+                sha256 = get_file_hash(os.path.join(output_dir, 'raw', instance))
+                pbar.update()
+                return sha256
+            except Exception as e:
+                pbar.update()
+                print(f"Error extracting for {instance}: {e}")
+                return None
+            
+        sha256s = executor.map(worker, metadata.index)
+        executor.shutdown(wait=True)
+
+    for k, sha256 in zip(metadata.index, sha256s):
+        if sha256 is not None:
+            if sha256 == metadata.loc[k, "sha256"]:
+                downloaded[sha256] = os.path.join('raw', k)
+            else:
+                print(f"Error downloading {k}: sha256s do not match")
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    file = os.path.join(output_dir, local_path)
+                    record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/datasets/ObjaverseXL.py

@@ -0,0 +1,92 @@
+import os
+import argparse
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+import objaverse.xl as oxl
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    parser.add_argument('--source', type=str, default='sketchfab',
+                        help='Data source to download annotations from (github, sketchfab)')
+
+
+def get_metadata(source, **kwargs):
+    if source == 'sketchfab':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_sketchfab.csv")
+    elif source == 'github':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_github.csv")
+    else:
+        raise ValueError(f"Invalid source: {source}")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    # download annotations
+    annotations = oxl.get_annotations()
+    annotations = annotations[annotations['sha256'].isin(metadata['sha256'].values)]
+    
+    # download and render objects
+    file_paths = oxl.download_objects(
+        annotations,
+        download_dir=os.path.join(output_dir, "raw"),
+        save_repo_format="zip",
+    )
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    for k, v in file_paths.items():
+        sha256 = metadata.loc[k, "sha256"]
+        downloaded[sha256] = os.path.relpath(v, output_dir)
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    import tempfile
+    import zipfile
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    if local_path.startswith('raw/github/repos/'):
+                        path_parts = local_path.split('/')
+                        file_name = os.path.join(*path_parts[5:])
+                        zip_file = os.path.join(output_dir, *path_parts[:5])
+                        with tempfile.TemporaryDirectory() as tmp_dir:
+                            with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+                                zip_ref.extractall(tmp_dir)
+                            file = os.path.join(tmp_dir, file_name)
+                            record = func(file, sha256)
+                    else:
+                        file = os.path.join(output_dir, local_path)
+                        record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/datasets/Toys4k.py

@@ -0,0 +1,92 @@
+import os
+import re
+import argparse
+import zipfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/Toys4k.csv")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(output_dir, exist_ok=True)
+
+    if not os.path.exists(os.path.join(output_dir, 'raw', 'toys4k_blend_files.zip')):
+        print("\033[93m")
+        print("Toys4k have to be downloaded manually")
+        print(f"Please download the toys4k_blend_files.zip file and place it in the {output_dir}/raw directory")
+        print("Visit https://github.com/rehg-lab/lowshot-shapebias/tree/main/toys4k for more information")
+        print("\033[0m")
+        raise FileNotFoundError("toys4k_blend_files.zip not found")
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    with zipfile.ZipFile(os.path.join(output_dir, 'raw', 'toys4k_blend_files.zip')) as zip_ref:
+        with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+            tqdm(total=len(metadata), desc="Extracting") as pbar:
+            def worker(instance: str) -> str:
+                try:
+                    zip_ref.extract(os.path.join('toys4k_blend_files', instance), os.path.join(output_dir, 'raw'))
+                    sha256 = get_file_hash(os.path.join(output_dir, 'raw/toys4k_blend_files', instance))
+                    pbar.update()
+                    return sha256
+                except Exception as e:
+                    pbar.update()
+                    print(f"Error extracting for {instance}: {e}")
+                    return None
+                
+            sha256s = executor.map(worker, metadata.index)
+            executor.shutdown(wait=True)
+
+    for k, sha256 in zip(metadata.index, sha256s):
+        if sha256 is not None:
+            if sha256 == metadata.loc[k, "sha256"]:
+                downloaded[sha256] = os.path.join("raw/toys4k_blend_files", k)
+            else:
+                print(f"Error downloading {k}: sha256s do not match")
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    # processing objects
+    records = []
+    max_workers = max_workers or os.cpu_count()
+    try:
+        with ThreadPoolExecutor(max_workers=max_workers) as executor, \
+            tqdm(total=len(metadata), desc=desc) as pbar:
+            def worker(metadatum):
+                try:
+                    local_path = metadatum['local_path']
+                    sha256 = metadatum['sha256']
+                    file = os.path.join(output_dir, local_path)
+                    record = func(file, sha256)
+                    if record is not None:
+                        records.append(record)
+                    pbar.update()
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+                    pbar.update()
+            
+            executor.map(worker, metadata)
+            executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    return pd.DataFrame.from_records(records)

dataset_toolkits/download.py

@@ -0,0 +1,52 @@
+import os
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(opt.output_dir, exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'local_path' in metadata.columns:
+            metadata = metadata[metadata['local_path'].isna()]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    downloaded = dataset_utils.download(metadata, **opt)
+    downloaded.to_csv(os.path.join(opt.output_dir, f'downloaded_{opt.rank}.csv'), index=False)

dataset_toolkits/encode_latent.py

@@ -0,0 +1,127 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import copy
+import json
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+
+import trellis.models as models
+import trellis.modules.sparse as sp
+
+
+torch.set_grad_enabled(False)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--feat_model', type=str, default='dinov2_vitl14_reg',
+                        help='Feature model')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/slat_enc_swin8_B_64l8_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str, default='results',
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str, default=None,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str, default=None,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.feat_model}_{opt.enc_pretrained.split("/")[-1]}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.feat_model}_{opt.enc_model}_{opt.ckpt}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    os.makedirs(os.path.join(opt.output_dir, 'latents', latent_name), exist_ok=True)
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        raise ValueError('metadata.csv not found')
+    if opt.instances is not None:
+        with open(opt.instances, 'r') as f:
+            sha256s = [line.strip() for line in f]
+        metadata = metadata[metadata['sha256'].isin(sha256s)]
+    else:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata[f'feature_{opt.feat_model}'] == True]
+        if f'latent_{latent_name}' in metadata.columns:
+            metadata = metadata[metadata[f'latent_{latent_name}'] == False]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # filter out objects that are already processed
+    sha256s = list(metadata['sha256'].values)
+    for sha256 in copy.copy(sha256s):
+        if os.path.exists(os.path.join(opt.output_dir, 'latents', latent_name, f'{sha256}.npz')):
+            records.append({'sha256': sha256, f'latent_{latent_name}': True})
+            sha256s.remove(sha256)
+
+    # encode latents
+    load_queue = Queue(maxsize=4)
+    try:
+        with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+            ThreadPoolExecutor(max_workers=32) as saver_executor:
+            def loader(sha256):
+                try:
+                    feats = np.load(os.path.join(opt.output_dir, 'features', opt.feat_model, f'{sha256}.npz'))
+                    load_queue.put((sha256, feats))
+                except Exception as e:
+                    print(f"Error loading features for {sha256}: {e}")
+            loader_executor.map(loader, sha256s)
+            
+            def saver(sha256, pack):
+                save_path = os.path.join(opt.output_dir, 'latents', latent_name, f'{sha256}.npz')
+                np.savez_compressed(save_path, **pack)
+                records.append({'sha256': sha256, f'latent_{latent_name}': True})
+                
+            for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
+                sha256, feats = load_queue.get()
+                feats = sp.SparseTensor(
+                    feats = torch.from_numpy(feats['patchtokens']).float(),
+                    coords = torch.cat([
+                        torch.zeros(feats['patchtokens'].shape[0], 1).int(),
+                        torch.from_numpy(feats['indices']).int(),
+                    ], dim=1),
+                ).cuda()
+                latent = encoder(feats, sample_posterior=False)
+                assert torch.isfinite(latent.feats).all(), "Non-finite latent"
+                pack = {
+                    'feats': latent.feats.cpu().numpy().astype(np.float32),
+                    'coords': latent.coords[:, 1:].cpu().numpy().astype(np.uint8),
+                }
+                saver_executor.submit(saver, sha256, pack)
+                
+            saver_executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.output_dir, f'latent_{latent_name}_{opt.rank}.csv'), index=False)

dataset_toolkits/encode_ss_latent.py

@@ -0,0 +1,128 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import copy
+import json
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+import utils3d
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+
+import trellis.models as models
+
+
+torch.set_grad_enabled(False)
+
+
+def get_voxels(instance):
+    position = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{instance}.ply'))[0]
+    coords = ((torch.tensor(position) + 0.5) * opt.resolution).int().contiguous()
+    ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long)
+    ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+    return ss
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str, default='results',
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str, default=None,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str, default=None,
+                        help='Checkpoint to load')
+    parser.add_argument('--resolution', type=int, default=64,
+                        help='Resolution')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model}_{opt.ckpt}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    os.makedirs(os.path.join(opt.output_dir, 'ss_latents', latent_name), exist_ok=True)
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        raise ValueError('metadata.csv not found')
+    if opt.instances is not None:
+        with open(opt.instances, 'r') as f:
+            instances = f.read().splitlines()
+        metadata = metadata[metadata['sha256'].isin(instances)]
+    else:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['voxelized'] == True]
+        if f'ss_latent_{latent_name}' in metadata.columns:
+            metadata = metadata[metadata[f'ss_latent_{latent_name}'] == False]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # filter out objects that are already processed
+    sha256s = list(metadata['sha256'].values)
+    for sha256 in copy.copy(sha256s):
+        if os.path.exists(os.path.join(opt.output_dir, 'ss_latents', latent_name, f'{sha256}.npz')):
+            records.append({'sha256': sha256, f'ss_latent_{latent_name}': True})
+            sha256s.remove(sha256)
+
+    # encode latents
+    load_queue = Queue(maxsize=4)
+    try:
+        with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+            ThreadPoolExecutor(max_workers=32) as saver_executor:
+            def loader(sha256):
+                try:
+                    ss = get_voxels(sha256)[None].float()
+                    load_queue.put((sha256, ss))
+                except Exception as e:
+                    print(f"Error loading features for {sha256}: {e}")
+            loader_executor.map(loader, sha256s)
+            
+            def saver(sha256, pack):
+                save_path = os.path.join(opt.output_dir, 'ss_latents', latent_name, f'{sha256}.npz')
+                np.savez_compressed(save_path, **pack)
+                records.append({'sha256': sha256, f'ss_latent_{latent_name}': True})
+                
+            for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
+                sha256, ss = load_queue.get()
+                ss = ss.cuda().float()
+                latent = encoder(ss, sample_posterior=False)
+                assert torch.isfinite(latent).all(), "Non-finite latent"
+                pack = {
+                    'mean': latent[0].cpu().numpy(),
+                }
+                saver_executor.submit(saver, sha256, pack)
+                
+            saver_executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.output_dir, f'ss_latent_{latent_name}_{opt.rank}.csv'), index=False)

dataset_toolkits/extract_feature.py

@@ -0,0 +1,179 @@
+import os
+import copy
+import sys
+import json
+import importlib
+import argparse
+import torch
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import utils3d
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from torchvision import transforms
+from PIL import Image
+
+
+torch.set_grad_enabled(False)
+
+
+def get_data(frames, sha256):
+    with ThreadPoolExecutor(max_workers=16) as executor:
+        def worker(view):
+            image_path = os.path.join(opt.output_dir, 'renders', sha256, view['file_path'])
+            try:
+                image = Image.open(image_path)
+            except:
+                print(f"Error loading image {image_path}")
+                return None
+            image = image.resize((518, 518), Image.Resampling.LANCZOS)
+            image = np.array(image).astype(np.float32) / 255
+            image = image[:, :, :3] * image[:, :, 3:]
+            image = torch.from_numpy(image).permute(2, 0, 1).float()
+
+            c2w = torch.tensor(view['transform_matrix'])
+            c2w[:3, 1:3] *= -1
+            extrinsics = torch.inverse(c2w)
+            fov = view['camera_angle_x']
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
+
+            return {
+                'image': image,
+                'extrinsics': extrinsics,
+                'intrinsics': intrinsics
+            }
+        
+        datas = executor.map(worker, frames)
+        for data in datas:
+            if data is not None:
+                yield data
+                
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--model', type=str, default='dinov2_vitl14_reg',
+                        help='Feature extraction model')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--batch_size', type=int, default=16)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+
+    feature_name = opt.model
+    os.makedirs(os.path.join(opt.output_dir, 'features', feature_name), exist_ok=True)
+
+    # load model
+    dinov2_model = torch.hub.load('facebookresearch/dinov2', opt.model)
+    dinov2_model.eval().cuda()
+    transform = transforms.Compose([
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    n_patch = 518 // 14
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        raise ValueError('metadata.csv not found')
+    if opt.instances is not None:
+        with open(opt.instances, 'r') as f:
+            instances = f.read().splitlines()
+        metadata = metadata[metadata['sha256'].isin(instances)]
+    else:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if f'feature_{feature_name}' in metadata.columns:
+            metadata = metadata[metadata[f'feature_{feature_name}'] == False]
+        metadata = metadata[metadata['voxelized'] == True]
+        metadata = metadata[metadata['rendered'] == True]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    sha256s = list(metadata['sha256'].values)
+    for sha256 in copy.copy(sha256s):
+        if os.path.exists(os.path.join(opt.output_dir, 'features', feature_name, f'{sha256}.npz')):
+            records.append({'sha256': sha256, f'feature_{feature_name}' : True})
+            sha256s.remove(sha256)
+
+    # extract features
+    load_queue = Queue(maxsize=4)
+    try:
+        with ThreadPoolExecutor(max_workers=8) as loader_executor, \
+            ThreadPoolExecutor(max_workers=8) as saver_executor:
+            def loader(sha256):
+                try:
+                    with open(os.path.join(opt.output_dir, 'renders', sha256, 'transforms.json'), 'r') as f:
+                        metadata = json.load(f)
+                    frames = metadata['frames']
+                    data = []
+                    for datum in get_data(frames, sha256):
+                        datum['image'] = transform(datum['image'])
+                        data.append(datum)
+                    positions = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0]
+                    load_queue.put((sha256, data, positions))
+                except Exception as e:
+                    print(f"Error loading data for {sha256}: {e}")
+
+            loader_executor.map(loader, sha256s)
+            
+            def saver(sha256, pack, patchtokens, uv):
+                pack['patchtokens'] = F.grid_sample(
+                    patchtokens,
+                    uv.unsqueeze(1),
+                    mode='bilinear',
+                    align_corners=False,
+                ).squeeze(2).permute(0, 2, 1).cpu().numpy()
+                pack['patchtokens'] = np.mean(pack['patchtokens'], axis=0).astype(np.float16)
+                save_path = os.path.join(opt.output_dir, 'features', feature_name, f'{sha256}.npz')
+                np.savez_compressed(save_path, **pack)
+                records.append({'sha256': sha256, f'feature_{feature_name}' : True})
+                
+            for _ in tqdm(range(len(sha256s)), desc="Extracting features"):
+                sha256, data, positions = load_queue.get()
+                positions = torch.from_numpy(positions).float().cuda()
+                indices = ((positions + 0.5) * 64).long()
+                assert torch.all(indices >= 0) and torch.all(indices < 64), "Some vertices are out of bounds"
+                n_views = len(data)
+                N = positions.shape[0]
+                pack = {
+                    'indices': indices.cpu().numpy().astype(np.uint8),
+                }
+                patchtokens_lst = []
+                uv_lst = []
+                for i in range(0, n_views, opt.batch_size):
+                    batch_data = data[i:i+opt.batch_size]
+                    bs = len(batch_data)
+                    batch_images = torch.stack([d['image'] for d in batch_data]).cuda()
+                    batch_extrinsics = torch.stack([d['extrinsics'] for d in batch_data]).cuda()
+                    batch_intrinsics = torch.stack([d['intrinsics'] for d in batch_data]).cuda()
+                    features = dinov2_model(batch_images, is_training=True)
+                    uv = utils3d.torch.project_cv(positions, batch_extrinsics, batch_intrinsics)[0] * 2 - 1
+                    patchtokens = features['x_prenorm'][:, dinov2_model.num_register_tokens + 1:].permute(0, 2, 1).reshape(bs, 1024, n_patch, n_patch)
+                    patchtokens_lst.append(patchtokens)
+                    uv_lst.append(uv)
+                patchtokens = torch.cat(patchtokens_lst, dim=0)
+                uv = torch.cat(uv_lst, dim=0)
+
+                # save features
+                saver_executor.submit(saver, sha256, pack, patchtokens, uv)
+                
+            saver_executor.shutdown(wait=True)
+    except:
+        print("Error happened during processing.")
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.output_dir, f'feature_{feature_name}_{opt.rank}.csv'), index=False)
+        

dataset_toolkits/render.py

@@ -0,0 +1,121 @@
+import os
+import json
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+import numpy as np
+from utils import sphere_hammersley_sequence
+
+
+BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = '/tmp'
+BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender'
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        os.system('sudo apt-get update')
+        os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6')
+        os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
+        os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
+
+
+def _render(file_path, sha256, output_dir, num_views):
+    output_folder = os.path.join(output_dir, 'renders', sha256)
+    
+    # Build camera {yaw, pitch, radius, fov}
+    yaws = []
+    pitchs = []
+    offset = (np.random.rand(), np.random.rand())
+    for i in range(num_views):
+        y, p = sphere_hammersley_sequence(i, num_views, offset)
+        yaws.append(y)
+        pitchs.append(p)
+    radius = [2] * num_views
+    fov = [40 / 180 * np.pi] * num_views
+    views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': f} for y, p, r, f in zip(yaws, pitchs, radius, fov)]
+    
+    args = [
+        BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'render.py'),
+        '--',
+        '--views', json.dumps(views),
+        '--object', os.path.expanduser(file_path),
+        '--resolution', '512',
+        '--output_folder', output_folder,
+        '--engine', 'CYCLES',
+        '--save_mesh',
+    ]
+    if file_path.endswith('.blend'):
+        args.insert(1, file_path)
+    
+    call(args, stdout=DEVNULL, stderr=DEVNULL)
+    
+    if os.path.exists(os.path.join(output_folder, 'transforms.json')):
+        return {'sha256': sha256, 'rendered': True}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--num_views', type=int, default=150,
+                        help='Number of views to render')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=8)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(os.path.join(opt.output_dir, 'renders'), exist_ok=True)
+    
+    # install blender
+    print('Checking blender...', flush=True)
+    _install_blender()
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    if opt.instances is None:
+        metadata = metadata[metadata['local_path'].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'rendered' in metadata.columns:
+            metadata = metadata[metadata['rendered'] == False]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    for sha256 in copy.copy(metadata['sha256'].values):
+        if os.path.exists(os.path.join(opt.output_dir, 'renders', sha256, 'transforms.json')):
+            records.append({'sha256': sha256, 'rendered': True})
+            metadata = metadata[metadata['sha256'] != sha256]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_render, output_dir=opt.output_dir, num_views=opt.num_views)
+    rendered = dataset_utils.foreach_instance(metadata, opt.output_dir, func, max_workers=opt.max_workers, desc='Rendering objects')
+    rendered = pd.concat([rendered, pd.DataFrame.from_records(records)])
+    rendered.to_csv(os.path.join(opt.output_dir, f'rendered_{opt.rank}.csv'), index=False)

dataset_toolkits/render_cond.py

@@ -0,0 +1,125 @@
+import os
+import json
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+import numpy as np
+from utils import sphere_hammersley_sequence
+
+
+BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = '/tmp'
+BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender'
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        os.system('sudo apt-get update')
+        os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6')
+        os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
+        os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
+
+
+def _render_cond(file_path, sha256, output_dir, num_views):
+    output_folder = os.path.join(output_dir, 'renders_cond', sha256)
+    
+    # Build camera {yaw, pitch, radius, fov}
+    yaws = []
+    pitchs = []
+    offset = (np.random.rand(), np.random.rand())
+    for i in range(num_views):
+        y, p = sphere_hammersley_sequence(i, num_views, offset)
+        yaws.append(y)
+        pitchs.append(p)
+    fov_min, fov_max = 10, 70
+    radius_min = np.sqrt(3) / 2 / np.sin(fov_max / 360 * np.pi)
+    radius_max = np.sqrt(3) / 2 / np.sin(fov_min / 360 * np.pi)
+    k_min = 1 / radius_max**2
+    k_max = 1 / radius_min**2
+    ks = np.random.uniform(k_min, k_max, (1000000,))
+    radius = [1 / np.sqrt(k) for k in ks]
+    fov = [2 * np.arcsin(np.sqrt(3) / 2 / r) for r in radius]
+    views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': f} for y, p, r, f in zip(yaws, pitchs, radius, fov)]
+    
+    args = [
+        BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'render.py'),
+        '--',
+        '--views', json.dumps(views),
+        '--object', os.path.expanduser(file_path),
+        '--output_folder', os.path.expanduser(output_folder),
+        '--resolution', '1024',
+    ]
+    if file_path.endswith('.blend'):
+        args.insert(1, file_path)
+    
+    call(args, stdout=DEVNULL)
+    
+    if os.path.exists(os.path.join(output_folder, 'transforms.json')):
+        return {'sha256': sha256, 'cond_rendered': True}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--num_views', type=int, default=24,
+                        help='Number of views to render')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=8)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(os.path.join(opt.output_dir, 'renders_cond'), exist_ok=True)
+    
+    # install blender
+    print('Checking blender...', flush=True)
+    _install_blender()
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    if opt.instances is None:
+        metadata = metadata[metadata['local_path'].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'cond_rendered' in metadata.columns:
+            metadata = metadata[metadata['cond_rendered'] == False]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    for sha256 in copy.copy(metadata['sha256'].values):
+        if os.path.exists(os.path.join(opt.output_dir, 'renders_cond', sha256, 'transforms.json')):
+            records.append({'sha256': sha256, 'cond_rendered': True})
+            metadata = metadata[metadata['sha256'] != sha256]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_render_cond, output_dir=opt.output_dir, num_views=opt.num_views)
+    cond_rendered = dataset_utils.foreach_instance(metadata, opt.output_dir, func, max_workers=opt.max_workers, desc='Rendering objects')
+    cond_rendered = pd.concat([cond_rendered, pd.DataFrame.from_records(records)])
+    cond_rendered.to_csv(os.path.join(opt.output_dir, f'cond_rendered_{opt.rank}.csv'), index=False)

dataset_toolkits/setup.sh

@@ -0,0 +1 @@
+pip install pillow imageio imageio-ffmpeg tqdm easydict opencv-python-headless pandas open3d objaverse huggingface_hub

dataset_toolkits/stat_latent.py

@@ -0,0 +1,66 @@
+import os
+import json
+import argparse
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--model', type=str, default='dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16',
+                        help='Latent model to use')
+    parser.add_argument('--num_samples', type=int, default=50000,
+                        help='Number of samples to use for calculating stats')
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+
+    # get file list
+    if os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    else:
+        raise ValueError('metadata.csv not found')
+    if opt.filter_low_aesthetic_score is not None:
+        metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+    metadata = metadata[metadata[f'latent_{opt.model}'] == True]
+    sha256s = metadata['sha256'].values
+    sha256s = np.random.choice(sha256s, min(opt.num_samples, len(sha256s)), replace=False)
+
+    # stats
+    means = []
+    mean2s = []
+    with ThreadPoolExecutor(max_workers=16) as executor, \
+        tqdm(total=len(sha256s), desc="Extracting features") as pbar:
+        def worker(sha256):
+            try:
+                feats = np.load(os.path.join(opt.output_dir, 'latents', opt.model, f'{sha256}.npz'))
+                feats = feats['feats']
+                means.append(feats.mean(axis=0))
+                mean2s.append((feats ** 2).mean(axis=0))
+                pbar.update()
+            except Exception as e:
+                print(f"Error extracting features for {sha256}: {e}")
+                pbar.update()
+
+        executor.map(worker, sha256s)
+        executor.shutdown(wait=True)
+
+    mean = np.array(means).mean(axis=0)
+    mean2 = np.array(mean2s).mean(axis=0)
+    std = np.sqrt(mean2 - mean ** 2)
+
+    print('mean:', mean)
+    print('std:', std)
+
+    with open(os.path.join(opt.output_dir, 'latents', opt.model, 'stats.json'), 'w') as f:
+        json.dump({
+            'mean': mean.tolist(),
+            'std': std.tolist(),
+        }, f, indent=4)
+        

dataset_toolkits/utils.py

@@ -0,0 +1,43 @@
+from typing import *
+import hashlib
+import numpy as np
+
+
+def get_file_hash(file: str) -> str:
+    sha256 = hashlib.sha256()
+    # Read the file from the path
+    with open(file, "rb") as f:
+        # Update the hash with the file content
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256.update(byte_block)
+    return sha256.hexdigest()
+
+# ===============LOW DISCREPANCY SEQUENCES================
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+def radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+def halton_sequence(dim, n):
+    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+def hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + halton_sequence(dim - 1, n)
+
+def sphere_hammersley_sequence(n, num_samples, offset=(0, 0)):
+    u, v = hammersley_sequence(2, n, num_samples)
+    u += offset[0] / num_samples
+    v += offset[1]
+    u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
+    theta = np.arccos(1 - 2 * u) - np.pi / 2
+    phi = v * 2 * np.pi
+    return [phi, theta]

dataset_toolkits/voxelize.py

@@ -0,0 +1,86 @@
+import os
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+import numpy as np
+import open3d as o3d
+import utils3d
+
+
+def _voxelize(file, sha256, output_dir):
+    mesh = o3d.io.read_triangle_mesh(os.path.join(output_dir, 'renders', sha256, 'mesh.ply'))
+    # clamp vertices to the range [-0.5, 0.5]
+    vertices = np.clip(np.asarray(mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6)
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
+    vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
+    assert np.all(vertices >= 0) and np.all(vertices < 64), "Some vertices are out of bounds"
+    vertices = (vertices + 0.5) / 64 - 0.5
+    utils3d.io.write_ply(os.path.join(output_dir, 'voxels', f'{sha256}.ply'), vertices)
+    return {'sha256': sha256, 'voxelized': True, 'num_voxels': len(vertices)}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output_dir', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--num_views', type=int, default=150,
+                        help='Number of views to render')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=None)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+
+    os.makedirs(os.path.join(opt.output_dir, 'voxels'), exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.output_dir, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.output_dir, 'metadata.csv'))
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'rendered' not in metadata.columns:
+            raise ValueError('metadata.csv does not have "rendered" column, please run "build_metadata.py" first')
+        metadata = metadata[metadata['rendered'] == True]
+        if 'voxelized' in metadata.columns:
+            metadata = metadata[metadata['voxelized'] == False]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    for sha256 in copy.copy(metadata['sha256'].values):
+        if os.path.exists(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply')):
+            pts = utils3d.io.read_ply(os.path.join(opt.output_dir, 'voxels', f'{sha256}.ply'))[0]
+            records.append({'sha256': sha256, 'voxelized': True, 'num_voxels': len(pts)})
+            metadata = metadata[metadata['sha256'] != sha256]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_voxelize, output_dir=opt.output_dir)
+    voxelized = dataset_utils.foreach_instance(metadata, opt.output_dir, func, max_workers=opt.max_workers, desc='Voxelizing')
+    voxelized = pd.concat([voxelized, pd.DataFrame.from_records(records)])
+    voxelized.to_csv(os.path.join(opt.output_dir, f'voxelized_{opt.rank}.csv'), index=False)