Update app.py

app.py

@@ -95,9 +95,14 @@ 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
-    torch.save(model.state_dict(), "lora_model.pth")
-    print("Model saved as lora_model.pth")
+    # 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 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}")
 
     print("Training completed.")