Update app.py

app.py

@@ -90,20 +90,6 @@ class Pix2PixDataset(torch.utils.data.Dataset):
 
 class UNetWrapper:
 
-    def push_to_hub(self, pth_name):
-        """Push model checkpoint and metadata to the Hugging Face Hub."""
-        try:
-            self.api.upload_file(
-                path_or_fileobj=pth_name,
-                path_in_repo=pth_name,
-                repo_id=self.repo_id,
-                token=self.token,
-                repo_type="model"
-            )
-            print(f"Model checkpoint successfully uploaded to {self.repo_id}")
-        except Exception as e:
-            print(f"Error uploading model: {e}")
-
 
             
 
@@ -146,52 +132,25 @@ class UNetWrapper:
         self.loss = checkpoint['loss']
         print(f"Checkpoint loaded: epoch {self.epoch}, loss: {self.loss}")
 
-
-    def push_to_hub(self):
+    def push_to_hub(self, pth_name):
+        """Push model checkpoint and metadata to the Hugging Face Hub."""
         try:
-            # Save model state and configuration
-            save_dict = {
-                'model_state_dict': self.model.state_dict(),
-                'model_config': {
-                    'big': isinstance(self.model, big_UNet),
-                    'img_size': 1024 if isinstance(self.model, big_UNet) else 256
-                },
-                'model_architecture': str(self.model),
-                'model_state':{
-                    'epoch': self.epoch,
-                    'loss': self.loss
-                }
-            }
-            
-            # Save model locally
-            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
-            try:
-                create_repo(
-                    repo_id=self.repo_id, 
-                    token=self.token,
-                    exist_ok=True
-                )
-            except Exception as e:
-                print(f"Repository creation note: {e}")
-            
-            """Push model checkpoint and metadata to the Hugging Face Hub."""
-            try:
-                self.api.upload_file(
-                    path_or_fileobj=pth_name,
-                    path_in_repo=pth_name,
-                    repo_id=self.repo_id,
-                    token=self.token,
-                    repo_type="model"
-                )
-                print(f"Model checkpoint successfully uploaded to {self.repo_id}")
-            except Exception as e:
-                print(f"Error uploading model: {e}")
+            self.api.upload_file(
+                path_or_fileobj=pth_name,
+                path_in_repo=pth_name,
+                repo_id=self.repo_id,
+                token=self.token,
+                repo_type="model"
+            )
+            print(f"Model checkpoint successfully uploaded to {self.repo_id}")
+        except Exception as e:
+            print(f"Error uploading model: {e}")
+
+        
+   
             
-            # Create and upload model card
-            model_card = f"""---
+        # Create and upload model card
+        model_card = f"""---
 tags:
 - unet
 - pix2pix
@@ -230,31 +189,31 @@ model.eval()
 ## Model Architecture
 
 {str(self.model)} """
-            rp(model_card)
-            
-            # Save and upload README
-            with open("README.md", "w") as f:
-                f.write(f"# Pix2Pix UNet Model\n\n"
-                        f"- **Image Size:** {save_dict['model_config']['img_size']}\n"
-                        f"- **Model Type:** {'big' if big else 'small'}_UNet ({save_dict['model_config']['img_size']})\n"
-                        f"## Model Architecture\n{str(self.model)}")
-            
-            self.api.upload_file(
-                path_or_fileobj="README.md",
-                path_in_repo="README.md",
-                repo_id=self.repo_id,
-                token=self.token,
-                repo_type="model"
-            )
-            
-            # Clean up local files
-            os.remove(pth_name)
-            os.remove("README.md")
-            
-            print(f"Model successfully uploaded to {self.repo_id}")
-            
-        except Exception as e:
-            print(f"Error uploading model: {e}")
+        rp(model_card)
+        
+        # Save and upload README
+        with open("README.md", "w") as f:
+            f.write(f"# Pix2Pix UNet Model\n\n"
+                    f"- **Image Size:** {save_dict['model_config']['img_size']}\n"
+                    f"- **Model Type:** {'big' if big else 'small'}_UNet ({save_dict['model_config']['img_size']})\n"
+                    f"## Model Architecture\n{str(self.model)}")
+        
+        self.api.upload_file(
+            path_or_fileobj="README.md",
+            path_in_repo="README.md",
+            repo_id=self.repo_id,
+            token=self.token,
+            repo_type="model"
+        )
+        
+        # Clean up local files
+        os.remove(pth_name)
+        os.remove("README.md")
+        
+        print(f"Model successfully uploaded to {self.repo_id}")
+        
+    except Exception as e:
+        print(f"Error uploading model: {e}")
 
 def prepare_input(image, device='cpu'):
     """Prepare image for inference"""