Add ETA to every training log line + epoch summary

train.py

@@ -6,6 +6,7 @@ Optimized for Colab free tier:
 - Auto-limits large datasets (WikiArt capped at 10K by default)
 - Latent pre-caching: train on pure tensors, no VAE during training
 - Gradient checkpointing + auto batch size = no OOM
+- ETA shown on every log line
 - All datasets pure parquet, open SDXL VAE (no login)
 """
 
@@ -28,7 +29,7 @@ DATASET_PRESETS = {
         "image_column": "image",
         "label_column": "",
         "num_classes": 0,
-        "max_default": 0,  # 0 = use all (~2.5K, small enough)
+        "max_default": 0,
         "description": "~2.5K cartoon/anime, unconditional, 181MB — fast",
     },
     "flowers": {
@@ -46,7 +47,7 @@ DATASET_PRESETS = {
         "image_column": "image",
         "label_column": "style",
         "num_classes": 0,
-        "max_default": 10000,  # Auto-cap: 105K is too many for Colab encoding
+        "max_default": 10000,
         "description": "~105K paintings with styles (auto-capped to 10K for speed)",
     },
     "art_painting": {
@@ -62,7 +63,6 @@ DATASET_PRESETS = {
 
 
 def auto_batch_size(model_size, image_size, gpu_mem_gb):
-    """Safe batch size for model + resolution + GPU."""
     param_mem = {"small": 0.66, "base": 1.68, "large": 3.35}
     base = param_mem.get(model_size, 1.0)
     act_per_sample = {"small": {256: 0.02, 512: 0.07},
@@ -78,6 +78,13 @@ def auto_batch_size(model_size, image_size, gpu_mem_gb):
     return max(1, bs)
 
 
+def _fmt_time(seconds):
+    """Format seconds into human readable string."""
+    if seconds < 60: return f"{seconds:.0f}s"
+    if seconds < 3600: return f"{seconds/60:.1f}m"
+    return f"{seconds/3600:.1f}h"
+
+
 @dataclass
 class TrainConfig:
     model_size: str = "small"
@@ -85,11 +92,11 @@ class TrainConfig:
     class_drop_prob: float = 0.1
     dataset_preset: str = "cartoon"
     image_size: int = 256
-    max_images: int = 0            # 0 = use dataset's default cap
+    max_images: int = 0
     vae_id: str = "madebyollin/sdxl-vae-fp16-fix"
     vae_scaling_factor: float = 0.13025
     latent_channels: int = 4
-    batch_size: int = 0            # 0 = auto
+    batch_size: int = 0
     gradient_accumulation_steps: int = 1
     learning_rate: float = 1e-4
     weight_decay: float = 0.01
@@ -175,7 +182,6 @@ def precache_latents(config, cache_path=None):
         transforms.CenterCrop(config.image_size), transforms.ToTensor(),
     ])
 
-    # Determine max images: user override > dataset default > all
     if config.max_images > 0:
         max_imgs = config.max_images
     elif preset.get("max_default", 0) > 0:
@@ -184,10 +190,9 @@ def precache_latents(config, cache_path=None):
     else:
         max_imgs = len(dataset)
     max_imgs = min(max_imgs, len(dataset))
-    print(f"  Encoding {max_imgs} of {len(dataset)} images")
 
-    # VAE encode batch size: bigger = faster. 64 for 256px, 32 for 512px
     encode_bs = 64 if config.image_size <= 256 else 32
+    print(f"  Encoding {max_imgs} images (batch={encode_bs})...")
 
     img_col, lbl_col = preset["image_column"], preset["label_column"]
     style_to_id = {}
@@ -220,7 +225,7 @@ def precache_latents(config, cache_path=None):
             speed = count / elapsed
             eta = (max_imgs - count) / speed if speed > 0 else 0
             if count % (encode_bs * 4) == 0:
-                print(f"  {count}/{max_imgs} ({speed:.0f} img/s, ~{eta:.0f}s left)")
+                print(f"  {count}/{max_imgs} | {speed:.0f} img/s | ETA {_fmt_time(eta)}")
 
     if batch_px:
         with torch.no_grad():
@@ -237,8 +242,7 @@ def precache_latents(config, cache_path=None):
         print(f"  {len(style_to_id)} style classes")
     torch.save(save_data, cache_path)
     mb = os.path.getsize(cache_path) / 1024**2
-    elapsed = time.time() - t0
-    print(f"Cached {count} latents -> {cache_path} ({mb:.0f}MB, {elapsed:.0f}s)")
+    print(f"Cached {count} latents -> {cache_path} ({mb:.0f}MB, {_fmt_time(time.time()-t0)})")
     del vae
     if torch.cuda.is_available(): torch.cuda.empty_cache()
     return cache_path
@@ -330,7 +334,7 @@ def train(config):
     scaler = GradScaler("cuda", enabled=config.mixed_precision and torch.cuda.is_available())
     fm = FlowMatchingScheduler(config.min_timestep, config.max_timestep)
     lat_size = config.image_size // 8
-    print(f"Steps: {total_steps}, Batch: {config.batch_size}")
+    print(f"Steps: {total_steps} | Batch: {config.batch_size} | Epochs: {config.num_epochs}")
 
     gs = 0; la = 0.0; vae = None; vae_loaded = False
     print(f"\nTraining!\n")
@@ -355,10 +359,14 @@ def train(config):
                 ema.update(model); gs += 1
                 if gs % config.log_every_n_steps == 0:
                     al = la / config.log_every_n_steps
+                    elapsed = time.time() - t_start
+                    sps = gs / max(elapsed, 1)
+                    remaining = (total_steps - gs) / sps if sps > 0 else 0
                     vram = torch.cuda.memory_allocated()/1024**3 if torch.cuda.is_available() else 0
-                    print(f"step={gs:>6d} | ep={epoch} | loss={al:.4f} | gn={gn:.2f} | "
-                          f"lr={opt.param_groups[0]['lr']:.2e} | vram={vram:.1f}G | "
-                          f"{gs/max(time.time()-t_start,1):.1f} st/s")
+                    pct = gs / total_steps * 100
+                    print(f"step={gs:>6d}/{total_steps} ({pct:.0f}%) | ep={epoch} | "
+                          f"loss={al:.4f} | gn={gn:.2f} | lr={opt.param_groups[0]['lr']:.2e} | "
+                          f"vram={vram:.1f}G | {sps:.1f} st/s | ETA {_fmt_time(remaining)}")
                     la = 0.0
                     if math.isnan(al) or al > 50: print("Diverged!"); return
                 if gs % config.sample_every_n_steps == 0:
@@ -380,8 +388,11 @@ def train(config):
                     torch.save({"model": model.state_dict(), "ema": ema.shadow,
                                 "optimizer": opt.state_dict(), "step": gs, "model_config": mcfg},
                                f"{config.output_dir}/checkpoints/step_{gs:07d}.pt")
-        print(f"Epoch {epoch} | {time.time()-et:.0f}s\n")
+        ep_time = time.time() - et
+        ep_eta = ep_time * (config.num_epochs - epoch - 1)
+        print(f"Epoch {epoch}/{config.num_epochs} done | {_fmt_time(ep_time)} | ETA {_fmt_time(ep_eta)}\n")
 
     final = f"{config.output_dir}/checkpoints/final.pt"
     torch.save({"model": model.state_dict(), "ema": ema.shadow, "model_config": mcfg, "step": gs}, final)
-    print(f"\nDone! {gs} steps, {(time.time()-t_start)/60:.1f}min -> {final}")
+    total_time = time.time() - t_start
+    print(f"\nDone! {gs} steps in {_fmt_time(total_time)} -> {final}")