nodes/smthemex-ComfyUI_TRELLIS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #38 from smthemex/Pr

Browse Source 
init
This commit is contained in:
smthemex
2025-04-05 17:14:18 +08:00
committed by GitHub
parent 5dbf87f99d 3a79a95d68
commit e491bafe9b
69 changed files with 153846 additions and 440 deletions
Download Patch File Download Diff File Expand all files Collapse all files
250327.json
-214
View File
@@ -1,214 +0,0 @@
{
"last_node_id": 14,
"last_link_id": 13,
"nodes": [
{
"id": 10,
"type": "Trellis_LoadModel",
"pos": [
604.9623413085938,
480.64385986328125
],
"size": [
315,
130
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "model",
"type": "MODEL_TRELLIS",
"links": [
10
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Trellis_LoadModel"
},
"widgets_values": [
"F:/test/ComfyUI/models/diffusers/JeffreyXiang/TRELLIS-image-large",
"dinov2_vitl14_reg4_pretrain.pth",
"xformers",
"auto"
]
},
{
"id": 12,
"type": "LoadImage",
"pos": [
583.5416870117188,
89.96366119384766
],
"size": [
315,
314
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
11
]
},
{
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"image (1).png",
"image"
]
},
{
"id": 11,
"type": "Trellis_Sampler",
"pos": [
974.942138671875,
181.10565185546875
],
"size": [
315,
438
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 11
},
{
"name": "model",
"type": "MODEL_TRELLIS",
"link": 10
}
],
"outputs": [
{
"name": "model_path",
"type": "STRING",
"links": [
13
],
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "Trellis_Sampler"
},
"widgets_values": [
581394181,
"randomize",
7.5,
12,
3,
12,
false,
512,
0.95,
"fast",
false,
"multidiffusion",
false,
false,
false,
false
]
},
{
"id": 14,
"type": "ShowText|pysssss",
"pos": [
1379.2635498046875,
272.084228515625
],
"size": [
315,
76
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "text",
"type": "STRING",
"link": 13,
"widget": {
"name": "text"
}
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": null,
"shape": 6
}
],
"properties": {
"Node name for S&R": "ShowText|pysssss"
},
"widgets_values": [
"",
"F:\\ComfyUI311\\ComfyUI\\output/fb690810-7fe2-40ea-981d-baadd7c5e581_0.glb"
]
}
],
"links": [
[
10,
10,
0,
11,
1,
"MODEL_TRELLIS"
],
[
11,
12,
0,
11,
0,
"IMAGE"
],
[
13,
11,
0,
14,
0,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 1,
"offset": [
-184.26351062065908,
40.35328013119408
]
}
},
"version": 0.4
}
README.md
+12 -21
View File
@@ -5,8 +5,9 @@
---
# Update
* 2025/03/27 kaolin support torch2.6 now;
# Update 2025/04/5
* support txt to 3D and meshto3D
* kaolin support torch2.6 now;
# 1. Installation
@@ -19,8 +20,8 @@ git clone https://github.com/smthemex/ComfyUI_TRELLIS.git
---
# 2. Requirements
本插件的测试环境是python3.11，torch2.5.1 cu124...
The testing environment for this node is Python 3.11, torch2.5.1 cu124...
本插件的测试环境是python3.11，torch2.6 cu126...
The testing environment for this node is Python 3.11, torch2.6 cu126...
```
pip install -r requirements.txt
@@ -113,33 +114,23 @@ if pre download ,fill local path in repo like this: x:/your/path/JeffreyXiang/TR
| ├── ss_flow_img_dit_L_16l8_fp16.json
| ├── ss_flow_img_dit_L_16l8_fp16.safetensors
```
* 3.2 dinov2
因为官方的代码每次加载dinov2都要连GitHub，所以我改成了离线版的，需要下载dinov2模型，[地址](https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth)
模型放在comfyUI/models/dinov2目录下
* 3.2 dinov2 and clip
因为官方的代码每次加载dinov2都要连GitHub，所以我改成了离线版的，需要下载dinov2模型，[地址](https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth) ,clip是常规的openai/clip-vit-large-patch14 模型，也为方便大陆用户改成了离线版
模型放在comfyUI/models/dinov2和clip目录下
```
├── ComfyUI/models/dinov2
| ├── dinov2_vitl14_reg4_pretrain.pth
├── ComfyUI/models/clip
| ├── clip_l.safetensors
```
---
# 4 Example
![](https://github.com/smthemex/ComfyUI_TRELLIS/blob/main/example1220.png)
![](https://github.com/smthemex/ComfyUI_TRELLIS/blob/main/exmaple.png)
![](https://github.com/smthemex/ComfyUI_TRELLIS/blob/main/batch_example.png
)
* img to 3D or txt to 3D
![](https://github.com/smthemex/ComfyUI_TRELLIS/blob/main/assets.png)
**previous update**
* 根据TRELLIS的更新，新增三视参考图渲染模式（示例图为实际效果，此模式速度更快）和高斯保存按钮，针对三视参考图新增图片加载节点（内置正方形裁切),三视图模式开启multi_image生效，否则是常规的单图模式（会输出三个结果）；
* 如果你输入的图片不是纯色背景，建议开启preprocess_image以获得最好的效果。（此次更新也修复了加载RGBA图片可能导致的变形错误）；
* 使用 [here](https://github.com/smthemex/ComfyUI_TRELLIS/issues/6) @planb788 方法，我制作了python3.11，torch2.5.1 cu124.的便携包，可以在[Google dirver](https://drive.google.com/file/d/174StpwP3D1qSD0RuhM1XwTene-5VMiL4/view?usp=drive_link) 或者 [夸克网盘](https://pan.quark.cn/s/8e07717bdff7)下载，注意，即便是便携包，也是需要配置VS和python的系统变量路径的；
* 增加批量渲染功能，注意过多图片可能会OOM；
* According to TRELLIS' update, a three view reference image rendering mode (the example image is the actual effect, which is faster) and a Gaussian save button have been added. For the three view reference image, a new image loading node (with built-in square cropping) has been added,Enabling multi_image in three view mode takes effect, otherwise it will be in regular single image mode (outputting three results);
* If the image you input is not a solid color background, it is recommended to enable 'preprocess_image' for the best effect. (This update also fixes deformation errors that may occur when loading RGBA images);
* Use [here](https://github.com/smthemex/ComfyUI_TRELLIS/issues/6) , I have created a portable package for ‘Python 3.11, Torch 2.5.1, and CU124’ using the @planb788 method, which can be found on [Google dirver](https://drive.google.com/file/d/174StpwP3D1qSD0RuhM1XwTene-5VMiL4/view?usp=drive_link) Or [夸克网盘](https://pan.quark.cn/s/8e07717bdff7) Download, note that even portable packages require configuring the system variable paths for VS and Python;
* Add batch rendering function, be aware that too many images may cause OOM
# 5 Citation
TRELLIS_node.py
+109 -38
View File
@@ -1,12 +1,15 @@
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import numpy as np
import gc
import os
import torch
import uuid
from .app import image_to_3d
from .trellis.pipelines import TrellisImageTo3DPipeline
from .app import image_to_3d,prepare_output
from .trellis.pipelines import TrellisImageTo3DPipeline,TrellisTextTo3DPipeline
from .utils import glb2obj_,obj2fbx_,tensor2imglist,pre_img
import folder_paths
@@ -32,7 +35,7 @@ class Trellis_LoadModel:
return {
"required": {
"repo": ("STRING", {"default": "JeffreyXiang/TRELLIS-image-large"}),
"dino": (["none"] + folder_paths.get_filename_list("dinov2"),),
"dino_clip": (["none"] + folder_paths.get_filename_list("dinov2")+folder_paths.get_filename_list("clip"),),
"attn_backend":(["xformers","flash-attn"],),
"spconv_algo":(["auto","flash-native"],),
}
@@ -43,7 +46,7 @@ class Trellis_LoadModel:
FUNCTION = "main_loader"
CATEGORY = "Trellis"
def main_loader(self, repo,dino,attn_backend,spconv_algo):
def main_loader(self, repo,dino_clip,attn_backend,spconv_algo):
if attn_backend=="xformers":
os.environ['ATTN_BACKEND'] = 'xformers'
else:
@@ -52,17 +55,33 @@ class Trellis_LoadModel:
os.environ['SPCONV_ALGO'] = 'auto'
else:
os.environ['SPCONV_ALGO'] = 'native'
if dino=="none":
raise "need choice dinov2 checkpoint"
TrellisImageTo3DPipeline.dino=folder_paths.get_full_path("dinov2", dino)
TrellisImageTo3DPipeline.dino_moudel=os.path.join(current_path,"facebookresearch/dinov2")
if repo:
model=TrellisImageTo3DPipeline.from_pretrained(repo)
infer_mode="txt" if "text" in repo else "img"
else:
raise "need fill repo"
return (model,)
if dino_clip=="none" :
raise "image mode need choice dinov2 or clip checkpoint"
if dino_clip!="none" :
if infer_mode=="img":
TrellisImageTo3DPipeline.dino=folder_paths.get_full_path("dinov2", dino_clip)
TrellisImageTo3DPipeline.dino_moudel=os.path.join(current_path,"facebookresearch/dinov2")
else:
TrellisTextTo3DPipeline.clip=folder_paths.get_full_path("clip", dino_clip)
TrellisTextTo3DPipeline.clip_moudel=os.path.join(current_path,"clip")
if repo:
infer_mode="txt" if "text" in repo else "img"
if infer_mode=="img":
model=TrellisImageTo3DPipeline.from_pretrained(repo)
else:
model=TrellisTextTo3DPipeline.from_pretrained(repo)
return ({"model": model,"infer_mode":infer_mode},)
class Trellis_Sampler:
@@ -71,8 +90,9 @@ class Trellis_Sampler:
return {
"required": {
"image": ("IMAGE",), # [B,H,W,C], C=3
"model": ("MODEL_TRELLIS",),
"prompt": ("STRING", {"multiline": True,"default":"Rugged, metallic texture with orange and white paint finish, suggesting a durable, industrial feel."}),
"mesh_path": ("STRING", {"default":"",}),
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
"cfg": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
"steps": ("INT", {"default": 12, "min": 1, "max": 50}),
@@ -87,8 +107,10 @@ class Trellis_Sampler:
"gaussians2ply": ("BOOLEAN", {"default": False},),
"covert2video": ("BOOLEAN", {"default": False},),
"glb2obj": ("BOOLEAN", {"default": False},),
"glb2fbx": ("BOOLEAN", {"default": False},),
}
"glb2fbx": ("BOOLEAN", {"default": False},), },
"optional":{ "image": ("IMAGE",), # [B,H,W,C], C=3
}
}
RETURN_TYPES = ("STRING", )
@@ -96,30 +118,79 @@ class Trellis_Sampler:
FUNCTION = "sampler_main"
CATEGORY = "Trellis"
def sampler_main(self, image, model, seed, cfg, steps,slat_cfg, slat_steps,preprocess_image,texture_size,mesh_simplify,mode,multi_image,multiimage_algo,gaussians2ply,covert2video,glb2obj,glb2fbx):
image_list,image_batch=tensor2imglist(image) #pil_list,batch
if multi_image and image_batch % 3 == 0:
print("********infer multi image,like Three views ******")
image_list=[image_list[i:i + 3] for i in range(0, len(image_list), 3)] #三等分列表
is_multiimage=True
else:
is_multiimage = False
trial_id = str(uuid.uuid4())
output_path = []
for i,img in enumerate(image_list):
model.cuda()
glb=image_to_3d(model,img,preprocess_image,covert2video,trial_id,seed,cfg,steps,slat_cfg,slat_steps,mesh_simplify,texture_size,mode,is_multiimage,gaussians2ply,multiimage_algo)
glb_path = f"{folder_paths.get_output_directory()}/{trial_id}_{i}.glb"
glb.export(glb_path)
output_path.append(glb_path)
model.cpu()
gc.collect()
torch.cuda.empty_cache()
print(f"glb save in {glb_path} ")
def sampler_main(self, model,prompt, mesh_path, seed, cfg, steps,slat_cfg, slat_steps,preprocess_image,texture_size,mesh_simplify,mode,multi_image,multiimage_algo,gaussians2ply,covert2video,glb2obj,glb2fbx,**kwargs):
infer_mode=model.get("infer_mode")
model=model.get("model")
if kwargs.get("image") is None and infer_mode=="img":
raise "when use img mode,need link a image"
if isinstance(kwargs.get("image"),torch.Tensor):
image_list,image_batch=tensor2imglist(kwargs.get("image")) #pil_list,batch
if multi_image and image_batch % 3 == 0:
print("********infer multi image,like Three views ******")
image_list=[image_list[i:i + 3] for i in range(0, len(image_list), 3)] #三等分列表
is_multiimage=True
else:
is_multiimage = False
else:
image_list,is_multiimage=None,False
trial_id = str(uuid.uuid4())
if infer_mode=="img":
output_path = []
for i,img in enumerate(image_list):
model.cuda()
glb_path=image_to_3d(model,img,preprocess_image,covert2video,trial_id,seed,cfg,steps,slat_cfg,slat_steps,mesh_simplify,texture_size,mode,is_multiimage,gaussians2ply,multiimage_algo)
output_path.append(glb_path)
else:
model.cuda()
if mesh_path:
if not mesh_path.endswith(".ply"):
raise "mesh_path must be .ply format"
import open3d as o3d
base_mesh = o3d.io.read_triangle_mesh(mesh_path)
outputs = model.run_variant(
base_mesh,
prompt=prompt,
seed=1,
# Optional parameters
slat_sampler_params={
"steps": slat_steps,
"cfg_strength": slat_cfg,
},
)
# outputs is a dictionary containing generated 3D assets in different formats:
# - outputs['gaussian']: a list of 3D Gaussians
# - outputs['radiance_field']: a list of radiance fields
# - outputs['mesh']: a list of meshes
else:
outputs = model.run(
prompt,
seed=seed,
formats=["gaussian", "mesh"],
sparse_structure_sampler_params={
"steps": steps,
"cfg_strength": cfg,
},
slat_sampler_params={
"steps": slat_steps,
"cfg_strength": slat_cfg,
},
)
output_path=prepare_output(outputs,covert2video,trial_id,False,gaussians2ply,mesh_simplify,texture_size,mode)
output_path=[output_path]
model.cpu()
gc.collect()
torch.cuda.empty_cache()
if glb2obj:
obj_paths=[]
for path in output_path:
__init__.py
+2 -1
View File
@@ -1,4 +1,5 @@
from .TRELLIS_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS",'WEB_DIRECTORY']
app.py
+73 -16
View File
@@ -1,13 +1,64 @@
import os.path
import random
from typing import *
import torch
import numpy as np
from .trellis.utils import render_utils, postprocessing_utils
import imageio
import folder_paths
from easydict import EasyDict as edict
from .trellis.pipelines import TrellisTextTo3DPipeline
from .trellis.representations import Gaussian, MeshExtractResult
from .trellis.utils import render_utils, postprocessing_utils
MAX_SEED = np.iinfo(np.int32).max
def image_to_3d(pipeline,image,preprocess_image:bool,covert2video:bool,trial_id: str, seed: int,ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int,mesh_simplify,texture_size,mode,is_multiimage,Gaussians2PLY,multiimage_algo):
def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
return {
'gaussian': {
**gs.init_params,
'_xyz': gs._xyz.cpu().numpy(),
'_features_dc': gs._features_dc.cpu().numpy(),
'_scaling': gs._scaling.cpu().numpy(),
'_rotation': gs._rotation.cpu().numpy(),
'_opacity': gs._opacity.cpu().numpy(),
},
'mesh': {
'vertices': mesh.vertices.cpu().numpy(),
'faces': mesh.faces.cpu().numpy(),
},
}
def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
gs = Gaussian(
aabb=state['gaussian']['aabb'],
sh_degree=state['gaussian']['sh_degree'],
mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
scaling_bias=state['gaussian']['scaling_bias'],
opacity_bias=state['gaussian']['opacity_bias'],
scaling_activation=state['gaussian']['scaling_activation'],
)
gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
mesh = edict(
vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
faces=torch.tensor(state['mesh']['faces'], device='cuda'),
)
return gs, mesh
def get_seed(randomize_seed: bool, seed: int) -> int:
"""
Get the random seed.
"""
return np.random.randint(0, MAX_SEED) if randomize_seed else seed
def image_to_3d(pipeline,image,preprocess_image:bool,covert2video:bool,trial_id: str, seed: int,ss_guidance_strength: float,
ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int,mesh_simplify,texture_size,mode,is_multiimage,Gaussians2PLY,multiimage_algo):
"""
Convert an image to a 3D model.
@@ -55,7 +106,11 @@ def image_to_3d(pipeline,image,preprocess_image:bool,covert2video:bool,trial_id:
},
mode=multiimage_algo,
)
return prepare_output(outputs,covert2video,trial_id,is_multiimage,Gaussians2PLY,mesh_simplify,texture_size,mode)
def prepare_output(outputs,covert2video,trial_id,is_multiimage,Gaussians2PLY,mesh_simplify,texture_size,mode):
if covert2video:
video_path = f"{trial_id}.mp4"
if is_multiimage:
@@ -72,20 +127,22 @@ def image_to_3d(pipeline,image,preprocess_image:bool,covert2video:bool,trial_id:
video = render_utils.render_video(outputs['mesh'][0])['normal']
imageio.mimsave(f"{trial_id}_mesh.mp4", video, fps=30)
state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
torch.cuda.empty_cache()
gs, mesh = unpack_state(state)
if Gaussians2PLY:
outputs['gaussian'][0].save_ply(f"{trial_id}.ply")
gs.save_ply(f"{trial_id}.ply")
glb = postprocessing_utils.to_glb(
outputs['gaussian'][0],
outputs['mesh'][0],
gs,
mesh,
# Optional parameters
simplify=mesh_simplify, # Ratio of triangles to remove in the simplification process
texture_size=texture_size, # Size of the texture used for the GLB
mode=mode,
uv_map=f"{folder_paths.get_output_directory()}/{trial_id}_uv_map.png",
)
return glb
uv_map=f"{folder_paths.get_output_directory()}/{trial_id}_uv_map.png",)
prefix = ''.join(random.choice("0123456789") for _ in range(5))
glb_path = f"{folder_paths.get_output_directory()}/{trial_id}_{prefix}.glb"
glb.export(glb_path)
print(f"glb save in {glb_path} ")
return glb_path
assets/example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 361 KiB

batch_example.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 312 KiB

clip/config.json
+171
View File
@@ -0,0 +1,171 @@
{
"_name_or_path": "clip-vit-large-patch14/",
"architectures": [
"CLIPModel"
],
"initializer_factor": 1.0,
"logit_scale_init_value": 2.6592,
"model_type": "clip",
"projection_dim": 768,
"text_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"bos_token_id": 0,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": 2,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "quick_gelu",
"hidden_size": 768,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 3072,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"max_position_embeddings": 77,
"min_length": 0,
"model_type": "clip_text_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 12,
"num_beam_groups": 1,
"num_beams": 1,
"num_hidden_layers": 12,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": 1,
"prefix": null,
"problem_type": null,
"projection_dim" : 768,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"task_specific_params": null,
"temperature": 1.0,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.16.0.dev0",
"use_bfloat16": false,
"vocab_size": 49408
},
"text_config_dict": {
"hidden_size": 768,
"intermediate_size": 3072,
"num_attention_heads": 12,
"num_hidden_layers": 12,
"projection_dim": 768
},
"torch_dtype": "float32",
"transformers_version": null,
"vision_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"bos_token_id": null,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": null,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "quick_gelu",
"hidden_size": 1024,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"image_size": 224,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 4096,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"min_length": 0,
"model_type": "clip_vision_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 16,
"num_beam_groups": 1,
"num_beams": 1,
"num_hidden_layers": 24,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": null,
"patch_size": 14,
"prefix": null,
"problem_type": null,
"projection_dim" : 768,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"task_specific_params": null,
"temperature": 1.0,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.16.0.dev0",
"use_bfloat16": false
},
"vision_config_dict": {
"hidden_size": 1024,
"intermediate_size": 4096,
"num_attention_heads": 16,
"num_hidden_layers": 24,
"patch_size": 14,
"projection_dim": 768
}
}
clip/merges.txt
+48895
View File
File diff suppressed because it is too large Load Diff
clip/preprocessor_config.json
+19
View File
@@ -0,0 +1,19 @@
{
"crop_size": 224,
"do_center_crop": true,
"do_normalize": true,
"do_resize": true,
"feature_extractor_type": "CLIPFeatureExtractor",
"image_mean": [
0.48145466,
0.4578275,
0.40821073
],
"image_std": [
0.26862954,
0.26130258,
0.27577711
],
"resample": 3,
"size": 224
}
clip/special_tokens_map.json
+1
View File
@@ -0,0 +1 @@
{"bos_token": {"content": "<|startoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "pad_token": "<|endoftext|>"}
clip/tokenizer.json
+98393
View File
File diff suppressed because it is too large Load Diff
clip/tokenizer_config.json
+34
View File
@@ -0,0 +1,34 @@
{
"unk_token": {
"content": "<|endoftext|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": true,
"__type": "AddedToken"
},
"bos_token": {
"content": "<|startoftext|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": true,
"__type": "AddedToken"
},
"eos_token": {
"content": "<|endoftext|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": true,
"__type": "AddedToken"
},
"pad_token": "<|endoftext|>",
"add_prefix_space": false,
"errors": "replace",
"do_lower_case": true,
"name_or_path": "openai/clip-vit-base-patch32",
"model_max_length": 77,
"special_tokens_map_file": "./special_tokens_map.json",
"tokenizer_class": "CLIPTokenizer"
}
clip/vocab.json
+1
View File
File diff suppressed because one or more lines are too long
configs/generation/slat_flow_img_dit_L_64l8p2_fp16.json
+102
View File
@@ -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": "JeffreyXiang/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
+101
View File
@@ -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": "JeffreyXiang/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
+101
View File
@@ -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": "JeffreyXiang/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
+101
View File
@@ -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": "JeffreyXiang/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
+70
View File
@@ -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": "JeffreyXiang/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
+69
View File
@@ -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": "JeffreyXiang/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
+69
View File
@@ -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": "JeffreyXiang/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
+70
View File
@@ -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": "JeffreyXiang/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
+73
View File
@@ -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
+71
View File
@@ -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
+105
View File
@@ -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
+65
View File
@@ -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
}
}
}
example.py
-54
View File
@@ -1,54 +0,0 @@
import os
# os.environ['ATTN_BACKEND'] = 'xformers' # Can be 'flash-attn' or 'xformers', default is 'flash-attn'
os.environ['SPCONV_ALGO'] = 'native' # Can be 'native' or 'auto', default is 'auto'.
# 'auto' is faster but will do benchmarking at the beginning.
# Recommended to set to 'native' if run only once.
import imageio
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import render_utils, postprocessing_utils
# Load a pipeline from a model folder or a Hugging Face model hub.
pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
pipeline.cuda()
# Load an image
image = Image.open("assets/example_image/T.png")
# Run the pipeline
outputs = pipeline.run(
image,
# Optional parameters
seed=1,
# sparse_structure_sampler_params={
# "steps": 12,
# "cfg_strength": 7.5,
# },
# slat_sampler_params={
# "steps": 12,
# "cfg_strength": 3,
# },
)
# outputs is a dictionary containing generated 3D assets in different formats:
# - outputs['gaussian']: a list of 3D Gaussians
# - outputs['radiance_field']: a list of radiance fields
# - outputs['mesh']: a list of meshes
# Render the outputs
video = render_utils.render_video(outputs['gaussian'][0])['color']
imageio.mimsave("sample_gs.mp4", video, fps=30)
video = render_utils.render_video(outputs['radiance_field'][0])['color']
imageio.mimsave("sample_rf.mp4", video, fps=30)
video = render_utils.render_video(outputs['mesh'][0])['normal']
imageio.mimsave("sample_mesh.mp4", video, fps=30)
# GLB files can be extracted from the outputs
glb = postprocessing_utils.to_glb(
outputs['gaussian'][0],
outputs['mesh'][0],
# Optional parameters
simplify=0.95, # Ratio of triangles to remove in the simplification process
texture_size=1024, # Size of the texture used for the GLB
)
glb.export("sample.glb")
example1220.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 364 KiB

exmaple.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 301 KiB

trellis/datasets/__init__.py
+58
View File
@@ -0,0 +1,58 @@
import importlib
__attributes = {
'SparseStructure': 'sparse_structure',
'SparseFeat2Render': 'sparse_feat2render',
'SLat2Render':'structured_latent2render',
'Slat2RenderGeo':'structured_latent2render',
'SparseStructureLatent': 'sparse_structure_latent',
'TextConditionedSparseStructureLatent': 'sparse_structure_latent',
'ImageConditionedSparseStructureLatent': 'sparse_structure_latent',
'SLat': 'structured_latent',
'TextConditionedSLat': 'structured_latent',
'ImageConditionedSLat': 'structured_latent',
}
__submodules = []
__all__ = list(__attributes.keys()) + __submodules
def __getattr__(name):
if name not in globals():
if name in __attributes:
module_name = __attributes[name]
module = importlib.import_module(f".{module_name}", __name__)
globals()[name] = getattr(module, name)
elif name in __submodules:
module = importlib.import_module(f".{name}", __name__)
globals()[name] = module
else:
raise AttributeError(f"module {__name__} has no attribute {name}")
return globals()[name]
# For Pylance
if __name__ == '__main__':
from .sparse_structure import SparseStructure
from .sparse_feat2render import SparseFeat2Render
from .structured_latent2render import (
SLat2Render,
Slat2RenderGeo,
)
from .sparse_structure_latent import (
SparseStructureLatent,
TextConditionedSparseStructureLatent,
ImageConditionedSparseStructureLatent,
)
from .structured_latent import (
SLat,
TextConditionedSLat,
ImageConditionedSLat,
)
trellis/datasets/components.py
+137
View File
@@ -0,0 +1,137 @@
from typing import *
from abc import abstractmethod
import os
import json
import torch
import numpy as np
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset
class StandardDatasetBase(Dataset):
"""
Base class for standard datasets.
Args:
roots (str): paths to the dataset
"""
def __init__(self,
roots: str,
):
super().__init__()
self.roots = roots.split(',')
self.instances = []
self.metadata = pd.DataFrame()
self._stats = {}
for root in self.roots:
key = os.path.basename(root)
self._stats[key] = {}
metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))
self._stats[key]['Total'] = len(metadata)
metadata, stats = self.filter_metadata(metadata)
self._stats[key].update(stats)
self.instances.extend([(root, sha256) for sha256 in metadata['sha256'].values])
metadata.set_index('sha256', inplace=True)
self.metadata = pd.concat([self.metadata, metadata])
@abstractmethod
def filter_metadata(self, metadata: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, int]]:
pass
@abstractmethod
def get_instance(self, root: str, instance: str) -> Dict[str, Any]:
pass
def __len__(self):
return len(self.instances)
def __getitem__(self, index) -> Dict[str, Any]:
try:
root, instance = self.instances[index]
return self.get_instance(root, instance)
except Exception as e:
print(e)
return self.__getitem__(np.random.randint(0, len(self)))
def __str__(self):
lines = []
lines.append(self.__class__.__name__)
lines.append(f' - Total instances: {len(self)}')
lines.append(f' - Sources:')
for key, stats in self._stats.items():
lines.append(f' - {key}:')
for k, v in stats.items():
lines.append(f' - {k}: {v}')
return '\n'.join(lines)
class TextConditionedMixin:
def __init__(self, roots, **kwargs):
super().__init__(roots, **kwargs)
self.captions = {}
for instance in self.instances:
sha256 = instance[1]
self.captions[sha256] = json.loads(self.metadata.loc[sha256]['captions'])
def filter_metadata(self, metadata):
metadata, stats = super().filter_metadata(metadata)
metadata = metadata[metadata['captions'].notna()]
stats['With captions'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
pack = super().get_instance(root, instance)
text = np.random.choice(self.captions[instance])
pack['cond'] = text
return pack
class ImageConditionedMixin:
def __init__(self, roots, *, image_size=518, **kwargs):
self.image_size = image_size
super().__init__(roots, **kwargs)
def filter_metadata(self, metadata):
metadata, stats = super().filter_metadata(metadata)
metadata = metadata[metadata[f'cond_rendered']]
stats['Cond rendered'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
pack = super().get_instance(root, instance)
image_root = os.path.join(root, 'renders_cond', instance)
with open(os.path.join(image_root, 'transforms.json')) as f:
metadata = json.load(f)
n_views = len(metadata['frames'])
view = np.random.randint(n_views)
metadata = metadata['frames'][view]
image_path = os.path.join(image_root, metadata['file_path'])
image = Image.open(image_path)
alpha = np.array(image.getchannel(3))
bbox = np.array(alpha).nonzero()
bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
aug_size_ratio = 1.2
aug_hsize = hsize * aug_size_ratio
aug_center_offset = [0, 0]
aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]]
aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)]
image = image.crop(aug_bbox)
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
alpha = image.getchannel(3)
image = image.convert('RGB')
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
alpha = torch.tensor(np.array(alpha)).float() / 255.0
image = image * alpha.unsqueeze(0)
pack['cond'] = image
return pack
trellis/datasets/sparse_feat2render.py
+134
View File
@@ -0,0 +1,134 @@
import os
from PIL import Image
import json
import numpy as np
import pandas as pd
import torch
import utils3d.torch
from ..modules.sparse.basic import SparseTensor
from .components import StandardDatasetBase
class SparseFeat2Render(StandardDatasetBase):
"""
SparseFeat2Render dataset.
Args:
roots (str): paths to the dataset
image_size (int): size of the image
model (str): model name
resolution (int): resolution of the data
min_aesthetic_score (float): minimum aesthetic score
max_num_voxels (int): maximum number of voxels
"""
def __init__(
self,
roots: str,
image_size: int,
model: str = 'dinov2_vitl14_reg',
resolution: int = 64,
min_aesthetic_score: float = 5.0,
max_num_voxels: int = 32768,
):
self.image_size = image_size
self.model = model
self.resolution = resolution
self.min_aesthetic_score = min_aesthetic_score
self.max_num_voxels = max_num_voxels
self.value_range = (0, 1)
super().__init__(roots)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'feature_{self.model}']]
stats['With features'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]
stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
return metadata, stats
def _get_image(self, root, instance):
with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:
metadata = json.load(f)
n_views = len(metadata['frames'])
view = np.random.randint(n_views)
metadata = metadata['frames'][view]
fov = metadata['camera_angle_x']
intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
c2w = torch.tensor(metadata['transform_matrix'])
c2w[:3, 1:3] *= -1
extrinsics = torch.inverse(c2w)
image_path = os.path.join(root, 'renders', instance, metadata['file_path'])
image = Image.open(image_path)
alpha = image.getchannel(3)
image = image.convert('RGB')
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
alpha = torch.tensor(np.array(alpha)).float() / 255.0
return {
'image': image,
'alpha': alpha,
'extrinsics': extrinsics,
'intrinsics': intrinsics,
}
def _get_feat(self, root, instance):
DATA_RESOLUTION = 64
feats_path = os.path.join(root, 'features', self.model, f'{instance}.npz')
feats = np.load(feats_path, allow_pickle=True)
coords = torch.tensor(feats['indices']).int()
feats = torch.tensor(feats['patchtokens']).float()
if self.resolution != DATA_RESOLUTION:
factor = DATA_RESOLUTION // self.resolution
coords = coords // factor
coords, idx = coords.unique(return_inverse=True, dim=0)
feats = torch.scatter_reduce(
torch.zeros(coords.shape[0], feats.shape[1], device=feats.device),
dim=0,
index=idx.unsqueeze(-1).expand(-1, feats.shape[1]),
src=feats,
reduce='mean'
)
return {
'coords': coords,
'feats': feats,
}
@torch.no_grad()
def visualize_sample(self, sample: dict):
return sample['image']
@staticmethod
def collate_fn(batch):
pack = {}
coords = []
for i, b in enumerate(batch):
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
coords = torch.cat(coords)
feats = torch.cat([b['feats'] for b in batch])
pack['feats'] = SparseTensor(
coords=coords,
feats=feats,
)
pack['image'] = torch.stack([b['image'] for b in batch])
pack['alpha'] = torch.stack([b['alpha'] for b in batch])
pack['extrinsics'] = torch.stack([b['extrinsics'] for b in batch])
pack['intrinsics'] = torch.stack([b['intrinsics'] for b in batch])
return pack
def get_instance(self, root, instance):
image = self._get_image(root, instance)
feat = self._get_feat(root, instance)
return {
**image,
**feat,
}
trellis/datasets/sparse_structure.py
+107
View File
@@ -0,0 +1,107 @@
import os
import json
from typing import Union
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
import utils3d
from .components import StandardDatasetBase
from ..representations.octree import DfsOctree as Octree
from ..renderers import OctreeRenderer
class SparseStructure(StandardDatasetBase):
"""
Sparse structure dataset
Args:
roots (str): path to the dataset
resolution (int): resolution of the voxel grid
min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
"""
def __init__(self,
roots,
resolution: int = 64,
min_aesthetic_score: float = 5.0,
):
self.resolution = resolution
self.min_aesthetic_score = min_aesthetic_score
self.value_range = (0, 1)
super().__init__(roots)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'voxelized']]
stats['Voxelized'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
position = utils3d.io.read_ply(os.path.join(root, 'voxels', f'{instance}.ply'))[0]
coords = ((torch.tensor(position) + 0.5) * self.resolution).int().contiguous()
ss = torch.zeros(1, self.resolution, self.resolution, self.resolution, dtype=torch.long)
ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
return {'ss': ss}
@torch.no_grad()
def visualize_sample(self, ss: Union[torch.Tensor, dict]):
ss = ss if isinstance(ss, torch.Tensor) else ss['ss']
renderer = OctreeRenderer()
renderer.rendering_options.resolution = 512
renderer.rendering_options.near = 0.8
renderer.rendering_options.far = 1.6
renderer.rendering_options.bg_color = (0, 0, 0)
renderer.rendering_options.ssaa = 4
renderer.pipe.primitive = 'voxel'
# Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
exts = []
ints = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(30)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
exts.append(extrinsics)
ints.append(intrinsics)
images = []
# Build each representation
ss = ss.cuda()
for i in range(ss.shape[0]):
representation = Octree(
depth=10,
aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
device='cuda',
primitive='voxel',
sh_degree=0,
primitive_config={'solid': True},
)
coords = torch.nonzero(ss[i, 0], as_tuple=False)
representation.position = coords.float() / self.resolution
representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')
image = torch.zeros(3, 1024, 1024).cuda()
tile = [2, 2]
for j, (ext, intr) in enumerate(zip(exts, ints)):
res = renderer.render(representation, ext, intr, colors_overwrite=representation.position)
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
images.append(image)
return torch.stack(images)
trellis/datasets/sparse_structure_latent.py
+189
View File
@@ -0,0 +1,189 @@
import os
import json
from typing import *
import numpy as np
import torch
import utils3d
from ..representations.octree import DfsOctree as Octree
from ..renderers import OctreeRenderer
from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin
from .. import models
from ..utils.dist_utils import read_file_dist
class SparseStructureLatentVisMixin:
def __init__(
self,
*args,
pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
ss_dec_path: Optional[str] = None,
ss_dec_ckpt: Optional[str] = None,
**kwargs
):
super().__init__(*args, **kwargs)
self.ss_dec = None
self.pretrained_ss_dec = pretrained_ss_dec
self.ss_dec_path = ss_dec_path
self.ss_dec_ckpt = ss_dec_ckpt
def _loading_ss_dec(self):
if self.ss_dec is not None:
return
if self.ss_dec_path is not None:
cfg = json.load(open(os.path.join(self.ss_dec_path, 'config.json'), 'r'))
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
ckpt_path = os.path.join(self.ss_dec_path, 'ckpts', f'decoder_{self.ss_dec_ckpt}.pt')
decoder.load_state_dict(torch.load(read_file_dist(ckpt_path), map_location='cpu', weights_only=True))
else:
decoder = models.from_pretrained(self.pretrained_ss_dec)
self.ss_dec = decoder.cuda().eval()
def _delete_ss_dec(self):
del self.ss_dec
self.ss_dec = None
@torch.no_grad()
def decode_latent(self, z, batch_size=4):
self._loading_ss_dec()
ss = []
if self.normalization is not None:
z = z * self.std.to(z.device) + self.mean.to(z.device)
for i in range(0, z.shape[0], batch_size):
ss.append(self.ss_dec(z[i:i+batch_size]))
ss = torch.cat(ss, dim=0)
self._delete_ss_dec()
return ss
@torch.no_grad()
def visualize_sample(self, x_0: Union[torch.Tensor, dict]):
x_0 = x_0 if isinstance(x_0, torch.Tensor) else x_0['x_0']
x_0 = self.decode_latent(x_0.cuda())
renderer = OctreeRenderer()
renderer.rendering_options.resolution = 512
renderer.rendering_options.near = 0.8
renderer.rendering_options.far = 1.6
renderer.rendering_options.bg_color = (0, 0, 0)
renderer.rendering_options.ssaa = 4
renderer.pipe.primitive = 'voxel'
# Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
exts = []
ints = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(30)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
exts.append(extrinsics)
ints.append(intrinsics)
images = []
# Build each representation
x_0 = x_0.cuda()
for i in range(x_0.shape[0]):
representation = Octree(
depth=10,
aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
device='cuda',
primitive='voxel',
sh_degree=0,
primitive_config={'solid': True},
)
coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False)
resolution = x_0.shape[-1]
representation.position = coords.float() / resolution
representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(resolution)), dtype=torch.uint8, device='cuda')
image = torch.zeros(3, 1024, 1024).cuda()
tile = [2, 2]
for j, (ext, intr) in enumerate(zip(exts, ints)):
res = renderer.render(representation, ext, intr, colors_overwrite=representation.position)
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
images.append(image)
return torch.stack(images)
class SparseStructureLatent(SparseStructureLatentVisMixin, StandardDatasetBase):
"""
Sparse structure latent dataset
Args:
roots (str): path to the dataset
latent_model (str): name of the latent model
min_aesthetic_score (float): minimum aesthetic score
normalization (dict): normalization stats
pretrained_ss_dec (str): name of the pretrained sparse structure decoder
ss_dec_path (str): path to the sparse structure decoder, if given, will override the pretrained_ss_dec
ss_dec_ckpt (str): name of the sparse structure decoder checkpoint
"""
def __init__(self,
roots: str,
*,
latent_model: str,
min_aesthetic_score: float = 5.0,
normalization: Optional[dict] = None,
pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
ss_dec_path: Optional[str] = None,
ss_dec_ckpt: Optional[str] = None,
):
self.latent_model = latent_model
self.min_aesthetic_score = min_aesthetic_score
self.normalization = normalization
self.value_range = (0, 1)
super().__init__(
roots,
pretrained_ss_dec=pretrained_ss_dec,
ss_dec_path=ss_dec_path,
ss_dec_ckpt=ss_dec_ckpt,
)
if self.normalization is not None:
self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'ss_latent_{self.latent_model}']]
stats['With sparse structure latents'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
latent = np.load(os.path.join(root, 'ss_latents', self.latent_model, f'{instance}.npz'))
z = torch.tensor(latent['mean']).float()
if self.normalization is not None:
z = (z - self.mean) / self.std
pack = {
'x_0': z,
}
return pack
class TextConditionedSparseStructureLatent(TextConditionedMixin, SparseStructureLatent):
"""
Text-conditioned sparse structure dataset
"""
pass
class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent):
"""
Image-conditioned sparse structure dataset
"""
pass
trellis/datasets/structured_latent.py
+218
View File
@@ -0,0 +1,218 @@
import json
import os
from typing import *
import numpy as np
import torch
import utils3d.torch
from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin
from ..modules.sparse.basic import SparseTensor
from .. import models
from ..utils.render_utils import get_renderer
from ..utils.dist_utils import read_file_dist
from ..utils.data_utils import load_balanced_group_indices
class SLatVisMixin:
def __init__(
self,
*args,
pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
slat_dec_path: Optional[str] = None,
slat_dec_ckpt: Optional[str] = None,
**kwargs
):
super().__init__(*args, **kwargs)
self.slat_dec = None
self.pretrained_slat_dec = pretrained_slat_dec
self.slat_dec_path = slat_dec_path
self.slat_dec_ckpt = slat_dec_ckpt
def _loading_slat_dec(self):
if self.slat_dec is not None:
return
if self.slat_dec_path is not None:
cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
decoder.load_state_dict(torch.load(read_file_dist(ckpt_path), map_location='cpu', weights_only=True))
else:
decoder = models.from_pretrained(self.pretrained_slat_dec)
self.slat_dec = decoder.cuda().eval()
def _delete_slat_dec(self):
del self.slat_dec
self.slat_dec = None
@torch.no_grad()
def decode_latent(self, z, batch_size=4):
self._loading_slat_dec()
reps = []
if self.normalization is not None:
z = z * self.std.to(z.device) + self.mean.to(z.device)
for i in range(0, z.shape[0], batch_size):
reps.append(self.slat_dec(z[i:i+batch_size]))
reps = sum(reps, [])
self._delete_slat_dec()
return reps
@torch.no_grad()
def visualize_sample(self, x_0: Union[SparseTensor, dict]):
x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
reps = self.decode_latent(x_0.cuda())
# Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
exts = []
ints = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(40)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
exts.append(extrinsics)
ints.append(intrinsics)
renderer = get_renderer(reps[0])
images = []
for representation in reps:
image = torch.zeros(3, 1024, 1024).cuda()
tile = [2, 2]
for j, (ext, intr) in enumerate(zip(exts, ints)):
res = renderer.render(representation, ext, intr)
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
images.append(image)
images = torch.stack(images)
return images
class SLat(SLatVisMixin, StandardDatasetBase):
"""
structured latent dataset
Args:
roots (str): path to the dataset
latent_model (str): name of the latent model
min_aesthetic_score (float): minimum aesthetic score
max_num_voxels (int): maximum number of voxels
normalization (dict): normalization stats
pretrained_slat_dec (str): name of the pretrained slat decoder
slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
slat_dec_ckpt (str): name of the slat decoder checkpoint
"""
def __init__(self,
roots: str,
*,
latent_model: str,
min_aesthetic_score: float = 5.0,
max_num_voxels: int = 32768,
normalization: Optional[dict] = None,
pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
slat_dec_path: Optional[str] = None,
slat_dec_ckpt: Optional[str] = None,
):
self.normalization = normalization
self.latent_model = latent_model
self.min_aesthetic_score = min_aesthetic_score
self.max_num_voxels = max_num_voxels
self.value_range = (0, 1)
super().__init__(
roots,
pretrained_slat_dec=pretrained_slat_dec,
slat_dec_path=slat_dec_path,
slat_dec_ckpt=slat_dec_ckpt,
)
self.loads = [self.metadata.loc[sha256, 'num_voxels'] for _, sha256 in self.instances]
if self.normalization is not None:
self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'latent_{self.latent_model}']]
stats['With latent'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]
stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz'))
coords = torch.tensor(data['coords']).int()
feats = torch.tensor(data['feats']).float()
if self.normalization is not None:
feats = (feats - self.mean) / self.std
return {
'coords': coords,
'feats': feats,
}
@staticmethod
def collate_fn(batch, split_size=None):
if split_size is None:
group_idx = [list(range(len(batch)))]
else:
group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
packs = []
for group in group_idx:
sub_batch = [batch[i] for i in group]
pack = {}
coords = []
feats = []
layout = []
start = 0
for i, b in enumerate(sub_batch):
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
feats.append(b['feats'])
layout.append(slice(start, start + b['coords'].shape[0]))
start += b['coords'].shape[0]
coords = torch.cat(coords)
feats = torch.cat(feats)
pack['x_0'] = SparseTensor(
coords=coords,
feats=feats,
)
pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
pack['x_0'].register_spatial_cache('layout', layout)
# collate other data
keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
for k in keys:
if isinstance(sub_batch[0][k], torch.Tensor):
pack[k] = torch.stack([b[k] for b in sub_batch])
elif isinstance(sub_batch[0][k], list):
pack[k] = sum([b[k] for b in sub_batch], [])
else:
pack[k] = [b[k] for b in sub_batch]
packs.append(pack)
if split_size is None:
return packs[0]
return packs
class TextConditionedSLat(TextConditionedMixin, SLat):
"""
Text conditioned structured latent dataset
"""
pass
class ImageConditionedSLat(ImageConditionedMixin, SLat):
"""
Image conditioned structured latent dataset
"""
pass
trellis/datasets/structured_latent2render.py
+160
View File
@@ -0,0 +1,160 @@
import os
from PIL import Image
import json
import numpy as np
import torch
import utils3d.torch
from ..modules.sparse.basic import SparseTensor
from .components import StandardDatasetBase
class SLat2Render(StandardDatasetBase):
"""
Dataset for Structured Latent and rendered images.
Args:
roots (str): paths to the dataset
image_size (int): size of the image
latent_model (str): latent model name
min_aesthetic_score (float): minimum aesthetic score
max_num_voxels (int): maximum number of voxels
"""
def __init__(
self,
roots: str,
image_size: int,
latent_model: str,
min_aesthetic_score: float = 5.0,
max_num_voxels: int = 32768,
):
self.image_size = image_size
self.latent_model = latent_model
self.min_aesthetic_score = min_aesthetic_score
self.max_num_voxels = max_num_voxels
self.value_range = (0, 1)
super().__init__(roots)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'latent_{self.latent_model}']]
stats['With latent'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]
stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
return metadata, stats
def _get_image(self, root, instance):
with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:
metadata = json.load(f)
n_views = len(metadata['frames'])
view = np.random.randint(n_views)
metadata = metadata['frames'][view]
fov = metadata['camera_angle_x']
intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
c2w = torch.tensor(metadata['transform_matrix'])
c2w[:3, 1:3] *= -1
extrinsics = torch.inverse(c2w)
image_path = os.path.join(root, 'renders', instance, metadata['file_path'])
image = Image.open(image_path)
alpha = image.getchannel(3)
image = image.convert('RGB')
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
alpha = torch.tensor(np.array(alpha)).float() / 255.0
return {
'image': image,
'alpha': alpha,
'extrinsics': extrinsics,
'intrinsics': intrinsics,
}
def _get_latent(self, root, instance):
data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz'))
coords = torch.tensor(data['coords']).int()
feats = torch.tensor(data['feats']).float()
return {
'coords': coords,
'feats': feats,
}
@torch.no_grad()
def visualize_sample(self, sample: dict):
return sample['image']
@staticmethod
def collate_fn(batch):
pack = {}
coords = []
for i, b in enumerate(batch):
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
coords = torch.cat(coords)
feats = torch.cat([b['feats'] for b in batch])
pack['latents'] = SparseTensor(
coords=coords,
feats=feats,
)
# collate other data
keys = [k for k in batch[0].keys() if k not in ['coords', 'feats']]
for k in keys:
if isinstance(batch[0][k], torch.Tensor):
pack[k] = torch.stack([b[k] for b in batch])
elif isinstance(batch[0][k], list):
pack[k] = sum([b[k] for b in batch], [])
else:
pack[k] = [b[k] for b in batch]
return pack
def get_instance(self, root, instance):
image = self._get_image(root, instance)
latent = self._get_latent(root, instance)
return {
**image,
**latent,
}
class Slat2RenderGeo(SLat2Render):
def __init__(
self,
roots: str,
image_size: int,
latent_model: str,
min_aesthetic_score: float = 5.0,
max_num_voxels: int = 32768,
):
super().__init__(
roots,
image_size,
latent_model,
min_aesthetic_score,
max_num_voxels,
)
def _get_geo(self, root, instance):
verts, face = utils3d.io.read_ply(os.path.join(root, 'renders', instance, 'mesh.ply'))
mesh = {
"vertices" : torch.from_numpy(verts),
"faces" : torch.from_numpy(face),
}
return {
"mesh" : mesh,
}
def get_instance(self, root, instance):
image = self._get_image(root, instance)
latent = self._get_latent(root, instance)
geo = self._get_geo(root, instance)
return {
**image,
**latent,
**geo,
}
trellis/models/__init__.py
+29 -3
View File
@@ -3,12 +3,20 @@ import importlib
__attributes = {
'SparseStructureEncoder': 'sparse_structure_vae',
'SparseStructureDecoder': 'sparse_structure_vae',
'SparseStructureFlowModel': 'sparse_structure_flow',
'SLatEncoder': 'structured_latent_vae',
'SLatGaussianDecoder': 'structured_latent_vae',
'SLatRadianceFieldDecoder': 'structured_latent_vae',
'SLatMeshDecoder': 'structured_latent_vae',
'ElasticSLatEncoder': 'structured_latent_vae',
'ElasticSLatGaussianDecoder': 'structured_latent_vae',
'ElasticSLatRadianceFieldDecoder': 'structured_latent_vae',
'ElasticSLatMeshDecoder': 'structured_latent_vae',
'SLatFlowModel': 'structured_latent_flow',
'ElasticSLatFlowModel': 'structured_latent_flow',
}
__submodules = []
@@ -64,7 +72,25 @@ def from_pretrained(path: str, **kwargs):
# For Pylance
if __name__ == '__main__':
from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
from .sparse_structure_vae import (
SparseStructureEncoder,
SparseStructureDecoder,
)
from .sparse_structure_flow import SparseStructureFlowModel
from .structured_latent_vae import SLatEncoder, SLatGaussianDecoder, SLatRadianceFieldDecoder, SLatMeshDecoder
from .structured_latent_flow import SLatFlowModel
from .structured_latent_vae import (
SLatEncoder,
SLatGaussianDecoder,
SLatRadianceFieldDecoder,
SLatMeshDecoder,
ElasticSLatEncoder,
ElasticSLatGaussianDecoder,
ElasticSLatRadianceFieldDecoder,
ElasticSLatMeshDecoder,
)
from .structured_latent_flow import (
SLatFlowModel,
ElasticSLatFlowModel,
)
trellis/models/sparse_elastic_mixin.py
+24
View File
@@ -0,0 +1,24 @@
from contextlib import contextmanager
from typing import *
import math
from ..modules import sparse as sp
from ..utils.elastic_utils import ElasticModuleMixin
class SparseTransformerElasticMixin(ElasticModuleMixin):
def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs):
return x.feats.shape[0]
@contextmanager
def with_mem_ratio(self, mem_ratio=1.0):
if mem_ratio == 1.0:
yield 1.0
return
num_blocks = len(self.blocks)
num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks)
exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks
for i in range(num_blocks):
self.blocks[i].use_checkpoint = i < num_checkpoint_blocks
yield exact_mem_ratio
for i in range(num_blocks):
self.blocks[i].use_checkpoint = False
trellis/models/structured_latent_flow.py
+50 -36
View File
@@ -9,6 +9,7 @@ from ..modules.norm import LayerNorm32
from ..modules import sparse as sp
from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
from .sparse_structure_flow import TimestepEmbedder
from .sparse_elastic_mixin import SparseTransformerElasticMixin
class SparseResBlock3d(nn.Module):
@@ -109,8 +110,9 @@ class SLatFlowModel(nn.Module):
self.qk_rms_norm_cross = qk_rms_norm_cross
self.dtype = torch.float16 if use_fp16 else torch.float32
assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
if self.io_block_channels is not None:
assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
self.t_embedder = TimestepEmbedder(model_channels)
if share_mod:
@@ -122,25 +124,27 @@ class SLatFlowModel(nn.Module):
if pe_mode == "ape":
self.pos_embedder = AbsolutePositionEmbedder(model_channels)
self.input_layer = sp.SparseLinear(in_channels, io_block_channels[0])
self.input_layer = sp.SparseLinear(in_channels, model_channels if io_block_channels is None else io_block_channels[0])
self.input_blocks = nn.ModuleList([])
for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
self.input_blocks.extend([
SparseResBlock3d(
chs,
model_channels,
out_channels=chs,
if io_block_channels is not None:
for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
self.input_blocks.extend([
SparseResBlock3d(
chs,
model_channels,
out_channels=chs,
)
for _ in range(num_io_res_blocks-1)
])
self.input_blocks.append(
SparseResBlock3d(
chs,
model_channels,
out_channels=next_chs,
downsample=True,
)
)
for _ in range(num_io_res_blocks-1)
])
self.input_blocks.append(
SparseResBlock3d(
chs,
model_channels,
out_channels=next_chs,
downsample=True,
)
)
self.blocks = nn.ModuleList([
ModulatedSparseTransformerCrossBlock(
@@ -159,24 +163,26 @@ class SLatFlowModel(nn.Module):
])
self.out_blocks = nn.ModuleList([])
for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
self.out_blocks.append(
SparseResBlock3d(
prev_chs * 2 if self.use_skip_connection else prev_chs,
model_channels,
out_channels=chs,
upsample=True,
if io_block_channels is not None:
for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
self.out_blocks.append(
SparseResBlock3d(
prev_chs * 2 if self.use_skip_connection else prev_chs,
model_channels,
out_channels=chs,
upsample=True,
)
)
)
self.out_blocks.extend([
SparseResBlock3d(
chs * 2 if self.use_skip_connection else chs,
model_channels,
out_channels=chs,
)
for _ in range(num_io_res_blocks-1)
])
self.out_layer = sp.SparseLinear(io_block_channels[0], out_channels)
self.out_blocks.extend([
SparseResBlock3d(
chs * 2 if self.use_skip_connection else chs,
model_channels,
out_channels=chs,
)
for _ in range(num_io_res_blocks-1)
])
self.out_layer = sp.SparseLinear(model_channels if io_block_channels is None else io_block_channels[0], out_channels)
self.initialize_weights()
if use_fp16:
@@ -260,3 +266,11 @@ class SLatFlowModel(nn.Module):
h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
h = self.out_layer(h.type(x.dtype))
return h
class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel):
"""
SLat Flow Model with elastic memory management.
Used for training with low VRAM.
"""
pass
trellis/models/structured_latent_vae/__init__.py
+4 -4
View File
@@ -1,4 +1,4 @@
from .encoder import SLatEncoder
from .decoder_gs import SLatGaussianDecoder
from .decoder_rf import SLatRadianceFieldDecoder
from .decoder_mesh import SLatMeshDecoder
from .encoder import SLatEncoder, ElasticSLatEncoder
from .decoder_gs import SLatGaussianDecoder, ElasticSLatGaussianDecoder
from .decoder_rf import SLatRadianceFieldDecoder, ElasticSLatRadianceFieldDecoder
from .decoder_mesh import SLatMeshDecoder, ElasticSLatMeshDecoder
trellis/models/structured_latent_vae/decoder_gs.py
+9
View File
@@ -6,6 +6,7 @@ from ...modules import sparse as sp
from ...utils.random_utils import hammersley_sequence
from .base import SparseTransformerBase
from ...representations import Gaussian
from ..sparse_elastic_mixin import SparseTransformerElasticMixin
class SLatGaussianDecoder(SparseTransformerBase):
@@ -120,3 +121,11 @@ class SLatGaussianDecoder(SparseTransformerBase):
h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
h = self.out_layer(h)
return self.to_representation(h)
class ElasticSLatGaussianDecoder(SparseTransformerElasticMixin, SLatGaussianDecoder):
"""
Slat VAE Gaussian decoder with elastic memory management.
Used for training with low VRAM.
"""
pass
trellis/models/structured_latent_vae/decoder_mesh.py
+9
View File
@@ -8,6 +8,7 @@ from ...modules import sparse as sp
from .base import SparseTransformerBase
from ...representations import MeshExtractResult
from ...representations.mesh import SparseFeatures2Mesh
from ..sparse_elastic_mixin import SparseTransformerElasticMixin
class SparseSubdivideBlock3d(nn.Module):
@@ -165,3 +166,11 @@ class SLatMeshDecoder(SparseTransformerBase):
h = h.type(x.dtype)
h = self.out_layer(h)
return self.to_representation(h)
class ElasticSLatMeshDecoder(SparseTransformerElasticMixin, SLatMeshDecoder):
"""
Slat VAE Mesh decoder with elastic memory management.
Used for training with low VRAM.
"""
pass
trellis/models/structured_latent_vae/decoder_rf.py
+9
View File
@@ -6,6 +6,7 @@ import numpy as np
from ...modules import sparse as sp
from .base import SparseTransformerBase
from ...representations import Strivec
from ..sparse_elastic_mixin import SparseTransformerElasticMixin
class SLatRadianceFieldDecoder(SparseTransformerBase):
@@ -102,3 +103,11 @@ class SLatRadianceFieldDecoder(SparseTransformerBase):
h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
h = self.out_layer(h)
return self.to_representation(h)
class ElasticSLatRadianceFieldDecoder(SparseTransformerElasticMixin, SLatRadianceFieldDecoder):
"""
Slat VAE Radiance Field Decoder with elastic memory management.
Used for training with low VRAM.
"""
pass
trellis/models/structured_latent_vae/encoder.py
+8
View File
@@ -4,6 +4,7 @@ import torch.nn as nn
import torch.nn.functional as F
from ...modules import sparse as sp
from .base import SparseTransformerBase
from ..sparse_elastic_mixin import SparseTransformerElasticMixin
class SLatEncoder(SparseTransformerBase):
@@ -70,3 +71,10 @@ class SLatEncoder(SparseTransformerBase):
return z, mean, logvar
else:
return z
class ElasticSLatEncoder(SparseTransformerElasticMixin, SLatEncoder):
"""
SLat VAE encoder with elastic memory management.
Used for training with low VRAM.
"""
trellis/pipelines/__init__.py
+1 -1
View File
@@ -1,6 +1,6 @@
from . import samplers
from .trellis_image_to_3d import TrellisImageTo3DPipeline
from .trellis_text_to_3d import TrellisTextTo3DPipeline
def from_pretrained(path: str):
"""
trellis/pipelines/base.py
+6 -4
View File
@@ -36,10 +36,12 @@ class Pipeline:
with open(config_file, 'r') as f:
args = json.load(f)['args']
_models = {
k: models.from_pretrained(f"{path}/{v}")
for k, v in args['models'].items()
}
_models = {}
for k, v in args['models'].items():
try:
_models[k] = models.from_pretrained(f"{path}/{v}")
except:
_models[k] = models.from_pretrained(v)
new_pipeline = Pipeline(_models)
new_pipeline._pretrained_args = args
trellis/pipelines/samplers/flow_euler.py
+2
View File
@@ -37,6 +37,8 @@ class FlowEulerSampler(Sampler):
def _inference_model(self, model, x_t, t, cond=None, **kwargs):
t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
if cond is not None and cond.shape[0] == 1 and x_t.shape[0] > 1:
cond = cond.repeat(x_t.shape[0], *([1] * (len(cond.shape) - 1)))
return model(x_t, t, cond, **kwargs)
def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):
trellis/pipelines/trellis_text_to_3d.py
+302
View File
@@ -0,0 +1,302 @@
from typing import *
import torch
import torch.nn as nn
import numpy as np
from transformers import CLIPTextModel, AutoTokenizer,CLIPTextConfig
from contextlib import nullcontext
from safetensors.torch import load_file
import open3d as o3d
from .base import Pipeline
from . import samplers
from ..modules import sparse as sp
class TrellisTextTo3DPipeline(Pipeline):
"""
Pipeline for inferring Trellis text-to-3D models.
Args:
models (dict[str, nn.Module]): The models to use in the pipeline.
sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
slat_sampler (samplers.Sampler): The sampler for the structured latent.
slat_normalization (dict): The normalization parameters for the structured latent.
text_cond_model (str): The name of the text conditioning model.
"""
def __init__(
self,
models: dict[str, nn.Module] = None,
sparse_structure_sampler: samplers.Sampler = None,
slat_sampler: samplers.Sampler = None,
slat_normalization: dict = None,
text_cond_model: str = None,
clip: str = None,
clip_moudel: str = "openai/clip-vit-large-patch14",
):
if models is None:
return
super().__init__(models)
self.sparse_structure_sampler = sparse_structure_sampler
self.slat_sampler = slat_sampler
self.sparse_structure_sampler_params = {}
self.slat_sampler_params = {}
self.slat_normalization = slat_normalization
self._init_text_cond_model(text_cond_model)
self.clip=clip
self.clip_moudel=clip_moudel
@staticmethod
def from_pretrained(path: str) -> "TrellisTextTo3DPipeline":
"""
Load a pretrained model.
Args:
path (str): The path to the model. Can be either local path or a Hugging Face repository.
"""
pipeline = super(TrellisTextTo3DPipeline, TrellisTextTo3DPipeline).from_pretrained(path)
new_pipeline = TrellisTextTo3DPipeline()
new_pipeline.__dict__ = pipeline.__dict__
args = pipeline._pretrained_args
new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args'])
new_pipeline.slat_sampler_params = args['slat_sampler']['params']
new_pipeline.slat_normalization = args['slat_normalization']
new_pipeline._init_text_cond_model(args['text_cond_model'])
return new_pipeline
def _init_text_cond_model(self, name: str):
"""
Initialize the text conditioning model.
"""
# load model
#model = CLIPTextModel.from_pretrained(name)
is_accelerate_available=False
try:
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
is_accelerate_available = True
except:
pass
ctx = init_empty_weights if is_accelerate_available else nullcontext
with ctx():
text_encoder_config=CLIPTextConfig.from_pretrained(self.clip_moudel, local_files_only=True)
model = CLIPTextModel(text_encoder_config)
tokenizer = AutoTokenizer.from_pretrained(self.clip_moudel)
clip_dict=load_file(self.clip)
model.load_state_dict(clip_dict,strict=False)
del clip_dict
torch.cuda.empty_cache()
model.eval()
self.clip_model = model.cuda()
self.text_cond_model = {
'model': self.clip_model,
'tokenizer': tokenizer,
}
self.text_cond_model['null_cond'] = self.encode_text([''])
@torch.no_grad()
def encode_text(self, text: List[str]) -> torch.Tensor:
"""
Encode the text.
"""
assert isinstance(text, list) and all(isinstance(t, str) for t in text), "text must be a list of strings"
encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
tokens = encoding['input_ids'].cuda()
embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
return embeddings
def get_cond(self, prompt: List[str]) -> dict:
"""
Get the conditioning information for the model.
Args:
prompt (List[str]): The text prompt.
Returns:
dict: The conditioning information
"""
cond = self.encode_text(prompt)
neg_cond = self.text_cond_model['null_cond']
self.clip_model.to('cpu')# unload clip
return {
'cond': cond,
'neg_cond': neg_cond,
}
def sample_sparse_structure(
self,
cond: dict,
num_samples: int = 1,
sampler_params: dict = {},
) -> torch.Tensor:
"""
Sample sparse structures with the given conditioning.
Args:
cond (dict): The conditioning information.
num_samples (int): The number of samples to generate.
sampler_params (dict): Additional parameters for the sampler.
"""
# Sample occupancy latent
flow_model = self.models['sparse_structure_flow_model']
reso = flow_model.resolution
noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device)
sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
z_s = self.sparse_structure_sampler.sample(
flow_model,
noise,
**cond,
**sampler_params,
verbose=True
).samples
# Decode occupancy latent
decoder = self.models['sparse_structure_decoder']
coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int()
return coords
def decode_slat(
self,
slat: sp.SparseTensor,
formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
) -> dict:
"""
Decode the structured latent.
Args:
slat (sp.SparseTensor): The structured latent.
formats (List[str]): The formats to decode the structured latent to.
Returns:
dict: The decoded structured latent.
"""
ret = {}
if 'mesh' in formats:
ret['mesh'] = self.models['slat_decoder_mesh'](slat)
if 'gaussian' in formats:
ret['gaussian'] = self.models['slat_decoder_gs'](slat)
if 'radiance_field' in formats:
ret['radiance_field'] = self.models['slat_decoder_rf'](slat)
return ret
def sample_slat(
self,
cond: dict,
coords: torch.Tensor,
sampler_params: dict = {},
) -> sp.SparseTensor:
"""
Sample structured latent with the given conditioning.
Args:
cond (dict): The conditioning information.
coords (torch.Tensor): The coordinates of the sparse structure.
sampler_params (dict): Additional parameters for the sampler.
"""
# Sample structured latent
flow_model = self.models['slat_flow_model']
noise = sp.SparseTensor(
feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
coords=coords,
)
sampler_params = {**self.slat_sampler_params, **sampler_params}
slat = self.slat_sampler.sample(
flow_model,
noise,
**cond,
**sampler_params,
verbose=True
).samples
std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device)
mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device)
slat = slat * std + mean
return slat
@torch.no_grad()
def run(
self,
prompt: str,
num_samples: int = 1,
seed: int = 42,
sparse_structure_sampler_params: dict = {},
slat_sampler_params: dict = {},
formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
) -> dict:
"""
Run the pipeline.
Args:
prompt (str): The text prompt.
num_samples (int): The number of samples to generate.
seed (int): The random seed.
sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
slat_sampler_params (dict): Additional parameters for the structured latent sampler.
formats (List[str]): The formats to decode the structured latent to.
"""
cond = self.get_cond([prompt])
torch.manual_seed(seed)
coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
slat = self.sample_slat(cond, coords, slat_sampler_params)
return self.decode_slat(slat, formats)
def voxelize(self, mesh: o3d.geometry.TriangleMesh) -> torch.Tensor:
"""
Voxelize a mesh.
Args:
mesh (o3d.geometry.TriangleMesh): The mesh to voxelize.
sha256 (str): The SHA256 hash of the mesh.
output_dir (str): The output directory.
"""
vertices = np.asarray(mesh.vertices)
aabb = np.stack([vertices.min(0), vertices.max(0)])
center = (aabb[0] + aabb[1]) / 2
scale = (aabb[1] - aabb[0]).max()
vertices = (vertices - center) / scale
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()])
return torch.tensor(vertices).int().cuda()
@torch.no_grad()
def run_variant(
self,
mesh: o3d.geometry.TriangleMesh,
prompt: str,
num_samples: int = 1,
seed: int = 42,
slat_sampler_params: dict = {},
formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
) -> dict:
"""
Run the pipeline for making variants of an asset.
Args:
mesh (o3d.geometry.TriangleMesh): The base mesh.
prompt (str): The text prompt.
num_samples (int): The number of samples to generate.
seed (int): The random seed
slat_sampler_params (dict): Additional parameters for the structured latent sampler.
formats (List[str]): The formats to decode the structured latent to.
"""
cond = self.get_cond([prompt])
coords = self.voxelize(mesh)
coords = torch.cat([
torch.arange(num_samples).repeat_interleave(coords.shape[0], 0)[:, None].int().cuda(),
coords.repeat(num_samples, 1)
], 1)
torch.manual_seed(seed)
slat = self.sample_slat(cond, coords, slat_sampler_params)
return self.decode_slat(slat, formats)
trellis/representations/octree/octree_dfs.py
+3 -18
View File
@@ -3,21 +3,6 @@ import torch.nn as nn
import torch.nn.functional as F
DEFAULT_TRIVEC_CONFIG = {
'dim': 8,
'rank': 8,
}
DEFAULT_VOXEL_CONFIG = {
'solid': False,
}
DEFAULT_DECOPOLY_CONFIG = {
'degree': 8,
'rank': 16,
}
class DfsOctree:
"""
Sparse Voxel Octree (SVO) implementation for PyTorch.
@@ -145,8 +130,8 @@ class DfsOctree:
@property
def get_density(self):
if self.primitive == 'voxel' and self.voxel_config['solid']:
return torch.full((self.position.shape[0], 1), 1000, dtype=torch.float32, device=self.device)
if self.primitive == 'voxel' and self.primitive_config.get('solid', False):
return torch.full((self.position.shape[0], 1), torch.finfo(torch.float32).max, dtype=torch.float32, device=self.device)
return self.density_activation(self.density)
@property
@@ -172,7 +157,7 @@ class DfsOctree:
return torch.cat([self.features_dc, self.features_ac], dim=-2)
def state_dict(self):
ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'trivec_config': self.trivec_config, 'voxel_config': self.voxel_config, 'primitive': self.primitive}
ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'primitive_config': self.primitive_config, 'primitive': self.primitive}
if hasattr(self, 'density_shift'):
ret['density_shift'] = self.density_shift
for data in set(self.data + self.param_names):
trellis/trainers/__init__.py
+63
View File
@@ -0,0 +1,63 @@
import importlib
__attributes = {
'BasicTrainer': 'basic',
'SparseStructureVaeTrainer': 'vae.sparse_structure_vae',
'SLatVaeGaussianTrainer': 'vae.structured_latent_vae_gaussian',
'SLatVaeRadianceFieldDecoderTrainer': 'vae.structured_latent_vae_rf_dec',
'SLatVaeMeshDecoderTrainer': 'vae.structured_latent_vae_mesh_dec',
'FlowMatchingTrainer': 'flow_matching.flow_matching',
'FlowMatchingCFGTrainer': 'flow_matching.flow_matching',
'TextConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
'ImageConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
'SparseFlowMatchingTrainer': 'flow_matching.sparse_flow_matching',
'SparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
'TextConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
'ImageConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
}
__submodules = []
__all__ = list(__attributes.keys()) + __submodules
def __getattr__(name):
if name not in globals():
if name in __attributes:
module_name = __attributes[name]
module = importlib.import_module(f".{module_name}", __name__)
globals()[name] = getattr(module, name)
elif name in __submodules:
module = importlib.import_module(f".{name}", __name__)
globals()[name] = module
else:
raise AttributeError(f"module {__name__} has no attribute {name}")
return globals()[name]
# For Pylance
if __name__ == '__main__':
from .basic import BasicTrainer
from .vae.sparse_structure_vae import SparseStructureVaeTrainer
from .vae.structured_latent_vae_gaussian import SLatVaeGaussianTrainer
from .vae.structured_latent_vae_rf_dec import SLatVaeRadianceFieldDecoderTrainer
from .vae.structured_latent_vae_mesh_dec import SLatVaeMeshDecoderTrainer
from .flow_matching.flow_matching import (
FlowMatchingTrainer,
FlowMatchingCFGTrainer,
TextConditionedFlowMatchingCFGTrainer,
ImageConditionedFlowMatchingCFGTrainer,
)
from .flow_matching.sparse_flow_matching import (
SparseFlowMatchingTrainer,
SparseFlowMatchingCFGTrainer,
TextConditionedSparseFlowMatchingCFGTrainer,
ImageConditionedSparseFlowMatchingCFGTrainer,
)
trellis/trainers/base.py
+451
View File
@@ -0,0 +1,451 @@
from abc import abstractmethod
import os
import time
import json
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
import numpy as np
from torchvision import utils
from torch.utils.tensorboard import SummaryWriter
from .utils import *
from ..utils.general_utils import *
from ..utils.data_utils import recursive_to_device, cycle, ResumableSampler
class Trainer:
"""
Base class for training.
"""
def __init__(self,
models,
dataset,
*,
output_dir,
load_dir,
step,
max_steps,
batch_size=None,
batch_size_per_gpu=None,
batch_split=None,
optimizer={},
lr_scheduler=None,
elastic=None,
grad_clip=None,
ema_rate=0.9999,
fp16_mode='inflat_all',
fp16_scale_growth=1e-3,
finetune_ckpt=None,
log_param_stats=False,
prefetch_data=True,
i_print=1000,
i_log=500,
i_sample=10000,
i_save=10000,
i_ddpcheck=10000,
**kwargs
):
assert batch_size is not None or batch_size_per_gpu is not None, 'Either batch_size or batch_size_per_gpu must be specified.'
self.models = models
self.dataset = dataset
self.batch_split = batch_split if batch_split is not None else 1
self.max_steps = max_steps
self.optimizer_config = optimizer
self.lr_scheduler_config = lr_scheduler
self.elastic_controller_config = elastic
self.grad_clip = grad_clip
self.ema_rate = [ema_rate] if isinstance(ema_rate, float) else ema_rate
self.fp16_mode = fp16_mode
self.fp16_scale_growth = fp16_scale_growth
self.log_param_stats = log_param_stats
self.prefetch_data = prefetch_data
if self.prefetch_data:
self._data_prefetched = None
self.output_dir = output_dir
self.i_print = i_print
self.i_log = i_log
self.i_sample = i_sample
self.i_save = i_save
self.i_ddpcheck = i_ddpcheck
if dist.is_initialized():
# Multi-GPU params
self.world_size = dist.get_world_size()
self.rank = dist.get_rank()
self.local_rank = dist.get_rank() % torch.cuda.device_count()
self.is_master = self.rank == 0
else:
# Single-GPU params
self.world_size = 1
self.rank = 0
self.local_rank = 0
self.is_master = True
self.batch_size = batch_size if batch_size_per_gpu is None else batch_size_per_gpu * self.world_size
self.batch_size_per_gpu = batch_size_per_gpu if batch_size_per_gpu is not None else batch_size // self.world_size
assert self.batch_size % self.world_size == 0, 'Batch size must be divisible by the number of GPUs.'
assert self.batch_size_per_gpu % self.batch_split == 0, 'Batch size per GPU must be divisible by batch split.'
self.init_models_and_more(**kwargs)
self.prepare_dataloader(**kwargs)
# Load checkpoint
self.step = 0
if load_dir is not None and step is not None:
self.load(load_dir, step)
elif finetune_ckpt is not None:
self.finetune_from(finetune_ckpt)
if self.is_master:
os.makedirs(os.path.join(self.output_dir, 'ckpts'), exist_ok=True)
os.makedirs(os.path.join(self.output_dir, 'samples'), exist_ok=True)
self.writer = SummaryWriter(os.path.join(self.output_dir, 'tb_logs'))
if self.world_size > 1:
self.check_ddp()
if self.is_master:
print('\n\nTrainer initialized.')
print(self)
@property
def device(self):
for _, model in self.models.items():
if hasattr(model, 'device'):
return model.device
return next(list(self.models.values())[0].parameters()).device
@abstractmethod
def init_models_and_more(self, **kwargs):
"""
Initialize models and more.
"""
pass
def prepare_dataloader(self, **kwargs):
"""
Prepare dataloader.
"""
self.data_sampler = ResumableSampler(
self.dataset,
shuffle=True,
)
self.dataloader = DataLoader(
self.dataset,
batch_size=self.batch_size_per_gpu,
num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
pin_memory=True,
drop_last=True,
persistent_workers=True,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
sampler=self.data_sampler,
)
self.data_iterator = cycle(self.dataloader)
@abstractmethod
def load(self, load_dir, step=0):
"""
Load a checkpoint.
Should be called by all processes.
"""
pass
@abstractmethod
def save(self):
"""
Save a checkpoint.
Should be called only by the rank 0 process.
"""
pass
@abstractmethod
def finetune_from(self, finetune_ckpt):
"""
Finetune from a checkpoint.
Should be called by all processes.
"""
pass
@abstractmethod
def run_snapshot(self, num_samples, batch_size=4, verbose=False, **kwargs):
"""
Run a snapshot of the model.
"""
pass
@torch.no_grad()
def visualize_sample(self, sample):
"""
Convert a sample to an image.
"""
if hasattr(self.dataset, 'visualize_sample'):
return self.dataset.visualize_sample(sample)
else:
return sample
@torch.no_grad()
def snapshot_dataset(self, num_samples=100):
"""
Sample images from the dataset.
"""
dataloader = torch.utils.data.DataLoader(
self.dataset,
batch_size=num_samples,
num_workers=0,
shuffle=True,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
data = next(iter(dataloader))
data = recursive_to_device(data, self.device)
vis = self.visualize_sample(data)
if isinstance(vis, dict):
save_cfg = [(f'dataset_{k}', v) for k, v in vis.items()]
else:
save_cfg = [('dataset', vis)]
for name, image in save_cfg:
utils.save_image(
image,
os.path.join(self.output_dir, 'samples', f'{name}.jpg'),
nrow=int(np.sqrt(num_samples)),
normalize=True,
value_range=self.dataset.value_range,
)
@torch.no_grad()
def snapshot(self, suffix=None, num_samples=64, batch_size=4, verbose=False):
"""
Sample images from the model.
NOTE: This function should be called by all processes.
"""
if self.is_master:
print(f'\nSampling {num_samples} images...', end='')
if suffix is None:
suffix = f'step{self.step:07d}'
# Assign tasks
num_samples_per_process = int(np.ceil(num_samples / self.world_size))
samples = self.run_snapshot(num_samples_per_process, batch_size=batch_size, verbose=verbose)
# Preprocess images
for key in list(samples.keys()):
if samples[key]['type'] == 'sample':
vis = self.visualize_sample(samples[key]['value'])
if isinstance(vis, dict):
for k, v in vis.items():
samples[f'{key}_{k}'] = {'value': v, 'type': 'image'}
del samples[key]
else:
samples[key] = {'value': vis, 'type': 'image'}
# Gather results
if self.world_size > 1:
for key in samples.keys():
samples[key]['value'] = samples[key]['value'].contiguous()
if self.is_master:
all_images = [torch.empty_like(samples[key]['value']) for _ in range(self.world_size)]
else:
all_images = []
dist.gather(samples[key]['value'], all_images, dst=0)
if self.is_master:
samples[key]['value'] = torch.cat(all_images, dim=0)[:num_samples]
# Save images
if self.is_master:
os.makedirs(os.path.join(self.output_dir, 'samples', suffix), exist_ok=True)
for key in samples.keys():
if samples[key]['type'] == 'image':
utils.save_image(
samples[key]['value'],
os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg'),
nrow=int(np.sqrt(num_samples)),
normalize=True,
value_range=self.dataset.value_range,
)
elif samples[key]['type'] == 'number':
min = samples[key]['value'].min()
max = samples[key]['value'].max()
images = (samples[key]['value'] - min) / (max - min)
images = utils.make_grid(
images,
nrow=int(np.sqrt(num_samples)),
normalize=False,
)
save_image_with_notes(
images,
os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg'),
notes=f'{key} min: {min}, max: {max}',
)
if self.is_master:
print(' Done.')
@abstractmethod
def update_ema(self):
"""
Update exponential moving average.
Should only be called by the rank 0 process.
"""
pass
@abstractmethod
def check_ddp(self):
"""
Check if DDP is working properly.
Should be called by all process.
"""
pass
@abstractmethod
def training_losses(**mb_data):
"""
Compute training losses.
"""
pass
def load_data(self):
"""
Load data.
"""
if self.prefetch_data:
if self._data_prefetched is None:
self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
data = self._data_prefetched
self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
else:
data = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
# if the data is a dict, we need to split it into multiple dicts with batch_size_per_gpu
if isinstance(data, dict):
if self.batch_split == 1:
data_list = [data]
else:
batch_size = list(data.values())[0].shape[0]
data_list = [
{k: v[i * batch_size // self.batch_split:(i + 1) * batch_size // self.batch_split] for k, v in data.items()}
for i in range(self.batch_split)
]
elif isinstance(data, list):
data_list = data
else:
raise ValueError('Data must be a dict or a list of dicts.')
return data_list
@abstractmethod
def run_step(self, data_list):
"""
Run a training step.
"""
pass
def run(self):
"""
Run training.
"""
if self.is_master:
print('\nStarting training...')
self.snapshot_dataset()
if self.step == 0:
self.snapshot(suffix='init')
else: # resume
self.snapshot(suffix=f'resume_step{self.step:07d}')
log = []
time_last_print = 0.0
time_elapsed = 0.0
while self.step < self.max_steps:
time_start = time.time()
data_list = self.load_data()
step_log = self.run_step(data_list)
time_end = time.time()
time_elapsed += time_end - time_start
self.step += 1
# Print progress
if self.is_master and self.step % self.i_print == 0:
speed = self.i_print / (time_elapsed - time_last_print) * 3600
columns = [
f'Step: {self.step}/{self.max_steps} ({self.step / self.max_steps * 100:.2f}%)',
f'Elapsed: {time_elapsed / 3600:.2f} h',
f'Speed: {speed:.2f} steps/h',
f'ETA: {(self.max_steps - self.step) / speed:.2f} h',
]
print(' | '.join([c.ljust(25) for c in columns]), flush=True)
time_last_print = time_elapsed
# Check ddp
if self.world_size > 1 and self.i_ddpcheck is not None and self.step % self.i_ddpcheck == 0:
self.check_ddp()
# Sample images
if self.step % self.i_sample == 0:
self.snapshot()
if self.is_master:
log.append((self.step, {}))
# Log time
log[-1][1]['time'] = {
'step': time_end - time_start,
'elapsed': time_elapsed,
}
# Log losses
if step_log is not None:
log[-1][1].update(step_log)
# Log scale
if self.fp16_mode == 'amp':
log[-1][1]['scale'] = self.scaler.get_scale()
elif self.fp16_mode == 'inflat_all':
log[-1][1]['log_scale'] = self.log_scale
# Save log
if self.step % self.i_log == 0:
## save to log file
log_str = '\n'.join([
f'{step}: {json.dumps(log)}' for step, log in log
])
with open(os.path.join(self.output_dir, 'log.txt'), 'a') as log_file:
log_file.write(log_str + '\n')
# show with mlflow
log_show = [l for _, l in log if not dict_any(l, lambda x: np.isnan(x))]
log_show = dict_reduce(log_show, lambda x: np.mean(x))
log_show = dict_flatten(log_show, sep='/')
for key, value in log_show.items():
self.writer.add_scalar(key, value, self.step)
log = []
# Save checkpoint
if self.step % self.i_save == 0:
self.save()
if self.is_master:
self.snapshot(suffix='final')
self.writer.close()
print('Training finished.')
def profile(self, wait=2, warmup=3, active=5):
"""
Profile the training loop.
"""
with torch.profiler.profile(
schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1),
on_trace_ready=torch.profiler.tensorboard_trace_handler(os.path.join(self.output_dir, 'profile')),
profile_memory=True,
with_stack=True,
) as prof:
for _ in range(wait + warmup + active):
self.run_step()
prof.step()
trellis/trainers/basic.py
+438
View File
@@ -0,0 +1,438 @@
import os
import copy
from functools import partial
from contextlib import nullcontext
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import numpy as np
from .utils import *
from .base import Trainer
from ..utils.general_utils import *
from ..utils.dist_utils import *
from ..utils import grad_clip_utils, elastic_utils
class BasicTrainer(Trainer):
"""
Trainer for basic training loop.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
"""
def __str__(self):
lines = []
lines.append(self.__class__.__name__)
lines.append(f' - Models:')
for name, model in self.models.items():
lines.append(f' - {name}: {model.__class__.__name__}')
lines.append(f' - Dataset: {indent(str(self.dataset), 2)}')
lines.append(f' - Dataloader:')
lines.append(f' - Sampler: {self.dataloader.sampler.__class__.__name__}')
lines.append(f' - Num workers: {self.dataloader.num_workers}')
lines.append(f' - Number of steps: {self.max_steps}')
lines.append(f' - Number of GPUs: {self.world_size}')
lines.append(f' - Batch size: {self.batch_size}')
lines.append(f' - Batch size per GPU: {self.batch_size_per_gpu}')
lines.append(f' - Batch split: {self.batch_split}')
lines.append(f' - Optimizer: {self.optimizer.__class__.__name__}')
lines.append(f' - Learning rate: {self.optimizer.param_groups[0]["lr"]}')
if self.lr_scheduler_config is not None:
lines.append(f' - LR scheduler: {self.lr_scheduler.__class__.__name__}')
if self.elastic_controller_config is not None:
lines.append(f' - Elastic memory: {indent(str(self.elastic_controller), 2)}')
if self.grad_clip is not None:
lines.append(f' - Gradient clip: {indent(str(self.grad_clip), 2)}')
lines.append(f' - EMA rate: {self.ema_rate}')
lines.append(f' - FP16 mode: {self.fp16_mode}')
return '\n'.join(lines)
def init_models_and_more(self, **kwargs):
"""
Initialize models and more.
"""
if self.world_size > 1:
# Prepare distributed data parallel
self.training_models = {
name: DDP(
model,
device_ids=[self.local_rank],
output_device=self.local_rank,
bucket_cap_mb=128,
find_unused_parameters=False
)
for name, model in self.models.items()
}
else:
self.training_models = self.models
# Build master params
self.model_params = sum(
[[p for p in model.parameters() if p.requires_grad] for model in self.models.values()]
, [])
if self.fp16_mode == 'amp':
self.master_params = self.model_params
self.scaler = torch.GradScaler() if self.fp16_mode == 'amp' else None
elif self.fp16_mode == 'inflat_all':
self.master_params = make_master_params(self.model_params)
self.fp16_scale_growth = self.fp16_scale_growth
self.log_scale = 20.0
elif self.fp16_mode is None:
self.master_params = self.model_params
else:
raise NotImplementedError(f'FP16 mode {self.fp16_mode} is not implemented.')
# Build EMA params
if self.is_master:
self.ema_params = [copy.deepcopy(self.master_params) for _ in self.ema_rate]
# Initialize optimizer
if hasattr(torch.optim, self.optimizer_config['name']):
self.optimizer = getattr(torch.optim, self.optimizer_config['name'])(self.master_params, **self.optimizer_config['args'])
else:
self.optimizer = globals()[self.optimizer_config['name']](self.master_params, **self.optimizer_config['args'])
# Initalize learning rate scheduler
if self.lr_scheduler_config is not None:
if hasattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name']):
self.lr_scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name'])(self.optimizer, **self.lr_scheduler_config['args'])
else:
self.lr_scheduler = globals()[self.lr_scheduler_config['name']](self.optimizer, **self.lr_scheduler_config['args'])
# Initialize elastic memory controller
if self.elastic_controller_config is not None:
assert any([isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)) for model in self.models.values()]), \
'No elastic module found in models, please inherit from ElasticModule or ElasticModuleMixin'
self.elastic_controller = getattr(elastic_utils, self.elastic_controller_config['name'])(**self.elastic_controller_config['args'])
for model in self.models.values():
if isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)):
model.register_memory_controller(self.elastic_controller)
# Initialize gradient clipper
if self.grad_clip is not None:
if isinstance(self.grad_clip, (float, int)):
self.grad_clip = float(self.grad_clip)
else:
self.grad_clip = getattr(grad_clip_utils, self.grad_clip['name'])(**self.grad_clip['args'])
def _master_params_to_state_dicts(self, master_params):
"""
Convert master params to dict of state_dicts.
"""
if self.fp16_mode == 'inflat_all':
master_params = unflatten_master_params(self.model_params, master_params)
state_dicts = {name: model.state_dict() for name, model in self.models.items()}
master_params_names = sum(
[[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
, [])
for i, (model_name, param_name) in enumerate(master_params_names):
state_dicts[model_name][param_name] = master_params[i]
return state_dicts
def _state_dicts_to_master_params(self, master_params, state_dicts):
"""
Convert a state_dict to master params.
"""
master_params_names = sum(
[[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
, [])
params = [state_dicts[name][param_name] for name, param_name in master_params_names]
if self.fp16_mode == 'inflat_all':
model_params_to_master_params(params, master_params)
else:
for i, param in enumerate(params):
master_params[i].data.copy_(param.data)
def load(self, load_dir, step=0):
"""
Load a checkpoint.
Should be called by all processes.
"""
if self.is_master:
print(f'\nLoading checkpoint from step {step}...', end='')
model_ckpts = {}
for name, model in self.models.items():
model_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')), map_location=self.device, weights_only=True)
model_ckpts[name] = model_ckpt
model.load_state_dict(model_ckpt)
if self.fp16_mode == 'inflat_all':
model.convert_to_fp16()
self._state_dicts_to_master_params(self.master_params, model_ckpts)
del model_ckpts
if self.is_master:
for i, ema_rate in enumerate(self.ema_rate):
ema_ckpts = {}
for name, model in self.models.items():
ema_ckpt = torch.load(os.path.join(load_dir, 'ckpts', f'{name}_ema{ema_rate}_step{step:07d}.pt'), map_location=self.device, weights_only=True)
ema_ckpts[name] = ema_ckpt
self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts)
del ema_ckpts
misc_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt')), map_location=torch.device('cpu'), weights_only=False)
self.optimizer.load_state_dict(misc_ckpt['optimizer'])
self.step = misc_ckpt['step']
self.data_sampler.load_state_dict(misc_ckpt['data_sampler'])
if self.fp16_mode == 'amp':
self.scaler.load_state_dict(misc_ckpt['scaler'])
elif self.fp16_mode == 'inflat_all':
self.log_scale = misc_ckpt['log_scale']
if self.lr_scheduler_config is not None:
self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler'])
if self.elastic_controller_config is not None:
self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller'])
if self.grad_clip is not None and not isinstance(self.grad_clip, float):
self.grad_clip.load_state_dict(misc_ckpt['grad_clip'])
del misc_ckpt
if self.world_size > 1:
dist.barrier()
if self.is_master:
print(' Done.')
if self.world_size > 1:
self.check_ddp()
def save(self):
"""
Save a checkpoint.
Should be called only by the rank 0 process.
"""
assert self.is_master, 'save() should be called only by the rank 0 process.'
print(f'\nSaving checkpoint at step {self.step}...', end='')
model_ckpts = self._master_params_to_state_dicts(self.master_params)
for name, model_ckpt in model_ckpts.items():
torch.save(model_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_step{self.step:07d}.pt'))
for i, ema_rate in enumerate(self.ema_rate):
ema_ckpts = self._master_params_to_state_dicts(self.ema_params[i])
for name, ema_ckpt in ema_ckpts.items():
torch.save(ema_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_ema{ema_rate}_step{self.step:07d}.pt'))
misc_ckpt = {
'optimizer': self.optimizer.state_dict(),
'step': self.step,
'data_sampler': self.data_sampler.state_dict(),
}
if self.fp16_mode == 'amp':
misc_ckpt['scaler'] = self.scaler.state_dict()
elif self.fp16_mode == 'inflat_all':
misc_ckpt['log_scale'] = self.log_scale
if self.lr_scheduler_config is not None:
misc_ckpt['lr_scheduler'] = self.lr_scheduler.state_dict()
if self.elastic_controller_config is not None:
misc_ckpt['elastic_controller'] = self.elastic_controller.state_dict()
if self.grad_clip is not None and not isinstance(self.grad_clip, float):
misc_ckpt['grad_clip'] = self.grad_clip.state_dict()
torch.save(misc_ckpt, os.path.join(self.output_dir, 'ckpts', f'misc_step{self.step:07d}.pt'))
print(' Done.')
def finetune_from(self, finetune_ckpt):
"""
Finetune from a checkpoint.
Should be called by all processes.
"""
if self.is_master:
print('\nFinetuning from:')
for name, path in finetune_ckpt.items():
print(f' - {name}: {path}')
model_ckpts = {}
for name, model in self.models.items():
model_state_dict = model.state_dict()
if name in finetune_ckpt:
model_ckpt = torch.load(read_file_dist(finetune_ckpt[name]), map_location=self.device, weights_only=True)
for k, v in model_ckpt.items():
if model_ckpt[k].shape != model_state_dict[k].shape:
if self.is_master:
print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.')
model_ckpt[k] = model_state_dict[k]
model_ckpts[name] = model_ckpt
model.load_state_dict(model_ckpt)
if self.fp16_mode == 'inflat_all':
model.convert_to_fp16()
else:
if self.is_master:
print(f'Warning: {name} not found in finetune_ckpt, skipped.')
model_ckpts[name] = model_state_dict
self._state_dicts_to_master_params(self.master_params, model_ckpts)
del model_ckpts
if self.world_size > 1:
dist.barrier()
if self.is_master:
print('Done.')
if self.world_size > 1:
self.check_ddp()
def update_ema(self):
"""
Update exponential moving average.
Should only be called by the rank 0 process.
"""
assert self.is_master, 'update_ema() should be called only by the rank 0 process.'
for i, ema_rate in enumerate(self.ema_rate):
for master_param, ema_param in zip(self.master_params, self.ema_params[i]):
ema_param.detach().mul_(ema_rate).add_(master_param, alpha=1.0 - ema_rate)
def check_ddp(self):
"""
Check if DDP is working properly.
Should be called by all process.
"""
if self.is_master:
print('\nPerforming DDP check...')
if self.is_master:
print('Checking if parameters are consistent across processes...')
dist.barrier()
try:
for p in self.master_params:
# split to avoid OOM
for i in range(0, p.numel(), 10000000):
sub_size = min(10000000, p.numel() - i)
sub_p = p.detach().view(-1)[i:i+sub_size]
# gather from all processes
sub_p_gather = [torch.empty_like(sub_p) for _ in range(self.world_size)]
dist.all_gather(sub_p_gather, sub_p)
# check if equal
assert all([torch.equal(sub_p, sub_p_gather[i]) for i in range(self.world_size)]), 'parameters are not consistent across processes'
except AssertionError as e:
if self.is_master:
print(f'\n\033[91mError: {e}\033[0m')
print('DDP check failed.')
raise e
dist.barrier()
if self.is_master:
print('Done.')
def run_step(self, data_list):
"""
Run a training step.
"""
step_log = {'loss': {}, 'status': {}}
amp_context = partial(torch.autocast, device_type='cuda') if self.fp16_mode == 'amp' else nullcontext
elastic_controller_context = self.elastic_controller.record if self.elastic_controller_config is not None else nullcontext
# Train
losses = []
statuses = []
elastic_controller_logs = []
zero_grad(self.model_params)
for i, mb_data in enumerate(data_list):
## sync at the end of each batch split
sync_contexts = [self.training_models[name].no_sync for name in self.training_models] if i != len(data_list) - 1 and self.world_size > 1 else [nullcontext]
with nested_contexts(*sync_contexts), elastic_controller_context():
with amp_context():
loss, status = self.training_losses(**mb_data)
l = loss['loss'] / len(data_list)
## backward
if self.fp16_mode == 'amp':
self.scaler.scale(l).backward()
elif self.fp16_mode == 'inflat_all':
scaled_l = l * (2 ** self.log_scale)
scaled_l.backward()
else:
l.backward()
## log
losses.append(dict_foreach(loss, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
statuses.append(dict_foreach(status, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
if self.elastic_controller_config is not None:
elastic_controller_logs.append(self.elastic_controller.log())
## gradient clip
if self.grad_clip is not None:
if self.fp16_mode == 'amp':
self.scaler.unscale_(self.optimizer)
elif self.fp16_mode == 'inflat_all':
model_grads_to_master_grads(self.model_params, self.master_params)
self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
if isinstance(self.grad_clip, float):
grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params, self.grad_clip)
else:
grad_norm = self.grad_clip(self.master_params)
if torch.isfinite(grad_norm):
statuses[-1]['grad_norm'] = grad_norm.item()
## step
if self.fp16_mode == 'amp':
prev_scale = self.scaler.get_scale()
self.scaler.step(self.optimizer)
self.scaler.update()
elif self.fp16_mode == 'inflat_all':
prev_scale = 2 ** self.log_scale
if not any(not p.grad.isfinite().all() for p in self.model_params):
if self.grad_clip is None:
model_grads_to_master_grads(self.model_params, self.master_params)
self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
self.optimizer.step()
master_params_to_model_params(self.model_params, self.master_params)
self.log_scale += self.fp16_scale_growth
else:
self.log_scale -= 1
else:
prev_scale = 1.0
if not any(not p.grad.isfinite().all() for p in self.model_params):
self.optimizer.step()
else:
print('\n\033[93mWarning: NaN detected in gradients. Skipping update.\033[0m')
## adjust learning rate
if self.lr_scheduler_config is not None:
statuses[-1]['lr'] = self.lr_scheduler.get_last_lr()[0]
self.lr_scheduler.step()
# Logs
step_log['loss'] = dict_reduce(losses, lambda x: np.mean(x))
step_log['status'] = dict_reduce(statuses, lambda x: np.mean(x), special_func={'min': lambda x: np.min(x), 'max': lambda x: np.max(x)})
if self.elastic_controller_config is not None:
step_log['elastic'] = dict_reduce(elastic_controller_logs, lambda x: np.mean(x))
if self.grad_clip is not None:
step_log['grad_clip'] = self.grad_clip if isinstance(self.grad_clip, float) else self.grad_clip.log()
# Check grad and norm of each param
if self.log_param_stats:
param_norms = {}
param_grads = {}
for name, param in self.backbone.named_parameters():
if param.requires_grad:
param_norms[name] = param.norm().item()
if param.grad is not None and torch.isfinite(param.grad).all():
param_grads[name] = param.grad.norm().item() / prev_scale
step_log['param_norms'] = param_norms
step_log['param_grads'] = param_grads
# Update exponential moving average
if self.is_master:
self.update_ema()
return step_log
trellis/trainers/flow_matching/flow_matching.py
+353
View File
@@ -0,0 +1,353 @@
from typing import *
import copy
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
from ..basic import BasicTrainer
from ...pipelines import samplers
from ...utils.general_utils import dict_reduce
from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
from .mixins.text_conditioned import TextConditionedMixin
from .mixins.image_conditioned import ImageConditionedMixin
class FlowMatchingTrainer(BasicTrainer):
"""
Trainer for diffusion model with flow matching objective.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
"""
def __init__(
self,
*args,
t_schedule: dict = {
'name': 'logitNormal',
'args': {
'mean': 0.0,
'std': 1.0,
}
},
sigma_min: float = 1e-5,
**kwargs
):
super().__init__(*args, **kwargs)
self.t_schedule = t_schedule
self.sigma_min = sigma_min
def diffuse(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Diffuse the data for a given number of diffusion steps.
In other words, sample from q(x_t | x_0).
Args:
x_0: The [N x C x ...] tensor of noiseless inputs.
t: The [N] tensor of diffusion steps [0-1].
noise: If specified, use this noise instead of generating new noise.
Returns:
x_t, the noisy version of x_0 under timestep t.
"""
if noise is None:
noise = torch.randn_like(x_0)
assert noise.shape == x_0.shape, "noise must have same shape as x_0"
t = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
x_t = (1 - t) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t) * noise
return x_t
def reverse_diffuse(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
"""
Get original image from noisy version under timestep t.
"""
assert noise.shape == x_t.shape, "noise must have same shape as x_t"
t = t.view(-1, *[1 for _ in range(len(x_t.shape) - 1)])
x_0 = (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * noise) / (1 - t)
return x_0
def get_v(self, x_0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""
Compute the velocity of the diffusion process at time t.
"""
return (1 - self.sigma_min) * noise - x_0
def get_cond(self, cond, **kwargs):
"""
Get the conditioning data.
"""
return cond
def get_inference_cond(self, cond, **kwargs):
"""
Get the conditioning data for inference.
"""
return {'cond': cond, **kwargs}
def get_sampler(self, **kwargs) -> samplers.FlowEulerSampler:
"""
Get the sampler for the diffusion process.
"""
return samplers.FlowEulerSampler(self.sigma_min)
def vis_cond(self, **kwargs):
"""
Visualize the conditioning data.
"""
return {}
def sample_t(self, batch_size: int) -> torch.Tensor:
"""
Sample timesteps.
"""
if self.t_schedule['name'] == 'uniform':
t = torch.rand(batch_size)
elif self.t_schedule['name'] == 'logitNormal':
mean = self.t_schedule['args']['mean']
std = self.t_schedule['args']['std']
t = torch.sigmoid(torch.randn(batch_size) * std + mean)
else:
raise ValueError(f"Unknown t_schedule: {self.t_schedule['name']}")
return t
def training_losses(
self,
x_0: torch.Tensor,
cond=None,
**kwargs
) -> Tuple[Dict, Dict]:
"""
Compute training losses for a single timestep.
Args:
x_0: The [N x C x ...] tensor of noiseless inputs.
cond: The [N x ...] tensor of additional conditions.
kwargs: Additional arguments to pass to the backbone.
Returns:
a dict with the key "loss" containing a tensor of shape [N].
may also contain other keys for different terms.
"""
noise = torch.randn_like(x_0)
t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
x_t = self.diffuse(x_0, t, noise=noise)
cond = self.get_cond(cond, **kwargs)
pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
assert pred.shape == noise.shape == x_0.shape
target = self.get_v(x_0, noise, t)
terms = edict()
terms["mse"] = F.mse_loss(pred, target)
terms["loss"] = terms["mse"]
# log loss with time bins
mse_per_instance = np.array([
F.mse_loss(pred[i], target[i]).item()
for i in range(x_0.shape[0])
])
time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
for i in range(10):
if (time_bin == i).sum() != 0:
terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
return terms, {}
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
sampler = self.get_sampler()
sample_gt = []
sample = []
cond_vis = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
data = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
noise = torch.randn_like(data['x_0'])
sample_gt.append(data['x_0'])
cond_vis.append(self.vis_cond(**data))
del data['x_0']
args = self.get_inference_cond(**data)
res = sampler.sample(
self.models['denoiser'],
noise=noise,
**args,
steps=50, cfg_strength=3.0, verbose=verbose,
)
sample.append(res.samples)
sample_gt = torch.cat(sample_gt, dim=0)
sample = torch.cat(sample, dim=0)
sample_dict = {
'sample_gt': {'value': sample_gt, 'type': 'sample'},
'sample': {'value': sample, 'type': 'sample'},
}
sample_dict.update(dict_reduce(cond_vis, None, {
'value': lambda x: torch.cat(x, dim=0),
'type': lambda x: x[0],
}))
return sample_dict
class FlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, FlowMatchingTrainer):
"""
Trainer for diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
"""
pass
class TextConditionedFlowMatchingCFGTrainer(TextConditionedMixin, FlowMatchingCFGTrainer):
"""
Trainer for text-conditioned diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
text_cond_model(str): Text conditioning model.
"""
pass
class ImageConditionedFlowMatchingCFGTrainer(ImageConditionedMixin, FlowMatchingCFGTrainer):
"""
Trainer for image-conditioned diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
image_cond_model (str): Image conditioning model.
"""
pass
trellis/trainers/flow_matching/mixins/classifier_free_guidance.py
+59
View File
@@ -0,0 +1,59 @@
import torch
import numpy as np
from ....utils.general_utils import dict_foreach
from ....pipelines import samplers
class ClassifierFreeGuidanceMixin:
def __init__(self, *args, p_uncond: float = 0.1, **kwargs):
super().__init__(*args, **kwargs)
self.p_uncond = p_uncond
def get_cond(self, cond, neg_cond=None, **kwargs):
"""
Get the conditioning data.
"""
assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance"
if self.p_uncond > 0:
# randomly drop the class label
def get_batch_size(cond):
if isinstance(cond, torch.Tensor):
return cond.shape[0]
elif isinstance(cond, list):
return len(cond)
else:
raise ValueError(f"Unsupported type of cond: {type(cond)}")
ref_cond = cond if not isinstance(cond, dict) else cond[list(cond.keys())[0]]
B = get_batch_size(ref_cond)
def select(cond, neg_cond, mask):
if isinstance(cond, torch.Tensor):
mask = torch.tensor(mask, device=cond.device).reshape(-1, *[1] * (cond.ndim - 1))
return torch.where(mask, neg_cond, cond)
elif isinstance(cond, list):
return [nc if m else c for c, nc, m in zip(cond, neg_cond, mask)]
else:
raise ValueError(f"Unsupported type of cond: {type(cond)}")
mask = list(np.random.rand(B) < self.p_uncond)
if not isinstance(cond, dict):
cond = select(cond, neg_cond, mask)
else:
cond = dict_foreach([cond, neg_cond], lambda x: select(x[0], x[1], mask))
return cond
def get_inference_cond(self, cond, neg_cond=None, **kwargs):
"""
Get the conditioning data for inference.
"""
assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance"
return {'cond': cond, 'neg_cond': neg_cond, **kwargs}
def get_sampler(self, **kwargs) -> samplers.FlowEulerCfgSampler:
"""
Get the sampler for the diffusion process.
"""
return samplers.FlowEulerCfgSampler(self.sigma_min)
trellis/trainers/flow_matching/mixins/image_conditioned.py
+93
View File
@@ -0,0 +1,93 @@
from typing import *
import torch
import torch.nn.functional as F
from torchvision import transforms
import numpy as np
from PIL import Image
from ....utils import dist_utils
class ImageConditionedMixin:
"""
Mixin for image-conditioned models.
Args:
image_cond_model: The image conditioning model.
"""
def __init__(self, *args, image_cond_model: str = 'dinov2_vitl14_reg', **kwargs):
super().__init__(*args, **kwargs)
self.image_cond_model_name = image_cond_model
self.image_cond_model = None # the model is init lazily
@staticmethod
def prepare_for_training(image_cond_model: str, **kwargs):
"""
Prepare for training.
"""
if hasattr(super(ImageConditionedMixin, ImageConditionedMixin), 'prepare_for_training'):
super(ImageConditionedMixin, ImageConditionedMixin).prepare_for_training(**kwargs)
# download the model
torch.hub.load('facebookresearch/dinov2', image_cond_model, pretrained=True)
def _init_image_cond_model(self):
"""
Initialize the image conditioning model.
"""
with dist_utils.local_master_first():
dinov2_model = torch.hub.load('facebookresearch/dinov2', self.image_cond_model_name, pretrained=True)
dinov2_model.eval().cuda()
transform = transforms.Compose([
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
self.image_cond_model = {
'model': dinov2_model,
'transform': transform,
}
@torch.no_grad()
def encode_image(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
"""
Encode the image.
"""
if isinstance(image, torch.Tensor):
assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
elif isinstance(image, list):
assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
image = [i.resize((518, 518), Image.LANCZOS) for i in image]
image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
image = torch.stack(image).cuda()
else:
raise ValueError(f"Unsupported type of image: {type(image)}")
if self.image_cond_model is None:
self._init_image_cond_model()
image = self.image_cond_model['transform'](image).cuda()
features = self.image_cond_model['model'](image, is_training=True)['x_prenorm']
patchtokens = F.layer_norm(features, features.shape[-1:])
return patchtokens
def get_cond(self, cond, **kwargs):
"""
Get the conditioning data.
"""
cond = self.encode_image(cond)
kwargs['neg_cond'] = torch.zeros_like(cond)
cond = super().get_cond(cond, **kwargs)
return cond
def get_inference_cond(self, cond, **kwargs):
"""
Get the conditioning data for inference.
"""
cond = self.encode_image(cond)
kwargs['neg_cond'] = torch.zeros_like(cond)
cond = super().get_inference_cond(cond, **kwargs)
return cond
def vis_cond(self, cond, **kwargs):
"""
Visualize the conditioning data.
"""
return {'image': {'value': cond, 'type': 'image'}}
trellis/trainers/flow_matching/mixins/text_conditioned.py
+68
View File
@@ -0,0 +1,68 @@
from typing import *
import os
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
import torch
from transformers import AutoTokenizer, CLIPTextModel
from ....utils import dist_utils
class TextConditionedMixin:
"""
Mixin for text-conditioned models.
Args:
text_cond_model: The text conditioning model.
"""
def __init__(self, *args, text_cond_model: str = 'openai/clip-vit-large-patch14', **kwargs):
super().__init__(*args, **kwargs)
self.text_cond_model_name = text_cond_model
self.text_cond_model = None # the model is init lazily
def _init_text_cond_model(self):
"""
Initialize the text conditioning model.
"""
# load model
with dist_utils.local_master_first():
model = CLIPTextModel.from_pretrained(self.text_cond_model_name)
tokenizer = AutoTokenizer.from_pretrained(self.text_cond_model_name)
model.eval()
model = model.cuda()
self.text_cond_model = {
'model': model,
'tokenizer': tokenizer,
}
self.text_cond_model['null_cond'] = self.encode_text([''])
@torch.no_grad()
def encode_text(self, text: List[str]) -> torch.Tensor:
"""
Encode the text.
"""
assert isinstance(text, list) and isinstance(text[0], str), "TextConditionedMixin only supports list of strings as cond"
if self.text_cond_model is None:
self._init_text_cond_model()
encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
tokens = encoding['input_ids'].cuda()
embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
return embeddings
def get_cond(self, cond, **kwargs):
"""
Get the conditioning data.
"""
cond = self.encode_text(cond)
kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
cond = super().get_cond(cond, **kwargs)
return cond
def get_inference_cond(self, cond, **kwargs):
"""
Get the conditioning data for inference.
"""
cond = self.encode_text(cond)
kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
cond = super().get_inference_cond(cond, **kwargs)
return cond
trellis/trainers/flow_matching/sparse_flow_matching.py
+286
View File
@@ -0,0 +1,286 @@
from typing import *
import os
import copy
import functools
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
from ...modules import sparse as sp
from ...utils.general_utils import dict_reduce
from ...utils.data_utils import cycle, BalancedResumableSampler
from .flow_matching import FlowMatchingTrainer
from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
from .mixins.text_conditioned import TextConditionedMixin
from .mixins.image_conditioned import ImageConditionedMixin
class SparseFlowMatchingTrainer(FlowMatchingTrainer):
"""
Trainer for sparse diffusion model with flow matching objective.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
"""
def prepare_dataloader(self, **kwargs):
"""
Prepare dataloader.
"""
self.data_sampler = BalancedResumableSampler(
self.dataset,
shuffle=True,
batch_size=self.batch_size_per_gpu,
)
self.dataloader = DataLoader(
self.dataset,
batch_size=self.batch_size_per_gpu,
num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
pin_memory=True,
drop_last=True,
persistent_workers=True,
collate_fn=functools.partial(self.dataset.collate_fn, split_size=self.batch_split),
sampler=self.data_sampler,
)
self.data_iterator = cycle(self.dataloader)
def training_losses(
self,
x_0: sp.SparseTensor,
cond=None,
**kwargs
) -> Tuple[Dict, Dict]:
"""
Compute training losses for a single timestep.
Args:
x_0: The [N x ... x C] sparse tensor of the inputs.
cond: The [N x ...] tensor of additional conditions.
kwargs: Additional arguments to pass to the backbone.
Returns:
a dict with the key "loss" containing a tensor of shape [N].
may also contain other keys for different terms.
"""
noise = x_0.replace(torch.randn_like(x_0.feats))
t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
x_t = self.diffuse(x_0, t, noise=noise)
cond = self.get_cond(cond, **kwargs)
pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
assert pred.shape == noise.shape == x_0.shape
target = self.get_v(x_0, noise, t)
terms = edict()
terms["mse"] = F.mse_loss(pred.feats, target.feats)
terms["loss"] = terms["mse"]
# log loss with time bins
mse_per_instance = np.array([
F.mse_loss(pred.feats[x_0.layout[i]], target.feats[x_0.layout[i]]).item()
for i in range(x_0.shape[0])
])
time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
for i in range(10):
if (time_bin == i).sum() != 0:
terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
return terms, {}
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
sampler = self.get_sampler()
sample_gt = []
sample = []
cond_vis = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
data = {k: v[:batch].cuda() if not isinstance(v, list) else v[:batch] for k, v in data.items()}
noise = data['x_0'].replace(torch.randn_like(data['x_0'].feats))
sample_gt.append(data['x_0'])
cond_vis.append(self.vis_cond(**data))
del data['x_0']
args = self.get_inference_cond(**data)
res = sampler.sample(
self.models['denoiser'],
noise=noise,
**args,
steps=50, cfg_strength=3.0, verbose=verbose,
)
sample.append(res.samples)
sample_gt = sp.sparse_cat(sample_gt)
sample = sp.sparse_cat(sample)
sample_dict = {
'sample_gt': {'value': sample_gt, 'type': 'sample'},
'sample': {'value': sample, 'type': 'sample'},
}
sample_dict.update(dict_reduce(cond_vis, None, {
'value': lambda x: torch.cat(x, dim=0),
'type': lambda x: x[0],
}))
return sample_dict
class SparseFlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, SparseFlowMatchingTrainer):
"""
Trainer for sparse diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
"""
pass
class TextConditionedSparseFlowMatchingCFGTrainer(TextConditionedMixin, SparseFlowMatchingCFGTrainer):
"""
Trainer for sparse text-conditioned diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
text_cond_model(str): Text conditioning model.
"""
pass
class ImageConditionedSparseFlowMatchingCFGTrainer(ImageConditionedMixin, SparseFlowMatchingCFGTrainer):
"""
Trainer for sparse image-conditioned diffusion model with flow matching objective and classifier-free guidance.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
t_schedule (dict): Time schedule for flow matching.
sigma_min (float): Minimum noise level.
p_uncond (float): Probability of dropping conditions.
image_cond_model (str): Image conditioning model.
"""
pass
trellis/trainers/utils.py
+77
View File
@@ -0,0 +1,77 @@
import torch.nn as nn
# FP16 utils
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
def make_master_params(model_params):
"""
Copy model parameters into a inflated tensor of full-precision parameters.
"""
master_params = _flatten_dense_tensors(
[param.detach().float() for param in model_params]
)
master_params = nn.Parameter(master_params)
master_params.requires_grad = True
return [master_params]
def unflatten_master_params(model_params, master_params):
"""
Unflatten the master parameters to look like model_params.
"""
return _unflatten_dense_tensors(master_params[0].detach(), model_params)
def model_params_to_master_params(model_params, master_params):
"""
Copy the model parameter data into the master parameters.
"""
master_params[0].detach().copy_(
_flatten_dense_tensors([param.detach().float() for param in model_params])
)
def master_params_to_model_params(model_params, master_params):
"""
Copy the master parameter data back into the model parameters.
"""
for param, master_param in zip(
model_params, _unflatten_dense_tensors(master_params[0].detach(), model_params)
):
param.detach().copy_(master_param)
def model_grads_to_master_grads(model_params, master_params):
"""
Copy the gradients from the model parameters into the master parameters
from make_master_params().
"""
master_params[0].grad = _flatten_dense_tensors(
[param.grad.data.detach().float() for param in model_params]
)
def zero_grad(model_params):
for param in model_params:
if param.grad is not None:
if param.grad.grad_fn is not None:
param.grad.detach_()
else:
param.grad.requires_grad_(False)
param.grad.zero_()
# LR Schedulers
from torch.optim.lr_scheduler import LambdaLR
class LinearWarmupLRScheduler(LambdaLR):
def __init__(self, optimizer, warmup_steps, last_epoch=-1):
self.warmup_steps = warmup_steps
super(LinearWarmupLRScheduler, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
def lr_lambda(self, current_step):
if current_step < self.warmup_steps:
return float(current_step + 1) / self.warmup_steps
return 1.0
trellis/trainers/vae/sparse_structure_vae.py
+130
View File
@@ -0,0 +1,130 @@
from typing import *
import copy
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from easydict import EasyDict as edict
from ..basic import BasicTrainer
class SparseStructureVaeTrainer(BasicTrainer):
"""
Trainer for Sparse Structure VAE.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
loss_type (str): Loss type. 'bce' for binary cross entropy, 'l1' for L1 loss, 'dice' for Dice loss.
lambda_kl (float): KL divergence loss weight.
"""
def __init__(
self,
*args,
loss_type='bce',
lambda_kl=1e-6,
**kwargs
):
super().__init__(*args, **kwargs)
self.loss_type = loss_type
self.lambda_kl = lambda_kl
def training_losses(
self,
ss: torch.Tensor,
**kwargs
) -> Tuple[Dict, Dict]:
"""
Compute training losses.
Args:
ss: The [N x 1 x H x W x D] tensor of binary sparse structure.
Returns:
a dict with the key "loss" containing a scalar tensor.
may also contain other keys for different terms.
"""
z, mean, logvar = self.training_models['encoder'](ss.float(), sample_posterior=True, return_raw=True)
logits = self.training_models['decoder'](z)
terms = edict(loss = 0.0)
if self.loss_type == 'bce':
terms["bce"] = F.binary_cross_entropy_with_logits(logits, ss.float(), reduction='mean')
terms["loss"] = terms["loss"] + terms["bce"]
elif self.loss_type == 'l1':
terms["l1"] = F.l1_loss(F.sigmoid(logits), ss.float(), reduction='mean')
terms["loss"] = terms["loss"] + terms["l1"]
elif self.loss_type == 'dice':
logits = F.sigmoid(logits)
terms["dice"] = 1 - (2 * (logits * ss.float()).sum() + 1) / (logits.sum() + ss.float().sum() + 1)
terms["loss"] = terms["loss"] + terms["dice"]
else:
raise ValueError(f'Invalid loss type {self.loss_type}')
terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
return terms, {}
@torch.no_grad()
def snapshot(self, suffix=None, num_samples=64, batch_size=1, verbose=False):
super().snapshot(suffix=suffix, num_samples=num_samples, batch_size=batch_size, verbose=verbose)
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
gts = []
recons = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
args = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
z = self.models['encoder'](args['ss'].float(), sample_posterior=False)
logits = self.models['decoder'](z)
recon = (logits > 0).long()
gts.append(args['ss'])
recons.append(recon)
sample_dict = {
'gt': {'value': torch.cat(gts, dim=0), 'type': 'sample'},
'recon': {'value': torch.cat(recons, dim=0), 'type': 'sample'},
}
return sample_dict
trellis/trainers/vae/structured_latent_vae_gaussian.py
+275
View File
@@ -0,0 +1,275 @@
from typing import *
import copy
import torch
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
import utils3d.torch
from ..basic import BasicTrainer
from ...representations import Gaussian
from ...renderers import GaussianRenderer
from ...modules.sparse import SparseTensor
from ...utils.loss_utils import l1_loss, l2_loss, ssim, lpips
class SLatVaeGaussianTrainer(BasicTrainer):
"""
Trainer for structured latent VAE.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
loss_type (str): Loss type. Can be 'l1', 'l2'
lambda_ssim (float): SSIM loss weight.
lambda_lpips (float): LPIPS loss weight.
lambda_kl (float): KL loss weight.
regularizations (dict): Regularization config.
"""
def __init__(
self,
*args,
loss_type: str = 'l1',
lambda_ssim: float = 0.2,
lambda_lpips: float = 0.2,
lambda_kl: float = 1e-6,
regularizations: Dict = {},
**kwargs
):
super().__init__(*args, **kwargs)
self.loss_type = loss_type
self.lambda_ssim = lambda_ssim
self.lambda_lpips = lambda_lpips
self.lambda_kl = lambda_kl
self.regularizations = regularizations
self._init_renderer()
def _init_renderer(self):
rendering_options = {"near" : 0.8,
"far" : 1.6,
"bg_color" : 'random'}
self.renderer = GaussianRenderer(rendering_options)
self.renderer.pipe.kernel_size = self.models['decoder'].rep_config['2d_filter_kernel_size']
def _render_batch(self, reps: List[Gaussian], extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
"""
Render a batch of representations.
Args:
reps: The dictionary of lists of representations.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
"""
ret = None
for i, representation in enumerate(reps):
render_pack = self.renderer.render(representation, extrinsics[i], intrinsics[i])
if ret is None:
ret = {k: [] for k in list(render_pack.keys()) + ['bg_color']}
for k, v in render_pack.items():
ret[k].append(v)
ret['bg_color'].append(self.renderer.bg_color)
for k, v in ret.items():
ret[k] = torch.stack(v, dim=0)
return ret
@torch.no_grad()
def _get_status(self, z: SparseTensor, reps: List[Gaussian]) -> Dict:
xyz = torch.cat([g.get_xyz for g in reps], dim=0)
xyz_base = (z.coords[:, 1:].float() + 0.5) / self.models['decoder'].resolution - 0.5
offset = xyz - xyz_base.unsqueeze(1).expand(-1, self.models['decoder'].rep_config['num_gaussians'], -1).reshape(-1, 3)
status = {
'xyz': xyz,
'offset': offset,
'scale': torch.cat([g.get_scaling for g in reps], dim=0),
'opacity': torch.cat([g.get_opacity for g in reps], dim=0),
}
for k in list(status.keys()):
status[k] = {
'mean': status[k].mean().item(),
'max': status[k].max().item(),
'min': status[k].min().item(),
}
return status
def _get_regularization_loss(self, reps: List[Gaussian]) -> Tuple[torch.Tensor, Dict]:
loss = 0.0
terms = {}
if 'lambda_vol' in self.regularizations:
scales = torch.cat([g.get_scaling for g in reps], dim=0) # [N x 3]
volume = torch.prod(scales, dim=1) # [N]
terms[f'reg_vol'] = volume.mean()
loss = loss + self.regularizations['lambda_vol'] * terms[f'reg_vol']
if 'lambda_opacity' in self.regularizations:
opacity = torch.cat([g.get_opacity for g in reps], dim=0)
terms[f'reg_opacity'] = (opacity - 1).pow(2).mean()
loss = loss + self.regularizations['lambda_opacity'] * terms[f'reg_opacity']
return loss, terms
def training_losses(
self,
feats: SparseTensor,
image: torch.Tensor,
alpha: torch.Tensor,
extrinsics: torch.Tensor,
intrinsics: torch.Tensor,
return_aux: bool = False,
**kwargs
) -> Tuple[Dict, Dict]:
"""
Compute training losses.
Args:
feats: The [N x * x C] sparse tensor of features.
image: The [N x 3 x H x W] tensor of images.
alpha: The [N x H x W] tensor of alpha channels.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
return_aux: Whether to return auxiliary information.
Returns:
a dict with the key "loss" containing a scalar tensor.
may also contain other keys for different terms.
"""
z, mean, logvar = self.training_models['encoder'](feats, sample_posterior=True, return_raw=True)
reps = self.training_models['decoder'](z)
self.renderer.rendering_options.resolution = image.shape[-1]
render_results = self._render_batch(reps, extrinsics, intrinsics)
terms = edict(loss = 0.0, rec = 0.0)
rec_image = render_results['color']
gt_image = image * alpha[:, None] + (1 - alpha[:, None]) * render_results['bg_color'][..., None, None]
if self.loss_type == 'l1':
terms["l1"] = l1_loss(rec_image, gt_image)
terms["rec"] = terms["rec"] + terms["l1"]
elif self.loss_type == 'l2':
terms["l2"] = l2_loss(rec_image, gt_image)
terms["rec"] = terms["rec"] + terms["l2"]
else:
raise ValueError(f"Invalid loss type: {self.loss_type}")
if self.lambda_ssim > 0:
terms["ssim"] = 1 - ssim(rec_image, gt_image)
terms["rec"] = terms["rec"] + self.lambda_ssim * terms["ssim"]
if self.lambda_lpips > 0:
terms["lpips"] = lpips(rec_image, gt_image)
terms["rec"] = terms["rec"] + self.lambda_lpips * terms["lpips"]
terms["loss"] = terms["loss"] + terms["rec"]
terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
reg_loss, reg_terms = self._get_regularization_loss(reps)
terms.update(reg_terms)
terms["loss"] = terms["loss"] + reg_loss
status = self._get_status(z, reps)
if return_aux:
return terms, status, {'rec_image': rec_image, 'gt_image': gt_image}
return terms, status
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
ret_dict = {}
gt_images = []
exts = []
ints = []
reps = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
args = {k: v[:batch].cuda() for k, v in data.items()}
gt_images.append(args['image'] * args['alpha'][:, None])
exts.append(args['extrinsics'])
ints.append(args['intrinsics'])
z = self.models['encoder'](args['feats'], sample_posterior=True, return_raw=False)
reps.extend(self.models['decoder'](z))
gt_images = torch.cat(gt_images, dim=0)
ret_dict.update({f'gt_image': {'value': gt_images, 'type': 'image'}})
# render single view
exts = torch.cat(exts, dim=0)
ints = torch.cat(ints, dim=0)
self.renderer.rendering_options.bg_color = (0, 0, 0)
self.renderer.rendering_options.resolution = gt_images.shape[-1]
render_results = self._render_batch(reps, exts, ints)
ret_dict.update({f'rec_image': {'value': render_results['color'], 'type': 'image'}})
# render multiview
self.renderer.rendering_options.resolution = 512
## Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
## render each view
miltiview_images = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(30)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
extrinsics = extrinsics.unsqueeze(0).expand(num_samples, -1, -1)
intrinsics = intrinsics.unsqueeze(0).expand(num_samples, -1, -1)
render_results = self._render_batch(reps, extrinsics, intrinsics)
miltiview_images.append(render_results['color'])
## Concatenate views
miltiview_images = torch.cat([
torch.cat(miltiview_images[:2], dim=-2),
torch.cat(miltiview_images[2:], dim=-2),
], dim=-1)
ret_dict.update({f'miltiview_image': {'value': miltiview_images, 'type': 'image'}})
self.renderer.rendering_options.bg_color = 'random'
return ret_dict
trellis/trainers/vae/structured_latent_vae_mesh_dec.py
+382
View File
@@ -0,0 +1,382 @@
from typing import *
import copy
import torch
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
import utils3d.torch
from ..basic import BasicTrainer
from ...representations import MeshExtractResult
from ...renderers import MeshRenderer
from ...modules.sparse import SparseTensor
from ...utils.loss_utils import l1_loss, smooth_l1_loss, ssim, lpips
from ...utils.data_utils import recursive_to_device
class SLatVaeMeshDecoderTrainer(BasicTrainer):
"""
Trainer for structured latent VAE Mesh Decoder.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
loss_type (str): Loss type. Can be 'l1', 'l2'
lambda_ssim (float): SSIM loss weight.
lambda_lpips (float): LPIPS loss weight.
"""
def __init__(
self,
*args,
depth_loss_type: str = 'l1',
lambda_depth: int = 1,
lambda_ssim: float = 0.2,
lambda_lpips: float = 0.2,
lambda_tsdf: float = 0.01,
lambda_color: float = 0.1,
**kwargs
):
super().__init__(*args, **kwargs)
self.depth_loss_type = depth_loss_type
self.lambda_depth = lambda_depth
self.lambda_ssim = lambda_ssim
self.lambda_lpips = lambda_lpips
self.lambda_tsdf = lambda_tsdf
self.lambda_color = lambda_color
self.use_color = self.lambda_color > 0
self._init_renderer()
def _init_renderer(self):
rendering_options = {"near" : 1,
"far" : 3}
self.renderer = MeshRenderer(rendering_options, device=self.device)
def _render_batch(self, reps: List[MeshExtractResult], extrinsics: torch.Tensor, intrinsics: torch.Tensor,
return_types=['mask', 'normal', 'depth']) -> Dict[str, torch.Tensor]:
"""
Render a batch of representations.
Args:
reps: The dictionary of lists of representations.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
return_types: vary in ['mask', 'normal', 'depth', 'normal_map', 'color']
Returns:
a dict with
reg_loss : [N] tensor of regularization losses
mask : [N x 1 x H x W] tensor of rendered masks
normal : [N x 3 x H x W] tensor of rendered normals
depth : [N x 1 x H x W] tensor of rendered depths
"""
ret = {k : [] for k in return_types}
for i, rep in enumerate(reps):
out_dict = self.renderer.render(rep, extrinsics[i], intrinsics[i], return_types=return_types)
for k in out_dict:
ret[k].append(out_dict[k][None] if k in ['mask', 'depth'] else out_dict[k])
for k in ret:
ret[k] = torch.stack(ret[k])
return ret
@staticmethod
def _tsdf_reg_loss(rep: MeshExtractResult, depth_map: torch.Tensor, extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
# Calculate tsdf
with torch.no_grad():
# Project points to camera and calculate pseudo-sdf as difference between gt depth and projected depth
projected_pts, pts_depth = utils3d.torch.project_cv(extrinsics=extrinsics, intrinsics=intrinsics, points=rep.tsdf_v)
projected_pts = (projected_pts - 0.5) * 2.0
depth_map_res = depth_map.shape[1]
gt_depth = torch.nn.functional.grid_sample(depth_map.reshape(1, 1, depth_map_res, depth_map_res),
projected_pts.reshape(1, 1, -1, 2), mode='bilinear', padding_mode='border', align_corners=True)
pseudo_sdf = gt_depth.flatten() - pts_depth.flatten()
# Truncate pseudo-sdf
delta = 1 / rep.res * 3.0
trunc_mask = pseudo_sdf > -delta
# Loss
gt_tsdf = pseudo_sdf[trunc_mask]
tsdf = rep.tsdf_s.flatten()[trunc_mask]
gt_tsdf = torch.clamp(gt_tsdf, -delta, delta)
return torch.mean((tsdf - gt_tsdf) ** 2)
def _calc_tsdf_loss(self, reps : list[MeshExtractResult], depth_maps, extrinsics, intrinsics) -> torch.Tensor:
tsdf_loss = 0.0
for i, rep in enumerate(reps):
tsdf_loss += self._tsdf_reg_loss(rep, depth_maps[i], extrinsics[i], intrinsics[i])
return tsdf_loss / len(reps)
@torch.no_grad()
def _flip_normal(self, normal: torch.Tensor, extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
"""
Flip normal to align with camera.
"""
normal = normal * 2.0 - 1.0
R = torch.zeros_like(extrinsics)
R[:, :3, :3] = extrinsics[:, :3, :3]
R[:, 3, 3] = 1.0
view_dir = utils3d.torch.unproject_cv(
utils3d.torch.image_uv(*normal.shape[-2:], device=self.device).reshape(1, -1, 2),
torch.ones(*normal.shape[-2:], device=self.device).reshape(1, -1),
R, intrinsics
).reshape(-1, *normal.shape[-2:], 3).permute(0, 3, 1, 2)
unflip = (normal * view_dir).sum(1, keepdim=True) < 0
normal *= unflip * 2.0 - 1.0
return (normal + 1.0) / 2.0
def _perceptual_loss(self, gt: torch.Tensor, pred: torch.Tensor, name: str) -> Dict[str, torch.Tensor]:
"""
Combination of L1, SSIM, and LPIPS loss.
"""
if gt.shape[1] != 3:
assert gt.shape[-1] == 3
gt = gt.permute(0, 3, 1, 2)
if pred.shape[1] != 3:
assert pred.shape[-1] == 3
pred = pred.permute(0, 3, 1, 2)
terms = {
f"{name}_loss" : l1_loss(gt, pred),
f"{name}_loss_ssim" : 1 - ssim(gt, pred),
f"{name}_loss_lpips" : lpips(gt, pred)
}
terms[f"{name}_loss_perceptual"] = terms[f"{name}_loss"] + terms[f"{name}_loss_ssim"] * self.lambda_ssim + terms[f"{name}_loss_lpips"] * self.lambda_lpips
return terms
def geometry_losses(
self,
reps: List[MeshExtractResult],
mesh: List[Dict],
normal_map: torch.Tensor,
extrinsics: torch.Tensor,
intrinsics: torch.Tensor,
):
with torch.no_grad():
gt_meshes = []
for i in range(len(reps)):
gt_mesh = MeshExtractResult(mesh[i]['vertices'].to(self.device), mesh[i]['faces'].to(self.device))
gt_meshes.append(gt_mesh)
target = self._render_batch(gt_meshes, extrinsics, intrinsics, return_types=['mask', 'depth', 'normal'])
target['normal'] = self._flip_normal(target['normal'], extrinsics, intrinsics)
terms = edict(geo_loss = 0.0)
if self.lambda_tsdf > 0:
tsdf_loss = self._calc_tsdf_loss(reps, target['depth'], extrinsics, intrinsics)
terms['tsdf_loss'] = tsdf_loss
terms['geo_loss'] += tsdf_loss * self.lambda_tsdf
return_types = ['mask', 'depth', 'normal', 'normal_map'] if self.use_color else ['mask', 'depth', 'normal']
buffer = self._render_batch(reps, extrinsics, intrinsics, return_types=return_types)
success_mask = torch.tensor([rep.success for rep in reps], device=self.device)
if success_mask.sum() != 0:
for k, v in buffer.items():
buffer[k] = v[success_mask]
for k, v in target.items():
target[k] = v[success_mask]
terms['mask_loss'] = l1_loss(buffer['mask'], target['mask'])
if self.depth_loss_type == 'l1':
terms['depth_loss'] = l1_loss(buffer['depth'] * target['mask'], target['depth'] * target['mask'])
elif self.depth_loss_type == 'smooth_l1':
terms['depth_loss'] = smooth_l1_loss(buffer['depth'] * target['mask'], target['depth'] * target['mask'], beta=1.0 / (2 * reps[0].res))
else:
raise ValueError(f"Unsupported depth loss type: {self.depth_loss_type}")
terms.update(self._perceptual_loss(buffer['normal'] * target['mask'], target['normal'] * target['mask'], 'normal'))
terms['geo_loss'] = terms['geo_loss'] + terms['mask_loss'] + terms['depth_loss'] * self.lambda_depth + terms['normal_loss_perceptual']
if self.use_color and normal_map is not None:
terms.update(self._perceptual_loss(normal_map[success_mask], buffer['normal_map'], 'normal_map'))
terms['geo_loss'] = terms['geo_loss'] + terms['normal_map_loss_perceptual'] * self.lambda_color
return terms
def color_losses(self, reps, image, alpha, extrinsics, intrinsics):
terms = edict(color_loss = torch.tensor(0.0, device=self.device))
buffer = self._render_batch(reps, extrinsics, intrinsics, return_types=['color'])
success_mask = torch.tensor([rep.success for rep in reps], device=self.device)
if success_mask.sum() != 0:
terms.update(self._perceptual_loss(image * alpha[:, None][success_mask], buffer['color'][success_mask], 'color'))
terms['color_loss'] = terms['color_loss'] + terms['color_loss_perceptual'] * self.lambda_color
return terms
def training_losses(
self,
latents: SparseTensor,
image: torch.Tensor,
alpha: torch.Tensor,
mesh: List[Dict],
extrinsics: torch.Tensor,
intrinsics: torch.Tensor,
normal_map: torch.Tensor = None,
) -> Tuple[Dict, Dict]:
"""
Compute training losses.
Args:
latents: The [N x * x C] sparse latents
image: The [N x 3 x H x W] tensor of images.
alpha: The [N x H x W] tensor of alpha channels.
mesh: The list of dictionaries of meshes.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
Returns:
a dict with the key "loss" containing a scalar tensor.
may also contain other keys for different terms.
"""
reps = self.training_models['decoder'](latents)
self.renderer.rendering_options.resolution = image.shape[-1]
terms = edict(loss = 0.0, rec = 0.0)
terms['reg_loss'] = sum([rep.reg_loss for rep in reps]) / len(reps)
terms['loss'] = terms['loss'] + terms['reg_loss']
geo_terms = self.geometry_losses(reps, mesh, normal_map, extrinsics, intrinsics)
terms.update(geo_terms)
terms['loss'] = terms['loss'] + terms['geo_loss']
if self.use_color:
color_terms = self.color_losses(reps, image, alpha, extrinsics, intrinsics)
terms.update(color_terms)
terms['loss'] = terms['loss'] + terms['color_loss']
return terms, {}
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
ret_dict = {}
gt_images = []
gt_normal_maps = []
gt_meshes = []
exts = []
ints = []
reps = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
args = recursive_to_device(data, 'cuda')
gt_images.append(args['image'] * args['alpha'][:, None])
if self.use_color and 'normal_map' in data:
gt_normal_maps.append(args['normal_map'])
gt_meshes.extend(args['mesh'])
exts.append(args['extrinsics'])
ints.append(args['intrinsics'])
reps.extend(self.models['decoder'](args['latents']))
gt_images = torch.cat(gt_images, dim=0)
ret_dict.update({f'gt_image': {'value': gt_images, 'type': 'image'}})
if self.use_color and gt_normal_maps:
gt_normal_maps = torch.cat(gt_normal_maps, dim=0)
ret_dict.update({f'gt_normal_map': {'value': gt_normal_maps, 'type': 'image'}})
# render single view
exts = torch.cat(exts, dim=0)
ints = torch.cat(ints, dim=0)
self.renderer.rendering_options.bg_color = (0, 0, 0)
self.renderer.rendering_options.resolution = gt_images.shape[-1]
gt_render_results = self._render_batch([
MeshExtractResult(vertices=mesh['vertices'].to(self.device), faces=mesh['faces'].to(self.device))
for mesh in gt_meshes
], exts, ints, return_types=['normal'])
ret_dict.update({f'gt_normal': {'value': self._flip_normal(gt_render_results['normal'], exts, ints), 'type': 'image'}})
return_types = ['normal']
if self.use_color:
return_types.append('color')
if 'normal_map' in data:
return_types.append('normal_map')
render_results = self._render_batch(reps, exts, ints, return_types=return_types)
ret_dict.update({f'rec_normal': {'value': render_results['normal'], 'type': 'image'}})
if 'color' in return_types:
ret_dict.update({f'rec_image': {'value': render_results['color'], 'type': 'image'}})
if 'normal_map' in return_types:
ret_dict.update({f'rec_normal_map': {'value': render_results['normal_map'], 'type': 'image'}})
# render multiview
self.renderer.rendering_options.resolution = 512
## Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
## render each view
multiview_normals = []
multiview_normal_maps = []
miltiview_images = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(30)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
extrinsics = extrinsics.unsqueeze(0).expand(num_samples, -1, -1)
intrinsics = intrinsics.unsqueeze(0).expand(num_samples, -1, -1)
render_results = self._render_batch(reps, extrinsics, intrinsics, return_types=return_types)
multiview_normals.append(render_results['normal'])
if 'color' in return_types:
miltiview_images.append(render_results['color'])
if 'normal_map' in return_types:
multiview_normal_maps.append(render_results['normal_map'])
## Concatenate views
multiview_normals = torch.cat([
torch.cat(multiview_normals[:2], dim=-2),
torch.cat(multiview_normals[2:], dim=-2),
], dim=-1)
ret_dict.update({f'multiview_normal': {'value': multiview_normals, 'type': 'image'}})
if 'color' in return_types:
miltiview_images = torch.cat([
torch.cat(miltiview_images[:2], dim=-2),
torch.cat(miltiview_images[2:], dim=-2),
], dim=-1)
ret_dict.update({f'multiview_image': {'value': miltiview_images, 'type': 'image'}})
if 'normal_map' in return_types:
multiview_normal_maps = torch.cat([
torch.cat(multiview_normal_maps[:2], dim=-2),
torch.cat(multiview_normal_maps[2:], dim=-2),
], dim=-1)
ret_dict.update({f'multiview_normal_map': {'value': multiview_normal_maps, 'type': 'image'}})
return ret_dict
trellis/trainers/vae/structured_latent_vae_rf_dec.py
+223
View File
@@ -0,0 +1,223 @@
from typing import *
import copy
import torch
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
import utils3d.torch
from ..basic import BasicTrainer
from ...representations import Strivec
from ...renderers import OctreeRenderer
from ...modules.sparse import SparseTensor
from ...utils.loss_utils import l1_loss, l2_loss, ssim, lpips
class SLatVaeRadianceFieldDecoderTrainer(BasicTrainer):
"""
Trainer for structured latent VAE Radiance Field Decoder.
Args:
models (dict[str, nn.Module]): Models to train.
dataset (torch.utils.data.Dataset): Dataset.
output_dir (str): Output directory.
load_dir (str): Load directory.
step (int): Step to load.
batch_size (int): Batch size.
batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
batch_split (int): Split batch with gradient accumulation.
max_steps (int): Max steps.
optimizer (dict): Optimizer config.
lr_scheduler (dict): Learning rate scheduler config.
elastic (dict): Elastic memory management config.
grad_clip (float or dict): Gradient clip config.
ema_rate (float or list): Exponential moving average rates.
fp16_mode (str): FP16 mode.
- None: No FP16.
- 'inflat_all': Hold a inflated fp32 master param for all params.
- 'amp': Automatic mixed precision.
fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
finetune_ckpt (dict): Finetune checkpoint.
log_param_stats (bool): Log parameter stats.
i_print (int): Print interval.
i_log (int): Log interval.
i_sample (int): Sample interval.
i_save (int): Save interval.
i_ddpcheck (int): DDP check interval.
loss_type (str): Loss type. Can be 'l1', 'l2'
lambda_ssim (float): SSIM loss weight.
lambda_lpips (float): LPIPS loss weight.
"""
def __init__(
self,
*args,
loss_type: str = 'l1',
lambda_ssim: float = 0.2,
lambda_lpips: float = 0.2,
**kwargs
):
super().__init__(*args, **kwargs)
self.loss_type = loss_type
self.lambda_ssim = lambda_ssim
self.lambda_lpips = lambda_lpips
self._init_renderer()
def _init_renderer(self):
rendering_options = {"near" : 0.8,
"far" : 1.6,
"bg_color" : 'random'}
self.renderer = OctreeRenderer(rendering_options)
self.renderer.pipe.primitive = 'trivec'
def _render_batch(self, reps: List[Strivec], extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
"""
Render a batch of representations.
Args:
reps: The dictionary of lists of representations.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
"""
ret = None
for i, representation in enumerate(reps):
render_pack = self.renderer.render(representation, extrinsics[i], intrinsics[i])
if ret is None:
ret = {k: [] for k in list(render_pack.keys()) + ['bg_color']}
for k, v in render_pack.items():
ret[k].append(v)
ret['bg_color'].append(self.renderer.bg_color)
for k, v in ret.items():
ret[k] = torch.stack(v, dim=0)
return ret
def training_losses(
self,
latents: SparseTensor,
image: torch.Tensor,
alpha: torch.Tensor,
extrinsics: torch.Tensor,
intrinsics: torch.Tensor,
return_aux: bool = False,
**kwargs
) -> Tuple[Dict, Dict]:
"""
Compute training losses.
Args:
latents: The [N x * x C] sparse latents
image: The [N x 3 x H x W] tensor of images.
alpha: The [N x H x W] tensor of alpha channels.
extrinsics: The [N x 4 x 4] tensor of extrinsics.
intrinsics: The [N x 3 x 3] tensor of intrinsics.
return_aux: Whether to return auxiliary information.
Returns:
a dict with the key "loss" containing a scalar tensor.
may also contain other keys for different terms.
"""
reps = self.training_models['decoder'](latents)
self.renderer.rendering_options.resolution = image.shape[-1]
render_results = self._render_batch(reps, extrinsics, intrinsics)
terms = edict(loss = 0.0, rec = 0.0)
rec_image = render_results['color']
gt_image = image * alpha[:, None] + (1 - alpha[:, None]) * render_results['bg_color'][..., None, None]
if self.loss_type == 'l1':
terms["l1"] = l1_loss(rec_image, gt_image)
terms["rec"] = terms["rec"] + terms["l1"]
elif self.loss_type == 'l2':
terms["l2"] = l2_loss(rec_image, gt_image)
terms["rec"] = terms["rec"] + terms["l2"]
else:
raise ValueError(f"Invalid loss type: {self.loss_type}")
if self.lambda_ssim > 0:
terms["ssim"] = 1 - ssim(rec_image, gt_image)
terms["rec"] = terms["rec"] + self.lambda_ssim * terms["ssim"]
if self.lambda_lpips > 0:
terms["lpips"] = lpips(rec_image, gt_image)
terms["rec"] = terms["rec"] + self.lambda_lpips * terms["lpips"]
terms["loss"] = terms["loss"] + terms["rec"]
if return_aux:
return terms, {}, {'rec_image': rec_image, 'gt_image': gt_image}
return terms, {}
@torch.no_grad()
def run_snapshot(
self,
num_samples: int,
batch_size: int,
verbose: bool = False,
) -> Dict:
dataloader = DataLoader(
copy.deepcopy(self.dataset),
batch_size=batch_size,
shuffle=True,
num_workers=0,
collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
)
# inference
ret_dict = {}
gt_images = []
exts = []
ints = []
reps = []
for i in range(0, num_samples, batch_size):
batch = min(batch_size, num_samples - i)
data = next(iter(dataloader))
args = {k: v[:batch].cuda() for k, v in data.items()}
gt_images.append(args['image'] * args['alpha'][:, None])
exts.append(args['extrinsics'])
ints.append(args['intrinsics'])
reps.extend(self.models['decoder'](args['latents']))
gt_images = torch.cat(gt_images, dim=0)
ret_dict.update({f'gt_image': {'value': gt_images, 'type': 'image'}})
# render single view
exts = torch.cat(exts, dim=0)
ints = torch.cat(ints, dim=0)
self.renderer.rendering_options.bg_color = (0, 0, 0)
self.renderer.rendering_options.resolution = gt_images.shape[-1]
render_results = self._render_batch(reps, exts, ints)
ret_dict.update({f'rec_image': {'value': render_results['color'], 'type': 'image'}})
# render multiview
self.renderer.rendering_options.resolution = 512
## Build camera
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
yaws = [y + yaws_offset for y in yaws]
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
## render each view
miltiview_images = []
for yaw, pitch in zip(yaws, pitch):
orig = torch.tensor([
np.sin(yaw) * np.cos(pitch),
np.cos(yaw) * np.cos(pitch),
np.sin(pitch),
]).float().cuda() * 2
fov = torch.deg2rad(torch.tensor(30)).cuda()
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
extrinsics = extrinsics.unsqueeze(0).expand(num_samples, -1, -1)
intrinsics = intrinsics.unsqueeze(0).expand(num_samples, -1, -1)
render_results = self._render_batch(reps, extrinsics, intrinsics)
miltiview_images.append(render_results['color'])
## Concatenate views
miltiview_images = torch.cat([
torch.cat(miltiview_images[:2], dim=-2),
torch.cat(miltiview_images[2:], dim=-2),
], dim=-1)
ret_dict.update({f'miltiview_image': {'value': miltiview_images, 'type': 'image'}})
self.renderer.rendering_options.bg_color = 'random'
return ret_dict
trellis/utils/data_utils.py
+226
View File
@@ -0,0 +1,226 @@
from typing import *
import math
import torch
import numpy as np
from torch.utils.data import Sampler, Dataset, DataLoader, DistributedSampler
import torch.distributed as dist
def recursive_to_device(
data: Any,
device: torch.device,
non_blocking: bool = False,
) -> Any:
"""
Recursively move all tensors in a data structure to a device.
"""
if hasattr(data, "to"):
return data.to(device, non_blocking=non_blocking)
elif isinstance(data, (list, tuple)):
return type(data)(recursive_to_device(d, device, non_blocking) for d in data)
elif isinstance(data, dict):
return {k: recursive_to_device(v, device, non_blocking) for k, v in data.items()}
else:
return data
def load_balanced_group_indices(
load: List[int],
num_groups: int,
equal_size: bool = False,
) -> List[List[int]]:
"""
Split indices into groups with balanced load.
"""
if equal_size:
group_size = len(load) // num_groups
indices = np.argsort(load)[::-1]
groups = [[] for _ in range(num_groups)]
group_load = np.zeros(num_groups)
for idx in indices:
min_group_idx = np.argmin(group_load)
groups[min_group_idx].append(idx)
if equal_size and len(groups[min_group_idx]) == group_size:
group_load[min_group_idx] = float('inf')
else:
group_load[min_group_idx] += load[idx]
return groups
def cycle(data_loader: DataLoader) -> Iterator:
while True:
for data in data_loader:
if isinstance(data_loader.sampler, ResumableSampler):
data_loader.sampler.idx += data_loader.batch_size # type: ignore[attr-defined]
yield data
if isinstance(data_loader.sampler, DistributedSampler):
data_loader.sampler.epoch += 1
if isinstance(data_loader.sampler, ResumableSampler):
data_loader.sampler.epoch += 1
data_loader.sampler.idx = 0
class ResumableSampler(Sampler):
"""
Distributed sampler that is resumable.
Args:
dataset: Dataset used for sampling.
rank (int, optional): Rank of the current process within :attr:`num_replicas`.
By default, :attr:`rank` is retrieved from the current distributed
group.
shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
indices.
seed (int, optional): random seed used to shuffle the sampler if
:attr:`shuffle=True`. This number should be identical across all
processes in the distributed group. Default: ``0``.
drop_last (bool, optional): if ``True``, then the sampler will drop the
tail of the data to make it evenly divisible across the number of
replicas. If ``False``, the sampler will add extra indices to make
the data evenly divisible across the replicas. Default: ``False``.
"""
def __init__(
self,
dataset: Dataset,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
) -> None:
self.dataset = dataset
self.epoch = 0
self.idx = 0
self.drop_last = drop_last
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
self.rank = dist.get_rank() if dist.is_initialized() else 0
# If the dataset length is evenly divisible by # of replicas, then there
# is no need to drop any data, since the dataset will be split equally.
if self.drop_last and len(self.dataset) % self.world_size != 0: # type: ignore[arg-type]
# Split to nearest available length that is evenly divisible.
# This is to ensure each rank receives the same amount of data when
# using this Sampler.
self.num_samples = math.ceil(
(len(self.dataset) - self.world_size) / self.world_size # type: ignore[arg-type]
)
else:
self.num_samples = math.ceil(len(self.dataset) / self.world_size) # type: ignore[arg-type]
self.total_size = self.num_samples * self.world_size
self.shuffle = shuffle
self.seed = seed
def __iter__(self) -> Iterator:
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type]
else:
indices = list(range(len(self.dataset))) # type: ignore[arg-type]
if not self.drop_last:
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
if padding_size <= len(indices):
indices += indices[:padding_size]
else:
indices += (indices * math.ceil(padding_size / len(indices)))[
:padding_size
]
else:
# remove tail of data to make it evenly divisible.
indices = indices[: self.total_size]
assert len(indices) == self.total_size
# subsample
indices = indices[self.rank : self.total_size : self.world_size]
# resume from previous state
indices = indices[self.idx:]
return iter(indices)
def __len__(self) -> int:
return self.num_samples
def state_dict(self) -> dict[str, int]:
return {
'epoch': self.epoch,
'idx': self.idx,
}
def load_state_dict(self, state_dict):
self.epoch = state_dict['epoch']
self.idx = state_dict['idx']
class BalancedResumableSampler(ResumableSampler):
"""
Distributed sampler that is resumable and balances the load among the processes.
Args:
dataset: Dataset used for sampling.
rank (int, optional): Rank of the current process within :attr:`num_replicas`.
By default, :attr:`rank` is retrieved from the current distributed
group.
shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
indices.
seed (int, optional): random seed used to shuffle the sampler if
:attr:`shuffle=True`. This number should be identical across all
processes in the distributed group. Default: ``0``.
drop_last (bool, optional): if ``True``, then the sampler will drop the
tail of the data to make it evenly divisible across the number of
replicas. If ``False``, the sampler will add extra indices to make
the data evenly divisible across the replicas. Default: ``False``.
"""
def __init__(
self,
dataset: Dataset,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
batch_size: int = 1,
) -> None:
assert hasattr(dataset, 'loads'), 'Dataset must have "loads" attribute to use BalancedResumableSampler'
super().__init__(dataset, shuffle, seed, drop_last)
self.batch_size = batch_size
self.loads = dataset.loads
def __iter__(self) -> Iterator:
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type]
else:
indices = list(range(len(self.dataset))) # type: ignore[arg-type]
if not self.drop_last:
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
if padding_size <= len(indices):
indices += indices[:padding_size]
else:
indices += (indices * math.ceil(padding_size / len(indices)))[
:padding_size
]
else:
# remove tail of data to make it evenly divisible.
indices = indices[: self.total_size]
assert len(indices) == self.total_size
# balance load among processes
num_batches = len(indices) // (self.batch_size * self.world_size)
balanced_indices = []
for i in range(num_batches):
start_idx = i * self.batch_size * self.world_size
end_idx = (i + 1) * self.batch_size * self.world_size
batch_indices = indices[start_idx:end_idx]
batch_loads = [self.loads[idx] for idx in batch_indices]
groups = load_balanced_group_indices(batch_loads, self.world_size, equal_size=True)
balanced_indices.extend([batch_indices[j] for j in groups[self.rank]])
# resume from previous state
indices = balanced_indices[self.idx:]
return iter(indices)
trellis/utils/dist_utils.py
+93
View File
@@ -0,0 +1,93 @@
import os
import io
from contextlib import contextmanager
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def setup_dist(rank, local_rank, world_size, master_addr, master_port):
os.environ['MASTER_ADDR'] = master_addr
os.environ['MASTER_PORT'] = master_port
os.environ['WORLD_SIZE'] = str(world_size)
os.environ['RANK'] = str(rank)
os.environ['LOCAL_RANK'] = str(local_rank)
torch.cuda.set_device(local_rank)
dist.init_process_group('nccl', rank=rank, world_size=world_size)
def read_file_dist(path):
"""
Read the binary file distributedly.
File is only read once by the rank 0 process and broadcasted to other processes.
Returns:
data (io.BytesIO): The binary data read from the file.
"""
if dist.is_initialized() and dist.get_world_size() > 1:
# read file
size = torch.LongTensor(1).cuda()
if dist.get_rank() == 0:
with open(path, 'rb') as f:
data = f.read()
data = torch.ByteTensor(
torch.UntypedStorage.from_buffer(data, dtype=torch.uint8)
).cuda()
size[0] = data.shape[0]
# broadcast size
dist.broadcast(size, src=0)
if dist.get_rank() != 0:
data = torch.ByteTensor(size[0].item()).cuda()
# broadcast data
dist.broadcast(data, src=0)
# convert to io.BytesIO
data = data.cpu().numpy().tobytes()
data = io.BytesIO(data)
return data
else:
with open(path, 'rb') as f:
data = f.read()
data = io.BytesIO(data)
return data
def unwrap_dist(model):
"""
Unwrap the model from distributed training.
"""
if isinstance(model, DDP):
return model.module
return model
@contextmanager
def master_first():
"""
A context manager that ensures master process executes first.
"""
if not dist.is_initialized():
yield
else:
if dist.get_rank() == 0:
yield
dist.barrier()
else:
dist.barrier()
yield
@contextmanager
def local_master_first():
"""
A context manager that ensures local master process executes first.
"""
if not dist.is_initialized():
yield
else:
if dist.get_rank() % torch.cuda.device_count() == 0:
yield
dist.barrier()
else:
dist.barrier()
yield
trellis/utils/elastic_utils.py
+228
View File
@@ -0,0 +1,228 @@
from abc import abstractmethod
from contextlib import contextmanager
from typing import Tuple
import torch
import torch.nn as nn
import numpy as np
class MemoryController:
"""
Base class for memory management during training.
"""
_last_input_size = None
_last_mem_ratio = []
@contextmanager
def record(self):
pass
def update_run_states(self, input_size=None, mem_ratio=None):
if self._last_input_size is None:
self._last_input_size = input_size
elif self._last_input_size!= input_size:
raise ValueError(f'Input size should not change for different ElasticModules.')
self._last_mem_ratio.append(mem_ratio)
@abstractmethod
def get_mem_ratio(self, input_size):
pass
@abstractmethod
def state_dict(self):
pass
@abstractmethod
def log(self):
pass
class LinearMemoryController(MemoryController):
"""
A simple controller for memory management during training.
The memory usage is modeled as a linear function of:
- the number of input parameters
- the ratio of memory the model use compared to the maximum usage (with no checkpointing)
memory_usage = k * input_size * mem_ratio + b
The controller keeps track of the memory usage and gives the
expected memory ratio to keep the memory usage under a target
"""
def __init__(
self,
buffer_size=1000,
update_every=500,
target_ratio=0.8,
available_memory=None,
max_mem_ratio_start=0.1,
params=None,
device=None
):
self.buffer_size = buffer_size
self.update_every = update_every
self.target_ratio = target_ratio
self.device = device or torch.cuda.current_device()
self.available_memory = available_memory or torch.cuda.get_device_properties(self.device).total_memory / 1024**3
self._memory = np.zeros(buffer_size, dtype=np.float32)
self._input_size = np.zeros(buffer_size, dtype=np.float32)
self._mem_ratio = np.zeros(buffer_size, dtype=np.float32)
self._buffer_ptr = 0
self._buffer_length = 0
self._params = tuple(params) if params is not None else (0.0, 0.0)
self._max_mem_ratio = max_mem_ratio_start
self.step = 0
def __repr__(self):
return f'LinearMemoryController(target_ratio={self.target_ratio}, available_memory={self.available_memory})'
def _add_sample(self, memory, input_size, mem_ratio):
self._memory[self._buffer_ptr] = memory
self._input_size[self._buffer_ptr] = input_size
self._mem_ratio[self._buffer_ptr] = mem_ratio
self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size
self._buffer_length = min(self._buffer_length + 1, self.buffer_size)
@contextmanager
def record(self):
torch.cuda.reset_peak_memory_stats(self.device)
self._last_input_size = None
self._last_mem_ratio = []
yield
self._last_memory = torch.cuda.max_memory_allocated(self.device) / 1024**3
self._last_mem_ratio = sum(self._last_mem_ratio) / len(self._last_mem_ratio)
self._add_sample(self._last_memory, self._last_input_size, self._last_mem_ratio)
self.step += 1
if self.step % self.update_every == 0:
self._max_mem_ratio = min(1.0, self._max_mem_ratio + 0.1)
self._fit_params()
def _fit_params(self):
memory_usage = self._memory[:self._buffer_length]
input_size = self._input_size[:self._buffer_length]
mem_ratio = self._mem_ratio[:self._buffer_length]
x = input_size * mem_ratio
y = memory_usage
k, b = np.polyfit(x, y, 1)
self._params = (k, b)
# self._visualize()
def _visualize(self):
import matplotlib.pyplot as plt
memory_usage = self._memory[:self._buffer_length]
input_size = self._input_size[:self._buffer_length]
mem_ratio = self._mem_ratio[:self._buffer_length]
k, b = self._params
plt.scatter(input_size * mem_ratio, memory_usage, c=mem_ratio, cmap='viridis')
x = np.array([0.0, 20000.0])
plt.plot(x, k * x + b, c='r')
plt.savefig(f'linear_memory_controller_{self.step}.png')
plt.cla()
def get_mem_ratio(self, input_size):
k, b = self._params
if k == 0: return np.random.rand() * self._max_mem_ratio
pred = (self.available_memory * self.target_ratio - b) / (k * input_size)
return min(self._max_mem_ratio, max(0.0, pred))
def state_dict(self):
return {
'params': self._params,
}
def load_state_dict(self, state_dict):
self._params = tuple(state_dict['params'])
def log(self):
return {
'params/k': self._params[0],
'params/b': self._params[1],
'memory': self._last_memory,
'input_size': self._last_input_size,
'mem_ratio': self._last_mem_ratio,
}
class ElasticModule(nn.Module):
"""
Module for training with elastic memory management.
"""
def __init__(self):
super().__init__()
self._memory_controller: MemoryController = None
@abstractmethod
def _get_input_size(self, *args, **kwargs) -> int:
"""
Get the size of the input data.
Returns:
int: The size of the input data.
"""
pass
@abstractmethod
def _forward_with_mem_ratio(self, *args, mem_ratio=0.0, **kwargs) -> Tuple[float, Tuple]:
"""
Forward with a given memory ratio.
"""
pass
def register_memory_controller(self, memory_controller: MemoryController):
self._memory_controller = memory_controller
def forward(self, *args, **kwargs):
if self._memory_controller is None or not torch.is_grad_enabled() or not self.training:
_, ret = self._forward_with_mem_ratio(*args, **kwargs)
else:
input_size = self._get_input_size(*args, **kwargs)
mem_ratio = self._memory_controller.get_mem_ratio(input_size)
mem_ratio, ret = self._forward_with_mem_ratio(*args, mem_ratio=mem_ratio, **kwargs)
self._memory_controller.update_run_states(input_size, mem_ratio)
return ret
class ElasticModuleMixin:
"""
Mixin for training with elastic memory management.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._memory_controller: MemoryController = None
@abstractmethod
def _get_input_size(self, *args, **kwargs) -> int:
"""
Get the size of the input data.
Returns:
int: The size of the input data.
"""
pass
@abstractmethod
@contextmanager
def with_mem_ratio(self, mem_ratio=1.0) -> float:
"""
Context manager for training with a reduced memory ratio compared to the full memory usage.
Returns:
float: The exact memory ratio used during the forward pass.
"""
pass
def register_memory_controller(self, memory_controller: MemoryController):
self._memory_controller = memory_controller
def forward(self, *args, **kwargs):
if self._memory_controller is None or not torch.is_grad_enabled() or not self.training:
ret = super().forward(*args, **kwargs)
else:
input_size = self._get_input_size(*args, **kwargs)
mem_ratio = self._memory_controller.get_mem_ratio(input_size)
with self.with_mem_ratio(mem_ratio) as exact_mem_ratio:
ret = super().forward(*args, **kwargs)
self._memory_controller.update_run_states(input_size, exact_mem_ratio)
return ret
trellis/utils/grad_clip_utils.py
+81
View File
@@ -0,0 +1,81 @@
from typing import *
import torch
import numpy as np
import torch.utils
class AdaptiveGradClipper:
"""
Adaptive gradient clipping for training.
"""
def __init__(
self,
max_norm=None,
clip_percentile=95.0,
buffer_size=1000,
):
self.max_norm = max_norm
self.clip_percentile = clip_percentile
self.buffer_size = buffer_size
self._grad_norm = np.zeros(buffer_size, dtype=np.float32)
self._max_norm = max_norm
self._buffer_ptr = 0
self._buffer_length = 0
def __repr__(self):
return f'AdaptiveGradClipper(max_norm={self.max_norm}, clip_percentile={self.clip_percentile})'
def state_dict(self):
return {
'grad_norm': self._grad_norm,
'max_norm': self._max_norm,
'buffer_ptr': self._buffer_ptr,
'buffer_length': self._buffer_length,
}
def load_state_dict(self, state_dict):
self._grad_norm = state_dict['grad_norm']
self._max_norm = state_dict['max_norm']
self._buffer_ptr = state_dict['buffer_ptr']
self._buffer_length = state_dict['buffer_length']
def log(self):
return {
'max_norm': self._max_norm,
}
def __call__(self, parameters, norm_type=2.0, error_if_nonfinite=False, foreach=None):
"""Clip the gradient norm of an iterable of parameters.
The norm is computed over all gradients together, as if they were
concatenated into a single vector. Gradients are modified in-place.
Args:
parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
single Tensor that will have gradients normalized
norm_type (float): type of the used p-norm. Can be ``'inf'`` for
infinity norm.
error_if_nonfinite (bool): if True, an error is thrown if the total
norm of the gradients from :attr:`parameters` is ``nan``,
``inf``, or ``-inf``. Default: False (will switch to True in the future)
foreach (bool): use the faster foreach-based implementation.
If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently
fall back to the slow implementation for other device types.
Default: ``None``
Returns:
Total norm of the parameter gradients (viewed as a single vector).
"""
max_norm = self._max_norm if self._max_norm is not None else float('inf')
grad_norm = torch.nn.utils.clip_grad_norm_(parameters, max_norm=max_norm, norm_type=norm_type, error_if_nonfinite=error_if_nonfinite, foreach=foreach)
if torch.isfinite(grad_norm):
self._grad_norm[self._buffer_ptr] = grad_norm
self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size
self._buffer_length = min(self._buffer_length + 1, self.buffer_size)
if self._buffer_length == self.buffer_size:
self._max_norm = np.percentile(self._grad_norm, self.clip_percentile)
self._max_norm = min(self._max_norm, self.max_norm) if self.max_norm is not None else self._max_norm
return grad_norm
trellis/utils/loss_utils.py
+92
View File
@@ -0,0 +1,92 @@
import torch
import torch.nn.functional as F
from torch.autograd import Variable
from math import exp
from lpips import LPIPS
def smooth_l1_loss(pred, target, beta=1.0):
diff = torch.abs(pred - target)
loss = torch.where(diff < beta, 0.5 * diff ** 2 / beta, diff - 0.5 * beta)
return loss.mean()
def l1_loss(network_output, gt):
return torch.abs((network_output - gt)).mean()
def l2_loss(network_output, gt):
return ((network_output - gt) ** 2).mean()
def gaussian(window_size, sigma):
gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
return gauss / gauss.sum()
def create_window(window_size, channel):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
return window
def psnr(img1, img2, max_val=1.0):
mse = F.mse_loss(img1, img2)
return 20 * torch.log10(max_val / torch.sqrt(mse))
def ssim(img1, img2, window_size=11, size_average=True):
channel = img1.size(-3)
window = create_window(window_size, channel)
if img1.is_cuda:
window = window.cuda(img1.get_device())
window = window.type_as(img1)
return _ssim(img1, img2, window, window_size, channel, size_average)
def _ssim(img1, img2, window, window_size, channel, size_average=True):
mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
C1 = 0.01 ** 2
C2 = 0.03 ** 2
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
if size_average:
return ssim_map.mean()
else:
return ssim_map.mean(1).mean(1).mean(1)
loss_fn_vgg = None
def lpips(img1, img2, value_range=(0, 1)):
global loss_fn_vgg
if loss_fn_vgg is None:
loss_fn_vgg = LPIPS(net='vgg').cuda().eval()
# normalize to [-1, 1]
img1 = (img1 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1
img2 = (img2 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1
return loss_fn_vgg(img1, img2).mean()
def normal_angle(pred, gt):
pred = pred * 2.0 - 1.0
gt = gt * 2.0 - 1.0
norms = pred.norm(dim=-1) * gt.norm(dim=-1)
cos_sim = (pred * gt).sum(-1) / (norms + 1e-9)
cos_sim = torch.clamp(cos_sim, -1.0, 1.0)
ang = torch.rad2deg(torch.acos(cos_sim[norms > 1e-9])).mean()
if ang.isnan():
return -1
return ang
trellis/utils/render_utils.py
+30 -30
View File
@@ -4,10 +4,7 @@ from tqdm import tqdm
import utils3d
from PIL import Image
from ..renderers import OctreeRenderer, MeshRenderer,GaussianRenderer
from ..renderers.gaussian_render import GaussianRenderer
# from ..renderers.mesh_renderer import MeshRenderer
# from ..renderers.octree_renderer import OctreeRenderer
from ..renderers import OctreeRenderer, GaussianRenderer, MeshRenderer
from ..representations import Octree, Gaussian, MeshExtractResult
from ..modules import sparse as sp
from .random_utils import sphere_hammersley_sequence
@@ -43,36 +40,44 @@ def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs):
return extrinsics, intrinsics
def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs):
def get_renderer(sample, **kwargs):
if isinstance(sample, Octree):
renderer = OctreeRenderer()
renderer.rendering_options.resolution = options.get('resolution', 512)
renderer.rendering_options.near = options.get('near', 0.8)
renderer.rendering_options.far = options.get('far', 1.6)
renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0))
renderer.rendering_options.ssaa = options.get('ssaa', 4)
renderer.rendering_options.resolution = kwargs.get('resolution', 512)
renderer.rendering_options.near = kwargs.get('near', 0.8)
renderer.rendering_options.far = kwargs.get('far', 1.6)
renderer.rendering_options.bg_color = kwargs.get('bg_color', (0, 0, 0))
renderer.rendering_options.ssaa = kwargs.get('ssaa', 4)
renderer.pipe.primitive = sample.primitive
elif isinstance(sample, Gaussian):
renderer = GaussianRenderer()
renderer.rendering_options.resolution = options.get('resolution', 512)
renderer.rendering_options.near = options.get('near', 0.8)
renderer.rendering_options.far = options.get('far', 1.6)
renderer.rendering_options.bg_color = options.get('bg_color', (0, 0, 0))
renderer.rendering_options.ssaa = options.get('ssaa', 1)
renderer.rendering_options.resolution = kwargs.get('resolution', 512)
renderer.rendering_options.near = kwargs.get('near', 0.8)
renderer.rendering_options.far = kwargs.get('far', 1.6)
renderer.rendering_options.bg_color = kwargs.get('bg_color', (0, 0, 0))
renderer.rendering_options.ssaa = kwargs.get('ssaa', 1)
renderer.pipe.kernel_size = kwargs.get('kernel_size', 0.1)
renderer.pipe.use_mip_gaussian = True
elif isinstance(sample, MeshExtractResult):
renderer = MeshRenderer()
renderer.rendering_options.resolution = options.get('resolution', 512)
renderer.rendering_options.near = options.get('near', 1)
renderer.rendering_options.far = options.get('far', 100)
renderer.rendering_options.ssaa = options.get('ssaa', 4)
renderer.rendering_options.resolution = kwargs.get('resolution', 512)
renderer.rendering_options.near = kwargs.get('near', 1)
renderer.rendering_options.far = kwargs.get('far', 100)
renderer.rendering_options.ssaa = kwargs.get('ssaa', 4)
else:
raise ValueError(f'Unsupported sample type: {type(sample)}')
return renderer
def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs):
renderer = get_renderer(sample, **options)
rets = {}
for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose):
if not isinstance(sample, MeshExtractResult):
if isinstance(sample, MeshExtractResult):
res = renderer.render(sample, extr, intr)
if 'normal' not in rets: rets['normal'] = []
rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
else:
res = renderer.render(sample, extr, intr, colors_overwrite=colors_overwrite)
if 'color' not in rets: rets['color'] = []
if 'depth' not in rets: rets['depth'] = []
@@ -83,10 +88,6 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N
rets['depth'].append(res['depth'].detach().cpu().numpy())
else:
rets['depth'].append(None)
else:
res = renderer.render(sample, extr, intr)
if 'normal' not in rets: rets['normal'] = []
rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
return rets
@@ -96,18 +97,17 @@ def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2
yaws = yaws.tolist()
pitch = pitch.tolist()
extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov)
return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
def render_multiview(sample, resolution=512, nviews=30,**kwargs):
def render_multiview(sample, resolution=512, nviews=30):
r = 2
fov = 40
bg_color=(0, 0, 0)
cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)]
yaws = [cam[0] for cam in cams]
pitchs = [cam[1] for cam in cams]
extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov) #计算相机内外参
res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov)
res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': (0, 0, 0)})
return res['color'], extrinsics, intrinsics
utils.py
+32
View File
@@ -11,7 +11,39 @@ import trimesh
from comfy.utils import common_upscale,ProgressBar
import folder_paths
current_path = os.path.dirname(os.path.abspath(__file__))
def convert_clip2diffuser(clip):
from transformers import (
CLIPTextModel,
CLIPTextConfig,
)
from contextlib import nullcontext
try:
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
is_accelerate_available = True
except:
pass
import comfy.model_management as mm
offload_device=mm.unet_offload_device()
text_encoder_config=CLIPTextConfig.from_pretrained(os.path.join(current_path, 'text_encoder'), local_files_only=True)
ctx = init_empty_weights if is_accelerate_available else nullcontext
with ctx():
text_encoder = CLIPTextModel(text_encoder_config)
text_encoder_sd = clip.get_sd()
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_open_clip_checkpoint,convert_ldm_clip_checkpoint
text_encoder = convert_ldm_clip_checkpoint(text_encoder_sd, text_encoder=text_encoder)
return text_encoder,text_encoder_2
#https://github.com/Abecid/realtime-3d-stylization-drawing/blob/97a687f148cec04e107f7fec8986c8b5615af1ec/gradio_app.py#L33
def glb2obj_(glb_path, obj_path):
print('Converting glb to obj')