meshgpt-pytorch - 基于注意力机制的先进3D网格生成框架

MeshGPT - Pytorch

在Pytorch中实现MeshGPT，使用注意力机制的最先进网格生成技术

还将添加文本条件，以实现最终的文本到3D资产转换

如果您有兴趣与他人合作复现这项工作，请加入

更新：Marcus已经训练并上传了一个可用模型到🤗 Huggingface！

致谢

感谢StabilityAI、A16Z开源AI资助计划和🤗 Huggingface的慷慨赞助，以及我的其他赞助者，让我能够独立开源当前的人工智能研究
感谢Einops让我的工作变得轻松
感谢Marcus进行初步代码审查（指出一些缺失的派生特征）以及运行首次成功的端到端实验
感谢Marcus 首次成功训练了一组基于标签条件的形状集合
感谢Quexi Ma发现了自动eos处理的多个错误
感谢Yingtian发现了空间标签平滑的位置高斯模糊中的一个错误
再次感谢Marcus进行实验验证，证明可以将系统从三角形扩展到四边形
感谢Marcus识别出文本条件的一个问题，并进行了所有导致问题解决的实验

安装

$ pip install meshgpt-pytorch

使用方法

import torch

from meshgpt_pytorch import (
    MeshAutoencoder,
    MeshTransformer
)

# 自动编码器

autoencoder = MeshAutoencoder(
    num_discrete_coors = 128
)

# 模拟输入

vertices = torch.randn((2, 121, 3))            # (批次, 顶点数, 坐标 (3))
faces = torch.randint(0, 121, (2, 64, 3))      # (批次, 面数, 顶点 (3))

# 确保对于可变长度的网格，faces用"-1"填充

# 前向传播faces

loss = autoencoder(
    vertices = vertices,
    faces = faces
)

loss.backward()

# 经过大量训练后...
# 您可以将上面的原始face数据传入transformer，以建模这个面顶点序列

transformer = MeshTransformer(
    autoencoder,
    dim = 512,
    max_seq_len = 768
)

loss = transformer(
    vertices = vertices,
    faces = faces
)

loss.backward()

# 在transformer经过大量训练后，您现在可以采样生成新的3D资产

faces_coordinates, face_mask = transformer.generate()

# (批次, 面数, 顶点 (3), 坐标 (3)), (批次, 面数)
# 现在对生成的3D资产进行后处理

对于基于文本条件的3D形状合成，只需在您的MeshTransformer上设置condition_on_text = True，然后将您的描述列表作为texts关键字参数传入

例如：

transformer = MeshTransformer(
    autoencoder,
    dim = 512,
    max_seq_len = 768,
    condition_on_text = True
)


loss = transformer(
    vertices = vertices,
    faces = faces,
    texts = ['一把高脚椅', '一个小茶壶'],
)

loss.backward()

# 在transformer经过大量训练后，您现在可以基于文本条件采样生成新的3D资产

faces_coordinates, face_mask = transformer.generate(
    texts = ['一张长桌'],
    cond_scale = 3.  # cond_scale > 1. 将启用无分类器引导 - 可以设置在3. - 10.之间的任何值
)

如果您想对网格进行标记化，以便在多模态transformer中使用，只需在自动编码器上调用.tokenize方法（或在自动编码器训练器实例上使用相同的方法来获取指数平滑模型）


mesh_token_ids = autoencoder.tokenize(
    vertices = vertices,
    faces = faces
)

# (批次, 面顶点数, 残差量化层)

类型检查

在项目根目录下运行

$ cp .env.sample .env

待办事项

引用

[引用内容保持不变]