Update app.py

app.py

@@ -1,7 +1,7 @@
 import torch
 from torch import nn, optim
-from torchvision import transforms, datasets, models
 from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms, datasets, models
 from PIL import Image
 import json
 import os
@@ -24,9 +24,9 @@ class ImageDescriptionDataset(Dataset):
         self.metadata = metadata
         self.image_names = list(metadata.keys())  # List of image filenames
         self.transform = transforms.Compose([
-            transforms.Resize((224, 224)),
+            transforms.Resize((512, 512)),
             transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
         ])
 
     def __len__(self):
@@ -35,76 +35,81 @@ class ImageDescriptionDataset(Dataset):
     def __getitem__(self, idx):
         image_name = self.image_names[idx]
         image_path = os.path.join(self.image_folder, image_name)
-        image = Image.open(image_path).convert("RGB")  # Open the image and convert to RGB
-        description = self.metadata[image_name]  # Get description for the image
-        image = self.transform(image)  # Apply transformations
+        image = Image.open(image_path).convert("RGB")
+        description = self.metadata[image_name]
+        image = self.transform(image)
         return image, description
 
-# LoRA Model Class (This is a placeholder, you'll need to implement the actual LoRA model)
+# LoRA Layer Implementation
+class LoRALayer(nn.Module):
+    def __init__(self, original_layer, rank=4):
+        super(LoRALayer, self).__init__()
+        self.original_layer = original_layer
+        self.rank = rank
+        self.lora_up = nn.Linear(original_layer.in_features, rank, bias=False)
+        self.lora_down = nn.Linear(rank, original_layer.out_features, bias=False)
+
+    def forward(self, x):
+        return self.original_layer(x) + self.lora_down(self.lora_up(x))
+
+# LoRA Model Class
 class LoRAModel(nn.Module):
     def __init__(self):
         super(LoRAModel, self).__init__()
-        self.backbone = models.resnet18(pretrained=True)  # Using a pre-trained ResNet18
-        
-        # Fixing the shape mismatch: Input size to the fc layer should match ResNet output
-        self.fc = nn.Linear(self.backbone.fc.in_features, 100)  # 100 is a placeholder for your output
-
-        # If you want to use LoRA, you will implement the low-rank adaptation mechanism here
+        self.backbone = models.resnet18(pretrained=True)  # Base model
+        self.backbone.fc = LoRALayer(self.backbone.fc)  # Replace the final layer with LoRA
 
     def forward(self, x):
-        x = self.backbone(x)  # Extract features using the ResNet18 backbone
-        x = self.fc(x)  # Apply the final fully connected layer
-        return x
+        return self.backbone(x)
 
-# Function to train LoRA
+# Training Function
 def train_lora(image_folder, metadata):
-    print("Starting training process...")
-    
-    # Create dataset and dataloaders
+    print("Starting LoRA training process...")
+
+    # Create dataset and dataloader
     dataset = ImageDescriptionDataset(image_folder, metadata)
     dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
-    
-    # Initialize model, loss, and optimizer
+
+    # Initialize model, loss function, and optimizer
     model = LoRAModel()
-    criterion = nn.CrossEntropyLoss()  # Placeholder loss function, can be adjusted
+    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
+    num_epochs = 5
     for epoch in range(num_epochs):
         print(f"Epoch {epoch + 1}/{num_epochs}")
         for batch_idx, (images, descriptions) in enumerate(dataloader):
-            # Here we would convert descriptions to a numerical format
-            # Since it's a placeholder, we use random labels for descriptions
-            labels = torch.randint(0, 100, (images.size(0),))  # Random labels as a placeholder
+            # Placeholder: Convert descriptions to labels
+            labels = torch.randint(0, 100, (images.size(0),))
 
             # 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
+
+            if batch_idx % 10 == 0:
                 print(f"Batch {batch_idx}, Loss: {loss.item()}")
 
-    print("Training completed.")
+    print("LoRA training completed.")
 
-# Gradio app function to load metadata and start training
+# Gradio App
 def start_training_gradio():
-    print("Preparing dataset...")
-    metadata = load_metadata(metadata_file)  # Load metadata
-    return train_lora(image_folder, metadata)
+    print("Loading metadata and preparing dataset...")
+    metadata = load_metadata(metadata_file)
+    train_lora(image_folder, metadata)
+    return "Training completed. Check the model outputs!"
 
-# Gradio interface
 demo = gr.Interface(
-    fn=start_training_gradio,  # Use the new function name here
+    fn=start_training_gradio,
     inputs=None,
     outputs="text",
-    title="Train LoRA on Your Dataset",
-    description="Click below to start training with the uploaded images and metadata."
+    title="Train LoRA Model",
+    description="Fine-tune a model using LoRA for consistent image generation."
 )
 
 demo.launch()