Update app.py

app.py

@@ -31,7 +31,7 @@ class Pix2PixDataset(torch.utils.data.Dataset):
         # 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
-        print(ds)
+        
         self.originals = [x for x in ds["train"] if x['label'] == 0]
         self.targets = [x for x in ds["train"] if x['label'] == 1]
         
@@ -62,7 +62,7 @@ class Pix2PixDataset(torch.utils.data.Dataset):
 def train_model(epochs):
     # Load the dataset
     ds = load_dataset(dataset_id)
-
+    print(f"ds{ds}")
     # Transform function to resize and convert to tensor
     transform = transforms.Compose([
         transforms.Resize((IMG_SIZE, IMG_SIZE)),
@@ -71,6 +71,7 @@ def train_model(epochs):
 
 
     dataset = Pix2PixDataset(ds)
+    print(f"dataset{dataset}")
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
     # Initialize model, loss function, and optimizer