Preview of phase 2-5 implementation; needs a full check

generator/src/data/dataset.py

@@ -51,17 +51,20 @@ class GeneratorDataset(Dataset):
         return img
 
 
-def get_transform(image_size: int, augment: bool = False) -> T.Compose:
+def get_transform(image_size: int, augment=False) -> T.Compose:
     """Build transform for generator training. Output is in [-1, 1].
 
-    augment=True adds horizontal flip + mild rotation + mild color jitter.
-    Use augment=False for validation / FID real-image sets.
+    augment=False    — no augmentation (for FID real-image sets)
+    augment="hflip"  — horizontal flip only (recommended for VAE/DDPM)
+    augment=True     — H-flip + rotation ±5° + mild color jitter (for GAN)
     """
     ops = [
         T.Resize(image_size),
         T.CenterCrop(image_size),
     ]
-    if augment:
+    if augment == "hflip":
+        ops.append(T.RandomHorizontalFlip(p=0.5))
+    elif augment:
         ops += [
             T.RandomHorizontalFlip(p=0.5),
             T.RandomRotation(degrees=5, interpolation=T.InterpolationMode.BILINEAR),