Update app.py

app.py

@@ -25,14 +25,20 @@ LR = 0.0002
 dataset_id = "K00B404/pix2pix_flux_set"
 model_repo_id = "K00B404/pix2pix_flux"
 
+# Create dataset and dataloader
 class Pix2PixDataset(torch.utils.data.Dataset):
     def __init__(self, ds):
+        # Filter dataset for 'original' and 'target' images
         self.originals = [x for x in ds["train"] if x['label'] == 'original']
         self.targets = [x for x in ds["train"] if x['label'] == 'target']
 
-        # Ensure original and target images match by their index
+        # Ensure the number of original and target images match
         assert len(self.originals) == len(self.targets), "Mismatch in number of original and target images."
 
+        # Debug: Print dataset size
+        print(f"Number of original images: {len(self.originals)}")
+        print(f"Number of target images: {len(self.targets)}")
+
     def __len__(self):
         return len(self.originals)