train.py

#
# For licensing see accompanying LICENSE file.
# Copyright (C) 2025 Apple Inc. All Rights Reserved.
#
#!/usr/bin/env python3
"""
Scalable Transformer Autoregressive Flow (STARFlow) Training Script

This script provides functionality for training transformer autoregressive flow models
with support for both image and video generation.

Usage:
    python train.py --model_config_path config.yaml --epochs 100
"""

import argparse
import builtins
import pathlib
import copy
import torch
import torch.nn.functional as F
import torchinfo
import torch.amp
import torch.utils
import torch.utils.data
import torchvision as tv
import numpy as np
import random
import transformer_flow
import utils
import time
import contextlib
import tqdm
import os
import gc
import sys
import wandb
import yaml
from typing import Dict, List, Optional, Tuple, Union

from dataset import read_tsv, aspect_ratio_to_image_size
from contextlib import nullcontext
from misc import print  # local_rank=0 print
from utils import simple_denoising, save_samples_unified, add_noise, encode_text, drop_label, load_model_config

# Set environment variables for local development
os.environ["PYTHONPATH"] = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) + ":" + os.environ.get("PYTHONPATH", "")
WANDB_API_KEY = os.environ.get("WANDB_API_KEY", None)


def self_denoise(model, samples, y, noise_std=0.1, lr=1, steps=1, disable_learnable_denoiser=False):
    if steps == 0:
        return samples
    
    outputs = []
    x = samples.clone()
    lr = noise_std ** 2 * lr
    with torch.enable_grad():
        x.requires_grad = True
        model.train()
        z, _, _, logdets = model(x, y)
        loss = model.get_loss(z, logdets)['loss'] * 65536
        grad = float(samples.numel()) / 65536 * torch.autograd.grad(loss, [x])[0]
        outputs += [(x - grad * lr).detach()]
    x = torch.cat(outputs, -1)
    return x








def main(args):
    # Load model configuration if provided
    if hasattr(args, 'model_config_path') and args.model_config_path:
        # Parse sys.argv to see which args were actually provided
        provided_args = set()
        for i, arg in enumerate(sys.argv[1:]):
            if arg.startswith('--'):
                arg_name = arg[2:].replace('-', '_')
                provided_args.add(arg_name)
        trainer_args = load_model_config(args.model_config_path)
        trainer_dict = vars(trainer_args)
        for k, v in vars(args).items():
            if k in provided_args:
                trainer_dict[k] = v
        args = argparse.Namespace(**trainer_dict)
    
    # global setup
    dist = utils.Distributed()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")   
    seed = args.train_seed if args.train_seed is not None else time.time_ns() % 2**32
    utils.set_random_seed(seed + dist.rank)
    print(f'set random seed {seed}')

    if dist.rank == 0 and WANDB_API_KEY is not None:
        job_name = f'{args.dataset}'
        if args.wandb_name is not None:
            wandb_names = args.wandb_name.split('+')
            if len(wandb_names) > 1:
                job_name += f'-{wandb_names[0]}-{getattr(args, wandb_names[1])}'
            else:
                job_name += f'-{wandb_names[0]}'

        wandb.login(key=WANDB_API_KEY)
        wandb.init(project="starflow", name=job_name, config=vars(args))
        wandb.run.save()
        wandb.run.log_code(os.path.dirname(os.path.realpath(__file__)))

    if args.use_pretrained_lm is not None:
        args.text = args.use_pretrained_lm  # need to match the text embedder

    print(f'{" Config ":-^80}')
    for k, v in sorted(vars(args).items()):
        print(f'{k:32s}: {v}')

    # dataset
    data_loader = utils.get_data(args, dist)
    total_num_images = data_loader.dataset.total_num_samples
    grad_accum_steps = max(args.acc, 1)
    num_batches_before_acc = len(data_loader)
    num_batches = num_batches_before_acc // grad_accum_steps
    
    print(f'{" Dataset Info ":-^80}')
    print(f'{num_batches} batches per epoch ({num_batches_before_acc} steps if consider {grad_accum_steps} accumulation), global batch size {args.batch_size} for {args.epochs} epochs')
    print(f'So it is {num_batches * args.batch_size:,} images per epoch')
    print(f'Target training on {args.batch_size * num_batches * args.epochs:,} images')
    print(f'Total {total_num_images:,} unique training examples')
    
    assert args.text is not None, "starflow needs text conditioning"
        
    # text encoder
    tokenizer, text_encoder = utils.setup_encoder(args, dist, device)
    text_encoder.requires_grad_(False)  # freeze text encoder

    # VAE & fixed noise
    if args.vae is not None:
        vae = utils.setup_vae(args, dist, device)
        vae.requires_grad_(False)  # freeze VAE
        args.img_size = args.img_size // vae.downsample_factor

    # main model
    model = utils.setup_transformer(args, dist, 
                                    txt_dim=text_encoder.config.hidden_size, 
                                    use_checkpoint=args.gradient_checkpoint,
                                    use_checkpoint_mlp=args.gradient_checkpoint_mlp).to(device)
    if dist.local_rank == 0:
        torchinfo.summary(model)

    # Load model before FSDP wrapping to support expansion
    if args.resume_path:
        print(f"Loading checkpoint from local path: {args.resume_path}")
        state_dict = torch.load(args.resume_path, map_location='cpu')
        model.load_state_dict(state_dict, strict=False)
        del state_dict; torch.cuda.empty_cache()
        epoch_start = args.resume_epoch if args.resume_epoch is not None else 0
    else:
        epoch_start = 0

    # setup for training
    model, model_ddp = utils.parallelize_model(args, model, dist, device)
    if args.text and args.fsdp_text_encoder:
        text_encoder = utils.parallelize_model(args, text_encoder, dist, device, [text_encoder.base_block_name])[1]
    trainable_params = [p for k, p in model_ddp.named_parameters() if p.requires_grad and not k.startswith('learnable_self_denoiser')]
    optimizer = torch.optim.AdamW(trainable_params, betas=(0.9, 0.95), lr=args.lr, weight_decay=1e-4)
    warmup_steps = args.warmup_steps if args.warmup_steps is not None else num_batches
    lr_schedule = utils.CosineLRSchedule(
        optimizer, warmup_steps, args.epochs * num_batches, args.min_lr, args.lr)
    if args.learnable_self_denoiser:
        denoiser_optimizer = torch.optim.AdamW(model_ddp.learnable_self_denoiser.parameters(), lr=1e-4, weight_decay=1e-4)
        denoiser_lr_schedule = utils.CosineLRSchedule(
            denoiser_optimizer, warmup_steps, args.epochs * num_batches, 1e-6, 1e-4)
    print('warmup_steps:', warmup_steps, 'num_batches:', num_batches, 'total steps:', args.epochs * num_batches)

    # Adjust learning rate schedule and counters if resuming    
    lr_schedule.counter += epoch_start * num_batches
    images_start = epoch_start * num_batches * args.batch_size


    if args.loss_scaling:
        scaler = torch.amp.GradScaler()
        if args.learnable_self_denoiser:
            denoiser_scaler = torch.amp.GradScaler()

    noise_std  = args.noise_std
    model_name = f'{args.patch_size}_{args.channels}_{args.blocks}_{args.layers_per_block}_{noise_std:.2f}'
    sample_dir: pathlib.Path = args.logdir / f'{args.dataset}_samples_{model_name}'
    model_ckpt_file = args.logdir / f'{args.dataset}_model_{model_name}.pth'
    opt_ckpt_file = args.logdir / f'{args.dataset}_opt_{model_name}.pth'
    if dist.local_rank == 0:
        sample_dir.mkdir(parents=True, exist_ok=True)

    print(f'{" Training ":-^80}')
    total_steps, total_images, total_training_time = epoch_start * num_batches, images_start, 0
    for epoch in range(epoch_start, args.epochs):
        metrics = utils.Metrics()
        for it, (x, y, meta) in enumerate(data_loader):
            if args.secondary_dataset is not None and random.random() < args.secondary_ratio:
                x, y, meta = next(data_loader.secondary_loader)  # load data from secondary dataset instead
            y_caption = copy.deepcopy(y)
            data_mode = 'image' if (x.dim() == 4) else 'video'
            
            start_time = time.time()
            x_aspect, video_mode = data_loader.dataset.get_batch_modes(x)
            x = x.to(device)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                # apply VAE over images
                if args.vae is not None:
                    with torch.no_grad():
                        if data_mode == 'video' and args.last_frame_cond:
                            x_last = x[:, -4:-3]  # use the last frame as additional condition
                            x = x[:, :-4]
                            x, x_last = vae.encode(x), vae.encode(x_last)
                            x = torch.cat([x_last, x], 1)
                            y = ["(last) " + desc for desc in y]

                        elif data_mode == 'video' and args.video_to_video:
                            x = torch.cat(x.chunk(2, dim=1)[::-1], 0)  # data is target:source
                            x = vae.encode(x)
                            x = torch.cat(x.chunk(2, dim=0), 1)
                            y = ["(v2v) " + desc for desc in y]

                        else:
                            x = vae.encode(x)

                # add noise to images  
                x, _ = add_noise(x, noise_std, args.noise_type)
                if data_mode == 'video' and args.drop_image > 0 and random.random() < args.drop_image:
                    x = x[:, 1:]
                    y = ["(extend) " + desc for desc in y]
                
                # Enable gradient computation for x
                x.requires_grad_(True)
                
                # Process labels/text based on model type
                with torch.no_grad():
                    y = encode_text(
                        text_encoder, tokenizer, 
                        drop_label(y, args.drop_label),
                        args.txt_size, device,
                        aspect_ratio=x_aspect if args.mix_aspect else None,
                        fps=meta.get('fps', None) if args.fps_cond else None,
                        noise_std=noise_std if args.cond_noise_level else None)

                # main training step
                needs_update = False # (it + 1) % grad_accum_steps == 0
                needs_zero_grad = it % grad_accum_steps == 0
                
                if needs_zero_grad:
                    optimizer.zero_grad()

                # main forward
                z, _, outputs, logdets = model_ddp(x, y)
                weights = noise_std / 0.3 if args.cond_noise_level else None
                loss_dict = model.get_loss(z, logdets, weights)
                loss = loss_dict['loss']
                if args.latent_norm_regularization > 0:
                    loss += args.latent_norm_regularization * sum([z.pow(2).mean() for z in outputs[:-1]])
                loss = loss / grad_accum_steps  # use gradient accumulation

                if dist.gather_concat(loss.view(1)).isnan().any():
                    if dist.local_rank == 0:
                        print('nan detected, skipping step')
                    continue
                
                with utils.sync_ctx(model_ddp, sync=needs_update) if grad_accum_steps > 1 else contextlib.nullcontext():
                    if args.loss_scaling:
                        scaler.scale(loss).backward()
                    else:
                        loss.backward()
                
                # Get gradient of x after backward pass
                if args.learnable_self_denoiser:
                    x_grad = x.grad.clone().detach()  # Clone to preserve the gradient
                    scale = (float(x.numel()) / scaler.get_scale()) if args.loss_scaling else float(x.numel())
                    score = x_grad * scale * grad_accum_steps * noise_std  # roughly std=1, similar to diffusion models
                    pred = model_ddp(x, y, denoiser=True)
                    loss_denoiser = F.mse_loss(pred, score, reduction='mean') / grad_accum_steps
                    loss_dict['loss_denoiser'] = loss_denoiser.item()
                    
                    with utils.sync_ctx(model_ddp, sync=needs_update) if grad_accum_steps > 1 else contextlib.nullcontext():
                        if args.loss_scaling:
                            denoiser_scaler.scale(loss_denoiser).backward()
                        else:
                            loss_denoiser.backward()
                                        
                # accumulate time
                total_training_time = total_training_time + (time.time() - start_time)
                if needs_update:
                    # Apply gradient clipping and monitor gradient norm
                    grad_norm = None
                    denoiser_grad_norm = None
                    skip_update = False
                    
                    if args.grad_clip > 0:
                        if args.loss_scaling:
                            scaler.unscale_(optimizer)
                        grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, args.grad_clip)
                        skip_update = grad_norm.item() > args.grad_clip if args.grad_skip and (total_steps > 100) else False
                        if args.learnable_self_denoiser:
                            if args.loss_scaling:
                                denoiser_scaler.unscale_(denoiser_optimizer)
                            denoiser_grad_norm = torch.nn.utils.clip_grad_norm_(model_ddp.learnable_self_denoiser.parameters(), args.grad_clip)
                            skip_update = skip_update or (denoiser_grad_norm.item() > args.grad_clip if args.grad_skip and (total_steps > 100) else False)
                    
                    if skip_update:
                        print(f'Skipping update due to large gradient norm {grad_norm.item():.4f} > {args.grad_clip:.4f}')
                        optimizer.zero_grad()
                        if args.learnable_self_denoiser:
                            denoiser_optimizer.zero_grad()
                    
                    if args.loss_scaling:
                        scaler.step(optimizer)
                        scaler.update()
                    else:
                        optimizer.step()

                    current_lr = lr_schedule.step() 
                    
                    if not skip_update:
                        metrics.update(loss_dict)

                    if args.learnable_self_denoiser:
                        denoiser_lr = denoiser_lr_schedule.step()
                        if args.loss_scaling:
                            denoiser_scaler.step(denoiser_optimizer)
                            denoiser_scaler.update()
                        else:
                            denoiser_optimizer.step()
                        if not skip_update:
                            metrics.update({'loss_denoiser': loss_denoiser.item()})

                    total_steps = total_steps + 1
                    total_images = total_images + args.batch_size
                        
                    # end of training step
                    if (it // grad_accum_steps) % 10 == 9:
                        speed = (total_images - images_start) / total_training_time
                        print(f"{total_steps:,} steps/{total_images:,} images ({speed:0.2f} samples/sec) - \t" + "\t".join(
                            ["{}: {:.4f}".format(k, v) for k, v in loss_dict.items()]))
                    
                        if dist.rank == 0 and WANDB_API_KEY is not None:
                            wandb_dict = {'speed': speed, 'steps': total_steps, 'lr': current_lr}
                            if grad_norm is not None:
                                wandb_dict['grad_norm'] = grad_norm.item()
                            if args.learnable_self_denoiser:
                                wandb_dict['denoiser_lr'] = denoiser_lr
                                if denoiser_grad_norm is not None:
                                    wandb_dict['denoiser_grad_norm'] = denoiser_grad_norm.item()
                            loss_dict.update(wandb_dict)
                            wandb.log(loss_dict, step=total_images)
            
            if args.dry_run:
                break

        # metrics_dict = {'lr': current_lr, **metrics.compute(dist)}

        # print metrics
        if False: # dist.local_rank == 0:
            metrics.print(metrics_dict, epoch + 1)
            print('\tLayer norm', ' '.join([f'{z.pow(2).mean():.4f}' for z in outputs]))
            print('\tLayer stdv', ' '.join([f'{z.std():.4f}' for z in outputs]))
            if dist.rank == 0 and WANDB_API_KEY is not None:
                wandb.log({f'epoch_{k}': v for k, v in metrics_dict.items()}, step=total_images)

        # save model and optimizer state
        if not args.dry_run:
            utils.save_model(args, dist, model, model_ckpt_file)
            if epoch % args.save_every == 0:  # save every 20 epochs
                utils.save_model(args, dist, model, str(model_ckpt_file) + f"_epoch{epoch+1:04d}")
            # utils.save_optimizer(args, dist, optimizer, lr_schedule, opt_ckpt_file)
        dist.barrier()

        # sample images (should i sample?)
        if args.sample_freq > 0 and (epoch % args.sample_freq == 0 or args.dry_run):
            model.eval()
            
            # Simple sampling using current batch data
            with torch.no_grad():
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    x_aspect = "16:9"
                    x_shape = aspect_ratio_to_image_size(
                        args.img_size * vae.downsample_factor, x_aspect,
                        multiple=vae.downsample_factor
                    )
                    if x.dim() == 5:
                        x_shape = (x.shape[0], 21, x.shape[2], x_shape[0] // vae.downsample_factor, x_shape[1] // vae.downsample_factor)
                    else:
                        x_shape = (x.shape[0], x.shape[1], x_shape[0] // vae.downsample_factor, x_shape[1] // vae.downsample_factor)
                    noise = torch.randn(*x_shape).to(device)
                    cfg = 3.5
                    y_caption = ["POV from the boat deck looking at a corgi wearing neon-pink sunglasses; wind noise feel, slight horizon bob, water droplets on lens occasionally, sun reflections flicker on the frames; natural lighting"]
                    y_caption = y_caption + [""] * len(y_caption)
                    sample_y = encode_text(
                        text_encoder, tokenizer, y_caption,
                        args.txt_size, device,
                        aspect_ratio=x_aspect if args.mix_aspect else None,
                        fps=meta.get('fps', [None])[0] if args.fps_cond else None,
                        noise_std=noise_std if args.cond_noise_level else None)
                
                    # Generate samples
                    samples = model(noise, sample_y, reverse=True, guidance=cfg, 
                                    jacobi=1 if noise.dim() == 5 else 0, verbose=True)
                    
                    # Apply self denoising if needed
                    sample_y = sample_y.chunk(2, dim=0)[0]  # Remove null captions for denoising
                    samples = simple_denoising(model, samples, sample_y, 
                        text_encoder, tokenizer, args, noise_std)
                    
                    # Decode with VAE if available
                    if args.vae is not None:
                        samples = vae.decode(samples)
                        
                    # Save samples using unified function
                    save_samples_unified(
                        samples=samples,
                        save_dir=sample_dir,
                        filename_prefix="train_samples",
                        epoch_or_iter=epoch+1,
                        fps=meta.get('fps', [16])[0],
                        dist=dist,
                        wandb_log=WANDB_API_KEY is not None,
                        wandb_step=total_images,
                        grid_arrangement="grid"  # Use simple grid for training
                    )
            model.train()

        if args.dry_run:
            break

    if dist.rank == 0 and WANDB_API_KEY is not None:
        wandb.finish()


def get_tarflow_parser():
    parser = argparse.ArgumentParser()

    # Model config path (same as sample.py)
    parser.add_argument('--model_config_path', default=None, type=str, help='path to YAML config file')

    # Dataset config
    parser.add_argument('--train_seed', default=None, type=int)
    parser.add_argument('--data', default='data', type=pathlib.Path)
    parser.add_argument('--logdir', default='./logs', type=pathlib.Path)
    parser.add_argument('--dataset', default='dummy', type=str)
    parser.add_argument('--wds', default=0, type=int)
    parser.add_argument('--mix_aspect', default=0, type=int)
    parser.add_argument('--img_size', default=32, type=int)
    parser.add_argument('--vid_size', default=None, type=str, help="num_frames:fps1:fps2 for video datasets. If None, image mode")
    parser.add_argument('--fps_cond', default=0, type=int, help="If 1, use fps from video dataset as condition")
    parser.add_argument('--no_flip', default=0, type=int)
    parser.add_argument('--caption_column', default='syn_detailed_description_w_caption', type=str,
                        help="If given 'folder', then extract caption from file name")
    
    # Optional Dadaset config
    parser.add_argument('--secondary_dataset', default=None, type=str, help="secondary dataset for training")
    parser.add_argument('--secondary_img_size', default=32, type=int, help="secondary dataset image size")
    parser.add_argument('--secondary_vid_size', default=None, type=str, help="secondary dataset video size")
    
    # VAE configuration
    parser.add_argument('--vae', default=None, type=str, help="pretrained VAE name")
    parser.add_argument('--vae_decoder_factor', default=1, type=float, help="VAE decoder scaling factor")
    parser.add_argument('--channel_size', default=3, type=int)
    parser.add_argument('--finetuned_vae', default=None, type=str)
    
    # Text encoder configuration
    parser.add_argument('--text', default=None, type=str, help="text encoder")
    parser.add_argument('--txt_size', default=0, type=int, help="maximum text length")
    
    # Model configuration
    parser.add_argument('--sos', default=0, type=int)
    parser.add_argument('--seq_order', default="R2L", type=str, choices=['R2L', 'L2R'])
    parser.add_argument('--patch_size', default=4, type=int)
    parser.add_argument('--channels', default=512, type=int)
    parser.add_argument('--top_block_channels', default=None, type=int)
    parser.add_argument('--blocks', default=4, type=int)
    parser.add_argument('--layers_per_block', default=8, type=int, nargs='*')
    parser.add_argument('--rope', default=0, type=int)
    parser.add_argument('--pt_seq_len', default=None, type=int)
    parser.add_argument('--adaln', default=0, type=int)
    parser.add_argument('--nvp', default=1, type=int)
    parser.add_argument('--use_softplus', default=0, type=int)
    parser.add_argument('--cond_top_only', default=0, type=int)
    parser.add_argument('--head_dim', default=64, type=int)
    parser.add_argument('--num_heads', default=None, type=int)
    parser.add_argument('--num_kv_heads', default=None, type=int)
    parser.add_argument('--use_swiglu', default=0, type=int)
    parser.add_argument('--use_qk_norm', default=0, type=int)
    parser.add_argument('--use_post_norm', default=0, type=int)
    parser.add_argument('--use_final_norm', default=0, type=int)
    parser.add_argument('--use_bias', default=1, type=int)
    parser.add_argument('--norm_type', default='layer_norm', type=str)
    parser.add_argument('--use_pretrained_lm', default=None, type=str, choices=['gemma3_4b', 'gemma3_1b', 'gemma2_2b'])
    parser.add_argument('--use_mm_attn', default=0, type=int)
    parser.add_argument('--soft_clip', default=0, type=float, help="soft clip the output values")
    parser.add_argument('--learnable_self_denoiser', default=0, type=int, help="Whether to use learnable self-denoiser")
    parser.add_argument('--conditional_denoiser', default=0, type=int, help="conditional denoiser")
    parser.add_argument('--noise_embed_denoiser', default=0, type=int, help="add noise embedding to the denoiser")
    parser.add_argument('--temporal_causal', default=0, type=int, help="Whether to use temporal causal model")
    parser.add_argument('--shallow_block_local', default=0, type=int, help="Whether to use local attention in shallow blocks")
    parser.add_argument('--denoiser_window', default=None, type=int, help="local window size for denoiser")
    parser.add_argument('--local_attn_window', default=None, type=int, help="Whether to use local attention")

    # Training configuration
    parser.add_argument('--noise_std', default=0.3, type=float)
    parser.add_argument('--noise_type', default='gaussian', choices=['gaussian', 'uniform'], type=str)
    parser.add_argument('--cond_noise_level', default=0, type=int, help="Whether to sample noise level as in diffusion models")
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--secondary_batch_size', default=128, type=int, help="only for secondary dataset")
    parser.add_argument('--secondary_ratio', default=0, type=float, help="a value between 0-1, ratio of using secondary data.")

    parser.add_argument('--acc', default=1, type=int)
    parser.add_argument('--fp8', default=0, type=int, help='Whether to use FP8 training')
    parser.add_argument('--use_8bit_adam', default=0, type=int, help='Whether to use 8-bit Adam optimizer')
    parser.add_argument('--epochs', default=1000, type=int)
    parser.add_argument('--epoch_length', default=50000, type=int)
    parser.add_argument('--lr', default=1e-4, type=float)
    parser.add_argument('--min_lr', default=1e-6, type=float)
    parser.add_argument('--drop_label', default=0, type=float)
    
    parser.add_argument('--drop_image', default=0, type=float)
    parser.add_argument('--last_frame_cond', default=0, type=int)
    parser.add_argument('--video_to_video', default=0, type=int)

    parser.add_argument('--resume_path', default=None, type=str)
    parser.add_argument('--resume_epoch', default=0, type=int) 
    parser.add_argument('--warmup_steps', default=None, type=int, help='Warmup steps for training')
    
    parser.add_argument('--fsdp', default=0, type=int)
    parser.add_argument('--fsdp_text_encoder', default=0, type=int)
    parser.add_argument('--gradient_checkpoint', default=0, type=int)
    parser.add_argument('--gradient_checkpoint_mlp', default=None, type=int)
    parser.add_argument('--compile', default=0, type=int, help='Whether to use torch.compile')
    parser.add_argument('--latent_norm_regularization', default=0, type=float, help='Regularization on latent norm, 1e-4 is a good value')
    parser.add_argument('--loss_scaling', default=1, type=int, help='Whether to use AMP')
    parser.add_argument('--grad_clip', default=0, type=float, help='Gradient clipping threshold, 0 to disable')
    parser.add_argument('--grad_skip', default=0, type=int, help='Skip gradient computation for the model')
    parser.add_argument('--dry_run', default=0, type=int, help='Dry run for quick tests')
    parser.add_argument('--wandb_name', default=None, type=str, help='Wandb name for the run')
    parser.add_argument('--save_every', default=20, type=int, help='Save model every N epochs')
    parser.add_argument('--sample_freq', default=1, type=int, help="sample every N epochs, 0 to disable")

    # Sampling configuration
    parser.add_argument('--cfg', default=0, type=float, nargs='+')
    parser.add_argument('--num_samples', default=4096, type=int)
    parser.add_argument('--sample_batch_size', default=256, type=int)
    return parser


if __name__ == '__main__':
    parser = get_tarflow_parser()    
    args = parser.parse_args()
    main(args)