Update README.md

README.md

@@ -1,7 +1,34 @@
+---
+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(
   (encoder): Sequential(
@@ -28,4 +55,4 @@ UNet(
     (8): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
     (9): Tanh()
   )
-)
+)