Update app.py

app.py

@@ -25,19 +25,6 @@ LR = 0.0002
 dataset_id = "K00B404/pix2pix_flux_set"
 model_repo_id = "K00B404/pix2pix_flux"
 
-# Training function
-def train_model(epochs):
-    # Load the dataset
-    ds = load_dataset(dataset_id)
-
-    # Transform function to resize and convert to tensor
-    transform = transforms.Compose([
-        transforms.Resize((IMG_SIZE, IMG_SIZE)),
-        transforms.ToTensor(),
-    ])
-
-    # Create dataset and dataloader
-    # Create dataset and dataloader
 class Pix2PixDataset(torch.utils.data.Dataset):
     def __init__(self, ds):
         self.originals = [x for x in ds["train"] if x['label'] == 'original']
@@ -60,6 +47,18 @@ class Pix2PixDataset(torch.utils.data.Dataset):
 
         # Return transformed original and target images
         return transform(original), transform(target)
+        
+# Training function
+def train_model(epochs):
+    # Load the dataset
+    ds = load_dataset(dataset_id)
+
+    # Transform function to resize and convert to tensor
+    transform = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+
 
     dataset = Pix2PixDataset(ds)
     dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)