{"id":31959536,"url":"https://github.com/huggingface/open-muse","last_synced_at":"2025-10-14T15:33:14.890Z","repository":{"id":153037806,"uuid":"602728970","full_name":"huggingface/open-muse","owner":"huggingface","description":"Open reproduction of MUSE for fast text2image generation. ","archived":false,"fork":false,"pushed_at":"2024-06-01T11:47:50.000Z","size":4538,"stargazers_count":356,"open_issues_count":41,"forks_count":31,"subscribers_count":38,"default_branch":"main","last_synced_at":"2025-09-05T04:27:46.767Z","etag":null,"topics":["cv","deep-learning","diffusion","generative-art","nlp","text2image","transformer"],"latest_commit_sha":null,"homepage":"https://huggingface.co/openMUSE","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/huggingface.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2023-02-16T20:25:03.000Z","updated_at":"2025-08-30T22:28:28.000Z","dependencies_parsed_at":null,"dependency_job_id":"4a1bfb60-97df-40dd-a3bf-6453d6526f81","html_url":"https://github.com/huggingface/open-muse","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/huggingface/open-muse","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/huggingface%2Fopen-muse","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/huggingface%2Fopen-muse/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/huggingface%2Fopen-muse/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/huggingface%2Fopen-muse/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/huggingface","download_url":"https://codeload.github.com/huggingface/open-muse/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/huggingface%2Fopen-muse/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279019320,"owners_count":26086711,"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","status":"online","status_checked_at":"2025-10-14T02:00:06.444Z","response_time":60,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"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":["cv","deep-learning","diffusion","generative-art","nlp","text2image","transformer"],"created_at":"2025-10-14T15:32:20.450Z","updated_at":"2025-10-14T15:33:14.881Z","avatar_url":"https://github.com/huggingface.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# open-muse\nAn open-reproduction effort to reproduce the transformer based [MUSE](https://muse-model.github.io/) model for fast text2image generation.\n\n## Demo\n\n👉 https://huggingface.co/spaces/openMUSE/MUSE\n\n## Goal\nThis repo is for reproduction of the [MUSE](https://arxiv.org/abs/2301.00704) model. The goal is to create a simple and scalable repo, to reproduce MUSE and build knowedge about VQ + transformers at scale.\nWe will use deduped LAION-2B + COYO-700M dataset for training.\n\nProject stages:\n1. Setup the codebase and train a class-conditional model on imagenet.\n2. Conduct text2image experiments on CC12M.\n3. Train improved VQGANs models.\n4. Train the full (base-256) model on LAION + COYO.\n5. Train the full (base-512) model on LAION + COYO.\n\nAll the artifacts of this project will be uploaded to the [openMUSE](https://huggingface.co/openMUSE) organization on the huggingface hub.\n\n## Usage\n\n### Installation\n\nFirst create a virtual environment and install the repo using:\n\n```bash\ngit clone https://github.com/huggingface/muse\ncd muse\npip install -e \".[extra]\"\n```\n\nYou'll need to install `PyTorch` and `torchvision` manually. We are using `torch==1.13.1` with `CUDA11.7` for training.\n\nFor distributed data parallel training we use `accelerate` library, although this may change in the future. For dataset loading, we use `webdataset` library. So the dataset should be in the `webdataset` format.\n\n### Models\n\nAt the momemnt we support following models:\n- `MaskGitTransformer` - The main transformer model from the paper.\n- `MaskGitVQGAN` - The VQGAN model from the [maskgit](https://github.com/google-research/maskgit) repo.\n- `VQGANModel` - The VQGAN model from the [taming transformers](https://github.com/CompVis/taming-transformers) repo.\n\nThe models are implemented under `muse` directory. All models implement the familiar `transformers` API. So you can use `from_pretrained` and `save_pretrained` methods to load and save the models. The model can be saved and loaded from the huggingface hub.\n\n#### VQGAN example:\n\n```python\nimport torch\nfrom torchvision import transforms\nfrom PIL import Image\nfrom muse import MaskGitVQGAN\n\n# Load the pre-trained vq model from the hub\nvq_model = MaskGitVQGAN.from_pretrained(\"openMUSE/maskgit-vqgan-imagenet-f16-256\")\n\n# encode and decode images using\nencode_transform = = transforms.Compose(\n    [\n        transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR),\n        transforms.CenterCrop(256),\n        transforms.ToTensor(),\n    ]\n)\nimage = Image.open(\"...\") #\npixel_values = encode_transform(image).unsqueeze(0)\nimage_tokens = vq_model.encode(pixel_values)\nrec_image = vq_model.decode(image_tokens)\n\n# Convert to PIL images\nrec_image = 2.0 * rec_image - 1.0\nrec_image = torch.clamp(rec_image, -1.0, 1.0)\nrec_image = (rec_image + 1.0) / 2.0\nrec_image *= 255.0\nrec_image = rec_image.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)\npil_images = [Image.fromarray(image) for image in rec_image]\n```\n\n#### MaskGitTransformer example for class-conditional generation:\n\n```python\nimport torch\nfrom muse import MaskGitTransformer, MaskGitVQGAN\nfrom muse.sampling import cosine_schedule\n\n# Load the pre-trained vq model from the hub\nvq_model = MaskGitVQGAN.from_pretrained(\"openMUSE/maskgit-vqgan-imagenet-f16-256\")\n\n# Initialize the MaskGitTransformer model\nmaskgit_model = MaskGitTransformer(\n    vocab_size=2025, #(1024 + 1000 + 1 = 2025 -\u003e Vq_tokens + Imagenet class ids + \u003cmask\u003e)\n    max_position_embeddings=257, # 256 + 1 for class token\n    hidden_size=512,\n    num_hidden_layers=8,\n    num_attention_heads=8,\n    intermediate_size=2048,\n    codebook_size=1024,\n    num_vq_tokens=256,\n    num_classes=1000,\n)\n\n# prepare the input batch\nimages = torch.randn(4, 3, 256, 256)\nclass_ids = torch.randint(0, 1000, (4,)) # random class ids\n# encode the images\nimage_tokens = vq_model.encode(images)\nbatch_size, seq_len = image_tokens.shape\n# Sample a random timestep for each image\ntimesteps = torch.rand(batch_size, device=image_tokens.device)\n# Sample a random mask probability for each image using timestep and cosine schedule\nmask_prob = cosine_schedule(timesteps)\nmask_prob = mask_prob.clip(min_masking_rate)\n# creat a random mask for each image\nnum_token_masked = (seq_len * mask_prob).round().clamp(min=1)\nbatch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)\nmask = batch_randperm \u003c num_token_masked.unsqueeze(-1)\n# mask images and create input and labels\ninput_ids = torch.where(mask, mask_id, image_tokens)\nlabels = torch.where(mask, image_tokens, -100)\n\n# shift the class ids by codebook size\nclass_ids = class_ids + vq_model.num_embeddings\n# prepend the class ids to the image tokens\ninput_ids = torch.cat([class_ids.unsqueeze(-1), input_ids], dim=-1)\n# prepend -100 to the labels as we don't want to predict the class ids\nlabels = torch.cat([-100 * torch.ones_like(class_ids).unsqueeze(-1), labels], dim=-1)\n\n# forward pass\nlogits, loss = maskgit_model(input_ids, labels=labels)\n\nloss.backward()\n\n# to generate images\nclass_ids = torch.randint(0, 1000, (4,)) # random class ids\ngenerated_tokens = maskgit_model.generate(class_ids=class_ids)\nrec_images = vq_model.decode(generated_tokens)\n```\n\n___Note___:\n- The vq model and transformer model are kept separate to be able to scale the transformer model independently. And we may pre-encode the images for faster training.\n- The masking is also done outside the model to be able to use different masking strategies without affecting the modeling code.\n\n## Basic explanation of MaskGit Generation Process\n\n1. Maskgits is a transformer that outputs logits given a sequence of tokens of both vq and class-conditioned label token\n\n2. The way the denoising process is done is to mask out with mask token ids and gradually denoise\n\n3. In the original implementation, this is done by first using a softmax on the last dim and randomly sampling as a categorical distribution. This will give our predicted tokens for each maskid. Then we get the probabilities for those tokens to be chosen. Finally, we get the topk highest confidence probabilities when gumbel*temp is added to it. Gumbel distribution is like a shifted normal distribution towards 0 which is used to model extreme events. So in extreme scenarios, we will like to see a different token being chosen from the default one\n\n4. For the lucidrian implementation, it first removes the highest-scoring (lowest probability) tokens by masking them with a given masking ratio. Then, except for the highest 10% of the logits that we get, we set it to -infinity so when we do the gumbel distribution on it, they will be ignored. Then update the input ids and the scores where the scores are just 1-the probability given by the softmax of the logits at the predicted ids interestingly\n\n## Training\nFor class-conditional imagenet we are using `accelerate` for DDP training and `webdataset` for data loading. The training script is available in `training/train_maskgit_imagenet.py`.\n\nWe use OmegaConf for configuration management. See `configs/template_config.yaml` for the configuration template. Below we explain the configuration parameters.\n\n```yaml\nwandb:\n  entity: ???\n\nexperiment:\n    name: ???\n    project: ???\n    output_dir: ???\n    max_train_examples: ???\n    save_every: 1000\n    eval_every: 500\n    generate_every: 1000\n    log_every: 50\n    log_grad_norm_every: 100\n    resume_from_checkpoint: latest\n\nmodel:\n    vq_model:\n        pretrained: \"openMUSE/maskgit-vqgan-imagenet-f16-256\"\n\n    transformer:\n        vocab_size: 2048 # (1024 + 1000 + 1 = 2025 -\u003e Vq + Imagenet + \u003cmask\u003e, use 2048 for even division by 8)\n        max_position_embeddings: 264 # (256 + 1 for class id, use 264 for even division by 8)\n        hidden_size: 768\n        num_hidden_layers: 12\n        num_attention_heads: 12\n        intermediate_size: 3072\n        codebook_size: 1024\n        num_vq_tokens: 256\n        num_classes: 1000\n        initializer_range: 0.02\n        layer_norm_eps: 1e-6\n        use_bias: False\n        use_normformer: True\n        use_encoder_layernorm: True\n        hidden_dropout: 0.0\n        attention_dropout: 0.0\n\n    gradient_checkpointing: True\n    enable_xformers_memory_efficient_attention: False\n\n\ndataset:\n    params:\n        train_shards_path_or_url: ???\n        eval_shards_path_or_url: ???\n        batch_size: ${training.batch_size}\n        shuffle_buffer_size: ???\n        num_workers: ???\n        resolution: 256\n        pin_memory: True\n        persistent_workers: True\n    preprocessing:\n        resolution: 256\n        center_crop: True\n        random_flip: False\n\noptimizer:\n    name: adamw # Can be adamw or lion or fused_adamw. Install apex for fused_adamw\n    params: # default adamw params\n        learning_rate: ???\n        scale_lr: False # scale learning rate by total batch size\n        beta1: 0.9\n        beta2: 0.999\n        weight_decay: 0.01\n        epsilon: 1e-8\n\nlr_scheduler:\n    scheduler: \"constant_with_warmup\"\n    params:\n        learning_rate: ${optimizer.params.learning_rate}\n        warmup_steps: 500\n\ntraining:\n    gradient_accumulation_steps: 1\n    batch_size: 128\n    mixed_precision: \"no\"\n    enable_tf32: True\n    use_ema: False\n    seed: 42\n    max_train_steps: ???\n    overfit_one_batch: False\n    min_masking_rate: 0.0\n    label_smoothing: 0.0\n    max_grad_norm: null\n```\n\nArguments with ??? are required.\n\n__wandb__:\n- `wandb.entity`: The wandb entity to use for logging.\n\n__experiment__:\n- `experiment.name`: The name of the experiment.\n- `experiment.project`: The wandb project to use for logging.\n- `experiment.output_dir`: The directory to save the checkpoints.\n- `experiment.max_train_examples`: The maximum number of training examples to use.\n- `experiment.save_every`: Save a checkpoint every `save_every` steps.\n- `experiment.eval_every`: Evaluate the model every `eval_every` steps.\n- `experiment.generate_every`: Generate images every `generate_every` steps.\n- `experiment.log_every`: Log the training metrics every `log_every` steps.\n- `log_grad_norm_every`: Log the gradient norm every `log_grad_norm_every` steps.\n- `experiment.resume_from_checkpoint`: The checkpoint to resume training from. Can be `latest` to resume from the latest checkpoint or path to a saved checkpoint. If `None` or the path does not exist then training starts from scratch.\n\n__model__:\n- `model.vq_model.pretrained`: The pretrained vq model to use. Can be a path to a saved checkpoint or a huggingface model name.\n- `model.transformer`: The transformer model configuration.\n- `model.gradient_checkpointing`: Enable gradient checkpointing for the transformer model.\n- `enable_xformers_memory_efficient_attention`: Enable memory efficient attention or flash attention for the transformer model. For flash attention we need to use `fp16` or `bf16`. [xformers](https://github.com/facebookresearch/xformers) needs to be installed for this to work.\n\n__dataset__:\n- `dataset.params.train_shards_path_or_url`: The path or url to the `webdataset` training shards.\n- `dataset.params.eval_shards_path_or_url`: The path or url to the `webdataset` evaluation shards.\n- `dataset.params.batch_size`: The batch size to use for training.\n- `dataset.params.shuffle_buffer_size`: The shuffle buffer size to use for training.\n- `dataset.params.num_workers`: The number of workers to use for data loading.\n- `dataset.params.resolution`: The resolution of the images to use for training.\n- `dataset.params.pin_memory`: Pin the memory for data loading.\n- `dataset.params.persistent_workers`: Use persistent workers for data loading.\n- `dataset.preprocessing.resolution`: The resolution of the images to use for preprocessing.\n- `dataset.preprocessing.center_crop`: Whether to center crop the images. If `False` then the images are randomly cropped to the `resolution`.\n- `dataset.preprocessing.random_flip`: Whether to randomly flip the images. If `False` then the images are not flipped.\n\n__optimizer__:\n- `optimizer.name`: The optimizer to use for training.\n- `optimizer.params`: The optimizer parameters.\n\n__lr_scheduler__:\n- `lr_scheduler.scheduler`: The learning rate scheduler to use for training.\n- `lr_scheduler.params`: The learning rate scheduler parameters.\n\n__training__:\n- `training.gradient_accumulation_steps`: The number of gradient accumulation steps to use for training.\n- `training.batch_size`: The batch size to use for training.\n- `training.mixed_precision`: The mixed precision mode to use for training. Can be `no`, `fp16` or `bf16`.\n- `training.enable_tf32`: Enable TF32 for training on Ampere GPUs.\n- `training.use_ema`: Enable EMA for training. Currently not supported.\n- `training.seed`: The seed to use for training.\n- `training.max_train_steps`: The maximum number of training steps.\n- `training.overfit_one_batch`: Whether to overfit one batch for debugging.\n- `training.min_masking_rate`: The minimum masking rate to use for training.\n- `training.label_smoothing`: The label smoothing value to use for training.\n- `max_grad_norm`: Max gradient norm.\n\n___Notes about training and dataset.___:\n\nWe randomly resample the shards (with replacement) and sample examples in buffer for training every time we resume/start the training run. This means our data loading is not determinitsic. We also don't do epoch based training but just using this for book keeping and being able to reuse the same training loop with other datasets/loaders.\n\n### Running experiments:\nSo far we are running experiments on single node. To launch a training run on a single node, run the following steps:\n\n1. Prepare the dataset in `webdataset` format. You can use the `scripts/convert_imagenet_to_wds.py` script to convert the imagenet dataset to `webdataset` format.\n2. First configure your training env using `accelerate config`.\n3. Create a `config.yaml` file for your experiment.\n4. Launch the training run using `accelerate launch`.\n\n```bash\naccelerate launch python -u training/train_maskgit_imagenet.py config=path/to/yaml/config\n```\n\nWith OmegaConf, commandline overrides are done in dot-notation format. E.g. if you want to override the dataset path, you would use the command `python -u train.py config=path/to/config dataset.params.path=path/to/dataset`.\n\nThe same command can be used to launch the training locally.\n\n## Steps\n### Setup the codebase and train a class-conditional model no imagenet.\n- [x] Setup repo-structure\n- [x] Add transformers and VQGAN model.\n- [x] Add a generation support for the model.\n- [x] Port the VQGAN from [maskgit](https://github.com/google-research/maskgit) repo for imagenet experiment.\n- [x] Finish and verify masking utils.\n- [ ] Add the masking arccos scheduling function from MUSE.\n- [x] Add EMA.\n- [x] Suport OmegaConf for training configuration.\n- [x] Add W\u0026B logging utils.\n- [x] Add WebDataset support. Not really needed for imagenet experiment but can work on this parallelly. (LAION is already available in this format so will be easier to use it).\n- [x] Add a training script for class conditional generation using imagenet.\n- [x] Make the codebase ready for the cluster training. Add SLURM scripts.\n\n### Conduct text2image experiments on CC12M.\n- [ ] Finish data loading, pre-processing utils.\n- [ ] Add CLIP and T5 support.\n- [ ] Add text2image training script.\n- [ ] Add eavluation scripts (FiD, CLIP score).\n- [ ] Train on CC12M. Here we could conduct different experiments:\n    - [ ] Train on CC12M with T5 conditioning.\n    - [ ] Train on CC12M with CLIP conditioning.\n    - [ ] Train on CC12M with CLIP + T5 conditioning (probably costly during training and experiments).\n    - [ ] Self conditioning from Bit Diffusion paper.\n- [ ] Collect different prompts for intermmediate evaluations (Can reuse the prompts for dalle-mini, parti-prompts).\n- [ ] Setup a space where people can play with the model and provide feedback, compare with other models etc.\n\n### Train improved VQGANs models.\n- [ ] Add training component models for VQGAN (EMA, discriminator, LPIPS etc).\n- [ ] VGQAN training script.\n\n\n### Misc tasks\n- [ ] Create a space for visualizing exploring dataset\n- [ ] Create a space where people can try to find their own images and can opt-out of the dataset.\n\n\n## Repo structure (WIP)\n```\n├── README.md\n├── configs                        -\u003e All training config files.\n│   └── template_config.yaml\n├── muse\n│   ├── __init__.py\n│   ├── data.py                    -\u003e All data related utils. Can create a data folder if needed.\n│   ├── logging.py                 -\u003e Misc logging utils.\n|   ├── lr_schedulers.py           -\u003e All lr scheduler related utils.\n│   ├── modeling_maskgit_vqgan.py  -\u003e VQGAN model from maskgit repo.\n│   ├── modeling_taming_vqgan.py   -\u003e VQGAN model from taming repo.\n│   └── modeling_transformer.py    -\u003e The main transformer model.\n│   ├── modeling_utils.py          -\u003e All model related utils, like save_pretrained, from_pretrained from hub etc\n│   ├── sampling.py                -\u003e Sampling/Generation utils.\n│   ├── training_utils.py          -\u003e Common training utils.\n├── pyproject.toml\n├── setup.cfg\n├── setup.py\n├── test.py\n└── training                       -\u003e All training scripts.\n    ├── __init__.py\n    ├── data.py                    -\u003e All data related utils. Can create a data folder if needed.\n    ├── optimizer.py               -\u003e All optimizer related utils and any new optimizer not available in PT.\n    ├── train_maskgit_imagenet.py\n    ├── train_muse.py\n    └── train_vqgan.py\n```\n\n## Acknowledgments\n\nThis project is hevily based on the following open-source repos. Thanks to all the authors for their amazing work.\n- [muse-maskgit-pytorch](https://github.com/lucidrains/muse-maskgit-pytorch) .  A big thanks to @lucidrains for this amazing work ❤️\n- [maskgit](https://github.com/google-research/maskgit) \n- [taming-transformers](https://github.com/CompVis/taming-transformers)\n- [open-clip](https://github.com/mlfoundations/open_clip)\n- [open-diffusion](https://github.com/mlfoundations/open-diffusion)\n- [dalle-mini](https://github.com/borisdayma/dalle-mini): ❤️\n- [transformers](https://github.com/huggingface/transformers)\n- [accelerate](https://github.com/huggingface/accelerate)\n- [diffusers](https://github.com/huggingface/diffusers)\n- [webdatset](https://github.com/webdataset/webdataset)\n\nAnd obivioulsy to PyTorch team for this amazing framework ❤️\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhuggingface%2Fopen-muse","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fhuggingface%2Fopen-muse","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhuggingface%2Fopen-muse/lists"}