Update app.py

app.py

@@ -67,8 +67,20 @@ class UNetWrapper:
         api = HfApi()
         # Create a repository if it doesn't exist
         create_repo(self.repo_id, exist_ok=True,token=self.token)
+        pth_name = 'model_weights.pth'
+        torch.save(self.model.state_dict(), pth_name)
+        from huggingface_hub import upload_file
+
+        upload_file(
+            path_or_fileobj=pth_name,
+            path_in_repo=pth_name,
+            repo_id=self.repo_id,,
+            token=self.token,
+        
+        )
+        #api.upload_file(repo_id=self.repo_id, path_in_repo=pth_name, path_or_fileobj=pth_name)
         # Push the model's state dict to the Hugging Face Hub
-        self.model.save_pretrained(self.repo_id,token=self.token)  # You may need to implement this method
+        #self.model.save_pretrained(self.repo_id,token=self.token)  # You may need to implement this method
 
 
 # Training function
@@ -94,7 +106,7 @@ def train_model(epochs):
 
     criterion = nn.L1Loss()
     optimizer = optim.Adam(model.parameters(), lr=LR)
-
+    output_text = []
     # Training loop
     for epoch in range(epochs):
         for i, (original, target) in enumerate(dataloader):
@@ -110,7 +122,12 @@ def train_model(epochs):
             optimizer.step()
 
             if i % 100 == 0:
-                print(f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.4f}")
+                status=f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}"
+                print(status)
+                output_text.append(status)
+                
+            # Here you could also use a delay to simulate training time
+            yield "\n".join(output_text)  # Send output to Gradio
 
     # Return trained model
     return model