Ensure train.py, dataset.py, and misc/ are included in Space

train.py

@@ -0,0 +1,555 @@
+#
+# 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)