Update app.py

app.py

@@ -1,12 +1,11 @@
 import torch
 from torch import nn, optim
 from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms, datasets, models
+from torch torchvision import transforms, datasets, models
 from PIL import Image
 import json
 import os
 import gradio as gr
-import shutil
 
 # Paths
 image_folder = "Images/"
@@ -96,27 +95,26 @@ def train_lora(image_folder, metadata):
             if batch_idx % 10 == 0:  # Log every 10 batches
                 print(f"Batch {batch_idx}, Loss: {loss.item()}")
 
-    # Save the trained model to /mnt/data/ for Hugging Face Space to access
-    save_path = '/mnt/data/lora_model.pth'
-    torch.save(model.state_dict(), save_path)
-    print(f"Model saved at {save_path}")
+    # Save the trained model
+    model_path = "lora_model.pth"
+    torch.save(model.state_dict(), model_path)
+    print(f"Model saved as {model_path}")
 
-    # Move the file to a location where we can access it for download
-    # Here, /mnt/data is directly accessible from the Hugging Face Space interface
-    print(f"Training completed. The model is saved and ready for download at {save_path}.")
-
-    return f"Training completed. Download the model from: [Download Model](sandbox:/mnt/data/lora_model.pth)"
+    print("Training completed.")
+    return model_path  # Return the path of the saved model
 
 # Gradio App
 def start_training_gradio():
     print("Loading metadata and preparing dataset...")
     metadata = load_metadata(metadata_file)
-    return train_lora(image_folder, metadata)
+    model_path = train_lora(image_folder, metadata)
+    return model_path  # This will return the model file path for download
 
+# Gradio interface
 demo = gr.Interface(
     fn=start_training_gradio,
     inputs=None,
-    outputs="text",
+    outputs=gr.File(),
     title="Train LoRA Model",
     description="Fine-tune a model using LoRA for consistent image generation."
 )