Upload README.md with huggingface_hub

README.md

@@ -1,36 +1,31 @@
 ---
-tags:
-- unet
-- pix2pix
-- pytorch
-library_name: pytorch
-license: wtfpl
-datasets:
-- K00B404/pix2pix_flux_set
-language:
-- en
-pipeline_tag: image-to-image
----
-# Pix2Pix UNet Model
-## Model Description
-Custom UNet model for Pix2Pix image translation.
-- **Image Size:** 1024
-- **Model Type:** big_UNet (1024)
-## Usage
-```python
-import torch
-from small_256_model import UNet as small_UNet
-from big_1024_model import UNet as big_UNet
-big = True
-# Load the model
-name='big_model_weights.pth' if big else 'small_model_weights.pth'
-checkpoint = torch.load(name)
-model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
-model.load_state_dict(checkpoint['model_state_dict'])
-model.eval()
-```
-## Model Architecture
-UNet(
+                        tags:
+                        - unet
+                        - pix2pix
+                        library_name: pytorch
+                        ---
+                        
+                        # Pix2Pix UNet Model
+                        
+                        ## Model Description
+                        Custom UNet model for Pix2Pix image translation.
+                        - Image Size: 256
+                        - Model Type: 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
+                        UNet(
   (encoder): Sequential(
     (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
     (1): ReLU(inplace=True)
@@ -38,21 +33,14 @@ UNet(
     (3): ReLU(inplace=True)
     (4): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
     (5): ReLU(inplace=True)
-    (6): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
-    (7): ReLU(inplace=True)
-    (8): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
-    (9): ReLU(inplace=True)
   )
   (decoder): Sequential(
-    (0): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
+    (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
     (1): ReLU(inplace=True)
-    (2): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
+    (2): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
     (3): ReLU(inplace=True)
-    (4): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
-    (5): ReLU(inplace=True)
-    (6): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
-    (7): ReLU(inplace=True)
-    (8): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
-    (9): Tanh()
+    (4): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
+    (5): Tanh()
   )
-) 
+)
+