Add v3 training script (self-contained, works with HF Jobs)

train_v3.py

@@ -0,0 +1,552 @@
+"""
+BokehFlow v3 Training Script
+Trains on RealBokeh_3MP dataset (timseizinger/RealBokeh_3MP)
+
+Self-contained — all model code is inline so this works as a standalone
+script in HF Jobs or any GPU environment.
+
+Usage:
+  # Quick test (200 scenes, 3 epochs)
+  VARIANT=small MAX_SCENES=200 EPOCHS=3 BATCH_SIZE=4 python train_v3.py
+
+  # Full training (all 3960 scenes, 10 epochs)
+  VARIANT=small EPOCHS=10 BATCH_SIZE=8 python train_v3.py
+
+Environment variables:
+  VARIANT:       nano/small/base (default: small)
+  MAX_SCENES:    limit scenes for testing (default: 0 = all)
+  EPOCHS:        number of epochs (default: 10)
+  BATCH_SIZE:    batch size (default: 4)
+  CROP_SIZE:     random crop size (default: 256)
+  LR:            learning rate (default: 2e-4)
+  HUB_MODEL_ID:  HF model repo to push to (default: asdf98/BokehFlow)
+
+Requirements:
+  pip install torch torchvision Pillow huggingface_hub trackio aiohttp
+"""
+
+import os, sys, time, json, math, random, glob
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+from dataclasses import dataclass
+
+
+# ===================================================================
+# Model (inline — identical to bokehflow_v3.py)
+# ===================================================================
+
+@dataclass
+class BokehFlowConfig:
+    variant: str = "small"
+    embed_dim: int = 96
+    depth_blocks: int = 6
+    bokeh_blocks: int = 6
+    fusion_every: int = 2
+    stem_channels: int = 48
+    patch_stride: int = 4
+    max_coc_radius: int = 31
+    num_depth_layers: int = 8
+    aperture_embed_dim: int = 64
+    dropout: float = 0.0
+    sensor_width_mm: float = 36.0
+    default_focal_mm: float = 50.0
+    default_fnumber: float = 2.0
+    default_focus_m: float = 2.0
+    ffn_expansion: int = 2
+    large_kernel: int = 7
+
+    def __post_init__(self):
+        if self.variant == "nano":
+            self.embed_dim = 48
+            self.depth_blocks = 4
+            self.bokeh_blocks = 4
+        elif self.variant == "small":
+            self.embed_dim = 96
+            self.depth_blocks = 6
+            self.bokeh_blocks = 6
+        elif self.variant == "base":
+            self.embed_dim = 192
+            self.depth_blocks = 8
+            self.bokeh_blocks = 8
+
+
+class GatedConvRecurrence(nn.Module):
+    def __init__(self, dim, kernel_size=7, ffn_expansion=2):
+        super().__init__()
+        k = kernel_size; p = k // 2
+        self.norm1 = nn.GroupNorm(8, dim)
+        self.dw1 = nn.Conv2d(dim, dim, k, padding=p, groups=dim, bias=False)
+        self.dw2 = nn.Conv2d(dim, dim, k, padding=p, groups=dim, bias=False)
+        self.pw = nn.Conv2d(dim, dim, 1, bias=False)
+        self.gate_proj = nn.Conv2d(dim, dim, 1, bias=True)
+        self.norm2 = nn.GroupNorm(8, dim)
+        h = dim * ffn_expansion
+        self.ffn = nn.Sequential(nn.Conv2d(dim, h, 1, bias=False), nn.GELU(), nn.Conv2d(h, dim, 1, bias=False))
+        nn.init.zeros_(self.pw.weight)
+        nn.init.zeros_(self.ffn[-1].weight)
+
+    def forward(self, x):
+        h = self.norm1(x)
+        spatial = self.dw2(F.silu(self.dw1(h)))
+        spatial = self.pw(spatial)
+        gate = torch.sigmoid(self.gate_proj(h))
+        x = x + spatial * gate
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class GatedConvRecurrenceWithACFM(GatedConvRecurrence):
+    def __init__(self, dim, kernel_size=7, ffn_expansion=2, aperture_embed_dim=64):
+        super().__init__(dim, kernel_size, ffn_expansion)
+        self.acfm = nn.Linear(aperture_embed_dim, dim * 2)
+        nn.init.zeros_(self.acfm.weight)
+        self.acfm.bias.data[:dim] = 1.0
+        self.acfm.bias.data[dim:] = 0.0
+
+    def forward(self, x, aperture_embed=None):
+        x = super().forward(x)
+        if aperture_embed is not None:
+            B, C, H, W = x.shape
+            ss = self.acfm(aperture_embed)
+            scale = ss[:, :C].view(B, C, 1, 1)
+            shift = ss[:, C:].view(B, C, 1, 1)
+            x = x * scale + shift
+        return x
+
+
+class ConvStem(nn.Module):
+    def __init__(self, in_ch=3, stem_ch=48, embed_dim=96):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_ch, stem_ch, 7, stride=2, padding=3, bias=False),
+            nn.GroupNorm(8, stem_ch), nn.GELU(),
+            nn.Conv2d(stem_ch, stem_ch, 3, stride=2, padding=1, groups=stem_ch, bias=False),
+            nn.Conv2d(stem_ch, embed_dim, 1, bias=False),
+            nn.GroupNorm(8, embed_dim), nn.GELU())
+    def forward(self, x): return self.net(x)
+
+
+class ApertureEncoder(nn.Module):
+    def __init__(self, embed_dim=64):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.Linear(3, embed_dim), nn.GELU(), nn.Linear(embed_dim, embed_dim), nn.GELU())
+        self.register_buffer('p_min', torch.tensor([1., 10., 0.1]))
+        self.register_buffer('p_max', torch.tensor([22., 200., 100.]))
+    def forward(self, f_number, focal_mm, focus_m):
+        p = torch.stack([f_number, focal_mm, focus_m], -1)
+        return self.mlp(((p - self.p_min) / (self.p_max - self.p_min + 1e-6)).clamp(0,1))
+
+
+class CrossFusion(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.gate_d = nn.Sequential(nn.Conv2d(d, d, 1), nn.Sigmoid())
+        self.gate_b = nn.Sequential(nn.Conv2d(d, d, 1), nn.Sigmoid())
+        self.proj_d = nn.Conv2d(d, d, 1, bias=False)
+        self.proj_b = nn.Conv2d(d, d, 1, bias=False)
+        nn.init.zeros_(self.proj_d.weight); nn.init.zeros_(self.proj_b.weight)
+    def forward(self, d_feat, b_feat):
+        return (d_feat + self.gate_d(b_feat) * self.proj_d(b_feat),
+                b_feat + self.gate_b(d_feat) * self.proj_b(d_feat))
+
+
+class DepthHead(nn.Module):
+    def __init__(self, dim=96):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, dim//2, 3, padding=1), nn.GELU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(dim//2, dim//4, 3, padding=1), nn.GELU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(dim//4, 1, 3, padding=1), nn.Softplus())
+    def forward(self, x): return self.net(x).clamp(max=100.0)
+
+
+class BokehHead(nn.Module):
+    def __init__(self, dim=96):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim, dim, 3, padding=1), nn.GELU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(dim, dim//2, 3, padding=1), nn.GELU(),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
+            nn.Conv2d(dim//2, 3, 3, padding=1))
+    def forward(self, x): return self.net(x)
+
+
+class PGCoC(nn.Module):
+    def __init__(self, sensor_width=36.0, max_radius=31, n_levels=5):
+        super().__init__()
+        self.sensor_width = sensor_width
+        self.max_radius = max_radius
+        self.n_levels = n_levels
+        self.kernels = nn.ParameterList()
+        for i in range(n_levels):
+            sigma = (i + 1) * max_radius / n_levels / 3.0
+            ks = int(sigma * 6) | 1; ks = max(ks, 3); ks = min(ks, 31)
+            k1d = torch.exp(-torch.arange(-(ks//2), ks//2+1).float()**2 / (2*sigma**2+1e-6))
+            k1d = k1d / k1d.sum()
+            k2d = k1d.unsqueeze(1) @ k1d.unsqueeze(0)
+            self.kernels.append(nn.Parameter(k2d.unsqueeze(0).unsqueeze(0), requires_grad=False))
+        self.refine = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.GELU(), nn.Conv2d(16, 3, 3, padding=1))
+
+    def _blur_at_level(self, image, kernel):
+        B, C, H, W = image.shape
+        k = kernel.expand(C, -1, -1, -1)
+        p = kernel.shape[-1] // 2
+        return F.conv2d(F.pad(image, [p]*4, mode='reflect'), k, groups=C)
+
+    def forward(self, image, depth, f_number, focal_mm, focus_m):
+        B, C, H, W = image.shape
+        f = focal_mm.view(-1,1,1,1); N = f_number.view(-1,1,1,1)
+        S1 = (focus_m.view(-1,1,1,1) * 1000).clamp(min=51)
+        D = (depth * 1000).clamp(min=100)
+        coc = (f**2 / (N * (S1 - f).clamp(min=1))) * (D - S1).abs() / D
+        coc_px = (coc * W / self.sensor_width / 2).clamp(0, self.max_radius)
+        coc_norm = coc_px / self.max_radius
+        blurred_levels = [self._blur_at_level(image, kernel) for kernel in self.kernels]
+        level_float = coc_norm * (self.n_levels - 1)
+        level_low = level_float.long().clamp(0, self.n_levels - 2)
+        level_frac = (level_float - level_low.float()).clamp(0, 1)
+        rendered = image.clone()
+        for lv in range(self.n_levels - 1):
+            mask = (level_low == lv).float()
+            if mask.sum() > 0:
+                interp = blurred_levels[lv] * (1 - level_frac) + blurred_levels[lv + 1] * level_frac
+                rendered = rendered * (1 - mask) + interp * mask
+        mask_top = (level_low >= self.n_levels - 2).float() * (level_frac > 0.99).float()
+        rendered = rendered * (1 - mask_top) + blurred_levels[-1] * mask_top
+        rendered = rendered + self.refine(rendered) * 0.1
+        return rendered, coc_px
+
+
+class BokehFlow(nn.Module):
+    def __init__(self, config=None):
+        super().__init__()
+        if config is None: config = BokehFlowConfig()
+        self.config = config; c = config
+        self.stem = ConvStem(3, c.stem_channels, c.embed_dim)
+        self.aperture_enc = ApertureEncoder(c.aperture_embed_dim)
+        self.depth_blocks = nn.ModuleList([
+            GatedConvRecurrence(c.embed_dim, c.large_kernel, c.ffn_expansion)
+            for _ in range(c.depth_blocks)])
+        self.bokeh_blocks = nn.ModuleList([
+            GatedConvRecurrenceWithACFM(c.embed_dim, c.large_kernel, c.ffn_expansion, c.aperture_embed_dim)
+            for _ in range(c.bokeh_blocks)])
+        n_fusions = max(c.depth_blocks, c.bokeh_blocks) // c.fusion_every
+        self.fusions = nn.ModuleList([CrossFusion(c.embed_dim) for _ in range(n_fusions)])
+        self.depth_head = DepthHead(c.embed_dim)
+        self.bokeh_head = BokehHead(c.embed_dim)
+        self.pgcoc = PGCoC(c.sensor_width_mm, c.max_coc_radius)
+        self.blend_w = nn.Parameter(torch.tensor(0.5))
+
+    def forward(self, image, f_number=None, focal_length_mm=None, focus_distance_m=None, **kw):
+        B = image.shape[0]; dev = image.device; c = self.config
+        if f_number is None: f_number = torch.full((B,), c.default_fnumber, device=dev)
+        if focal_length_mm is None: focal_length_mm = torch.full((B,), c.default_focal_mm, device=dev)
+        if focus_distance_m is None: focus_distance_m = torch.full((B,), c.default_focus_m, device=dev)
+        ae = self.aperture_enc(f_number, focal_length_mm, focus_distance_m)
+        feat = self.stem(image)
+        d_feat = feat; b_feat = feat; fi = 0
+        n_blk = max(c.depth_blocks, c.bokeh_blocks)
+        for i in range(n_blk):
+            if i < c.depth_blocks: d_feat = self.depth_blocks[i](d_feat)
+            if i < c.bokeh_blocks: b_feat = self.bokeh_blocks[i](b_feat, ae)
+            if (i+1) % c.fusion_every == 0 and fi < len(self.fusions):
+                d_feat, b_feat = self.fusions[fi](d_feat, b_feat); fi += 1
+        depth = self.depth_head(d_feat)
+        if depth.shape[2:] != image.shape[2:]:
+            depth = F.interpolate(depth, image.shape[2:], mode='bilinear', align_corners=False)
+        physics_bokeh, coc_map = self.pgcoc(image, depth, f_number, focal_length_mm, focus_distance_m)
+        learned_bokeh = self.bokeh_head(b_feat)
+        if learned_bokeh.shape[2:] != image.shape[2:]:
+            learned_bokeh = F.interpolate(learned_bokeh, image.shape[2:], mode='bilinear', align_corners=False)
+        w = torch.sigmoid(self.blend_w)
+        bokeh = (w * physics_bokeh + (1-w) * (image + learned_bokeh)).clamp(0, 1)
+        return {'bokeh': bokeh, 'depth': depth, 'coc_map': coc_map}
+
+
+class BokehFlowLoss(nn.Module):
+    def forward(self, pred, targets):
+        bp, bg = pred['bokeh'], targets['bokeh_gt']
+        l1 = F.l1_loss(bp, bg)
+        C1, C2 = 0.01**2, 0.03**2
+        mu_p = F.avg_pool2d(bp, 11, 1, 5); mu_g = F.avg_pool2d(bg, 11, 1, 5)
+        mu_pp = mu_p*mu_p; mu_gg = mu_g*mu_g; mu_pg = mu_p*mu_g
+        sig_pp = F.avg_pool2d(bp*bp, 11, 1, 5) - mu_pp
+        sig_gg = F.avg_pool2d(bg*bg, 11, 1, 5) - mu_gg
+        sig_pg = F.avg_pool2d(bp*bg, 11, 1, 5) - mu_pg
+        ssim_map = ((2*mu_pg+C1)*(2*sig_pg+C2)) / ((mu_pp+mu_gg+C1)*(sig_pp+sig_gg+C2))
+        ssim_loss = 1.0 - ssim_map.mean()
+        return {'total': l1 + ssim_loss, 'l1': l1.detach(), 'ssim': ssim_loss.detach()}
+
+
+# ===================================================================
+# Dataset
+# ===================================================================
+
+class RealBokehDataset(Dataset):
+    """Loads from local disk after snapshot_download."""
+    def __init__(self, root, crop_size=256, split='train', target_fstop='2.0'):
+        self.crop = crop_size
+        self.pairs = []
+        in_dir = Path(root) / split / 'in'
+        gt_dir = Path(root) / split / 'gt'
+        meta_dir = Path(root) / split / 'metadata'
+
+        for in_path in sorted(in_dir.glob('*_f22.JPG')):
+            sid = in_path.stem.split('_')[0]
+            gt_path = gt_dir / sid / f'{sid}_f{target_fstop}.JPG'
+            meta_path = meta_dir / f'{sid}.json'
+            if gt_path.exists():
+                meta = {}
+                if meta_path.exists():
+                    with open(meta_path) as f:
+                        meta = json.load(f)
+                self.pairs.append((str(in_path), str(gt_path), meta))
+
+        print(f"RealBokehDataset: {len(self.pairs)} pairs found (target f/{target_fstop})")
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def __getitem__(self, idx):
+        from PIL import Image
+        import torchvision.transforms.functional as TF
+
+        in_path, gt_path, meta = self.pairs[idx]
+        inp = Image.open(in_path).convert('RGB')
+        gt = Image.open(gt_path).convert('RGB')
+
+        # Resize to manageable size first, then crop
+        short = min(inp.size)
+        if short > 512:
+            scale = 512.0 / short
+            new_w = int(inp.size[0] * scale)
+            new_h = int(inp.size[1] * scale)
+            inp = inp.resize((new_w, new_h), Image.LANCZOS)
+            gt = gt.resize((new_w, new_h), Image.LANCZOS)
+
+        inp = TF.to_tensor(inp)
+        gt = TF.to_tensor(gt)
+
+        # Random crop
+        _, h, w = inp.shape
+        cs = self.crop
+        if h >= cs and w >= cs:
+            i = random.randint(0, h - cs)
+            j = random.randint(0, w - cs)
+            inp = inp[:, i:i+cs, j:j+cs]
+            gt = gt[:, i:i+cs, j:j+cs]
+        else:
+            inp = F.interpolate(inp.unsqueeze(0), (cs, cs), mode='bilinear', align_corners=False).squeeze(0)
+            gt = F.interpolate(gt.unsqueeze(0), (cs, cs), mode='bilinear', align_corners=False).squeeze(0)
+
+        # Random horizontal flip
+        if random.random() > 0.5:
+            inp = inp.flip(-1)
+            gt = gt.flip(-1)
+
+        focal = meta.get('focal_length', 50.0)
+        focus = meta.get('focus_plane_distance', 2.0)
+        fnum = 2.0
+
+        return {
+            'image': inp,
+            'bokeh_gt': gt,
+            'f_number': torch.tensor(fnum, dtype=torch.float32),
+            'focal_length_mm': torch.tensor(float(focal), dtype=torch.float32),
+            'focus_distance_m': torch.tensor(float(focus), dtype=torch.float32),
+        }
+
+
+# ===================================================================
+# Data download
+# ===================================================================
+
+def download_realbokeh(max_scenes=None):
+    """Download RealBokeh_3MP using snapshot_download with exact patterns."""
+    from huggingface_hub import snapshot_download
+
+    data_dir = '/tmp/realbokeh'
+    check_file = Path(data_dir) / 'train' / 'in' / '1_f22.JPG'
+    if check_file.exists():
+        n = len(list(Path(data_dir).glob('train/in/*_f22.JPG')))
+        print(f"Data already cached: {n} scenes")
+        return data_dir
+
+    print("Fetching metadata to build download list...")
+    snapshot_download(
+        'timseizinger/RealBokeh_3MP',
+        repo_type='dataset',
+        local_dir=data_dir,
+        allow_patterns=['train/metadata/*.json'],
+    )
+
+    meta_dir = Path(data_dir) / 'train' / 'metadata'
+    scene_ids = sorted([p.stem for p in meta_dir.glob('*.json')], key=lambda x: int(x))
+
+    if max_scenes:
+        scene_ids = scene_ids[:max_scenes]
+
+    print(f"Found {len(scene_ids)} scenes. Downloading input + f/2.0 GT images...")
+
+    patterns = []
+    for sid in scene_ids:
+        patterns.append(f'train/in/{sid}_f22.JPG')
+        patterns.append(f'train/gt/{sid}/{sid}_f2.0.JPG')
+
+    t0 = time.time()
+    snapshot_download(
+        'timseizinger/RealBokeh_3MP',
+        repo_type='dataset',
+        local_dir=data_dir,
+        allow_patterns=patterns,
+    )
+    elapsed = time.time() - t0
+    n_found = len(list(Path(data_dir).glob('train/in/*_f22.JPG')))
+    print(f"Downloaded {n_found} scenes in {elapsed:.0f}s")
+    return data_dir
+
+
+# ===================================================================
+# Training loop
+# ===================================================================
+
+def train():
+    import trackio
+
+    VARIANT = os.environ.get('VARIANT', 'small')
+    MAX_SCENES = int(os.environ.get('MAX_SCENES', '0')) or None
+    EPOCHS = int(os.environ.get('EPOCHS', '10'))
+    BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '4'))
+    CROP_SIZE = int(os.environ.get('CROP_SIZE', '256'))
+    LR = float(os.environ.get('LR', '2e-4'))
+    HUB_MODEL_ID = os.environ.get('HUB_MODEL_ID', 'asdf98/BokehFlow')
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"Device: {device}")
+    if device == 'cuda':
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+
+    trackio.init(project="bokehflow", name=f"v3-{VARIANT}-e{EPOCHS}-lr{LR}")
+
+    data_dir = download_realbokeh(max_scenes=MAX_SCENES)
+
+    ds = RealBokehDataset(data_dir, crop_size=CROP_SIZE)
+    dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4,
+                    pin_memory=True, drop_last=True, persistent_workers=True)
+    print(f"Batches per epoch: {len(dl)}")
+
+    config = BokehFlowConfig(variant=VARIANT)
+    model = BokehFlow(config).to(device)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Model: BokehFlow-{VARIANT}, {n_params:,} params")
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)
+    total_steps = EPOCHS * len(dl)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, total_steps, eta_min=LR/20)
+    loss_fn = BokehFlowLoss()
+
+    scaler = torch.amp.GradScaler('cuda', enabled=(device == 'cuda'))
+
+    global_step = 0
+    best_loss = float('inf')
+
+    for epoch in range(EPOCHS):
+        model.train()
+        epoch_loss = 0.0
+        t_epoch = time.time()
+
+        for batch_idx, batch in enumerate(dl):
+            t_step = time.time()
+            image = batch['image'].to(device)
+            bokeh_gt = batch['bokeh_gt'].to(device)
+            f_number = batch['f_number'].to(device)
+            focal_mm = batch['focal_length_mm'].to(device)
+            focus_m = batch['focus_distance_m'].to(device)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            with torch.amp.autocast('cuda', enabled=(device == 'cuda')):
+                out = model(image, f_number, focal_mm, focus_m)
+                losses = loss_fn(out, {'bokeh_gt': bokeh_gt})
+                loss = losses['total']
+
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+
+            epoch_loss += loss.item()
+            global_step += 1
+            step_time = time.time() - t_step
+
+            if global_step % 10 == 0 or batch_idx == 0:
+                lr_now = scheduler.get_last_lr()[0]
+                print(f"Ep {epoch+1}/{EPOCHS} [{batch_idx+1}/{len(dl)}] "
+                      f"loss={loss.item():.4f} l1={losses['l1'].item():.4f} "
+                      f"ssim={losses['ssim'].item():.4f} lr={lr_now:.2e} "
+                      f"step={step_time*1000:.0f}ms")
+                trackio.log({
+                    "loss": loss.item(),
+                    "l1": losses['l1'].item(),
+                    "ssim_loss": losses['ssim'].item(),
+                    "lr": lr_now,
+                    "step_ms": step_time * 1000,
+                    "epoch": epoch + 1,
+                })
+
+            if device == 'cuda' and global_step == 1:
+                vram = torch.cuda.max_memory_allocated() / 1e9
+                print(f"Peak VRAM after 1st step: {vram:.2f} GB")
+                trackio.log({"peak_vram_gb": vram})
+
+        epoch_time = time.time() - t_epoch
+        avg_loss = epoch_loss / len(dl)
+        print(f"Epoch {epoch+1}/{EPOCHS} done in {epoch_time:.0f}s, avg_loss={avg_loss:.4f}")
+        trackio.log({"epoch_avg_loss": avg_loss, "epoch_time_s": epoch_time})
+
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save({
+                'model_state_dict': model.state_dict(),
+                'config': config.__dict__,
+                'epoch': epoch + 1,
+                'loss': avg_loss,
+            }, '/tmp/bokehflow_best.pt')
+            print(f"  Saved best model (loss={avg_loss:.4f})")
+
+    # Push to hub
+    print("\nPushing model to Hub...")
+    from huggingface_hub import HfApi
+    api = HfApi()
+
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'config': config.__dict__,
+        'epoch': EPOCHS,
+        'loss': avg_loss,
+    }, '/tmp/bokehflow_final.pt')
+
+    for fname in ['bokehflow_best.pt', 'bokehflow_final.pt']:
+        fpath = f'/tmp/{fname}'
+        if os.path.exists(fpath):
+            api.upload_file(
+                path_or_fileobj=fpath,
+                path_in_repo=f'checkpoints/{fname}',
+                repo_id=HUB_MODEL_ID,
+            )
+            print(f"  Uploaded {fname}")
+
+    print(f"\nTraining complete! Model: https://huggingface.co/{HUB_MODEL_ID}")
+
+
+if __name__ == '__main__':
+    train()