Update app.py

app.py

@@ -60,29 +60,94 @@ class UNetWrapper:
     def __init__(self, unet_model, repo_id):
         self.model = unet_model
         self.repo_id = repo_id
-        self.token = os.getenv('HF_WRITE')
+        self.token = os.getenv('HF_WRITE')  # Make sure this environment variable is set
+        self.api = HfApi()
 
     def push_to_hub(self):
-        # Initialize the Hugging Face API
-        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,
-            repo_type="model"
-        
-        )
-        #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
-
+        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)
+            }
+            
+            # Save model locally
+            pth_name = '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}")
+            
+            # Upload the model file
+            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"
+            )
+            
+            # Create and upload model card
+            model_card = f"""---
+                        tags:
+                        - unet
+                        - pix2pix
+                        library_name: pytorch
+                        ---
+                        
+                        # 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)"}
+                        
+                        ## Usage
+                        
+                        ```python
+                        import torch
+                        from small_256_model import UNet as small_UNet
+                        from big_1024_model import UNet as big_UNet
+                        
+                        # Load the model
+                        checkpoint = torch.load('model_weights.pth')
+                        model = big_UNet() if checkpoint['model_config']['big'] else small_UNet()
+                        model.load_state_dict(checkpoint['model_state_dict'])
+                        model.eval()
+                        Model Architecture
+                        {str(self.model)}
+                    """
+            # Save and upload README
+            with open("README.md", "w") as f:
+                f.write(model_card)
+            
+            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}")
 
 # Training function
 def train_model(epochs):
@@ -134,12 +199,9 @@ def train_model(epochs):
     return model
 
 # Push model to Hugging Face Hub
-def push_model_to_hub(model, repo_name):
-    # Usage example
-    model_wrapper = UNetWrapper(model, model_repo_id)
-    model_wrapper.push_to_hub()
-    # Push the model to the Hugging Face hub
-    #model.push_to_hub(repo_name)
+def push_model_to_hub(model, repo_id):
+    wrapper = UNetWrapper(model, repo_id)
+    wrapper.push_to_hub()
 
 # Gradio interface function
 def gradio_train(epochs):