Update app.py

app.py

@@ -15,7 +15,7 @@ from big_1024_model import UNet as big_UNet
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-big = False if device == torch.device('cpu') else True
+big = True if device == torch.device('cpu') else False
 
 # Parameters
 IMG_SIZE = 1024 if big else 256
@@ -190,7 +190,11 @@ def run_inference(image):
     output = output.cpu().squeeze(0).permute(1, 2, 0).numpy()
     output = ((output - output.min()) / (output.max() - output.min()) * 255).astype(np.uint8)
     return output
-
+    
+def to_hub(model):
+    wrapper = UNetWrapper(model, model_repo_id)
+    wrapper.push_to_hub()
+    
 def train_model(epochs):
     """Training function"""
     global global_model
@@ -223,15 +227,17 @@ def train_model(epochs):
                 status = f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}"
                 print(status)
                 output_text.append(status)
-    
+        
+        to_hub(model)
+        
     global_model = model
     return model, "\n".join(output_text)
 
+
 def gradio_train(epochs):
     """Gradio training interface function"""
     model, training_log = train_model(int(epochs))
-    wrapper = UNetWrapper(model, model_repo_id)
-    wrapper.push_to_hub()
+    to_hub(model)
     return f"{training_log}\n\nModel trained for {epochs} epochs and pushed to {model_repo_id}"
 
 def gradio_inference(input_image):