Upload train_vae.py with huggingface_hub

train_vae.py

@@ -0,0 +1,729 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch.amp import autocast, GradScaler
+from torch.utils.tensorboard import SummaryWriter
+import numpy as np
+from PIL import Image
+from glob import glob
+from time import time
+import argparse
+import logging
+import os
+import json
+
+from models import AutoencoderKL, DiT_models
+from custom_dataset import StyleTransferDataset, create_style_transfer_dataloader
+
+
+#################################################################################
+#                             VAE loss function                                 #
+#################################################################################
+
+class SSIMLoss(nn.Module):
+    def __init__(self, window_size=11, size_average=True):
+        super(SSIMLoss, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = 1
+        self.window = self.create_window(window_size, self.channel)
+    
+    def gaussian(self, window_size, sigma):
+        gauss = torch.Tensor([np.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+        return gauss/gauss.sum()
+    
+    def create_window(self, window_size, channel):
+        _1D_window = self.gaussian(window_size, 1.5).unsqueeze(1)
+        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+        return window
+    
+    def _ssim(self, img1, img2, window, window_size, channel, size_average=True):
+        mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel)
+        mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel)
+        
+        mu1_sq = mu1.pow(2)
+        mu2_sq = mu2.pow(2)
+        mu1_mu2 = mu1 * mu2
+        
+        sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1_sq
+        sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2_sq
+        sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1_mu2
+        
+        C1 = 0.01**2
+        C2 = 0.03**2
+        
+        ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+        
+        if size_average:
+            return ssim_map.mean()
+        else:
+            return ssim_map.mean(1).mean(1).mean(1)
+    
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+        
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = self.create_window(self.window_size, channel)
+            
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+            
+            self.window = window
+            self.channel = channel
+        
+        ssim_val = self._ssim(img1, img2, window, self.window_size, channel, self.size_average)
+        return 1 - ssim_val  
+
+
+class VAELoss(nn.Module):
+    def __init__(
+        self,
+        kl_weight=1e-6,
+        l1_weight=1.0,
+        ssim_weight=1.0,
+    ):
+        super().__init__()
+        self.kl_weight = kl_weight
+        self.l1_weight = l1_weight
+        self.ssim_weight = ssim_weight
+
+        self.l1_loss = nn.L1Loss()
+        self.ssim_loss = SSIMLoss()
+    
+    def forward(self, recon, target, posterior):
+
+        l1_loss = self.l1_loss(recon, target)
+        
+        # Convert from [-1, 1] to [0, 1] for SSIM calculation
+        # SSIM constants (C1, C2) are designed for [0, 1] range
+        recon_01 = (recon + 1.0) / 2.0
+        target_01 = (target + 1.0) / 2.0
+        ssim_loss = self.ssim_loss(recon_01, target_01)
+        
+        kl_loss = posterior.kl().mean()
+        
+        total_loss = (
+            self.l1_weight * l1_loss + 
+            self.ssim_weight * ssim_loss + 
+            self.kl_weight * kl_loss
+        )
+        
+        return {
+            'total_loss': total_loss,
+            'l1_loss': self.l1_weight * l1_loss ,
+            'ssim_loss': self.ssim_weight * ssim_loss,
+            'kl_loss': self.kl_weight * kl_loss,
+        }
+
+
+#################################################################################
+#                             Training Helper Functions                         #
+#################################################################################
+
+def create_logger(experiment_dir):
+    if experiment_dir is not None:
+        logging.basicConfig(
+            level=logging.INFO,
+            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[
+                logging.StreamHandler(),
+                logging.FileHandler(f"{experiment_dir}/log.txt")
+            ]
+        )
+        logger = logging.getLogger(__name__)
+    else:
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def cleanup():
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def get_lr_scheduler(optimizer, args, steps_per_epoch):
+    if args.lr_scheduler == 'none':
+        return None
+    
+    total_steps = args.epochs * steps_per_epoch
+    warmup_steps = args.warmup_epochs * steps_per_epoch
+    
+    if args.lr_scheduler == 'linear':
+        # Warmup + Linear Decay
+        def lr_lambda(current_step):
+            if current_step < warmup_steps:
+                return float(current_step) / float(max(1, warmup_steps))
+            else:
+                progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+                return max(0.0, 1.0 - progress)
+        
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    
+    elif args.lr_scheduler == 'cosine':
+        # Warmup + Cosine Decay
+        def lr_lambda(current_step):
+            if current_step < warmup_steps:
+                return float(current_step) / float(max(1, warmup_steps))
+            else:
+                progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+                return max(0.0, 0.5 * (1.0 + np.cos(np.pi * progress)))
+        
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    
+    elif args.lr_scheduler == 'constant':
+        # Warmup + Constant
+        def lr_lambda(current_step):
+            if current_step < warmup_steps:
+                return float(current_step) / float(max(1, warmup_steps))
+            else:
+                return 1.0
+        
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    
+    else:
+        raise ValueError(f"Unknown lr_scheduler: {args.lr_scheduler}")
+
+
+@torch.no_grad()
+def save_samples(vae, dataloader, device, save_dir, num_samples=8, patches_per_image=4, patch_size=512, is_conditional=False, use_fp16=False, vae_domain1=False, vae_domain2=False, multiscale=False, multiscale_levels=None):
+    vae.eval()
+    os.makedirs(save_dir, exist_ok=True)
+    
+    if multiscale_levels is None:
+        multiscale_levels = [32, 64, 128, 256, 512]
+    
+    saved_count = 0
+    for a_large_images, b_large_images, paths in dataloader:
+        a_large_images = a_large_images.to(device)
+        b_large_images = b_large_images.to(device)
+        # multiscale_sizes=None disables multiscale, otherwise enables it
+        # If multiscale_levels is provided, automatically enable multiscale
+        if multiscale:
+            multiscale_sizes_param = multiscale_levels if multiscale_levels is not None else [32, 64, 128, 256, 512]
+        elif multiscale_levels is not None:
+            # User provided multiscale_levels without multiscale flag, auto-enable
+            multiscale_sizes_param = multiscale_levels
+        else:
+            multiscale_sizes_param = None
+        if vae_domain1:
+            b_images, _, pos_info = StyleTransferDataset.crop_patches_from_large_images_with_pos(
+                a_large_images,
+                b_large_images,
+                patch_size=patch_size,
+                patches_per_image=patches_per_image,
+                width_norm=15000.0,
+                height_norm=20000.0,
+                multiscale_sizes=multiscale_sizes_param
+            )
+        elif vae_domain2:
+            _, b_images, pos_info = StyleTransferDataset.crop_patches_from_large_images_with_pos(
+                a_large_images,
+                b_large_images,
+                patch_size=patch_size,
+                patches_per_image=patches_per_image,
+                width_norm=15000.0,
+                height_norm=20000.0,
+                multiscale_sizes=multiscale_sizes_param
+            )
+        
+        del a_large_images, b_large_images
+        torch.cuda.empty_cache()
+        
+        batch = b_images[:num_samples - saved_count]
+        pos_batch = pos_info[:num_samples - saved_count] if pos_info is not None else None
+        
+        # Use autocast for inference if fp16 is enabled
+        with autocast('cuda', enabled=use_fp16):
+            if is_conditional and pos_batch is not None:
+                recon, _ = vae(batch, sample_posterior=False, pos_context=pos_batch)
+            else:
+                recon, _ = vae(batch, sample_posterior=False)
+        
+        if vae_domain2:
+            batch_ori = (batch[:, 0:1, :, :] + 1.0) / 2.0 
+            recon_ori = (recon[:, 0:1, :, :] + 1.0) / 2.0 
+            
+            # Z-score denormalization for retar: (normalized * std + mean)
+            # batch_retar = (batch[:, 1:2, :, :] + 1.0) / 2.0 * 9
+            # recon_retar = (recon[:, 1:2, :, :] + 1.0) / 2.0  * 9
+            batch_retar = batch[:, 1:2, :, :] * 11.41 + 5.61
+            recon_retar = recon[:, 1:2, :, :] * 11.41 + 5.61
+            
+            for i in range(batch.shape[0]):
+                orig_ori = (batch_ori[i, 0].cpu().numpy() * 255).astype(np.uint8)
+                recon_ori_img = (recon_ori[i, 0].cpu().numpy() * 255).astype(np.uint8)
+                
+                # Clip retar values to [0, 90] range and scale to [0, 255] for display
+                orig_retar = (batch_retar[i, 0].cpu().numpy() * 255 / 9).astype(np.uint8)
+                recon_retar_img = (recon_retar[i, 0].cpu().numpy() * 255 / 9).astype(np.uint8)
+
+                orig_ori_pil = Image.fromarray(orig_ori)
+                recon_ori_pil = Image.fromarray(recon_ori_img)
+                orig_retar_pil = Image.fromarray(orig_retar)
+                recon_retar_pil = Image.fromarray(recon_retar_img)
+
+                combined = Image.new('L', (orig_ori_pil.width * 4, orig_ori_pil.height))
+                combined.paste(orig_ori_pil, (0, 0))
+                combined.paste(recon_ori_pil, (orig_ori_pil.width, 0))
+                combined.paste(orig_retar_pil, (orig_ori_pil.width * 2, 0))
+                combined.paste(recon_retar_pil, (orig_ori_pil.width * 3, 0))
+
+        elif vae_domain1:
+            batch_a = (batch[:, 0:1, :, :] + 1.0) / 2.0 
+            recon_a = (recon[:, 0:1, :, :] + 1.0) / 2.0 
+
+            for i in range(batch.shape[0]):
+                orig_a = (batch_a[i, 0].cpu().numpy() * 255).astype(np.uint8)
+                recon_a_img = (recon_a[i, 0].cpu().numpy() * 255).astype(np.uint8)
+
+                orig_a_pil = Image.fromarray(orig_a)
+                recon_a_pil = Image.fromarray(recon_a_img)
+
+                combined = Image.new('L', (orig_a_pil.width * 2, orig_a_pil.height))
+                combined.paste(orig_a_pil, (0, 0))
+                combined.paste(recon_a_pil, (orig_a_pil.width, 0))
+                
+
+        combined.save(f"{save_dir}/sample_{saved_count}.png")
+        saved_count += 1
+        if saved_count >= num_samples:
+            break
+    
+    vae.train()
+
+
+
+def main(args):
+    assert torch.cuda.is_available(), "Training requires at least one GPU."
+    
+    # Setup DDP
+    rank = int(os.environ.get("RANK", 0))
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    
+    torch.cuda.set_device(local_rank)
+    device = local_rank
+    
+    dist.init_process_group("nccl")
+    
+    seed = args.global_seed * world_size + rank
+    torch.manual_seed(seed)
+    print(f"Starting rank={rank}, local_rank={local_rank}, seed={seed}, world_size={world_size}.")
+    
+    # Setup experiment folder
+    is_master = (rank == 0)
+    if is_master:
+        os.makedirs(args.results_dir, exist_ok=True)
+        experiment_index = len(glob(f"{args.results_dir}/*"))
+        model_name = args.vae_model if args.vae_model else "VAE-Custom"
+        experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_name}"
+        checkpoint_dir = f"{experiment_dir}/checkpoints"
+        sample_dir = f"{experiment_dir}/samples"
+        tensorboard_dir = f"{experiment_dir}/tensorboard"
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        os.makedirs(sample_dir, exist_ok=True)
+        os.makedirs(tensorboard_dir, exist_ok=True)
+        logger = create_logger(experiment_dir)
+        logger.info(f"Experiment directory created at {experiment_dir}")
+        
+        writer = SummaryWriter(tensorboard_dir)
+        logger.info(f"TensorBoard logs will be saved to {tensorboard_dir}")
+    else:
+        logger = create_logger(None)
+        sample_dir = None
+        writer = None
+    
+    if args.vae_model:
+        if is_master:
+            logger.info(f"Creating VAE model: {args.vae_model}")
+        if args.vae_model not in DiT_models:
+            raise ValueError(f"Unknown VAE model: {args.vae_model}. Available: {[k for k in DiT_models.keys() if k.startswith('VAE')]}")
+        
+        vae_fn = DiT_models[args.vae_model]
+        vae = vae_fn(
+            in_channels=args.in_channels,
+            out_ch=args.out_channels,
+            resolution=args.image_size,
+        ).to(device)
+        if is_master:
+            logger.info(f"Using predefined VAE model: {args.vae_model}")
+
+    else:
+        if is_master:
+            logger.info("Creating VAE model with custom parameters")
+        vae = AutoencoderKL(
+            embed_dim=args.embed_dim,
+            in_channels=args.in_channels,
+            out_ch=args.out_channels,
+            ch=args.ch,
+            ch_mult=tuple(args.ch_mult),
+            num_res_blocks=args.num_res_blocks,
+            attn_resolutions=args.attn_resolutions,
+            dropout=args.dropout,
+            resolution=args.image_size,
+            z_channels=args.z_channels,
+            double_z=args.double_z,
+            use_mid_attn=False, 
+        ).to(device)
+        if is_master:
+            logger.info(f"VAE middle attention: DISABLED (saves ~68GB memory)")
+    
+    if is_master:
+        logger.info(f"VAE Parameters: {sum(p.numel() for p in vae.parameters()):,}")
+    
+    is_conditional_vae = hasattr(vae, 'condition_net')
+    if is_master:
+        if is_conditional_vae:
+            logger.info("✓ Using Conditional VAE with position information")
+        else:
+            logger.info("Using standard VAE (no position conditioning)")
+    
+    
+    vae = DDP(vae, device_ids=[device], find_unused_parameters=True)
+    if is_master:
+        logger.info("Using find_unused_parameters=True to handle attention layers")
+    
+    opt = torch.optim.AdamW(vae.parameters(), lr=args.learning_rate, weight_decay=0.0)
+    criterion = VAELoss(
+        kl_weight=args.kl_weight,
+        l1_weight=args.l1_weight,
+        ssim_weight=args.ssim_weight,
+    ).to(device)
+    
+
+    scaler = GradScaler('cuda', enabled=args.fp16)
+    actual_batch_size = int(args.global_batch_size // dist.get_world_size())
+    world_size = dist.get_world_size()
+    loader = create_style_transfer_dataloader(
+        pairing_json_path=args.data_path,
+        batch_size=actual_batch_size,
+        patch_size=args.image_size,
+        patches_per_image=args.patches_per_image,
+        num_workers=args.num_workers,
+        shuffle=True,
+        drop_last=True,
+        device=device,
+        distributed=(world_size > 1), 
+        rank=rank,
+        world_size=world_size
+    )
+    
+    if is_master:
+        logger.info(f"Dataset contains {len(loader.dataset):,} large images")
+        logger.info(f"Global batch size: {args.global_batch_size}, Actual batch size: {actual_batch_size}")
+        logger.info(f"Patches per image: {args.patches_per_image}")
+    
+    steps_per_epoch = len(loader)
+    scheduler = get_lr_scheduler(opt, args, steps_per_epoch)
+    if is_master:
+        if scheduler:
+            logger.info(f"Using LR scheduler: {args.lr_scheduler} with {args.warmup_epochs} warmup epochs")
+            logger.info(f"Total steps: {args.epochs * steps_per_epoch}, Warmup steps: {args.warmup_epochs * steps_per_epoch}")
+        else:
+            logger.info("No LR scheduler (constant learning rate)")
+    
+    train_steps = 0
+    start_epoch = 0
+    
+    if args.resume:
+        if is_master:
+            logger.info(f"Resuming from checkpoint: {args.resume}")
+        checkpoint = torch.load(args.resume, map_location=f"cuda:{device}", weights_only=False)
+        vae.module.load_state_dict(checkpoint["vae"], strict=False)
+        if is_master:
+            logger.info(f"Note: using strict=False to ignore unexpected keys (e.g., old attention weights)")
+        opt.load_state_dict(checkpoint["opt"])
+        train_steps = checkpoint.get("train_steps", 0)
+
+        if args.start_epoch is not None:
+            start_epoch = args.start_epoch
+            if is_master:
+                logger.info(f"Using manually specified start epoch: {start_epoch}")
+        else:
+            start_epoch = checkpoint.get("epoch", 0)
+            if is_master:
+                if "epoch" in checkpoint:
+                    logger.info(f"Loaded epoch from checkpoint: {start_epoch}")
+                else:
+                    logger.info(f"No epoch info in checkpoint, starting from epoch 0")
+        
+        if scheduler and "scheduler" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler"])
+            if is_master:
+                logger.info(f"Resumed scheduler from step {train_steps}")
+        
+        if is_master:
+            logger.info(f"Resumed from epoch {start_epoch}, step {train_steps}")
+    elif args.start_epoch is not None:
+        start_epoch = args.start_epoch
+        if is_master:
+            logger.info(f"Starting from manually specified epoch: {start_epoch} (without resume)")
+    
+
+    vae.train()
+    running_loss = 0
+    running_l1_loss = 0
+    running_ssim_loss = 0
+    running_kl_loss = 0
+    log_steps = 0
+    start_time = time()
+    
+    if is_master:
+        logger.info(f"Training for {args.epochs} epochs (from epoch {start_epoch} to {args.epochs})...")
+    
+    for epoch in range(start_epoch, args.epochs):
+        if is_master:
+            logger.info(f"Beginning epoch {epoch}...")
+        
+        for a_large_images, b_large_images, paths in loader:
+            # Use non_blocking transfer to overlap data loading with computation
+            a_large_images = a_large_images.to(device, non_blocking=True)
+            b_large_images = b_large_images.to(device, non_blocking=True)
+
+            # Crop patches on GPU (already on GPU from .to(device))
+            if args.multiscale:
+                multiscale_sizes = getattr(args, 'multiscale_levels', [32, 64, 128, 256, 512])
+            elif hasattr(args, 'multiscale_levels') and args.multiscale_levels is not None:
+                multiscale_sizes = args.multiscale_levels
+            else:
+                multiscale_sizes = None
+            if args.vae_domain2:
+                _, b_images, pos_info = StyleTransferDataset.crop_patches_from_large_images_with_pos(
+                    a_large_images, b_large_images,
+                    patch_size=args.image_size,
+                    patches_per_image=args.patches_per_image,
+                    width_norm=15000.0,
+                    height_norm=20000.0,
+                    multiscale_sizes=multiscale_sizes
+                )
+            elif args.vae_domain1:
+                b_images, _, pos_info = StyleTransferDataset.crop_patches_from_large_images_with_pos(
+                    a_large_images,
+                    b_large_images,
+                    patch_size=args.image_size,
+                    patches_per_image=args.patches_per_image,
+                    width_norm=15000.0,
+                    height_norm=20000.0,
+                    multiscale_sizes=multiscale_sizes
+                )
+
+            images = b_images
+            
+            with autocast('cuda', enabled=args.fp16):
+                if is_conditional_vae and pos_info is not None:
+                    recon, posterior = vae(images, sample_posterior=True, pos_context=pos_info)
+                else:
+                    recon, posterior = vae(images, sample_posterior=True)
+                losses = criterion(recon, images, posterior)
+                loss = losses['total_loss']
+
+            # Extract loss values without blocking
+            loss_val = loss.item()
+            l1_loss_val = losses['l1_loss'].item()
+            ssim_loss_val = losses['ssim_loss'].item()
+            kl_loss_val = losses['kl_loss'].item()
+            
+            opt.zero_grad()
+            if args.fp16:
+                scaler.scale(loss).backward()
+                scaler.step(opt)
+                scaler.update()
+            else:
+                loss.backward()
+                opt.step()
+
+            # Accumulate losses locally (no sync needed)
+            running_loss += loss_val
+            running_l1_loss += l1_loss_val
+            running_ssim_loss += ssim_loss_val
+            running_kl_loss += kl_loss_val
+            log_steps += 1
+            train_steps += 1
+            
+            if scheduler:
+                scheduler.step()
+            
+            # Only sync and log periodically (reduces GPU-CPU synchronization overhead)
+            if train_steps % args.log_every == 0:
+                # Synchronize only when logging (not every step)
+                torch.cuda.synchronize()
+                end_time = time()
+                steps_per_sec = log_steps / (end_time - start_time)
+                
+                # Compute local averages
+                local_avg_loss = running_loss / log_steps
+                local_avg_l1 = running_l1_loss / log_steps
+                local_avg_ssim = running_ssim_loss / log_steps
+                local_avg_kl = running_kl_loss / log_steps
+                
+                avg_loss_tensor = torch.tensor(local_avg_loss, device=device)
+                avg_l1_tensor = torch.tensor(local_avg_l1, device=device)
+                avg_ssim_tensor = torch.tensor(local_avg_ssim, device=device)
+                avg_kl_tensor = torch.tensor(local_avg_kl, device=device)
+                
+                dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
+                dist.all_reduce(avg_l1_tensor, op=dist.ReduceOp.SUM)
+                dist.all_reduce(avg_ssim_tensor, op=dist.ReduceOp.SUM)
+                dist.all_reduce(avg_kl_tensor, op=dist.ReduceOp.SUM)
+                
+                avg_loss = avg_loss_tensor.item() / dist.get_world_size()
+                avg_l1 = avg_l1_tensor.item() / dist.get_world_size()
+                avg_ssim = avg_ssim_tensor.item() / dist.get_world_size()
+                avg_kl = avg_kl_tensor.item() / dist.get_world_size()
+                
+                if is_master:
+                    if writer is not None:
+                        writer.add_scalar('Loss/total', avg_loss, train_steps)
+                        writer.add_scalar('Loss/l1', avg_l1, train_steps)
+                        writer.add_scalar('Loss/ssim', avg_ssim, train_steps)
+                        writer.add_scalar('Loss/kl', avg_kl, train_steps)
+                        writer.add_scalar('Training/steps_per_sec', steps_per_sec, train_steps)
+                        writer.add_scalar('Training/learning_rate', opt.param_groups[0]['lr'], train_steps)
+                    
+                    current_lr = opt.param_groups[0]['lr']
+                    logger.info(
+                        f"(step={train_steps:07d}) "
+                        f"Loss: {avg_loss:.4f} | "
+                        f"L1: {avg_l1:.4f} | "
+                        f"SSIM: {avg_ssim:.4f} | "
+                        f"KL: {avg_kl:.6f} | "
+                        f"LR: {current_lr:.2e} | "
+                        f"Steps/Sec: {steps_per_sec:.2f}"
+                    )
+                
+                running_loss = 0
+                running_l1_loss = 0
+                running_ssim_loss = 0
+                running_kl_loss = 0
+                log_steps = 0
+                start_time = time()
+
+            if args.sample_every > 0 and train_steps % args.sample_every == 0 and train_steps > 0:
+                if is_master:
+                    sample_subdir = f"{sample_dir}/step_{train_steps:07d}"
+                    # Calculate multiscale_sizes using same logic as training loop
+                    if args.multiscale:
+                        sample_multiscale_levels = getattr(args, 'multiscale_levels', [32, 64, 128, 256, 512])
+                    elif hasattr(args, 'multiscale_levels') and args.multiscale_levels is not None:
+                        sample_multiscale_levels = args.multiscale_levels
+                    else:
+                        sample_multiscale_levels = None
+                    
+                    save_samples(vae, loader, device, sample_subdir, 
+                               num_samples=args.vis_num_samples, 
+                               patches_per_image=args.patches_per_image,
+                               patch_size=args.image_size,
+                               is_conditional=is_conditional_vae,
+                               use_fp16=args.fp16,
+                               vae_domain1=args.vae_domain1,
+                               vae_domain2=args.vae_domain2,
+                               multiscale=(sample_multiscale_levels is not None),
+                               multiscale_levels=sample_multiscale_levels)
+                    logger.info(f"Saved {args.vis_num_samples} samples to {sample_subdir}")
+                    
+                    if writer is not None:
+                        sample_image_path = f"{sample_subdir}/sample_0.png"
+                        if os.path.exists(sample_image_path):
+                            sample_img = Image.open(sample_image_path)
+                            sample_img_array = np.array(sample_img)
+                            if len(sample_img_array.shape) == 2:
+                                sample_img_tensor = torch.from_numpy(sample_img_array).unsqueeze(0).float() / 255.0
+                            else:
+                                sample_img_tensor = torch.from_numpy(sample_img_array).permute(2, 0, 1).float() / 255.0
+                            writer.add_image('Samples/reconstruction', sample_img_tensor, train_steps)
+                dist.barrier()
+            
+            if train_steps % args.ckpt_every == 0 and train_steps > 0:
+                if is_master:
+                    checkpoint = {
+                        "vae": vae.module.state_dict(),
+                        "opt": opt.state_dict(),
+                        "train_steps": train_steps,
+                        "epoch": epoch,
+                        "args": args
+                    }
+                    if scheduler:
+                        checkpoint["scheduler"] = scheduler.state_dict()
+                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+                    torch.save(checkpoint, checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path} (epoch {epoch}, step {train_steps})")
+                dist.barrier()
+    
+    if is_master:
+        logger.info("Done!")
+ 
+    if is_master and writer is not None:
+        writer.close()
+        logger.info("TensorBoard writer closed.")
+    
+    cleanup()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    
+    # data parameters
+    parser.add_argument("--data-path", type=str, required=True, help="Path to pairing JSON file")
+    parser.add_argument("--image-size", type=int, default=256, help="Patch size")
+    parser.add_argument("--patches-per-image", type=int, default=4, help="Number of patches to crop from each large image")
+    parser.add_argument("--multiscale", action="store_true", help="Enable multiscale training: randomly crop from specified sizes and resize to image-size")
+    parser.add_argument("--multiscale-levels", type=int, nargs='+', default=[32, 64, 128, 256, 512],
+                       help="Multiscale crop sizes (default: 32 64 128 256 512). Example: --multiscale-levels 128 256 512")
+    parser.add_argument("--results-dir", type=str, default="results_vae")
+    
+    # VAE model selection (either use predefined model or custom parameters)
+    parser.add_argument("--vae-model", type=str, default=None,
+                       help="Predefined VAE model name (e.g., VAE-KL-f8, VAE-KL-f16). If specified, overrides custom architecture parameters.")
+    
+    # VAE architecture parameters (when --vae-model is not specified)
+    parser.add_argument("--embed-dim", type=int, default=4, help="Latent embedding dimension")
+    parser.add_argument("--z-channels", type=int, default=4, help="Number of latent channels")
+    parser.add_argument("--in-channels", type=int, default=3, help="Number of input channels")
+    parser.add_argument("--out-channels", type=int, default=3, help="Number of output channels")
+    parser.add_argument("--ch", type=int, default=128, help="Base channel count")
+    parser.add_argument("--ch-mult", type=int, nargs="+", default=[1, 2, 4, 4], help="Channel multipliers")
+    parser.add_argument("--num-res-blocks", type=int, default=2, help="Number of residual blocks per level")
+    parser.add_argument("--attn-resolutions", type=int, nargs="*", default=[], help="Resolutions at which to apply attention")
+    parser.add_argument("--dropout", type=float, default=0.0, help="Dropout rate")
+    parser.add_argument("--double-z", action="store_true", default=True, help="Double z for mean and variance")
+    
+    # loss parameters
+    parser.add_argument("--kl-weight", type=float, default=1e-6, help="Weight for KL divergence loss")
+    parser.add_argument("--l1-weight", type=float, default=1.0, help="Weight for L1 reconstruction loss")
+    parser.add_argument("--ssim-weight", type=float, default=1.0, help="Weight for SSIM reconstruction loss")
+    
+    # training parameters
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--global-batch-size", type=int, default=4)
+    parser.add_argument("--learning-rate", type=float, default=4.5e-6)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--ckpt-every", type=int, default=5000)
+    parser.add_argument("--sample-every", type=int, default=1000, help="Save reconstruction samples every N steps")
+    parser.add_argument("--vis_num-samples", type=int, default=8, help="Number of reconstruction samples to save")
+    parser.add_argument("--fp16", action="store_true", help="Use mixed precision training")
+    parser.add_argument("--resume", type=str, default=None, help="Path to checkpoint to resume from")
+    parser.add_argument("--start-epoch", type=int, default=None, help="Manually specify starting epoch (useful for old checkpoints without epoch info)")
+    parser.add_argument("--vae-domain1", action="store_true", help="use domain1 for training")
+    parser.add_argument("--vae-domain2", action="store_true", help="use domain2 for training")
+
+    # scheduler parameters
+    parser.add_argument("--lr-scheduler", type=str, default="linear", 
+                       choices=["none", "linear", "cosine", "constant"],
+                       help="Learning rate scheduler type")
+    parser.add_argument("--warmup-epochs", type=int, default=0, 
+                       help="Number of warmup epochs (linear warmup from 0 to initial lr)")
+    
+    args = parser.parse_args()
+    main(args)