{"id":26770796,"url":"https://github.com/vvvm23/vqvae-2","last_synced_at":"2025-06-27T00:35:58.730Z","repository":{"id":48157578,"uuid":"350845841","full_name":"vvvm23/vqvae-2","owner":"vvvm23","description":"PyTorch implementation of VQ-VAE-2 from \"Generating Diverse High-Fidelity Images with VQ-VAE-2\"","archived":false,"fork":false,"pushed_at":"2023-02-15T08:50:24.000Z","size":83480,"stargazers_count":152,"open_issues_count":3,"forks_count":20,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-15T18:08:50.848Z","etag":null,"topics":["generative-model","pytorch","vae","vq-vae","vq-vae-2"],"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/vvvm23.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}},"created_at":"2021-03-23T20:18:37.000Z","updated_at":"2025-04-07T19:06:44.000Z","dependencies_parsed_at":"2023-02-18T23:01:53.466Z","dependency_job_id":null,"html_url":"https://github.com/vvvm23/vqvae-2","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/vvvm23%2Fvqvae-2","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/vvvm23%2Fvqvae-2/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/vvvm23%2Fvqvae-2/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/vvvm23%2Fvqvae-2/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/vvvm23","download_url":"https://codeload.github.com/vvvm23/vqvae-2/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":249125955,"owners_count":21216705,"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":["generative-model","pytorch","vae","vq-vae","vq-vae-2"],"created_at":"2025-03-28T23:16:38.224Z","updated_at":"2025-04-15T18:09:00.670Z","avatar_url":"https://github.com/vvvm23.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Generating Diverse High-Fidelity Images with VQ-VAE-2 [Work in Progress]\nPyTorch implementation of Hierarchical, Vector Quantized, Variational Autoencoders (VQ-VAE-2) \nfrom the paper \"Generating Diverse High-Fidelity Images with VQ-VAE-2\"\n\nOriginal paper can be found [here](https://arxiv.org/abs/1906.00446)\n\nVector Quantizing layer based off implementation by [@rosinality](https://github.com/rosinality) \nfound [here](https://github.com/rosinality/vq-vae-2-pytorch).\n\nAiming for a focus on supporting an arbitrary number of VQ-VAE \"levels\". Most implementations in \nPyTorch typically only use 2 which is limiting at higher resolutions. This repository contains \ncheckpoints for a 3-level and 5-level VQ-VAE-2, trained on FFHQ1024.\n\nThis project will not only contain the VQ-VAE-2 architecture, but also an example autoregressive \nprior and latent dataset extraction.\n\n\u003e This project is very much Work-in-Progress.\n\u003e VQ-VAE-2 model is mostly complete. PixelSnail prior models are still experimental\n\u003e and most definitely do not work.\n\n## Usage\n### VQ-VAE-2 Usage\nRun VQ-VAE-2 training using the config `task_name` found in `hps.py`. Defaults to `cifar10`:\n```\npython main-vqvae.py --task task_name\n```\n\nEvaluate VQ-VAE-2 from parameters `state_dict_path` on task `task_name`. Defaults to `cifar10`:\n```\npython main-vqvae.py --task task_name --load-path state_dict_path --evaluate\n```\n\nOther useful flags:\n```\n--no-save       # disables saving of files during training\n--cpu           # do not use GPU\n--batch-size    # overrides batch size in cfg.py, useful for evaluating on larger batch size\n--no-tqdm       # disable tqdm status bars\n--no-save       # disables saving of files\n--no-amp        # disables using native AMP (Automatic Mixed Precision) operations\n--save-jpg      # save all images as jpg instead of png, useful for extreme resolutions\n```\n\n### Latent Dataset Generation\nRun latent dataset generation using VQ-VAE-2 saved at `path` that was trained on task `task_name`. Defaults to `cifar10`:\n```\npython main-latents.py path --task task_name\n```\nResult is saved in `latent-data` directory.\n\nOther useful flags:\n```\n--cpu           # do not use GPU\n--batch-size    # overrides batch size in cfg.py, useful for evaluating on larger batch size\n--no-tqdm       # disable tqdm status bars\n--no-save       # disables saving of files\n--no-amp        # disables using native AMP (Automatic Mixed Precision) operations\n```\n\n### Discrete Prior Usage\nRun level `level` PixelSnail discrete prior training using the config `task_name` found in `hps.py` using latent dataset saved at path `latent_dataset.pt` and VQ-VAE `vqvae_path` to dequantize conditioning variables. Defaults to `cifar10`:\n```\npython main-pixelsnail.py latent_dataset.pt vqvae_path.pt level --task task_name\n```\n\nOther useful flags:\n```\n--cpu           # do not use GPU\n--load-path     # resume from saved state on disk\n--batch-size    # overrides batch size in cfg.py, useful for evaluating on larger batch size\n--save-jpg      # save all images as jpg instead of png, useful for extreme resolutions\n--no-tqdm       # disable tqdm status bars\n--no-save       # disables saving of files\n```\n\n### Sample Generation \nRun sampling script on trained VQ-VAE-2 and PixelSnail priors using the config `task_name` (default `cifar10`) found in `hps.py`.\nThe first positional argument is the path to the VQ-VAE-2 checkpoint. The remaining `L` positional arguments are the PixelSnail \nprior checkpoints from level `0` to `L`.\n```\npython main-sample.py vq_vae_path.pt pixelsnail_0_path.pt pixel_snail_1_path.pt ... --task task_name\n```\n\nOther useful flags:\n```\n--cpu           # do not use GPU\n--batch-size    # overrides batch size in cfg.py, useful for evaluating on larger batch size\n--nb-samples    # number of samples to generate. defaults to 1.\n--no-tqdm       # disable tqdm status bars\n--no-save       # disables saving of files\n--no-amp        # disables using native AMP (Automatic Mixed Precision) operations\n--save-jpg      # save all images as jpg instead of png, useful for extreme resolutions\n--temperature   # controls softmax temperature during sampling\n```\n\n## Modifications\n- Replacing residual layers with ReZero layers.\n\n## Samples\n*Reconstructions from FFHQ1024 using a 3-level VQ-VAE-2*\n![Reconstructions from FFHQ1024 using a 3-level VQ-VAE-2](recon-example.png)\n\n## Checkpoints\n[FFHQ1024 - 3-level VQ-VAE-2](ffhq1024-state-dict-0017.pt)\n\n[FFHQ1024 - 5-level VQ-VAE-2](ffhq1024-large-state-dict-0010.pt)\n\n### Roadmap\n- [X] Server mode (no fancy printing)\n- [X] Experiment directories (containing logs / checkpoints / etc)\n- [X] Accumulated gradient training (for larger batch sizes on limited resources)\n- [X] Samples and checkpoints on FFHQ1024\n- [X] Latent dataset generation\n- [ ] Autoregressive prior models / training scripts\n- [X] Full system sampling\n- [ ] Prettier outputs\n- [ ] Output logging\n\n### Citations\n**Generating Diverse High-Fidelity Images with VQ-VAE-2**\n```\n@misc{razavi2019generating,\n      title={Generating Diverse High-Fidelity Images with VQ-VAE-2}, \n      author={Ali Razavi and Aaron van den Oord and Oriol Vinyals},\n      year={2019},\n      eprint={1906.00446},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG}\n}\n```\n\n**PixelSNAIL: An Improved Autoregressive Generative Model**\n```\n@misc{chen2017pixelsnail,\n      title={PixelSNAIL: An Improved Autoregressive Generative Model}, \n      author={Xi Chen and Nikhil Mishra and Mostafa Rohaninejad and Pieter Abbeel},\n      year={2017},\n      eprint={1712.09763},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG}\n}\n```\n\n**ReZero is All You Need: Fast Convergence at Large Depth**\n```\n@misc{bachlechner2020rezero,\n      title={ReZero is All You Need: Fast Convergence at Large Depth}, \n      author={Thomas Bachlechner and Bodhisattwa Prasad Majumder and Huanru Henry Mao and Garrison W. Cottrell and Julian McAuley},\n      year={2020},\n      eprint={2003.04887},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fvvvm23%2Fvqvae-2","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fvvvm23%2Fvqvae-2","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fvvvm23%2Fvqvae-2/lists"}