{"id":13595172,"url":"https://github.com/lucidrains/soundstorm-pytorch","last_synced_at":"2025-05-14T03:10:45.962Z","repository":{"id":166254272,"uuid":"641729782","full_name":"lucidrains/soundstorm-pytorch","owner":"lucidrains","description":"Implementation of SoundStorm, Efficient Parallel Audio Generation from Google Deepmind, in Pytorch","archived":false,"fork":false,"pushed_at":"2025-04-24T13:40:30.000Z","size":270,"stargazers_count":1492,"open_issues_count":7,"forks_count":92,"subscribers_count":50,"default_branch":"main","last_synced_at":"2025-04-24T14:03:37.902Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","audio-generation","deep-learning","non-autoregressive","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-05-17T03:46:22.000Z","updated_at":"2025-04-24T13:40:30.000Z","dependencies_parsed_at":"2024-09-25T00:40:34.727Z","dependency_job_id":"3a5e0e0a-7658-474d-b417-1cb84db3078f","html_url":"https://github.com/lucidrains/soundstorm-pytorch","commit_stats":{"total_commits":85,"total_committers":6,"mean_commits":"14.166666666666666","dds":"0.11764705882352944","last_synced_commit":"602b61694d2460f1b162ef442c6fe856a525fdb5"},"previous_names":[],"tags_count":47,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fsoundstorm-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fsoundstorm-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fsoundstorm-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fsoundstorm-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/soundstorm-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254059512,"owners_count":22007769,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanism","audio-generation","deep-learning","non-autoregressive","transformers"],"created_at":"2024-08-01T16:01:45.278Z","updated_at":"2025-05-14T03:10:40.944Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python","HarmonyOS"],"sub_categories":["Windows Manager"],"readme":"\u003cimg src=\"./soundstorm.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Soundstorm - Pytorch\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2305.09636\"\u003eSoundStorm\u003c/a\u003e, Efficient Parallel Audio Generation from Google Deepmind, in Pytorch.\n\nThey basically applied \u003ca href=\"https://arxiv.org/abs/2202.04200\"\u003eMaskGiT\u003c/a\u003e to the residual vector quantized codes from \u003ca href=\"https://github.com/lucidrains/audiolm-pytorch#soundstream--encodec\"\u003eSoundstream\u003c/a\u003e. The transformer architecture they chose to use is one that fits well with the audio domain, named \u003ca href=\"https://arxiv.org/abs/2005.08100\"\u003eConformer\u003c/a\u003e\n\n\u003ca href=\"https://google-research.github.io/seanet/soundstorm/examples/\"\u003eProject Page\u003c/a\u003e\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStability\u003c/a\u003e and \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for their generous sponsorships to work on and open source cutting edge artificial intelligence research\n\n- \u003ca href=\"https://github.com/lucasnewman\"\u003eLucas Newman\u003c/a\u003e for numerous contributions, including the initial training code, acoustic prompting logic, per-level quantizer decoding!\n\n- \u003ca href=\"https://huggingface.co/docs/accelerate/index\"\u003e🤗 Accelerate\u003c/a\u003e for providing a simple and powerful solution for training\n\n- \u003ca href=\"https://einops.rocks/\"\u003eEinops\u003c/a\u003e for the indispensable abstraction that makes building neural networks fun, easy, and uplifting\n\n- \u003ca href=\"https://github.com/stevenhillis\"\u003eSteven Hillis\u003c/a\u003e for submitting the correct masking strategy and for verifying that the repository works! 🙏\n\n- \u003ca href=\"https://github.com/lucasnewman\"\u003eLucas Newman\u003c/a\u003e for basically training a small working Soundstorm with models across multiple repositories, showing it all works end-to-end. Models include \u003ca href=\"https://github.com/lucidrains/audiolm-pytorch\"\u003eSoundStream\u003c/a\u003e, \u003ca href=\"https://github.com/lucidrains/spear-tts-pytorch\"\u003eText-to-Semantic T5\u003c/a\u003e, and finally the SoundStorm transformer here.\n\n- \u003ca href=\"https://github.com/Jiang-Stan\"\u003e@Jiang-Stan\u003c/a\u003e for identifying a critical bug in the iterative demasking!\n\n## Install\n\n```bash\n$ pip install soundstorm-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom soundstorm_pytorch import SoundStorm, ConformerWrapper\n\nconformer = ConformerWrapper(\n    codebook_size = 1024,\n    num_quantizers = 12,\n    conformer = dict(\n        dim = 512,\n        depth = 2\n    ),\n)\n\nmodel = SoundStorm(\n    conformer,\n    steps = 18,          # 18 steps, as in original maskgit paper\n    schedule = 'cosine'  # currently the best schedule is cosine\n)\n\n# get your pre-encoded codebook ids from the soundstream from a lot of raw audio\n\ncodes = torch.randint(0, 1024, (2, 1024, 12)) # (batch, seq, num residual VQ)\n\n# do the below in a loop for a ton of data\n\nloss, _ = model(codes)\nloss.backward()\n\n# model can now generate in 18 steps. ~2 seconds sounds reasonable\n\ngenerated = model.generate(1024, batch_size = 2) # (2, 1024)\n```\n\nTo directly train on raw audio, you need to pass in your pretrained `SoundStream` into `SoundStorm`. You can train your own `SoundStream` at \u003ca href=\"https://github.com/lucidrains/audiolm-pytorch#soundstream--encodec\"\u003eaudiolm-pytorch\u003c/a\u003e.\n\n```python\nimport torch\nfrom soundstorm_pytorch import SoundStorm, ConformerWrapper, Conformer, SoundStream\n\nconformer = ConformerWrapper(\n    codebook_size = 1024,\n    num_quantizers = 12,\n    conformer = dict(\n        dim = 512,\n        depth = 2\n    ),\n)\n\nsoundstream = SoundStream(\n    codebook_size = 1024,\n    rq_num_quantizers = 12,\n    attn_window_size = 128,\n    attn_depth = 2\n)\n\nmodel = SoundStorm(\n    conformer,\n    soundstream = soundstream   # pass in the soundstream\n)\n\n# find as much audio you'd like the model to learn\n\naudio = torch.randn(2, 10080)\n\n# course it through the model and take a gazillion tiny steps\n\nloss, _ = model(audio)\nloss.backward()\n\n# and now you can generate state-of-the-art speech\n\ngenerated_audio = model.generate(seconds = 30, batch_size = 2)  # generate 30 seconds of audio (it will calculate the length in seconds based off the sampling frequency and cumulative downsamples in the soundstream passed in above)\n```\n\nComplete text-to-speech will rely on a trained `TextToSemantic` encoder / decoder transformer. You will then load the weights and pass it into the `SoundStorm` as `spear_tts_text_to_semantic`\n\nThis is a work-in-progress, as `spear-tts-pytorch` only has the model architecture complete, and not the pretraining + pseudo-labeling + backtranslation logic.\n\n```python\nfrom spear_tts_pytorch import TextToSemantic\n\ntext_to_semantic = TextToSemantic(\n    dim = 512,\n    source_depth = 12,\n    target_depth = 12,\n    num_text_token_ids = 50000,\n    num_semantic_token_ids = 20000,\n    use_openai_tokenizer = True\n)\n\n# load the trained text-to-semantic transformer\n\ntext_to_semantic.load('/path/to/trained/model.pt')\n\n# pass it into the soundstorm\n\nmodel = SoundStorm(\n    conformer,\n    soundstream = soundstream,\n    spear_tts_text_to_semantic = text_to_semantic\n).cuda()\n\n# and now you can generate state-of-the-art speech\n\ngenerated_speech = model.generate(\n    texts = [\n        'the rain in spain stays mainly in the plain',\n        'the quick brown fox jumps over the lazy dog'\n    ]\n) # (2, n) - raw waveform decoded from soundstream\n```\n\n## Todo\n\n- [x] integrate soundstream\n- [x] when generating, and length can be defined in seconds (takes into sampling freq etc)\n- [x] make sure grouped rvq is supported. concat embeddings rather than sum across group dimension\n- [x] just copy conformer over and redo shaw's relative positional embedding with rotary embedding. nobody uses shaw anymore.\n- [x] default flash attention to true\n- [x] remove batchnorm, and just use layernorm, but after the swish (as in normformer paper)\n- [x] trainer with accelerate - thanks to @lucasnewman\n- [x] allow for variable lengthed sequence training and generation, by passing in `mask` at `forward` and `generate`\n\n- [ ] option to return list of audio files when generating\n- [ ] turn it into a command line tool\n- [ ] add cross attention and adaptive layernorm conditioning\n\n## Citations\n\n```bibtex\n@misc{borsos2023soundstorm,\n    title   = {SoundStorm: Efficient Parallel Audio Generation}, \n    author  = {Zalán Borsos and Matt Sharifi and Damien Vincent and Eugene Kharitonov and Neil Zeghidour and Marco Tagliasacchi},\n    year    = {2023},\n    eprint  = {2305.09636},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.SD}\n}\n```\n\n```bibtex\n@inproceedings{dao2022flashattention,\n    title   = {Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},\n    author  = {Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\\'e}, Christopher},\n    booktitle = {Advances in Neural Information Processing Systems},\n    year    = {2022}\n}\n```\n\n```bibtex\n@article{Chang2022MaskGITMG,\n    title   = {MaskGIT: Masked Generative Image Transformer},\n    author  = {Huiwen Chang and Han Zhang and Lu Jiang and Ce Liu and William T. Freeman},\n    journal = {2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n    year    = {2022},\n    pages   = {11305-11315}\n}\n```\n\n```bibtex\n@article{Lezama2022ImprovedMI,\n    title   = {Improved Masked Image Generation with Token-Critic},\n    author  = {Jos{\\'e} Lezama and Huiwen Chang and Lu Jiang and Irfan Essa},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2209.04439}\n}\n```\n\n```bibtex\n@inproceedings{Nijkamp2021SCRIPTSP,\n    title   = {SCRIPT: Self-Critic PreTraining of Transformers},\n    author  = {Erik Nijkamp and Bo Pang and Ying Nian Wu and Caiming Xiong},\n    booktitle = {North American Chapter of the Association for Computational Linguistics},\n    year    = {2021}\n}\n```\n\n```bibtex\n@inproceedings{rogozhnikov2022einops,\n    title   = {Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation},\n    author  = {Alex Rogozhnikov},\n    booktitle = {International Conference on Learning Representations},\n    year    = {2022},\n    url     = {https://openreview.net/forum?id=oapKSVM2bcj}\n}\n```\n\n```bibtex\n@misc{su2021roformer,\n    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding},\n    author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},\n    year    = {2021},\n    eprint  = {2104.09864},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n```\n\n```bibtex\n@inproceedings{Zhou2024ValueRL,\n    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},\n    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},\n    year    = {2024},\n    url     = {https://api.semanticscholar.org/CorpusID:273532030}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fsoundstorm-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fsoundstorm-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fsoundstorm-pytorch/lists"}