Update app.py

app.py

@@ -110,34 +110,36 @@ class UNetWrapper:
             
             # Create and upload model card
             model_card = f"""---
-                        tags:
-                        - unet
-                        - pix2pix
-                        library_name: pytorch
-                        ---
-                        
-                        # Pix2Pix UNet Model
-                        
-                        ## Model Description
-                        Custom UNet model for Pix2Pix image translation.
-                        - Image Size: {1024 if isinstance(self.model, big_UNet) else 256}
-                        - Model Type: {"Big (1024)" if isinstance(self.model, big_UNet) else "Small (256)"}
-                        
-                        ## Usage
-                        
-                        ```python
-                        import torch
-                        from small_256_model import UNet as small_UNet
-                        from big_1024_model import UNet as big_UNet
-                        
-                        # Load the model
-                        checkpoint = torch.load('model_weights.pth')
-                        model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
-                        model.load_state_dict(checkpoint['model_state_dict'])
-                        model.eval()
-                        Model Architecture
-                        {str(self.model)}
-                    """
+tags:
+- unet
+- pix2pix
+- pytorch
+library_name: pytorch
+---
+
+# Pix2Pix UNet Model
+
+## Model Description
+Custom UNet model for Pix2Pix image translation.
+- **Image Size:** {1024 if isinstance(self.model, big_UNet) else 256}
+- **Model Type:** {"Big (1024)" if isinstance(self.model, big_UNet) else "Small (256)"}
+
+## Usage
+
+```python
+import torch
+from small_256_model import UNet as small_UNet
+from big_1024_model import UNet as big_UNet
+
+# Load the model
+checkpoint = torch.load('model_weights.pth')
+model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
+model.load_state_dict(checkpoint['model_state_dict'])
+model.eval()
+
+Model Architecture
+
+{str(self.model)} """
             # Save and upload README
             with open("README.md", "w") as f:
                 f.write(model_card)