Update app.py

app.py

@@ -28,14 +28,10 @@ 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
-        #https://huggingface.co/datasets/K00B404/pix2pix_flux_set/viewer/default/train?f[label][value]=0
-        #https://huggingface.co/datasets/K00B404/pix2pix_flux_set/viewer/default/train?f[label][value]=1
-        
+        # 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]
-        
-        
+
         # Ensure the number of original and target images match
         assert len(self.originals) == len(self.targets), "Mismatch in number of original and target images."
 
@@ -47,13 +43,13 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         return len(self.originals)
 
     def __getitem__(self, idx):
-        # Load original and target images for the given index
+        # 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 = Image.open(original_img).convert('RGB')
-        target = Image.open(target_img).convert('RGB')
+        original = original_img.convert('RGB')
+        target = target_img.convert('RGB')
 
         # Return transformed original and target images
         return transform(original), transform(target)