Update app.py

app.py

@@ -66,10 +66,99 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         target = target_img.convert('RGB')
         return self.transform(original), self.transform(target)
 
-# UNetWrapper class remains the same
 class UNetWrapper:
-    # ... [Previous UNetWrapper implementation remains unchanged]
-    pass
+    def __init__(self, unet_model, repo_id):
+        self.model = unet_model
+        self.repo_id = repo_id
+        self.token = os.getenv('NEW_TOKEN') # Make sure this environment variable is set
+        self.api = HfApi(token=os.getenv('NEW_TOKEN'))
+
+    def push_to_hub(self):
+        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}")
+            
 
 def prepare_input(image, device='cpu'):
     """Prepare image for inference"""