{"id":15601062,"url":"https://github.com/lucidrains/parti-pytorch","last_synced_at":"2025-04-12T18:45:17.943Z","repository":{"id":38079683,"uuid":"506404806","full_name":"lucidrains/parti-pytorch","owner":"lucidrains","description":"Implementation of Parti, Google's pure attention-based text-to-image neural network, in Pytorch","archived":false,"fork":false,"pushed_at":"2023-12-08T14:48:21.000Z","size":347,"stargazers_count":530,"open_issues_count":4,"forks_count":24,"subscribers_count":13,"default_branch":"main","last_synced_at":"2025-04-04T00:53:34.710Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","text-to-image","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":"2022-06-22T21:01:28.000Z","updated_at":"2025-04-01T07:44:16.000Z","dependencies_parsed_at":"2023-11-16T20:10:34.084Z","dependency_job_id":"b6605bec-23a9-4cdd-9470-6f9f1cb855f0","html_url":"https://github.com/lucidrains/parti-pytorch","commit_stats":{"total_commits":57,"total_committers":2,"mean_commits":28.5,"dds":0.01754385964912286,"last_synced_commit":"5075148e3cbe788e5e4b96985c82208013f024b7"},"previous_names":[],"tags_count":23,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fparti-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fparti-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fparti-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fparti-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/parti-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248617341,"owners_count":21134190,"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","deep-learning","text-to-image","transformers"],"created_at":"2024-10-03T02:13:31.226Z","updated_at":"2025-04-12T18:45:17.915Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./parti.jpeg\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Parti - Pytorch\n\nImplementation of \u003ca href=\"https://parti.research.google/\"\u003eParti\u003c/a\u003e, Google's pure attention-based text-to-image neural network, in Pytorch. \u003ca src=\"https://sites.research.google/parti/\"\u003eProject Page\u003c/a\u003e\n\nThis repository also contains working training code for \u003ca href=\"https://ai.googleblog.com/2022/05/vector-quantized-image-modeling-with.html\"\u003eViT VQGan VAE\u003c/a\u003e. It also contains some additional modifications for faster training from vision transformers literature.\n\n\u003ca href=\"https://www.youtube.com/watch?v=qS-iYnp00uc\"\u003eYannic Kilcher\u003c/a\u003e\n\nPlease join \u003ca href=\"https://discord.gg/xBPBXfcFHd\"\u003e\u003cimg alt=\"Join us on Discord\" src=\"https://img.shields.io/discord/823813159592001537?color=5865F2\u0026logo=discord\u0026logoColor=white\"\u003e\u003c/a\u003e if you are interested in helping out with the replication with the \u003ca href=\"https://laion.ai/\"\u003eLAION\u003c/a\u003e community\n\n## Install\n\n```bash\n$ pip install parti-pytorch\n```\n\n## Usage\n\nFirst you will need to train your Transformer VQ-GAN VAE\n\n```python\nfrom parti_pytorch import VitVQGanVAE, VQGanVAETrainer\n\nvit_vae = VitVQGanVAE(\n    dim = 256,               # dimensions\n    image_size = 256,        # target image size\n    patch_size = 16,         # size of the patches in the image attending to each other\n    num_layers = 3           # number of layers\n).cuda()\n\ntrainer = VQGanVAETrainer(\n    vit_vae,\n    folder = '/path/to/your/images',\n    num_train_steps = 100000,\n    lr = 3e-4,\n    batch_size = 4,\n    grad_accum_every = 8,\n    amp = True\n)\n\ntrainer.train()\n```\n\nThen\n\n```python\nimport torch\nfrom parti_pytorch import Parti, VitVQGanVAE\n\n# first instantiate your ViT VQGan VAE\n# a VQGan VAE made of transformers\n\nvit_vae = VitVQGanVAE(\n    dim = 256,               # dimensions\n    image_size = 256,        # target image size\n    patch_size = 16,         # size of the patches in the image attending to each other\n    num_layers = 3           # number of layers\n).cuda()\n\nvit_vae.load_state_dict(torch.load(f'/path/to/vae.pt')) # you will want to load the exponentially moving averaged VAE\n\n# then you plugin the ViT VqGan VAE into your Parti as so\n\nparti = Parti(\n    vae = vit_vae,            # vit vqgan vae\n    dim = 512,                # model dimension\n    depth = 8,                # depth\n    dim_head = 64,            # attention head dimension\n    heads = 8,                # attention heads\n    dropout = 0.,             # dropout\n    cond_drop_prob = 0.25,    # conditional dropout, for classifier free guidance\n    ff_mult = 4,              # feedforward expansion factor\n    t5_name = 't5-large',     # name of your T5\n)\n\n# ready your training text and images\n\ntexts = [\n    'a child screaming at finding a worm within a half-eaten apple',\n    'lizard running across the desert on two feet',\n    'waking up to a psychedelic landscape',\n    'seashells sparkling in the shallow waters'\n]\n\nimages = torch.randn(4, 3, 256, 256).cuda()\n\n# feed it into your parti instance, with return_loss set to True\n\nloss = parti(\n    texts = texts,\n    images = images,\n    return_loss = True\n)\n\nloss.backward()\n\n# do this for a long time on much data\n# then...\n\nimages = parti.generate(texts = [\n    'a whale breaching from afar',\n    'young girl blowing out candles on her birthday cake',\n    'fireworks with blue and green sparkles'\n], cond_scale = 3., return_pil_images = True) # conditioning scale for classifier free guidance\n\n# List[PILImages] (256 x 256 RGB)\n```\n\nRealistically, when scaling up, you'll want to pre-encode your text into tokens and their respective mask\n\n```python\nfrom parti_pytorch.t5 import t5_encode_text\n\nimages = torch.randn(4, 3, 256, 256).cuda()\n\ntext_token_embeds, text_mask = t5_encode_text([\n    'a child screaming at finding a worm within a half-eaten apple',\n    'lizard running across the desert on two feet',\n    'waking up to a psychedelic landscape',\n    'seashells sparkling in the shallow waters'\n], name = 't5-large', output_device = images.device)\n\n# store somewhere, then load with the dataloader\n\nloss = parti(\n    text_token_embeds = text_token_embeds,\n    text_mask = text_mask,\n    images = images,\n    return_loss = True\n)\n\nloss.backward()\n```\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStabilityAI\u003c/a\u003e for the sponsorship, as well as my other sponsors, for affording me the independence to open source artificial intelligence.\n\n- \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for the transformers library and the ease for encoding text with T5 language model\n\n## Todo\n\n- [x] add 2d relative positional bias to parti autoregressive transformer\n- [x] cite all techniques adopted from vision transformer literature in vit vqgan if they work\n- [x] get working vit vqgan-vae trainer code, as discriminator needs to be trained\n- [x] use crossformer embed layer for initial convolution in discriminator\n- [ ] preencoding of text with designated t5\n- [ ] training code for parti\n- [ ] inference caching\n- [ ] automatic filtering with Coca https://github.com/lucidrains/CoCa-pytorch\n- [ ] bring in the super-resoluting convolutional net mentioned in the paper, with training code\n- [ ] initialize 2d rel pos bias in conv-like pattern\n- [ ] consider a small nerf-like MLP at the end of vit-vqgan, similar to https://arxiv.org/abs/2107.04589\n\n## Citations\n\n```bibtex\n@inproceedings{Yu2022Pathways\n    title   = {Pathways Autoregressive Text-to-Image Model},\n    author  = {Jiahui Yu*, Yuanzhong Xu†, Jing Yu Koh†, Thang Luong†, Gunjan Baid†, Zirui Wang†, Vijay Vasudevan†, Alexander Ku†, Yinfei Yang, Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge†, Yonghui Wu*},\n    year    = {2022}\n}\n```\n\n```bibtex\n@article{Shleifer2021NormFormerIT,\n    title   = {NormFormer: Improved Transformer Pretraining with Extra Normalization},\n    author  = {Sam Shleifer and Jason Weston and Myle Ott},\n    journal = {ArXiv},\n    year    = {2021},\n    volume  = {abs/2110.09456}\n}\n```\n\n```bibtex\n@article{Sankararaman2022BayesFormerTW,\n    title   = {BayesFormer: Transformer with Uncertainty Estimation},\n    author  = {Karthik Abinav Sankararaman and Sinong Wang and Han Fang},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2206.00826}\n}\n```\n\n```bibtex\n@article{Lee2021VisionTF,\n    title   = {Vision Transformer for Small-Size Datasets},\n    author  = {Seung Hoon Lee and Seunghyun Lee and Byung Cheol Song},\n    journal = {ArXiv},\n    year    = {2021},\n    volume  = {abs/2112.13492}\n}\n```\n\n```bibtex\n@article{Chu2021DoWR,\n    title   = {Do We Really Need Explicit Position Encodings for Vision Transformers?},\n    author  = {Xiangxiang Chu and Bo Zhang and Zhi Tian and Xiaolin Wei and Huaxia Xia},\n    journal = {ArXiv},\n    year    = {2021},\n    volume  = {abs/2102.10882}\n}\n```\n\n```bibtex\n@article{So2021PrimerSF,\n    title   = {Primer: Searching for Efficient Transformers for Language Modeling},\n    author  = {David R. So and Wojciech Ma'nke and Hanxiao Liu and Zihang Dai and Noam M. Shazeer and Quoc V. Le},\n    journal = {ArXiv},\n    year    = {2021},\n    volume  = {abs/2109.08668}\n}\n```\n\n```bibtex\n@inproceedings{Wang2021CrossFormerAV,\n    title   = {CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention},\n    author  = {Wenxiao Wang and Lulian Yao and Long Chen and Binbin Lin and Deng Cai and Xiaofei He and Wei Liu},\n    year    = {2021}\n}\n```\n\n```bibtex\n@misc{mentzer2023finite,\n    title   = {Finite Scalar Quantization: VQ-VAE Made Simple},\n    author  = {Fabian Mentzer and David Minnen and Eirikur Agustsson and Michael Tschannen},\n    year    = {2023},\n    eprint  = {2309.15505},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n\n```bibtex\n@misc{yu2023language,\n    title   = {Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation},\n    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},\n    year    = {2023},\n    eprint  = {2310.05737},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fparti-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fparti-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fparti-pytorch/lists"}