base on Implementation of MeshGPT, SOTA Mesh generation using Attention, in Pytorch <img src="./meshgpt.png" width="450px"></img> ## MeshGPT - Pytorch Implementation of <a href="https://arxiv.org/abs/2311.15475">MeshGPT</a>, SOTA Mesh generation using Attention, in Pytorch Will also add text conditioning, for eventual text-to-3d asset Please join <a href="https://discord.gg/xBPBXfcFHd"><img alt="Join us on Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a> if you are interested in collaborating with others to replicate this work Update: <a href="https://github.com/MarcusLoppe">Marcus</a> has trained and uploaded <a href="https://huggingface.co/MarcusLoren/MeshGPT-preview">a working model</a> to 🤗 Huggingface! ## Appreciation - <a href="https://stability.ai/">StabilityAI</a>, <a href="https://a16z.com/supporting-the-open-source-ai-community/">A16Z Open Source AI Grant Program</a>, and <a href="https://huggingface.co/">🤗 Huggingface</a> for the generous sponsorships, as well as my other sponsors, for affording me the independence to open source current artificial intelligence research - <a href="https://github.com/arogozhnikov/einops">Einops</a> for making my life easy - <a href="https://github.com/MarcusLoppe">Marcus</a> for the initial code review (pointing out some missing derived features) as well as running the first successful end-to-end experiments - <a href="https://github.com/MarcusLoppe">Marcus</a> for the <a href="https://github.com/lucidrains/meshgpt-pytorch/issues/18#issuecomment-1859214710">first successful training</a> of a collection of shapes conditioned on labels - <a href="https://github.com/qixuema">Quexi Ma</a> for finding numerous bugs with automatic eos handling - <a href="https://github.com/thuliu-yt16">Yingtian</a> for finding a bug with the gaussian blurring of the positions for spatial label smoothing - <a href="https://github.com/MarcusLoppe">Marcus</a> yet again for running the experiments to validate that it is possible to extend the system from triangles to <a href="https://github.com/lucidrains/meshgpt-pytorch/issues/54#issuecomment-1906789076">quads</a> - <a href="https://github.com/MarcusLoppe">Marcus</a> for identifying <a href="https://github.com/lucidrains/meshgpt-pytorch/issues/80">an issue</a> with text conditioning and for running all the experiments that led to it being resolved ## Install ```bash $ pip install meshgpt-pytorch ``` ## Usage ```python import torch from meshgpt_pytorch import ( MeshAutoencoder, MeshTransformer ) # autoencoder autoencoder = MeshAutoencoder( num_discrete_coors = 128 ) # mock inputs vertices = torch.randn((2, 121, 3)) # (batch, num vertices, coor (3)) faces = torch.randint(0, 121, (2, 64, 3)) # (batch, num faces, vertices (3)) # make sure faces are padded with `-1` for variable lengthed meshes # forward in the faces loss = autoencoder( vertices = vertices, faces = faces ) loss.backward() # after much training... # you can pass in the raw face data above to train a transformer to model this sequence of face vertices transformer = MeshTransformer( autoencoder, dim = 512, max_seq_len = 768 ) loss = transformer( vertices = vertices, faces = faces ) loss.backward() # after much training of transformer, you can now sample novel 3d assets faces_coordinates, face_mask = transformer.generate() # (batch, num faces, vertices (3), coordinates (3)), (batch, num faces) # now post process for the generated 3d asset ``` For <a href="https://www.youtube.com/watch?v=NXX0dKw4SjI">text-conditioned 3d shape synthesis</a>, simply set `condition_on_text = True` on your `MeshTransformer`, and then pass in your list of descriptions as the `texts` keyword argument ex. ```python transformer = MeshTransformer( autoencoder, dim = 512, max_seq_len = 768, condition_on_text = True ) loss = transformer( vertices = vertices, faces = faces, texts = ['a high chair', 'a small teapot'], ) loss.backward() # after much training of transformer, you can now sample novel 3d assets conditioned on text faces_coordinates, face_mask = transformer.generate( texts = ['a long table'], cond_scale = 8., # a cond_scale > 1. will enable classifier free guidance - can be placed anywhere from 3. - 10. remove_parallel_component = True # from https://arxiv.org/abs/2410.02416 ) ``` If you want to tokenize meshes, for use in your multimodal transformer, simply invoke `.tokenize` on your autoencoder (or same method on autoencoder trainer instance for the exponentially smoothed model) ```python mesh_token_ids = autoencoder.tokenize( vertices = vertices, faces = faces ) # (batch, num face vertices, residual quantized layer) ``` ## Typecheck At the project root, run ```bash $ cp .env.sample .env ``` ## Todo - [x] autoencoder - [x] encoder sageconv with torch geometric - [x] proper scatter mean accounting for padding for meaning the vertices and RVQ the vertices before gathering back for decoder - [x] complete decoder and reconstruction loss + commitment loss - [x] handle variable lengthed faces - [x] add option to use residual LFQ, latest quantization development that scales code utilization - [x] xcit linear attention in encoder and decoder - [x] figure out how to auto-derive `face_edges` directly from faces and vertices - [x] embed any derived values (area, angles, etc) from the vertices before sage convs - [ ] add an extra graph conv stage in the encoder, where vertices are enriched with their connected vertex neighbors, before aggregating into faces. make optional - [ ] allow for encoder to noise the vertices, so autoencoder is a bit denoising. consider conditioning decoder on noise level, if varying - [ ] transformer - [x] properly mask out eos logit during generation - [x] make sure it trains - [x] take care of sos token automatically - [x] take care of eos token automatically if sequence length or mask is passed in - [x] handle variable lengthed faces - [x] on forwards - [x] on generation, do all eos logic + substitute everything after eos with pad id - [x] generation + cache kv - [x] trainer wrapper with hf accelerate - [x] autoencoder - take care of ema - [x] transformer - [x] text conditioning using own CFG library - [x] complete preliminary text conditioning - [x] make sure CFG library can support passing in arguments to the two separate calls when cond scaling (as well as aggregating their outputs) - [ ] polish up the magic dataset decorator and see if it can be moved to CFG library - [x] hierarchical transformers (using the RQ transformer) - [x] fix caching in simple gateloop layer in other repo - [x] local attention - [x] fix kv caching for two-staged hierarchical transformer - 7x faster now, and faster than original non-hierarchical transformer - [x] fix caching for gateloop layers - [x] allow for customization of model dimensions of fine vs coarse attention network - [x] figure out if autoencoder is really necessary - it is necessary, ablations are in the paper - [x] when mesh discretizer is passed in, one can inject inter-face attention with the relative distance - [x] additional embeddings (angles, area, normal), can also be appended before coarse transformer attention - [ ] make transformer efficient - [ ] reversible networks - [ ] speculative decoding option - [ ] spend a day on documentation ## Citations ```bibtex @inproceedings{Siddiqui2023MeshGPTGT, title = {MeshGPT: Generating Triangle Meshes with Decoder-Only Transformers}, author = {Yawar Siddiqui and Antonio Alliegro and Alexey Artemov and Tatiana Tommasi and Daniele Sirigatti and Vladislav Rosov and Angela Dai and Matthias Nie{\ss}ner}, year = {2023}, url = {https://api.semanticscholar.org/CorpusID:265457242} } ``` ```bibtex @inproceedings{dao2022flashattention, title = {Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness}, author = {Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher}, booktitle = {Advances in Neural Information Processing Systems}, year = {2022} } ``` ```bibtex @inproceedings{Leviathan2022FastIF, title = {Fast Inference from Transformers via Speculative Decoding}, author = {Yaniv Leviathan and Matan Kalman and Y. Matias}, booktitle = {International Conference on Machine Learning}, year = {2022}, url = {https://api.semanticscholar.org/CorpusID:254096365} } ``` ```bibtex @misc{yu2023language, title = {Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation}, author = {Lijun Yu and José Lezama and Nitesh B. Gundavarapu and Luca Versari and Kihyuk Sohn and David Minnen and Yong Cheng and Agrim Gupta and Xiuye Gu and Alexander G. Hauptmann and Boqing Gong and Ming-Hsuan Yang and Irfan Essa and David A. Ross and Lu Jiang}, year = {2023}, eprint = {2310.05737}, archivePrefix = {arXiv}, primaryClass = {cs.CV} } ``` ```bibtex @article{Lee2022AutoregressiveIG, title = {Autoregressive Image Generation using Residual Quantization}, author = {Doyup Lee and Chiheon Kim and Saehoon Kim and Minsu Cho and Wook-Shin Han}, journal = {2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, year = {2022}, pages = {11513-11522}, url = {https://api.semanticscholar.org/CorpusID:247244535} } ``` ```bibtex @inproceedings{Katsch2023GateLoopFD, title = {GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling}, author = {Tobias Katsch}, year = {2023}, url = {https://api.semanticscholar.org/CorpusID:265018962} } ``` ", Assign "at most 3 tags" to the expected json: {"id":"5634","tags":[]} "only from the tags list I provide: [{"id":77,"name":"3d"},{"id":89,"name":"agent"},{"id":17,"name":"ai"},{"id":54,"name":"algorithm"},{"id":24,"name":"api"},{"id":44,"name":"authentication"},{"id":3,"name":"aws"},{"id":27,"name":"backend"},{"id":60,"name":"benchmark"},{"id":72,"name":"best-practices"},{"id":39,"name":"bitcoin"},{"id":37,"name":"blockchain"},{"id":1,"name":"blog"},{"id":45,"name":"bundler"},{"id":58,"name":"cache"},{"id":21,"name":"chat"},{"id":49,"name":"cicd"},{"id":4,"name":"cli"},{"id":64,"name":"cloud-native"},{"id":48,"name":"cms"},{"id":61,"name":"compiler"},{"id":68,"name":"containerization"},{"id":92,"name":"crm"},{"id":34,"name":"data"},{"id":47,"name":"database"},{"id":8,"name":"declarative-gui "},{"id":9,"name":"deploy-tool"},{"id":53,"name":"desktop-app"},{"id":6,"name":"dev-exp-lib"},{"id":59,"name":"dev-tool"},{"id":13,"name":"ecommerce"},{"id":26,"name":"editor"},{"id":66,"name":"emulator"},{"id":62,"name":"filesystem"},{"id":80,"name":"finance"},{"id":15,"name":"firmware"},{"id":73,"name":"for-fun"},{"id":2,"name":"framework"},{"id":11,"name":"frontend"},{"id":22,"name":"game"},{"id":81,"name":"game-engine "},{"id":23,"name":"graphql"},{"id":84,"name":"gui"},{"id":91,"name":"http"},{"id":5,"name":"http-client"},{"id":51,"name":"iac"},{"id":30,"name":"ide"},{"id":78,"name":"iot"},{"id":40,"name":"json"},{"id":83,"name":"julian"},{"id":38,"name":"k8s"},{"id":31,"name":"language"},{"id":10,"name":"learning-resource"},{"id":33,"name":"lib"},{"id":41,"name":"linter"},{"id":28,"name":"lms"},{"id":16,"name":"logging"},{"id":76,"name":"low-code"},{"id":90,"name":"message-queue"},{"id":42,"name":"mobile-app"},{"id":18,"name":"monitoring"},{"id":36,"name":"networking"},{"id":7,"name":"node-version"},{"id":55,"name":"nosql"},{"id":57,"name":"observability"},{"id":46,"name":"orm"},{"id":52,"name":"os"},{"id":14,"name":"parser"},{"id":74,"name":"react"},{"id":82,"name":"real-time"},{"id":56,"name":"robot"},{"id":65,"name":"runtime"},{"id":32,"name":"sdk"},{"id":71,"name":"search"},{"id":63,"name":"secrets"},{"id":25,"name":"security"},{"id":85,"name":"server"},{"id":86,"name":"serverless"},{"id":70,"name":"storage"},{"id":75,"name":"system-design"},{"id":79,"name":"terminal"},{"id":29,"name":"testing"},{"id":12,"name":"ui"},{"id":50,"name":"ux"},{"id":88,"name":"video"},{"id":20,"name":"web-app"},{"id":35,"name":"web-server"},{"id":43,"name":"webassembly"},{"id":69,"name":"workflow"},{"id":87,"name":"yaml"}]" returns me the "expected json"