Update app.py

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 import json
 import os
 import gradio as gr
+import shutil
 
 # Paths
 image_folder = "Images/"
@@ -95,23 +96,22 @@ def train_lora(image_folder, metadata):
             if batch_idx % 10 == 0:  # Log every 10 batches
                 print(f"Batch {batch_idx}, Loss: {loss.item()}")
 
-    # Define the folder where the model will be saved
-    save_folder = "models"
-    os.makedirs(save_folder, exist_ok=True)  # Create the folder if it doesn't exist
+    # 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 in the specified folder
-    model_save_path = os.path.join(save_folder, "lora_model.pth")
-    torch.save(model.state_dict(), model_save_path)
-    print(f"Model saved at {model_save_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}.")
 
-    print("Training completed.")
+    return f"Training completed. Download the model from: [Download Model](sandbox:/mnt/data/lora_model.pth)"
 
 # Gradio App
 def start_training_gradio():
     print("Loading metadata and preparing dataset...")
     metadata = load_metadata(metadata_file)
-    train_lora(image_folder, metadata)
-    return "Training completed. Check the model outputs!"
+    return train_lora(image_folder, metadata)
 
 demo = gr.Interface(
     fn=start_training_gradio,