Update app.py

app.py

@@ -5,6 +5,8 @@ from torch.utils.data import DataLoader
 from torchvision import transforms
 from datasets import load_dataset
 from huggingface_hub import Repository
+from huggingface_hub import HfApi, HfFolder, Repository, create_repo
+
 import gradio as gr
 from PIL import Image
 import os
@@ -53,7 +55,21 @@ class Pix2PixDataset(torch.utils.data.Dataset):
 
         # Apply the necessary transforms
         return self.transform(original), self.transform(target)
-        
+
+class UNetWrapper:
+    def __init__(self, unet_model, repo_id):
+        self.model = unet_model
+        self.repo_id = repo_id
+
+    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)
+        # Push the model's state dict to the Hugging Face Hub
+        self.model.save_pretrained(self.repo_id)  # You may need to implement this method
+
+
 # Training function
 def train_model(epochs):
     # Load the dataset
@@ -100,8 +116,11 @@ def train_model(epochs):
 
 # 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)
+    #model.push_to_hub(repo_name)
 
 # Gradio interface function
 def gradio_train(epochs):