Update app.py

app.py

@@ -29,11 +29,11 @@ model_repo_id = "K00B404/pix2pix_flux"
 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']
-        print(f"nr images{len(ds["train"])}")
-        print(f"nr 170{ds["train"][170]}")
-        print(f"nr 171{ds["train"][171]}")
+        #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
+        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."