Update app.py

app.py

@@ -75,31 +75,31 @@ def train_lora(image_folder, metadata):
     criterion = nn.CrossEntropyLoss()  # Update this if your task changes
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-# Training loop
-num_epochs = 5  # Adjust the number of epochs based on your needs
-for epoch in range(num_epochs):
-    print(f"Epoch {epoch + 1}/{num_epochs}")
-    for batch_idx, (images, descriptions) in enumerate(dataloader):
-        # Convert descriptions to a numerical format (if applicable)
-        labels = torch.randint(0, 100, (images.size(0),))  # Placeholder labels
-
-        # Forward pass
-        outputs = model(images)
-        loss = criterion(outputs, labels)
-        
-        # Backward pass
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        
-        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")
-
-print("Training completed.")
+    # Training loop
+    num_epochs = 5  # Adjust the number of epochs based on your needs
+    for epoch in range(num_epochs):
+        print(f"Epoch {epoch + 1}/{num_epochs}")
+        for batch_idx, (images, descriptions) in enumerate(dataloader):
+            # Convert descriptions to a numerical format (if applicable)
+            labels = torch.randint(0, 100, (images.size(0),))  # Placeholder labels
+
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            # Backward pass
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            
+            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")
+
+    print("Training completed.")
 
 # Gradio App
 def start_training_gradio():