Update app.py

app.py

@@ -27,7 +27,7 @@ model_repo_id = "K00B404/pix2pix_flux"
 
 # Create dataset and dataloader
 class Pix2PixDataset(torch.utils.data.Dataset):
-    def __init__(self, ds):
+    def __init__(self, ds, transform):
         # Filter dataset for 'original' (label = 0) and 'target' (label = 1) images
         self.originals = [x for x in ds["train"] if x['label'] == 0]
         self.targets = [x for x in ds["train"] if x['label'] == 1]
@@ -39,35 +39,34 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         print(f"Number of original images: {len(self.originals)}")
         print(f"Number of target images: {len(self.targets)}")
 
+        self.transform = transform  # Store the transform
+
     def __len__(self):
         return len(self.originals)
 
     def __getitem__(self, idx):
-        # Directly use the 'image' object without loading via Image.open()
         original_img = self.originals[idx]['image']
         target_img = self.targets[idx]['image']
 
-        # Apply the necessary transforms
-        original = original_img.convert('RGB')
-        target = target_img.convert('RGB')
+        original = original_img.convert('RGB')  # Convert to RGB if needed
+        target = target_img.convert('RGB')      # Convert to RGB if needed
 
-        # Return transformed original and target images
-        return transform(original), transform(target)
+        # Apply the necessary transforms
+        return self.transform(original), self.transform(target)
         
 # Training function
 def train_model(epochs):
     # Load the dataset
     ds = load_dataset(dataset_id)
     print(f"ds{ds}")
-    # Transform function to resize and convert to tensor
+    # Create the transform function outside of the dataset class
     transform = transforms.Compose([
         transforms.Resize((IMG_SIZE, IMG_SIZE)),
         transforms.ToTensor(),
     ])
-
-
-    dataset = Pix2PixDataset(ds)
-    print(f"dataset{dataset}")
+    
+    # Create dataset and dataloader
+    dataset = Pix2PixDataset(ds, transform)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
     # Initialize model, loss function, and optimizer