Update README.md

README.md

@@ -2,8 +2,70 @@
 tags:
 - model_hub_mixin
 - pytorch_model_hub_mixin
+license: mit
 ---
 
+# ImageGenerationTAU: Autoencoder for MNIST Image Generation
+
+## Model Details
+
+- **Model Architecture:** Convolutional Autoencoder  
+- **Framework:** PyTorch  
+- **Input Shape:** (1, 28, 28) (Grayscale MNIST Images)  
+- **Latent Dimension:** User-defined (`hidden_dim`)  
+- **Dataset:** [MNIST Handwritten Digits](http://yann.lecun.com/exdb/mnist/)  
+
+## Model Description
+
+The **ImageGenerationTAU** model is a **convolutional autoencoder** designed for **image generation and feature extraction** from MNIST. It consists of:
+- An **encoder** that compresses the input image into a **low-dimensional representation**.
+- A **decoder** that reconstructs the original image from the compressed representation.
+
+This model can be used for **image denoising, feature learning, and generative tasks**.
+
+## Training Details
+
+- **Loss Function:** Smooth L1 Loss  
+- **Optimizer:** Adam  
+- **Batch Size:** 512  
+- **Number of Epochs:** TBD  
+- **Regularization:** Batch Normalization  
+
+### Model Architecture
+
+```python
+class ImageGenerationTAU(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, hidden_dim):
+        super(ImageGenerationTAU, self).__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.ReLU(),
+            nn.BatchNorm2d(64),
+            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.ReLU(),
+            nn.BatchNorm2d(32),
+            nn.Flatten(),
+            nn.Linear(32 * 7 * 7, hidden_dim),
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear(hidden_dim, 32 * 7 * 7),
+            nn.ReLU(),
+            nn.Unflatten(1, (32, 7, 7)),
+            nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2),
+            nn.ReLU(),
+            nn.BatchNorm2d(64),
+            nn.ConvTranspose2d(64, 1, kernel_size=2, stride=2),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+```
+
 This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
 - Library: [More Information Needed]
 - Docs: [More Information Needed]