{"id":21135238,"url":"https://github.com/cloneofsimo/vqgan-training","last_synced_at":"2025-04-13T02:21:49.731Z","repository":{"id":258409398,"uuid":"866123675","full_name":"cloneofsimo/vqgan-training","owner":"cloneofsimo","description":"Train VAE like a boss","archived":false,"fork":false,"pushed_at":"2024-10-21T20:10:08.000Z","size":32283,"stargazers_count":270,"open_issues_count":5,"forks_count":12,"subscribers_count":7,"default_branch":"main","last_synced_at":"2025-04-04T04:41:13.502Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cloneofsimo.png","metadata":{"files":{"readme":"README.hf.md","changelog":null,"contributing":null,"funding":null,"license":null,"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":"2024-10-01T17:29:25.000Z","updated_at":"2025-03-14T10:15:14.000Z","dependencies_parsed_at":"2024-12-20T15:10:31.482Z","dependency_job_id":"c77ca127-5f91-4653-926f-fb8ebfaf8d52","html_url":"https://github.com/cloneofsimo/vqgan-training","commit_stats":null,"previous_names":["cloneofsimo/vqgan-training"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fvqgan-training","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fvqgan-training/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fvqgan-training/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fvqgan-training/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cloneofsimo","download_url":"https://codeload.github.com/cloneofsimo/vqgan-training/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248654701,"owners_count":21140348,"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":[],"created_at":"2024-11-20T06:46:50.304Z","updated_at":"2025-04-13T02:21:49.702Z","avatar_url":"https://github.com/cloneofsimo.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"---\nlicense: apache-2.0\n---\n\n# Equivariant 16ch, f8 VAE\n\n\u003cvideo controls autoplay src=\"https://cdn-uploads.huggingface.co/production/uploads/6311151c64939fabc00c8436/6DQGRWvQvDXp2xQlvwvwU.mp4\"\u003e\u003c/video\u003e\n\nAuraEquiVAE is a novel autoencoder that addresses multiple problems of existing conventional VAEs. First, unlike traditional VAEs that have significantly small log-variance, this model admits large noise to the latent space.\nAdditionally, unlike traditional VAEs, the latent space is equivariant under `Z_2 X Z_2` group operations (Horizontal / Vertical flip).\n\nTo understand the equivariance, we apply suitable group actions to both the latent space globally and locally. The latent is represented as `Z = (z_1, ..., z_n)`, and we perform a global permutation group action `g_global` on the tuples such that `g_global` is isomorphic to the `Z_2 x Z_2` group.\nWe also apply a local action `g_local` to individual `z_i` elements such that `g_local` is also isomorphic to the `Z_2 x Z_2` group.\n\nIn our specific case, `g_global` corresponds to flips, while `g_local` corresponds to sign flips on specific latent dimensions. Changing 2 channels for sign flips for both horizontal and vertical directions was chosen empirically.\n\nThe model has been trained using the approach described in [Mastering VAE Training](https://github.com/cloneofsimo/vqgan-training), where detailed explanations for the training process can be found.\n\n## How to use\n\nTo use the weights, copy paste the [VAE](https://github.com/cloneofsimo/vqgan-training/blob/03e04401cf49fe55be612d1f568be0110aa0fad1/ae.py) implementation.\n\n```python\nfrom ae import VAE\nimport torch\nfrom PIL import Image\n\nvae = VAE(\n    resolution=256,\n    in_channels=3,\n    ch=256,\n    out_ch=3,\n    ch_mult=[1, 2, 4, 4],\n    num_res_blocks=2,\n    z_ch\n).cuda().bfloat16()\n\nfrom safetensors.torch import load_file\nstate_dict = load_file(\"./vae_epoch_3_step_49501_bf16.pt\")\nvae.load_state_dict(state_dict)\n\nimgpath = 'contents/lavender.jpg'\n\nimg_orig = Image.open(imgpath).convert(\"RGB\")\noffset = 128\nW = 768\nimg_orig = img_orig.crop((offset, offset, W + offset, W + offset))\nimg = transforms.ToTensor()(img_orig).unsqueeze(0).cuda()\nimg = (img - 0.5) / 0.5\n\nwith torch.no_grad():\n    z = vae.encoder(img)\n    z = z.clamp(-8.0, 8.0) # this is latent!!\n\n# flip horizontal\nz = torch.flip(z, [-1]) # this corresponds to g_global\nz[:, -4:-2] = -z[:, -4:-2] # this corresponds to g_local\n\n# flip vertical\nz = torch.flip(z, [-2])\nz[:, -2:] = -z[:, -2:]\n\n\nwith torch.no_grad():\n    decz = vae.decoder(z) # this is image!\n\ndecimg = ((decz + 1) / 2).clamp(0, 1).squeeze(0).cpu().float().numpy().transpose(1, 2, 0)\ndecimg = (decimg * 255).astype('uint8')\ndecimg = Image.fromarray(decimg) # PIL image.\n\n```\n\n## Citation\n\nIf you find this model useful, please cite:\n\n```\n@misc{Training VQGAN and VAE, with detailed explanation,\n  author = {Simo Ryu},\n  title = {Training VQGAN and VAE, with detailed explanation},\n  year = {2024},\n  publisher = {GitHub},\n  journal = {GitHub repository},\n  howpublished = {\\url{https://github.com/cloneofsimo/vqgan-training}},\n}\n```\n\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fvqgan-training","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcloneofsimo%2Fvqgan-training","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fvqgan-training/lists"}