Update app.py

app.py

@@ -148,19 +148,18 @@ class UNetWrapper:
             
         except Exception as e:
             print(f"Error uploading model: {e}")
-
+            
 # Training function
 def train_model(epochs):
     # Load the dataset
     ds = load_dataset(dataset_id)
     print(f"ds{ds}")
-    # Create the transform function outside of the dataset class
+    
     transform = transforms.Compose([
         transforms.Resize((IMG_SIZE, IMG_SIZE)),
         transforms.ToTensor(),
     ])
     
-    # Create dataset and dataloader
     dataset = Pix2PixDataset(ds, transform)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
@@ -173,8 +172,10 @@ def train_model(epochs):
     criterion = nn.L1Loss()
     optimizer = optim.Adam(model.parameters(), lr=LR)
     output_text = []
+    
     # Training loop
     for epoch in range(epochs):
+        model.train()
         for i, (original, target) in enumerate(dataloader):
             original, target = original.to(device), target.to(device)
             optimizer.zero_grad()
@@ -188,36 +189,26 @@ def train_model(epochs):
             optimizer.step()
 
             if i % 100 == 0:
-                status=f"Epoch [{epoch}/{epochs}], Step [{i}/{len(dataloader)}], Loss: {loss.item():.8f}"
+                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
 
-# Push model to Hugging Face Hub
-def push_model_to_hub(model, repo_id):
-    wrapper = UNetWrapper(model, repo_id)
-    wrapper.push_to_hub()
+    return model, "\n".join(output_text)
 
 # Gradio interface function
 def gradio_train(epochs):
-    model = train_model(int(epochs))
+    model, training_log = train_model(int(epochs))
     push_model_to_hub(model, model_repo_id)
-    return f"Model trained for {epochs} epochs on the {dataset_id} dataset and pushed to Hugging Face Hub {model_repo_id} repository."
+    return f"{training_log}\n\nModel trained for {epochs} epochs on the {dataset_id} dataset and pushed to Hugging Face Hub {model_repo_id} repository."
 
 # Gradio Interface
 gr_interface = gr.Interface(
     fn=gradio_train,
     inputs=gr.Number(label="Number of Epochs"),
-    outputs="text",
+    outputs=gr.Textbox(label="Training Progress", lines=10),
     title="Pix2Pix Model Training",
     description="Train the Pix2Pix model and push it to the Hugging Face Hub repository."
 )
-
 if __name__ == '__main__':
     # Create or clone the repository if necessary
     repo = Repository(local_dir=model_repo_id, clone_from=model_repo_id)