Update README.md

README.md

@@ -23,13 +23,13 @@ pipeline_tag: text-to-image
 
 <!-- Provide a quick summary of what the model is/does. -->
 
-This model is based in [Muse](https://muse-model.github.io/) and trained using the code hosted on [ZeroCool940711/muse-maskgit-pytorch](https://github.com/ZeroCool940711/muse-maskgit-pytorch), which is based on [`lucidrains/muse-maskgit-pytorch`](https://github.com/lucidrains/muse-maskgit-pytorch).
+This model is based in [Muse](https://muse-model.github.io/) and trained using the code hosted on [Sygil-Dev/muse-maskgit-pytorch](https://github.com/Sygil-Dev/muse-maskgit-pytorch), which is based on [`lucidrains/muse-maskgit-pytorch`](https://github.com/lucidrains/muse-maskgit-pytorch) and a collaboration between the [Sygil-Dev](https://github.com/Sygil-Dev) and [ShoukanLabs](https://github.com/ShoukanLabs) teams.
 
 # Model Details
-This model is a new model trained from scratch based on [Muse](https://muse-model.github.io/), trained on the [Imaginary Network Expanded Dataset](https://github.com/Sygil-Dev/INE-dataset), with the big advantage of allowing the use of multiple namespaces (labeled tags) to control various parts of the final generation.
+This model is a new model trained from scratch based on [Muse](https://muse-model.github.io/), trained on a subset of the [Imaginary Network Expanded Dataset](https://github.com/Sygil-Dev/INE-dataset), with the big advantage of allowing the use of multiple namespaces (labeled tags) to control various parts of the final generation.
 The use of namespaces (eg. “species:seal” or “studio:dc”) stops the model from misinterpreting a seal as the singer Seal, or DC Comics as Washington DC. 
 
-Note: As of right now, only the first VAE and MaskGit has been trained, we still need to train the Super Resolution VAE for the model to be usable even tho we might be able to reuse the first VAE depending on the quality of it once the training progresses more.
+Note: As of right now, only the first VAE and MaskGit has been trained with different configuration, we are trying to find the best balance between quality, performance and vram usage so Muse can be used on all kind of devices, we still need to train the Super Resolution VAE for the model to be usable even tho we might be able to reuse the first VAE depending on the quality of it once the training progresses more.
 
 If you find my work useful, please consider supporting me on [GitHub Sponsors](https://github.com/sponsors/ZeroCool940711)! 
 
@@ -44,11 +44,12 @@ This model is still in its infancy and it's meant to be constantly updated and t
     - [vae.sygil_muse_v0.1.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.sygil_muse_v0.1.pt): Trained from scratch for 3.0M steps with **dim: 128** and **vq_codebook_size: 256**.
     - [maskgit.sygil_muse_v0.1.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/maskgit.sygil_muse_v0.1.pt): Maskgit trained from the VAE for 3.46M steps
     - [vae.sygil_muse_v0.5.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.sygil_muse_v0.5.pt): Trained from scratch for 1.99M steps with **dim: 128** and **vq_codebook_size: 8192**.
+    - [vae.sygil_muse_v0.7.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.sygil_muse_v0.7.pt): Trained from scratch for 440K steps with **dim: 128**, **vq_codebook_dim 4096** and **vq_codebook_size: 16384**.
+    - [maskgit.sygil_muse_v0.7.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/maskgit.sygil_muse_v0.7.pt): Maskgit trained from the **vae.sygil_muse_v0.7.pt** VAE for 108K steps.
   - #### Beta:
-    - [vae.263000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.263000.pt): Trained from scratch for 263K steps and higher **vq_codebook_dim** and **vq_codebook_size** than before.
-    - [maskgit.23739.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/maskgit.23739.pt): Maskgit trained from the VAE for 23K steps.
+    - [vae.190000.pt](https://huggingface.co/Sygil/Sygil-Muse/blob/main/vae.190000.pt): Trained from scratch for 190K steps with **dim: 32**,**vq_codebook_dim: 8192** and **vq_codebook_size: 8192**.
   
-  Note: Checkpoints under the Beta section are updated daily or at least 3-4 times a week. While the beta checkpoints can be used as they are only the latest version is kept on the repo and the older checkpoints are removed when a new one
+  Note: Checkpoints under the Beta section are updated daily or at least 3-4 times a week. While the beta checkpoints can be used as they are, only the latest version is kept on the repo and the older checkpoints are removed when a new one
   is uploaded to keep the repo clean.
   
 ## Training
@@ -60,26 +61,26 @@ The model was trained on the following dataset:
 **Hardware and others**
 - **Hardware:** 1 x Nvidia RTX 3050 GPU
 - **Hours Trained:** NaN.
-- **Gradient Accumulations**: 20
+- **Gradient Accumulations**: 10
 - **Batch:** 1
 - **Learning Rate:** 1e-5
 - **Learning Rate Scheduler:** `constant_with_warmup`
-- **Scheduler Power:** 0.5
+- **Scheduler Power:** 1.0
 - **Optimizer:** Adam
 - **Warmup Steps:** 10,000
 - **Number of Cycles:** 1
 - **Resolution/Image Size**: First trained at a resolution of 64x64, then increased to 256x256 and then to 512x512. Check the notes down below for more details on this.
-- **Dimension:** 128
-- **vq_codebook_dim:** 4096
-- **vq_codebook_size:** 16384
+- **Dimension:** 32
+- **vq_codebook_dim:** 8192
+- **vq_codebook_size:** 8192
 - **heads:** 8
 - **depth:** 4
 - **Random Crop:** True
-- **Total Training Steps:** 263,000
+- **Total Training Steps:** 190,000
 
   Note: On Muse we can change the image_size or resolution at any time without having to train the model from scratch again, this allows us to first train the model at low resolution using the same `dim` and `vq_codebook_size` to train faster and then we can increase the `image_size` and use a higher resolution once the model has trained enough.
 
-Developed by: [ZeroCool](https://github.com/ZeroCool940711) at [Sygil-Dev](https://github.com/Sygil-Dev/)
+Developed by: [ZeroCool](https://github.com/ZeroCool940711) at [Sygil-Dev](https://github.com/Sygil-Dev/).
 
 # License
 This model is open access and available to all, with a CreativeML Open RAIL++-M License further specifying rights and usage.