Update app.py

app.py

@@ -31,8 +31,9 @@ global_model = None
 def load_model():
     """Load the model at startup"""
     global global_model
+    weights_name = 'big_model_weights.pth' if big else 'small_model_weights.pth'
     try:
-        checkpoint = torch.load('model_weights.pth', map_location=device)
+        checkpoint = torch.load(weights_name, map_location=device)
         model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
         model.load_state_dict(checkpoint['model_state_dict'])
         model.to(device)
@@ -86,7 +87,7 @@ class UNetWrapper:
             }
             
             # Save model locally
-            pth_name = 'model_weights.pth'
+            pth_name = 'big_model_weights.pth' if big else  'small_model_weights.pth'
             torch.save(save_dict, pth_name)
             
             # Create repo if it doesn't exist
@@ -115,14 +116,20 @@ tags:
 - 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 if isinstance(self.model, big_UNet) else 256}
-- **Model Type:** {"Big (1024)" if isinstance(self.model, big_UNet) else "Small (256)"}
+- **Image Size:** 1024
+- **Model Type:** Big (1024)
 
 ## Usage
 
@@ -130,9 +137,10 @@ Custom UNet model for Pix2Pix image translation.
 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
-checkpoint = torch.load('model_weights.pth')
+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()