transfusion-pytorch


Nametransfusion-pytorch JSON
Version 0.3.4 PyPI version JSON
download
home_pageNone
SummaryTransfusion in Pytorch
upload_time2024-11-05 16:23:48
maintainerNone
docs_urlNone
authorNone
requires_python>=3.8
licenseMIT License Copyright (c) 2024 Phil Wang Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
keywords artificial intelligence attention mechanism deep learning rectified flow transformers
VCS
bugtrack_url
requirements No requirements were recorded.
Travis-CI No Travis.
coveralls test coverage No coveralls.
            <img src="./transfusion.png" width="400px"></img>

## Transfusion - Pytorch

Pytorch implementation of [Transfusion](https://www.arxiv.org/abs/2408.11039), "Predict the Next Token and Diffuse Images with One Multi-Modal Model", from MetaAI.

In this repo, we will substitute diffusion with flow matching given the success of Flux from Black Forest Labs (but will keep the original paper title given Transflow does not have the same ring). This repository will also attempt to extend to any number of modalities.

## Install

```bash
$ pip install transfusion-pytorch
```

## Usage

One modality, say images

```python
from torch import randint, randn
from transfusion_pytorch import Transfusion

model = Transfusion(
    num_text_tokens = 256,
    dim_latent = 384,
    modality_default_shape = (4,),  # fallback, in the case the language model did not produce a valid modality shape
    transformer = dict(
        dim = 512,
        depth = 8
    )
)

# any torch.long is text, torch.float is modalities

text_and_images = [
    [randint(0, 256, (16,)), randn(4, 384), randint(0, 256, (8,)), randn(6, 384)],
    [randint(0, 256, (16,)), randn(7, 384), randint(0, 256, (5,)), randn(2, 384), randint(0, 256, (9,))]
]

loss = model(text_and_images)

loss.backward()

# after much training

one_multimodal_sample = model.sample()
```

Multiple different modalities

```python
from torch import randint, randn
from transfusion_pytorch import Transfusion

model = Transfusion(
    num_text_tokens = 256,
    dim_latent = (384, 192),                 # specify multiple latent dimensions
    modality_default_shape = ((4,), (2,)),   # default shapes for first and second modality
    transformer = dict(
        dim = 512,
        depth = 8
    )
)

# then for the Tensors of type float, you can pass a tuple[int, Tensor] and specify the modality index in the first position

# any torch.long is text, torch.float is modalities

text_images_and_audio = [
    [randint(0, 256, (16,)), (0, randn(4, 384)), randint(0, 256, (8,)), (1, randn(6, 192))],
    [randint(0, 256, (16,)), randn(7, 384), randint(0, 256, (5,)), (1, randn(2, 192)), randint(0, 256, (9,))]
]

loss = model(text_images_and_audio)

loss.backward()

# after much training

one_multimodal_sample = model.sample()
```

Automatically taking care of encoding and decoding of images

```python
import torch
from torch import nn, randint, randn
from transfusion_pytorch import Transfusion, print_modality_sample

mock_encoder = nn.Conv2d(3, 384, 3, padding = 1)
mock_decoder = nn.Conv2d(384, 3, 3, padding = 1)

model = Transfusion(
    num_text_tokens = 12,
    dim_latent = 384,
    channel_first_latent = True,
    modality_default_shape = (4, 4),
    modality_encoder = mock_encoder,
    modality_decoder = mock_decoder,
    transformer = dict(
        dim = 512,
        depth = 8
    )
)

text_and_images = [
    [
        randint(0, 12, (16,)),  # 16 text tokens
        randn(3, 8, 8),         # (8 x 8) 3 channeled image
        randint(0, 12, (8,)),   # 8 text tokens
        randn(3, 7, 7)          # (7 x 7) 3 channeled image
    ],
    [
        randint(0, 12, (16,)),  # 16 text tokens
        randn(3, 8, 5),         # (8 x 5) 3 channeled image
        randint(0, 12, (5,)),   # 5 text tokens
        randn(3, 2, 16),        # (2 x 16) 3 channeled image
        randint(0, 12, (9,))    # 9 text tokens
    ]
]

loss = model(text_and_images)

loss.backward()

# after much training

one_multimodal_sample = model.sample()

print_modality_sample(one_multimodal_sample)
```

To pretrain on language first, just pass in your text as type `Int['batch seq']`

```python
import torch
from transfusion_pytorch import Transfusion

model = Transfusion(
    num_text_tokens = 256,
    dim_latent = 384,
    transformer = dict(
        dim = 512,
        depth = 8,
    )
).cuda()

text = torch.randint(0, 256, (2, 1024)).cuda()

loss = model(text)
loss.backward()

# after much training

sampled = model.generate_text_only(text[:, :1], 1024)
```

## Todo

- [ ] use N-dimensional alibi with flex attention (configure for only certain amount of heads) for relative positions for any modality
- [ ] test out modality only training on oxford flowers
- [ ] given findings in pi-zero robotics foundation model, add mixture of experts for both attention and feedforward as options

## Citations

```bibtex
@inproceedings{Zhou2024TransfusionPT,
    title  = {Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model},
    author = {Chunting Zhou and Lili Yu and Arun Babu and Kushal Tirumala and Michihiro Yasunaga and Leonid Shamis and Jacob Kahn and Xuezhe Ma and Luke Zettlemoyer and Omer Levy},
    year   = {2024},
    url    = {https://api.semanticscholar.org/CorpusID:271909855}
}
```

```bibtex
@misc{Rubin2024,
    author  = {Ohad Rubin},
    url     = {https://medium.com/@ohadrubin/exploring-weight-decay-in-layer-normalization-challenges-and-a-reparameterization-solution-ad4d12c24950}
}
```

```bibtex
@article{Nguyen2024MinPS,
    title   = {Min P Sampling: Balancing Creativity and Coherence at High Temperature},
    author  = {Minh Nguyen and Andrew Baker and Andreas Kirsch and Clement Neo},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2407.01082},
    url     = {https://api.semanticscholar.org/CorpusID:270870613}
}
```

```bibtex
@article{Bao2022AllAW,
    title   = {All are Worth Words: A ViT Backbone for Diffusion Models},
    author  = {Fan Bao and Shen Nie and Kaiwen Xue and Yue Cao and Chongxuan Li and Hang Su and Jun Zhu},
    journal = {2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year    = {2022},
    pages   = {22669-22679},
    url     = {https://api.semanticscholar.org/CorpusID:253581703}
}
```

```bibtex
@inproceedings{Zhao2024MonoFormerOT,
    title     = {MonoFormer: One Transformer for Both Diffusion and Autoregression},
    author    = {Chuyang Zhao and Yuxing Song and Wenhao Wang and Haocheng Feng and Errui Ding and Yifan Sun and Xinyan Xiao and Jingdong Wang},
    year      = {2024},
    url       = {https://api.semanticscholar.org/CorpusID:272832492}
}
```

```bibtex
@article{Yang2024ConsistencyFM,
    title   = {Consistency Flow Matching: Defining Straight Flows with Velocity Consistency},
    author  = {Ling Yang and Zixiang Zhang and Zhilong Zhang and Xingchao Liu and Minkai Xu and Wentao Zhang and Chenlin Meng and Stefano Ermon and Bin Cui},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2407.02398},
    url     = {https://api.semanticscholar.org/CorpusID:270878436}
}
```

```bibtex
@inproceedings{Zhou2024ValueRL,
    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},
    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273532030}
}
```

```bibtex
@inproceedings{Yao2024FasterDiTTF,
    title   = {FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification},
    author  = {Jingfeng Yao and Wang Cheng and Wenyu Liu and Xinggang Wang},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273346237}
}
```

            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "transfusion-pytorch",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.8",
    "maintainer_email": null,
    "keywords": "artificial intelligence, attention mechanism, deep learning, rectified flow, transformers",
    "author": null,
    "author_email": "Phil Wang <lucidrains@gmail.com>",
    "download_url": "https://files.pythonhosted.org/packages/ca/55/6d915d7225bdd944a7647b3282b4ffa346e2fb920b17338d252217caf801/transfusion_pytorch-0.3.4.tar.gz",
    "platform": null,
    "description": "<img src=\"./transfusion.png\" width=\"400px\"></img>\n\n## Transfusion - Pytorch\n\nPytorch implementation of [Transfusion](https://www.arxiv.org/abs/2408.11039), \"Predict the Next Token and Diffuse Images with One Multi-Modal Model\", from MetaAI.\n\nIn this repo, we will substitute diffusion with flow matching given the success of Flux from Black Forest Labs (but will keep the original paper title given Transflow does not have the same ring). This repository will also attempt to extend to any number of modalities.\n\n## Install\n\n```bash\n$ pip install transfusion-pytorch\n```\n\n## Usage\n\nOne modality, say images\n\n```python\nfrom torch import randint, randn\nfrom transfusion_pytorch import Transfusion\n\nmodel = Transfusion(\n    num_text_tokens = 256,\n    dim_latent = 384,\n    modality_default_shape = (4,),  # fallback, in the case the language model did not produce a valid modality shape\n    transformer = dict(\n        dim = 512,\n        depth = 8\n    )\n)\n\n# any torch.long is text, torch.float is modalities\n\ntext_and_images = [\n    [randint(0, 256, (16,)), randn(4, 384), randint(0, 256, (8,)), randn(6, 384)],\n    [randint(0, 256, (16,)), randn(7, 384), randint(0, 256, (5,)), randn(2, 384), randint(0, 256, (9,))]\n]\n\nloss = model(text_and_images)\n\nloss.backward()\n\n# after much training\n\none_multimodal_sample = model.sample()\n```\n\nMultiple different modalities\n\n```python\nfrom torch import randint, randn\nfrom transfusion_pytorch import Transfusion\n\nmodel = Transfusion(\n    num_text_tokens = 256,\n    dim_latent = (384, 192),                 # specify multiple latent dimensions\n    modality_default_shape = ((4,), (2,)),   # default shapes for first and second modality\n    transformer = dict(\n        dim = 512,\n        depth = 8\n    )\n)\n\n# then for the Tensors of type float, you can pass a tuple[int, Tensor] and specify the modality index in the first position\n\n# any torch.long is text, torch.float is modalities\n\ntext_images_and_audio = [\n    [randint(0, 256, (16,)), (0, randn(4, 384)), randint(0, 256, (8,)), (1, randn(6, 192))],\n    [randint(0, 256, (16,)), randn(7, 384), randint(0, 256, (5,)), (1, randn(2, 192)), randint(0, 256, (9,))]\n]\n\nloss = model(text_images_and_audio)\n\nloss.backward()\n\n# after much training\n\none_multimodal_sample = model.sample()\n```\n\nAutomatically taking care of encoding and decoding of images\n\n```python\nimport torch\nfrom torch import nn, randint, randn\nfrom transfusion_pytorch import Transfusion, print_modality_sample\n\nmock_encoder = nn.Conv2d(3, 384, 3, padding = 1)\nmock_decoder = nn.Conv2d(384, 3, 3, padding = 1)\n\nmodel = Transfusion(\n    num_text_tokens = 12,\n    dim_latent = 384,\n    channel_first_latent = True,\n    modality_default_shape = (4, 4),\n    modality_encoder = mock_encoder,\n    modality_decoder = mock_decoder,\n    transformer = dict(\n        dim = 512,\n        depth = 8\n    )\n)\n\ntext_and_images = [\n    [\n        randint(0, 12, (16,)),  # 16 text tokens\n        randn(3, 8, 8),         # (8 x 8) 3 channeled image\n        randint(0, 12, (8,)),   # 8 text tokens\n        randn(3, 7, 7)          # (7 x 7) 3 channeled image\n    ],\n    [\n        randint(0, 12, (16,)),  # 16 text tokens\n        randn(3, 8, 5),         # (8 x 5) 3 channeled image\n        randint(0, 12, (5,)),   # 5 text tokens\n        randn(3, 2, 16),        # (2 x 16) 3 channeled image\n        randint(0, 12, (9,))    # 9 text tokens\n    ]\n]\n\nloss = model(text_and_images)\n\nloss.backward()\n\n# after much training\n\none_multimodal_sample = model.sample()\n\nprint_modality_sample(one_multimodal_sample)\n```\n\nTo pretrain on language first, just pass in your text as type `Int['batch seq']`\n\n```python\nimport torch\nfrom transfusion_pytorch import Transfusion\n\nmodel = Transfusion(\n    num_text_tokens = 256,\n    dim_latent = 384,\n    transformer = dict(\n        dim = 512,\n        depth = 8,\n    )\n).cuda()\n\ntext = torch.randint(0, 256, (2, 1024)).cuda()\n\nloss = model(text)\nloss.backward()\n\n# after much training\n\nsampled = model.generate_text_only(text[:, :1], 1024)\n```\n\n## Todo\n\n- [ ] use N-dimensional alibi with flex attention (configure for only certain amount of heads) for relative positions for any modality\n- [ ] test out modality only training on oxford flowers\n- [ ] given findings in pi-zero robotics foundation model, add mixture of experts for both attention and feedforward as options\n\n## Citations\n\n```bibtex\n@inproceedings{Zhou2024TransfusionPT,\n    title  = {Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model},\n    author = {Chunting Zhou and Lili Yu and Arun Babu and Kushal Tirumala and Michihiro Yasunaga and Leonid Shamis and Jacob Kahn and Xuezhe Ma and Luke Zettlemoyer and Omer Levy},\n    year   = {2024},\n    url    = {https://api.semanticscholar.org/CorpusID:271909855}\n}\n```\n\n```bibtex\n@misc{Rubin2024,\n    author  = {Ohad Rubin},\n    url     = {https://medium.com/@ohadrubin/exploring-weight-decay-in-layer-normalization-challenges-and-a-reparameterization-solution-ad4d12c24950}\n}\n```\n\n```bibtex\n@article{Nguyen2024MinPS,\n    title   = {Min P Sampling: Balancing Creativity and Coherence at High Temperature},\n    author  = {Minh Nguyen and Andrew Baker and Andreas Kirsch and Clement Neo},\n    journal = {ArXiv},\n    year    = {2024},\n    volume  = {abs/2407.01082},\n    url     = {https://api.semanticscholar.org/CorpusID:270870613}\n}\n```\n\n```bibtex\n@article{Bao2022AllAW,\n    title   = {All are Worth Words: A ViT Backbone for Diffusion Models},\n    author  = {Fan Bao and Shen Nie and Kaiwen Xue and Yue Cao and Chongxuan Li and Hang Su and Jun Zhu},\n    journal = {2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n    year    = {2022},\n    pages   = {22669-22679},\n    url     = {https://api.semanticscholar.org/CorpusID:253581703}\n}\n```\n\n```bibtex\n@inproceedings{Zhao2024MonoFormerOT,\n    title     = {MonoFormer: One Transformer for Both Diffusion and Autoregression},\n    author    = {Chuyang Zhao and Yuxing Song and Wenhao Wang and Haocheng Feng and Errui Ding and Yifan Sun and Xinyan Xiao and Jingdong Wang},\n    year      = {2024},\n    url       = {https://api.semanticscholar.org/CorpusID:272832492}\n}\n```\n\n```bibtex\n@article{Yang2024ConsistencyFM,\n    title   = {Consistency Flow Matching: Defining Straight Flows with Velocity Consistency},\n    author  = {Ling Yang and Zixiang Zhang and Zhilong Zhang and Xingchao Liu and Minkai Xu and Wentao Zhang and Chenlin Meng and Stefano Ermon and Bin Cui},\n    journal = {ArXiv},\n    year    = {2024},\n    volume  = {abs/2407.02398},\n    url     = {https://api.semanticscholar.org/CorpusID:270878436}\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\n```bibtex\n@inproceedings{Yao2024FasterDiTTF,\n    title   = {FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification},\n    author  = {Jingfeng Yao and Wang Cheng and Wenyu Liu and Xinggang Wang},\n    year    = {2024},\n    url     = {https://api.semanticscholar.org/CorpusID:273346237}\n}\n```\n",
    "bugtrack_url": null,
    "license": "MIT License  Copyright (c) 2024 Phil Wang  Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:  The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.  THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.",
    "summary": "Transfusion in Pytorch",
    "version": "0.3.4",
    "project_urls": {
        "Homepage": "https://pypi.org/project/transfusion-pytorch/",
        "Repository": "https://github.com/lucidrains/transfusion-pytorch"
    },
    "split_keywords": [
        "artificial intelligence",
        " attention mechanism",
        " deep learning",
        " rectified flow",
        " transformers"
    ],
    "urls": [
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "1015afe3839e13b520a8f75956d4a8390df8560fe1d051d29fbe5dd16cdc7475",
                "md5": "d96559c4cefca901a9cbe134c630f7bf",
                "sha256": "24198997555c26cac44732ab699b6cd43d70dab24e942cca355d0a99d28b5b39"
            },
            "downloads": -1,
            "filename": "transfusion_pytorch-0.3.4-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "d96559c4cefca901a9cbe134c630f7bf",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.8",
            "size": 21554,
            "upload_time": "2024-11-05T16:23:44",
            "upload_time_iso_8601": "2024-11-05T16:23:44.598924Z",
            "url": "https://files.pythonhosted.org/packages/10/15/afe3839e13b520a8f75956d4a8390df8560fe1d051d29fbe5dd16cdc7475/transfusion_pytorch-0.3.4-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "ca556d915d7225bdd944a7647b3282b4ffa346e2fb920b17338d252217caf801",
                "md5": "cd1ba628bcc3ec2cba540540ec00c98e",
                "sha256": "a905b43d7963935d374aa4f40ac8ae320eefc0c92a23fee270b46ecf63fef07b"
            },
            "downloads": -1,
            "filename": "transfusion_pytorch-0.3.4.tar.gz",
            "has_sig": false,
            "md5_digest": "cd1ba628bcc3ec2cba540540ec00c98e",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.8",
            "size": 359594,
            "upload_time": "2024-11-05T16:23:48",
            "upload_time_iso_8601": "2024-11-05T16:23:48.617333Z",
            "url": "https://files.pythonhosted.org/packages/ca/55/6d915d7225bdd944a7647b3282b4ffa346e2fb920b17338d252217caf801/transfusion_pytorch-0.3.4.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2024-11-05 16:23:48",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "lucidrains",
    "github_project": "transfusion-pytorch",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": true,
    "lcname": "transfusion-pytorch"
}
        
Elapsed time: 0.44571s