train.py

import itertools
import math
import os
import random
import subprocess
import sys
from pathlib import Path
from typing import Iterable, Iterator, List, NamedTuple
os.environ.setdefault("WANDB_MODE", "disabled")
import wandb
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision.utils
from einops import rearrange
from lightning import Fabric
from wandb.integration.lightning.fabric import WandbLogger
from dataset import ImageFolderDataset
from gan_loss import GANLoss
from lpips import LPIPS
from models import ARModel, Encoder, Decoder, Linear
from tqdm import tqdm
from omegaconf import OmegaConf
from torch.utils.data import DataLoader
import copy
import torch.distributed as dist
from glob import glob
try:
    from thop import profile
    from thop.vision.basic_hooks import count_linear
except ImportError:
    profile = None
    count_linear = None


torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False

# Phase utils
class Target(NamedTuple):
    DO_AE: bool = False
    DO_L2: bool = False
    DO_L1: bool = False
    DO_LPIPS: bool = False
    DO_GAN_G: bool = False
    DO_GAN_D: bool = False
    DO_PRIOR_ENC: bool = False
    DO_PRIOR_AR: bool = False

class Phase(NamedTuple):
    num_steps: int
    targets: List[Target]       
    internal_steps: List[int]   

# Parse phases such as "1000:DO_L2-DO_PRIOR_AR:1 7000:DO_PRIOR_AR:1".
def parse_phases(phases_str):
    phases = [] # list of Phase objects
    for phase_str in phases_str.split(' '):
        num_steps, targets_str, internal_steps_str = phase_str.split(':')
        num_steps = int(num_steps)
        targets = [Target(**{k: True for obj in target_str.split(',') for k in obj.split('-') }) for target_str in targets_str.split(',')]
        internal_steps = [int(step) for step in internal_steps_str.split(',')]
        phases.append(Phase(num_steps, targets, internal_steps))
    return phases

def get_phase(global_step, phases, phase_step_accum, config):
    target=None
    for phase_idx, phase_step in enumerate(phase_step_accum):
        if global_step <= phase_step:
            internel_step = (global_step - phase_step_accum[phase_idx-1]) if phase_idx > 0 else global_step
            internel_accumulate = list(itertools.accumulate(phases[phase_idx].internal_steps))
            internel_step = internel_step % internel_accumulate[-1]

            for inner_idx in range(len(internel_accumulate)):
                if internel_step < internel_accumulate[inner_idx]:
                    target = phases[phase_idx].targets[inner_idx]
                    break
            if target is not None:
                break

    DO_AE = any([target.DO_L2, target.DO_L1, target.DO_LPIPS, target.DO_GAN_G])

    # assert inconsistency configuration
    assert not (DO_AE and  not config.train_ae), f"train_ae={config.train_ae} is inconsistent with DO_AE={DO_AE} in phases"
    assert not (target.DO_PRIOR_AR and  not config.train_ar), f"train_ar={config.train_ar} is inconsistent with DO_PRIOR_AR={target.DO_PRIOR_AR} in phases"
    assert not ((target.DO_GAN_G or target.DO_GAN_D) and  not config.use_gan_loss), f"use_gan_loss={config.use_gan_loss} is inconsistent with DO_GAN_G or DO_GAN_D={(target.DO_GAN_G or target.DO_GAN_D)} in phases"
    assert not (target.DO_PRIOR_ENC and  not config.use_prior_loss), f"use_prior_loss={config.use_prior_loss} is inconsistent with DO_PRIOR_ENC={target.DO_PRIOR_ENC} in phases"
    assert not (target.DO_LPIPS and  not config.use_lpips_loss), f"use_lpips_loss={config.use_lpips_loss} is inconsistent with DO_LPIPS={target.DO_LPIPS} in phases"

    # update target with global_step
    target = Target(DO_L1=target.DO_L1, 
                    DO_L2=target.DO_L2, 
                    DO_LPIPS=target.DO_LPIPS, 
                    DO_GAN_G=target.DO_GAN_G and (global_step >= config.gan_start),
                    DO_GAN_D=target.DO_GAN_D and (global_step >= config.gan_start),
                    DO_PRIOR_ENC=target.DO_PRIOR_ENC and (global_step >= config.prior_start),
                    DO_PRIOR_AR=target.DO_PRIOR_AR,
                    DO_AE=DO_AE)
    return phase_idx, inner_idx, target, internel_step

class InfiniteIterator(Iterator):
    def __init__(self, iterable: Iterable):
        self.iterable = iterable
        self._it = iter(iterable)

    def __iter__(self):
        return self

    def __next__(self):
        try:
            return next(self._it)
        except StopIteration:
            self._it = iter(self.iterable)
            return next(self._it)

# Image Processing Utils
def patchify(x, patch_size):
    x = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
    return x

def unpatchify(x, image_size, patch_size):
    x = rearrange(x, 'b (h w) (p1 p2 c) -> b c (h p1) (w p2)', p1=patch_size, p2=patch_size, h=image_size//patch_size, w=image_size//patch_size)
    return x

def img_unint8_to_norm(x):
    return x.float() / 127.5 - 1.0

def img_denormalize(x):
    return x.clamp(-1,1) * 0.5 + 0.5

def img_norm_to_uint8(x):
    return torch.clamp(127.5 * x + 128.0, 0, 255).byte()

# -----------------------------
# Tail-dropout + noise-query utils
# -----------------------------

def _get_decoder_context_type(dec: nn.Module) -> str:
    # `dec` is Decoder, `dec.dec` is Transformer
    ctx = getattr(getattr(dec, "dec", None), "context_type", None)
    return ctx if ctx is not None else "none"

def _get_decoder_query_len(dec: nn.Module) -> int:
    return int(getattr(getattr(dec, "dec", None), "max_seq_len", 0))

def _t_to_k(t: torch.Tensor, K: int, t2k: float = 1.0) -> torch.Tensor:
    """
    Map t in [0,1] to discrete k in [0, K-1].
    """
    if K <= 1:
        return torch.zeros_like(t, dtype=torch.long)
    t_tmp = (t2k * t).clamp(0.0, 1.0)
    k = torch.floor(t_tmp * K).to(torch.long).clamp(0, K - 1)
    return k

class DiTiScheduler:
    """
    Piecewise linear time scheduler from Selftok, adapted for wARTok.
    """
    def __init__(
        self,
        n_timesteps: int,
        K: int,
        stages: str | List[int],
        k_per_stage: str | List[int],
        reverse_t: bool = False,
    ):
        self.n_timesteps = n_timesteps
        self.K = K
        self.reverse_t = reverse_t

        if isinstance(stages, str):
            stages = [int(s.strip()) for s in stages.split(',')]
        if isinstance(k_per_stage, str):
            k_per_stage = [int(k.strip()) for k in k_per_stage.split(',')]
        self.stages = stages
        
        # Scaling logic: if sum(k_per_stage) != K, scale it
        total_k = sum(k_per_stage)
        if total_k != K:
            print(f"[DiTiScheduler] Warning: sum(k_per_stage)={total_k} does not match K={K}. Scaling k_per_stage...")
            scale_factor = K / total_k
            scaled_k = [int(k * scale_factor) for k in k_per_stage]
            # Fix rounding errors to ensure sum is exactly K
            diff = K - sum(scaled_k)
            if diff != 0:
                scaled_k[-1] += diff
            k_per_stage = scaled_k
            print(f"[DiTiScheduler] Scaled k_per_stage: {k_per_stage}")
            
        self.k_per_stage = k_per_stage
        
        self.t_to_idx = torch.zeros(n_timesteps).long()
        
        # Precompute lookup table
        current_stage = 0
        sum_indices = 0
        # stages should be cumulative, e.g., 200, 400, 600...
        
        for t in range(n_timesteps):
            # Check if we moved to next stage
            if current_stage < len(self.stages) and t >= self.stages[current_stage]:
                sum_indices += self.k_per_stage[current_stage]
                current_stage += 1
            
            if current_stage >= len(self.stages):
                 idx = self.K - 1
            else:
                prev_stage_t = self.stages[current_stage - 1] if current_stage > 0 else 0
                current_stage_t = self.stages[current_stage]
                
                current_steps = float(current_stage_t - prev_stage_t)
                current_k = float(self.k_per_stage[current_stage])
                
                t_adj = t - prev_stage_t
                idx = int(float(t_adj) / current_steps * current_k + sum_indices)
            
            self.t_to_idx[t] = idx

    def to_indices(self, t: torch.Tensor) -> torch.Tensor:
        """
        Convert t in [0, 1] to indices [0, K-1] using the lookup table.
        """
        if self.reverse_t:
            t = 1.0-t
        device = t.device
        # Map [0, 1] to [0, n_timesteps]
        t_idx = (t * self.n_timesteps).long().clamp(0, self.n_timesteps - 1)
        # Move the lookup table to the input device.
        return self.t_to_idx.to(device)[t_idx].clamp(0, self.K - 1)


def build_taildrop_attn_mask(
    *,
    t: torch.Tensor,
    context_len: int,
    query_len: int,
    context_type: str,
    t2k: float = 1.0,
    device=None,
    diti_scheduler: DiTiScheduler = None,
) -> torch.Tensor | None:
    """
    Build a boolean attention mask (True for keep, False for masked).

    - cross_attn: returns (B, 1, query_len, context_len)
    - concat:     returns (B, 1, total_len, total_len) where total_len = query_len + context_len
    """
    device = device or t.device
    B = t.shape[0]
    if context_len <= 0:
        return None

    if diti_scheduler is not None:
        k = diti_scheduler.to_indices(t) # (B,)
    else:
        k = _t_to_k(t, context_len, t2k=t2k)  # (B,)
        
    keep = (torch.arange(context_len, device=device)[None, :] <= k[:, None])  # (B, context_len)


    if context_type == "cross_attn":
        # (B, 1, query_len, context_len)
        return keep[:, None, :].unsqueeze(1).expand(B, 1, query_len, context_len)

    if context_type == "concat":
        total_len = query_len + context_len
        keep_total = torch.cat(
            [torch.ones((B, query_len), device=device, dtype=torch.bool), keep],
            dim=1,
        )  # (B, total_len)
        # (B, 1, total_len, total_len)
        return keep_total[:, None, :].unsqueeze(1).expand(B, 1, total_len, total_len)

    return None

def make_noise_image_patch_query(x: torch.Tensor, t: torch.Tensor, query_len: int) -> torch.Tensor:
    """
    Replace dummy query with a noised version of the (clean) image patches x.
    x: (B, L, D)
    t: (B,)
    Returns noised patches, noise, and velocity, each of shape (B, L, D).
    """
    B, L, D = x.shape
    noise = torch.randn_like(x)
    t = t.view(B, 1, 1).to(x.dtype)
    x_noised = (1.0 - t) * x + t * noise
    return x_noised, noise, noise-x

# Seed utils
def seed_everything(seed):
    # manually call these to ensure reproducibility
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    # fabric seed everything
    Fabric.seed_everything(seed, workers=True)

def worker_init_fn(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def make_worker_init_fn(base_seed: int):
    base_seed = int(base_seed) % (2**32)

    def _init(worker_id: int):
        rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
        worker_seed = (base_seed + worker_id + 1000 * rank) % (2**32)
        np.random.seed(worker_seed)
        random.seed(worker_seed)
        torch.manual_seed(worker_seed)

    return _init


# FID Utils                    
def adm_fid_evaluator(sample_cached_path, gt_cache_path, config, fabric):
    if not os.path.exists(gt_cache_path):
        raise FileNotFoundError(f"Ground-truth cache not found: {gt_cache_path}")
    if not os.path.exists(sample_cached_path):
        raise FileNotFoundError(f"Sample cache not found: {sample_cached_path}")

    fid_script = str(Path(__file__).with_name('calc_fid.py'))
    env = os.environ.copy()
    cmd = [sys.executable,fid_script,"--ref_batch",gt_cache_path,"--sample_batch",sample_cached_path,"--batch_size",str(config.eval_batch_size),]
    fabric.print(f"Running FID evaluation via {fid_script}...")

    FID = None
    process = subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True,bufsize=1,env=env,)
    assert process.stdout is not None
    for line in process.stdout:
        line = line.rstrip("\n")
        if line:
            fabric.print(line, flush=True)
        if line.startswith("FID_RESULT:"):
            try:
                FID = float(line.split("FID_RESULT:")[1].strip())
            except ValueError:
                pass
    retcode = process.wait()
    if retcode != 0 or FID is None or not math.isfinite(FID):
        raise RuntimeError(f"FID evaluation failed: exit={retcode}, result={FID}")

    return FID


# Training Utils
@torch.no_grad()
@torch._dynamo.disable
def ema_update(model, ema_model, ema_rate):
    if model is None or ema_model is None:
        return
    
    for p,ema_p in zip(model.parameters(),ema_model.parameters()):
            ema_p.copy_(p.detach().lerp(ema_p, ema_rate))
    


@torch.no_grad()
@torch._dynamo.disable
def sampling(enc, dec, ar_model,  bz, class_label=None, temperature=1.0, topK=0, topP=0., cfg=0.0, cfg_schedule=None, cfg_power=None, cache_kv=False, config=None,t=None,query_tokens=None, **params):
    token_ids = ar_model.sampling(bz, class_label, temperature, topK, topP, cfg, cfg_schedule, cfg_power, cache_kv)
    quant = enc.quantizer.get_codes_w_indices(token_ids)
    decoded = dec(quant, class_label, attn_mask=None, query_tokens=query_tokens,t=t)
    return decoded

@torch.no_grad()
@torch._dynamo.disable
def reconstruction(enc, dec, x, labels, return_idx: bool = False, query_tokens=None,t=None):
    quant, idx, _, _ = enc(x,labels, training=False)
    decoded = dec(quant, labels, attn_mask=None, query_tokens=query_tokens, t=t)
    return decoded, idx

@torch.no_grad()
@torch._dynamo.disable
def euler_sample_loop(enc, dec, ar_model,  bz, class_label=None, temperature=1.0, topK=None, topP=None, cfg=1.0, cfg_schedule=None, cfg_power=None, cache_kv=False, config=None, diti_scheduler=None):
    """
    Euler solver for Flow Matching with dynamic token masking (DiTi).
    """
    if config is None:
        raise ValueError("config is required for euler_sample_loop")
    # Derive device / shape from config to avoid relying on outer-scope globals.
    device = next(dec.parameters()).device
    patch_dim = config.patch_dim
    token_ids = ar_model.sampling(bz, class_label, temperature, topK, topP, cfg, cfg_schedule, cfg_power, cache_kv)
    quant = enc.quantizer.get_codes_w_indices(token_ids)
    query_len = config.x_len

    num_steps = 25
    x = torch.randn((bz, query_len, patch_dim), device=device)
    context_len = quant.shape[1]
    timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
    dt = 1.0 / num_steps
    
    for i in range(num_steps):
        t_curr = timesteps[i]
        t_tensor = torch.full((bz,), t_curr, device=device)
        
        attn_bias = None
        if config.use_tail_dropout and config.tied_timestep:
            if diti_scheduler is not None:
                attn_bias = build_taildrop_attn_mask(
                    t=t_tensor, context_len=context_len, query_len=query_len,
                    context_type=config.context_type, device=device, diti_scheduler=diti_scheduler,
                )
            else:
                attn_bias = build_taildrop_attn_mask(
                    t=t_tensor, context_len=context_len, query_len=query_len,
                    context_type=config.context_type, t2k=config.tail_dropout_t2k, device=device,
                )
        decoded = dec(quant, class_label, attn_mask=attn_bias, query_tokens=x, t=t_tensor)
        if config.l2_v_target:
            v_pred = decoded
            x = x - v_pred * dt
        else:
            x0_pred = decoded
            t_frac = dt/t_tensor
            t_frac = t_frac[:,None,None]
            x = x * (1-t_frac) +  x0_pred * t_frac
    return x

@torch.no_grad()
@torch._dynamo.disable
def euler_recon_loop(enc, dec, x, class_label=None, config=None, diti_scheduler=None):
    """
    Euler solver for Flow Matching with dynamic token masking (DiTi).
    """
    if config is None:
        raise ValueError("config is required for euler_recon_loop")
    # Derive device / shape from config to avoid relying on outer-scope globals.
    device = next(dec.parameters()).device
    B = x.shape[0]

    quant, idx, _, _ = enc(x,class_label, training=False)
    context_len = quant.shape[1]
    query_len = x.shape[1]

    num_steps = 25
    timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
    dt = 1.0 / num_steps
    xt = torch.randn_like(x)
    for i in range(num_steps):
        t_tensor = torch.full((B,), timesteps[i], device=device)
        attn_bias = None
        if config.use_tail_dropout and config.tied_timestep:
            if diti_scheduler is not None:
                attn_bias = build_taildrop_attn_mask(
                    t=t_tensor, context_len=context_len, query_len=query_len,
                    context_type=config.context_type, device=device, diti_scheduler=diti_scheduler,
                )
            else:
                attn_bias = build_taildrop_attn_mask(
                    t=t_tensor, context_len=context_len, query_len=query_len,
                    context_type=config.context_type, t2k=config.tail_dropout_t2k, device=device,
                )
        decoded = dec(quant, class_label, attn_mask=attn_bias, query_tokens=xt, t=t_tensor)
        if config.l2_v_target:
            v_pred = decoded
            xt = xt - v_pred * dt
        else:
            x0_pred = decoded
            t_frac = dt/t_tensor
            t_frac = t_frac[:,None,None]
            xt = xt * (1-t_frac) +  x0_pred * t_frac
    return xt, idx

def toogle_train_eval(model, train=True):
    if model is None: return
    if train:
        model.train()
    else:
        model.eval()

def toogle_require_grad(model, grads=True):
    if model is None:
        return
    model.requires_grad_(grads)

def zero_nan_gradients(model):
    for name, param in model.named_parameters():
        if param.grad is not None:
            param.grad.nan_to_num_(nan=0.0, posinf=1e5, neginf=-1e5)

# Model Prams and GFLOPS Utils
def print_model_info(fabric,model, inputs, name="Model"):
    if model is None: return
    if profile is None:
        fabric.print(f"[{name}] thop not installed, skipping GFLOPS calculation.")
        params = sum(p.numel() for p in model.parameters())
        fabric.print(f"[{name}] Params: {params/1e6:.2f}M")
        return

    # Move inputs to device
    device = next(model.parameters()).device
    inputs = tuple(i.to(device) if isinstance(i, torch.Tensor) else i for i in inputs)

    try:
        model_copy = copy.deepcopy(model)
        model_copy.eval()
        custom_ops = {}
        if count_linear is not None:
            custom_ops[Linear] = count_linear

        macs, params = profile(model_copy, inputs=inputs, custom_ops=custom_ops, verbose=False)
        gflops = macs * 2 / 1e9
        fabric.print(f"[{name}] Params: {params/1e6:.2f}M, GFLOPS: {gflops:.2f}")
        del model_copy
    except Exception as e:
        fabric.print(f"[{name}] Failed to calculate GFLOPS: {e}")
        params = sum(p.numel() for p in model.parameters())
        fabric.print(f"[{name}] Params: {params/1e6:.2f}M")
        
def print_model_stats(fabric,config, enc, dec, ar_model):
    x_len = int(config.x_len)
    patch_dim = int(config.patch_dim)
    num_classes = int(config.num_classes)
    z_len = int(config.z_len)
    z_dim = int(config.z_dim)
    codebook_size = int(config.codebook_size)
    
    x_dummy = torch.randn(1, x_len, patch_dim)
    label_dummy = torch.zeros(1, num_classes)
    
    print_model_info(fabric,enc, (x_dummy, label_dummy), "Encoder")
    
    z_dummy = torch.randn(1, z_len, z_dim)
    # If Decoder uses noise query, it requires query_tokens (and optionally t) at forward time.
    query_type = getattr(getattr(dec, "dec", None), "query_type", "learnable")
    if query_type == "noise":
        query_len = int(getattr(config, "x_len", z_len))
        query_dummy = torch.randn(1, query_len, patch_dim)
        t_dummy = torch.zeros(1)
        print_model_info(fabric, dec, (z_dummy, label_dummy, None, query_dummy, t_dummy), "Decoder")
    else:
        print_model_info(fabric,dec, (z_dummy, label_dummy), "Decoder")
    
    if ar_model is not None:
        idx_dummy = torch.randint(0, codebook_size, (1, z_len))
        print_model_info(fabric,ar_model, (idx_dummy, label_dummy), "ARModel")


def train(config):
    if config.use_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    # Parse devices
    devices_value = config.devices
    if isinstance(devices_value, str) and "," in devices_value:
        devices_for_fabric = [int(d.strip()) for d in devices_value.split(',')]
        num_devices = len(devices_for_fabric)
    else:
        try:
            num_devices = int(devices_value)
            devices_for_fabric = num_devices
        except (TypeError, ValueError):
            raise ValueError(f"Invalid value for devices: {devices_value}. Expected an integer or comma-separated string like '0,1,2'.")

    if config.batch_size % num_devices or config.batch_size < num_devices:
        raise ValueError("Global batch_size must be divisible by the number of devices")
    if config.eval_data_dir is not None and not Path(config.eval_fid_ref_path).is_file():
        raise FileNotFoundError("Prepare eval_fid_ref_path before training, or set eval_data_dir=null to disable periodic evaluation")

    # prepare seed and save dir
    if config.seed is not None:
        seed = int(config.seed) if config.seed is not None else 0
    seed_everything(seed)
    dl_generator = torch.Generator()
    dl_generator.manual_seed(seed)
    worker_init = make_worker_init_fn(seed)

    if "EXPERIMENT_SAVE_DIR" in os.environ:
        save_dir = os.environ["EXPERIMENT_SAVE_DIR"]
    else:
        experiment_index = len(glob(f"{config.save_dir}/*"))
        save_dir = config.save_dir+f"/{experiment_index:03d}"
        os.environ["EXPERIMENT_SAVE_DIR"] = save_dir
    logger = WandbLogger(project=config.wandb_project, save_dir=save_dir, name=config.wandb_name)

    # initialize fabric
    fabric = Fabric(accelerator="gpu", devices=devices_for_fabric, precision=config.precision, loggers=logger)
    fabric.launch()
    fabric.print(config)
    fabric.print(f"Global seed set to {config.seed}")

    # Initialize DiTi Scheduler if config provided
    diti_scheduler = None
    if getattr(config, "time_stages", None) and getattr(config, "k_per_stage", None):
        fabric.print("Initializing DiTi Scheduler for piecewise linear time-to-k mapping...")
        diti_scheduler = DiTiScheduler(
            n_timesteps=1000, 
            K=int(config.z_len),
            stages=config.time_stages,
            k_per_stage=config.k_per_stage,
            reverse_t=config.reverse_t
        )

    # Models and optimizers
    enc = Encoder(config).to(fabric.device) # always initialize encoder
    dec = Decoder(config).to(fabric.device) 
    if not config.train_ae:
        enc = enc.requires_grad_(False)
        dec = dec.requires_grad_(False)
    ar_model = ARModel(config).to(fabric.device) if config.train_ar else None 
    gan_loss = GANLoss(**config.GANLoss).to(fabric.device) if config.use_gan_loss else None
    perceptual_loss = LPIPS().to(fabric.device).eval().requires_grad_(False) if config.use_lpips_loss else None


    # Load Ckpt
    if config.ae_ckpt_path != "":
        fabric.print(f'resume ae from ckpt {config.ae_ckpt_path}')

        ckpt_states = fabric.load(config.ae_ckpt_path)
        if enc is not None and ckpt_states.get('enc', None) is not None:
            enc.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['enc'].items() })
            fabric.print("loaded enc")
        if dec is not None and ckpt_states.get('dec', None) is not None:
            dec.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['dec'].items() })
            fabric.print("loaded dec")

        if gan_loss is not None and ckpt_states.get('gan_loss', None) is not None:
            gan_loss.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['gan_loss'].items() })

    if config.ar_ckpt_path!="":
        fabric.print(f'resume ar from ckpt {config.ar_ckpt_path}')
        ar_ckpt_states = fabric.load(config.ar_ckpt_path)
        if config.load_ar_model and ar_model is not None and ar_ckpt_states.get('ar_model', None) is not None:
            ar_model.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['ar_model'].items() } )
            fabric.print("loaded ar")

    # for ema
    ae_ema_rate = math.pow(0.5,1/config.ae_ema_halflife) if config.ae_ema_halflife>0 else 0.
    ar_ema_rate = math.pow(0.5,1/config.ar_ema_halflife) if config.ar_ema_halflife>0 else 0.
    enc_ema = copy.deepcopy(enc).requires_grad_(False).eval() if ae_ema_rate>0 else None
    dec_ema = copy.deepcopy(dec).requires_grad_(False).eval() if ae_ema_rate>0 else None
    freeze_ar_model = copy.deepcopy(ar_model).requires_grad_(False).eval() if config.train_ar else None
    ar_model_ema = copy.deepcopy(ar_model).requires_grad_(False).eval() if config.train_ar and ar_ema_rate>0 else None
    
    if config.ae_ckpt_path != "":
        fabric.print(f'resume ae ema from ckpt {config.ae_ckpt_path}')
        ckpt_states = fabric.load(config.ae_ckpt_path)
        if enc_ema is not None and ckpt_states.get('enc_ema', None) is not None:
            enc_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['enc_ema'].items() })
        if dec_ema is not None and ckpt_states.get('dec_ema', None) is not None:
            dec_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ckpt_states['dec_ema'].items() })
        print("loaded ae ema")

    if config.ar_ckpt_path!="" and config.load_ar_model_ema and ar_model_ema is not None and ar_ckpt_states.get('ar_model_ema', None) is not None:
        ar_model_ema.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['ar_model_ema'].items() })
        print("loaded ar ema")
    if config.ar_ckpt_path!="" and config.load_ar_model_ema and freeze_ar_model is not None and ar_ckpt_states.get('freeze_ar_model', None) is not None:
        freeze_ar_model.load_state_dict({k.split('_orig_mod.')[-1]: v for k, v in ar_ckpt_states['freeze_ar_model'].items() })
        print("loaded frozen ar")

    # for inference
    enc_raw = enc
    dec_raw = dec
    ar_model_raw = ar_model
    enc_ema_raw = enc_ema
    dec_ema_raw = dec_ema
    ar_model_ema_raw = ar_model_ema

    if fabric.is_global_zero:
        fabric.print("Calculating model stats...")
        print_model_stats(fabric,config, enc_raw, dec_raw, ar_model_raw)
        root_dir = Path(save_dir)
        root_dir.mkdir(exist_ok=True, parents=True)
        OmegaConf.save(config, root_dir / "config.yaml", resolve=True)
        img_dir  = Path(save_dir+ "/images")
        img_dir.mkdir(exist_ok=True, parents=True)
        ckpt_dir  = Path(save_dir+ "/ckpts")
        ckpt_dir.mkdir(exist_ok=True, parents=True)
    
    fabric.barrier()

    if config.torch_compile:
        if enc is not None and config.compile_ae:
            enc = torch.compile(enc)
        if enc_ema is not None and config.compile_ae:
            enc_ema = torch.compile(enc_ema)
        if dec is not None and config.compile_ae:
            dec = torch.compile(dec)
        if dec_ema is not None and config.compile_ae:
            dec_ema = torch.compile(dec_ema)
        if ar_model is not None and config.compile_ar:
            ar_model = torch.compile(ar_model)
        if ar_model_ema is not None and config.compile_ar:
            ar_model_ema = torch.compile(ar_model_ema)
        if freeze_ar_model is not None and config.compile_ar:
            freeze_ar_model= torch.compile(freeze_ar_model)
        if gan_loss is not None and config.compile_gan:
            gan_loss = torch.compile(gan_loss)

    opt_enc = torch.optim.AdamW(enc.parameters() , lr=config.lr_enc, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_enc) if config.train_ae else None
    opt_ar = torch.optim.AdamW(ar_model.parameters(), lr=config.lr_ar, betas=(config.ar_beta1, config.ar_beta2), weight_decay=config.weight_decay_ar) if config.train_ar else None
    opt_dec = torch.optim.AdamW(dec.parameters(), lr=config.lr_dec, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_dec) if config.train_ae else None
    opt_gan_loss = torch.optim.AdamW(gan_loss.parameters(), lr=config.lr_gan_loss, betas=(config.beta1, config.beta2), weight_decay=config.weight_decay_gan) if config.use_gan_loss else None
    if  config.resume_optimizer:
        if config.train_ae:       
            fabric.print("resumed AE optimizer")
            opt_enc.load_state_dict(ckpt_states['opt_enc'])
            opt_dec.load_state_dict(ckpt_states['opt_dec'])
        if config.train_ar:
            fabric.print("resumed AR optimizer")
            ar_optimizer_states = ar_ckpt_states if config.ar_ckpt_path else ckpt_states
            opt_ar.load_state_dict(ar_optimizer_states['opt_ar'])
        if config.use_gan_loss:
            fabric.print("resumed GAN optimizer")
            opt_gan_loss.load_state_dict(ckpt_states['opt_gan_loss'])

    # Datasets
    dataset = ImageFolderDataset(path=config['data_dir'],
        resolution=config.image_size,
        use_label=config.use_label,
        max_size=None,
        xflip=config.xflip,
        crop_type=config.crop_type,
        deterministic_crop=bool(getattr(config, "deterministic_crop", False)),
        crop_seed=int(config.seed) if config.seed is not None else 0,
    )
    train_loader = DataLoader(
        dataset,
        batch_size=config.batch_size//num_devices,
        shuffle=True,
        num_workers=config.num_workers,
        pin_memory=True,
        drop_last=True,
        persistent_workers=config.num_workers > 0,
        worker_init_fn=worker_init,
        generator=dl_generator,
    )
    eval_loader=None
    if config.eval_data_dir is not None:
        eval_dataset = ImageFolderDataset(path=config['eval_data_dir'],
            resolution=config.image_size,
            use_label=config.use_label,
            max_size=None,
            xflip=False,
            crop_type=config.eval_crop_type,
            deterministic_crop=bool(getattr(config, "deterministic_crop", False)),
            crop_seed=int(config.seed) if config.seed is not None else 0,
        )
        eval_loader = DataLoader(
            eval_dataset,
            batch_size=config.batch_size//num_devices,
            shuffle=False,
            pin_memory=True,
            drop_last=False,
            worker_init_fn=worker_init,
            generator=dl_generator,
        )
    
    # Setup Model, Optimizer, DataLoader with Fabric
    if config.train_ae:
        enc, opt_enc = fabric.setup(enc, opt_enc)
        dec, opt_dec = fabric.setup(dec, opt_dec)
        if enc_ema is not None:
            enc_ema = enc_ema.to(fabric.device)
        if dec_ema is not None:
            dec_ema = dec_ema.to(fabric.device)
    if config.train_ar:
        ar_model, opt_ar = fabric.setup(ar_model, opt_ar)
        if ar_model_ema is not None:
            ar_model_ema = ar_model_ema.to(fabric.device)
        if freeze_ar_model is not None:
            freeze_ar_model = freeze_ar_model.to(fabric.device)
    if config.use_gan_loss:
        gan_loss, opt_gan_loss = fabric.setup(gan_loss, opt_gan_loss)

    train_loader = fabric.setup_dataloaders(train_loader)
    train_loader = InfiniteIterator(train_loader)
    if eval_loader is not None:
        eval_loader = fabric.setup_dataloaders(eval_loader)


    if config.train_ar:
        ar_model.mark_forward_method('sampling')

    global_step = 0

    phases = parse_phases(config.phases)
    fabric.print("Training Phases: ", phases)
    total_phase_steps = sum(phase.num_steps for phase in phases)
    total_phase = len(phases)

    prev_phase_idx=-1
    phase_step_accum = list(itertools.accumulate([phase.num_steps for phase in phases]))
    pbar = tqdm(total=total_phase_steps, disable=not fabric.is_global_zero, dynamic_ncols=True, file=sys.stdout)
    pbar.set_description(f"Total Steps {total_phase_steps}")
    data_buffer = []
    while global_step < total_phase_steps:
        # Assign gradient and optimizer based on current phase and target
        phase_idx, inner_idx, target, internel_step = get_phase(global_step, phases, phase_step_accum, config)
        if not any(target):
            phase_idx, inner_idx, target, internel_step = get_phase(global_step+1, phases, phase_step_accum, config)
        if phase_idx != prev_phase_idx:
            if config.lr_scheduler != 'none':
                scheduler_enc = torch.optim.lr_scheduler.CosineAnnealingLR(opt_enc, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_enc_min)  if config.train_ae else None
                scheduler_dec = torch.optim.lr_scheduler.CosineAnnealingLR(opt_dec, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_dec_min)  if config.train_ae else None
                scheduler_ar = torch.optim.lr_scheduler.CosineAnnealingLR(opt_ar, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_ar_min)  if config.train_ar else None
                scheduler_gan_loss = torch.optim.lr_scheduler.CosineAnnealingLR(opt_gan_loss, T_max=phases[phase_idx].num_steps/len(phases[phase_idx].internal_steps),eta_min=config.lr_gan_loss_min)  if config.use_gan_loss else None
            else:
                scheduler_enc = None
                scheduler_dec = None
                scheduler_ar = None
                scheduler_gan_loss = None


        enc_grad = target.DO_AE or target.DO_PRIOR_ENC
        dec_grad = target.DO_AE
        prior_grad = target.DO_PRIOR_AR
        disc_grad = target.DO_GAN_D 
        toogle_require_grad(enc, enc_grad)
        toogle_require_grad(dec, dec_grad)
        toogle_require_grad(ar_model, prior_grad)
        toogle_require_grad(gan_loss, disc_grad)
        opts = []
        if enc_grad:
            opts.append((enc, opt_enc, scheduler_enc))
        if dec_grad:
            opts.append((dec, opt_dec, scheduler_dec))
        if prior_grad:
            opts.append((ar_model, opt_ar, scheduler_ar))
        if disc_grad:
            opts.append((gan_loss, opt_gan_loss, scheduler_gan_loss))

        # input
        # Clear buffer at start of cycle
        if internel_step == 0:
            data_buffer = []

        if inner_idx == 0:
            batch = next(train_loader)
            data_buffer.append(batch)
        else:
            if len(data_buffer) > 0:
                batch = data_buffer.pop(0)
            else:
                batch = next(train_loader)

        imgs, labels = batch if isinstance(batch, (list, tuple)) and len(batch) == 2 else (batch, None)
        imgs = img_unint8_to_norm(imgs)

        # no labels or drop labels
        uncond_labels = F.one_hot(torch.full((imgs.shape[0],), config.num_classes - 1, device=fabric.device, dtype=torch.long), num_classes=config.num_classes).float()

        if labels is None or len(labels)==0 or not config.use_label:
            labels = uncond_labels
        elif config.label_drop_prob > 0:
            labels = torch.where(torch.rand(imgs.shape[0], device=fabric.device) < config.label_drop_prob, uncond_labels, labels)

        x = patchify(imgs, config.patch_size)

        
        # forward, build need computation graph for gradient calculation
        with torch.set_grad_enabled(enc_grad):
            quant, idx, one_hot, vqloss = enc(x, labels, training=enc_grad)
        
        if target.DO_AE:
            
            # For Diffusion/FM decoder
            t_fm,xt,noise,v = None,None,None,None
            if config.use_noise_query:
                t_fm = torch.rand((x.shape[0],), device=fabric.device) 
                xt, noise, v = make_noise_image_patch_query(x, t_fm, config.z_len) 

            # For Tail Dropout
            attn_bias = None
            if config.use_tail_dropout and global_step>config.tail_dropout_start:
                if config.tied_timestep:
                    t_tail_drop = t_fm
                else:
                    t_tail_drop = torch.rand((x.shape[0],), device=fabric.device)
                    no_drop_t = torch.ones_like(t_tail_drop)
                    t_tail_drop = torch.where(torch.rand_like(t_tail_drop)<=config.tail_dropout_p,t_tail_drop,no_drop_t)
                attn_bias = build_taildrop_attn_mask(t=t_tail_drop,context_len=int(quant.shape[1]),query_len=config.x_len,
                                    context_type=config.context_type,t2k=config.tail_dropout_t2k,
                                    device=fabric.device,
                                    diti_scheduler=diti_scheduler,
                                                    ) if config.use_tail_dropout else None

            decoded = dec(quant, labels, attn_mask=attn_bias, query_tokens=xt, t=t_fm)

            v_hat = decoded if config.l2_v_target else None
            x_hat = xt - decoded * t_fm[:,None,None] if config.l2_v_target else decoded

            imgs_hat = unpatchify(x_hat, config.image_size, config.patch_size)
            
            l2_target = v if config.l2_v_target else x
            l2_input = v_hat if config.l2_v_target else x_hat

        if target.DO_PRIOR_ENC or target.DO_PRIOR_AR:
            with torch.set_grad_enabled(prior_grad):
                ar_logits = ar_model(idx,labels)

        if target.DO_PRIOR_ENC:
            with torch.no_grad():
                ema_ar_logits = freeze_ar_model(idx,labels)
                with torch.amp.autocast('cuda',enabled=False):
                    ar_logits_detach = ar_logits.detach().float()
                    temp_low = config.prior_ar_temp_low
                    temp_high = config.prior_ar_temp_high
                    steps = config.z_len
                    if config.prior_ar_temp_schedule == 'cosine':
                        t = torch.linspace(0, 1, steps, device=fabric.device)
                        temp = temp_low + 0.5 * (temp_high - temp_low) * (1 - torch.cos(math.pi * t))
                    elif config.prior_ar_temp_schedule == 'linear':
                        temp = torch.linspace(temp_low, temp_high, steps, device=fabric.device)
                    elif config.prior_ar_temp_schedule == 'constant':
                        temp = torch.full((steps,), temp_low, device=fabric.device)
                    else:
                        raise NotImplementedError(f"Unknown schedule: {config.prior_ar_temp_schedule}")

                    b, n = ema_ar_logits.shape[0], ema_ar_logits.shape[1]
                    temp = temp.view(1, n, 1).expand(b, n, 1)

                    ema_ar_logits = ema_ar_logits.detach().float() / temp
                    prior_ema_ar_loss = F.cross_entropy(ema_ar_logits.reshape(-1, ema_ar_logits.shape[-1]).float(), idx.reshape(-1)) 

                    ema_ar_log_probs= torch.log_softmax(ema_ar_logits, dim=-1)
                    ar_log_probs = torch.log_softmax(ar_logits_detach, dim=-1)
                    entropy_ar_log_probs = -(ar_log_probs.exp() * ar_log_probs).sum(dim=-1)
                    entropy_ema_ar_log_probs = -(ema_ar_log_probs.exp() * ema_ar_log_probs).sum(dim=-1)
                    log_diff = ar_log_probs - ema_ar_log_probs


        # caculate the loss
        l2_loss = config.l2_weight * F.mse_loss(l2_input, l2_target) if target.DO_L2 else 0.
        l1_loss = config.l1_weight * F.l1_loss(x_hat, x) if target.DO_L1 else 0.
        lpips_loss = config.lpips_weight * perceptual_loss(imgs_hat, imgs).mean() if target.DO_LPIPS else 0.
        ae_loss = l2_loss + l1_loss + lpips_loss

        gan_G_loss, gan_G_loss_dict = gan_loss(imgs_hat, imgs, global_step=global_step, loss='G') if target.DO_GAN_G else (0., {})
        gan_D_loss, gan_D_loss_dict = gan_loss(imgs_hat, imgs, global_step=global_step, loss='D') if target.DO_GAN_D else (0., {})
        gan_G_loss_weight = (config.gan_G_weight * min(1., gan_loss.calculate_adaptive_weight(ae_loss,gan_G_loss,dec.dec.out.weight))) if target.DO_GAN_G and config.disc_adaptive_weight else config.gan_G_weight
        gan_G_loss = gan_G_loss * gan_G_loss_weight
        gan_D_loss = gan_D_loss * config.gan_G_weight
        correct_tokens = (ar_logits.argmax(dim=-1) == idx).detach().float().sum().item() if target.DO_PRIOR_AR else 0.
        total_tokens = torch.tensor(idx.numel()).item() if target.DO_PRIOR_AR else 0.
        correct_token_rate = correct_tokens / total_tokens if target.DO_PRIOR_AR else 0.
        prior_enc_loss = config.prior_enc_weight * (one_hot * log_diff.detach()).mean() if target.DO_PRIOR_ENC else 0.
        prior_ar_loss = F.cross_entropy(ar_logits.reshape(-1, ar_logits.shape[-1]).float(), idx.reshape(-1)) if target.DO_PRIOR_AR else 0.

        loss = ae_loss + prior_enc_loss + prior_ar_loss + gan_G_loss + gan_D_loss + vqloss 

        # backward and optimization
        fabric.backward(loss)
        
        for model, opt, scheduler in opts:
            if model is not None and opt is not None:
                zero_nan_gradients(model)
                if config.grad_clip>0:
                    fabric.clip_gradients(model, opt, max_norm=config.grad_clip)
                opt.step()
                opt.zero_grad(set_to_none=True)
            if scheduler is not None:
                scheduler.step() 
        if config.train_ae:
            if enc_grad:
                ema_update(enc, enc_ema, ae_ema_rate)
            if dec_grad:
                ema_update(dec, dec_ema, ae_ema_rate)
            if config.use_gan_loss and disc_grad:
                ema_update(gan_loss, gan_loss_ema, ae_ema_rate)
        if config.train_ar and prior_grad:
            ema_update(ar_model, ar_model_ema, ar_ema_rate)

        #visualization
        if (global_step + 1) % config.visualize_freq == 0:
            grid = x[:config.visualize_img_num]
            if not config.ema_reconstruction:
                if config.train_ae:
                    toogle_train_eval(enc_raw,train=False) 
                    toogle_train_eval(dec_raw,train=False) 
                eval_enc = enc_raw
                eval_dec = dec_raw
            else:
                eval_enc = enc_ema_raw
                eval_dec = dec_ema_raw
            if not config.ema_sampling:
                if config.train_ar:
                    toogle_train_eval(ar_model_raw,train=False) 
                eval_ar = ar_model_raw
            else:
                eval_ar = ar_model_ema_raw
            with torch.no_grad():
                noise_t = torch.ones((x.shape[0],), device=fabric.device) if config.use_noise_query else None
                if config.train_ae:
                    noise_recon = torch.randn_like(x) if config.use_noise_query else None
                    if config.use_noise_query and config.diffusion_decoder:
                        all_recon_images,_ = euler_recon_loop(eval_enc, eval_dec, x, labels, config=config, diti_scheduler=diti_scheduler)
                    else:
                        decoded_recon,_ = reconstruction(eval_enc, eval_dec, x, labels, query_tokens=noise_recon, t=noise_t)
                        all_recon_images = decoded_recon if not config.l2_v_target else noise_recon-decoded_recon
                    grid = torch.cat([grid, all_recon_images[:config.visualize_img_num]], dim=0)

                if config.train_ar:
                    noise_sample = torch.randn_like(x) if config.use_noise_query else None
                    if config.use_noise_query and config.diffusion_decoder:
                        all_sample_images = euler_sample_loop(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
                                            class_label=labels,temperature=config.temperature, 
                                            topK=config.topK, topP=config.topP, 
                                            cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
                                            cache_kv=config.cache_kv, config=config, diti_scheduler=diti_scheduler)
                    else:
                        decoded_sample = sampling(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
                                                class_label=labels,temperature=config.temperature, 
                                                topK=config.topK, topP=config.topP, 
                                                cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
                                                cache_kv=config.cache_kv, config=config, t=noise_t ,query_tokens=noise_sample)
                        all_sample_images = decoded_sample if not config.l2_v_target else noise_sample-decoded_sample
                    grid = torch.cat([grid, all_sample_images[:config.visualize_img_num]], dim=0)
                grid = img_denormalize(unpatchify(grid, config.image_size, config.patch_size))
                grid = torchvision.utils.make_grid(grid, nrow=config.visualize_img_num, normalize=False)
                fabric.log_dict({"visualization":wandb.Image(grid),'global_step':global_step+1})
            if not config.ema_reconstruction and config.train_ae:
                toogle_train_eval(enc_raw,train=True) 
                toogle_train_eval(dec_raw,train=True) 
            if not config.ema_sampling and config.train_ar:
                toogle_train_eval(ar_model_raw,train=True) 
           
        # eval and save ckpt
        rFID = 0.0
        gFID = 0.0
        if eval_loader is not None and (global_step + 1) % config.eval_freq == 0:
             
            sample_cached_path = os.path.join(save_dir, "sample_images.npz")
            recon_cached_path = os.path.join(save_dir, "recon_images.npz")
            gt_cache_path = config.eval_fid_ref_path

            samples_buf = []
            recons_buf = []
            if not config.ema_reconstruction:
                if config.train_ae:
                    toogle_train_eval(enc_raw,train=False) 
                    toogle_train_eval(dec_raw,train=False) 
                eval_enc = enc_raw
                eval_dec = dec_raw
            else:
                eval_enc = enc_ema_raw
                eval_dec = dec_ema_raw
            
            if not config.ema_sampling:
                if config.train_ar:
                    toogle_train_eval(ar_model_raw, train=False)
                eval_ar = ar_model_raw
            else:
                eval_ar = ar_model_ema_raw


            with torch.no_grad():

                fabric.barrier()
                for i, batch in enumerate(tqdm(eval_loader,disable=not fabric.is_global_zero,dynamic_ncols=True,file=sys.stdout,desc="Evaluating")):
                    imgs, labels = batch if isinstance(batch, (list, tuple)) and len(batch) == 2 else (batch, None)
                    x = img_unint8_to_norm(imgs)
                    x = patchify(x, config.patch_size)
                    uncond_labels = F.one_hot(torch.full((imgs.shape[0],), config.num_classes - 1, device=fabric.device, dtype=torch.long), num_classes=config.num_classes).float()
                    noise_t = torch.ones((labels.shape[0],), device=fabric.device) if config.use_noise_query else None

                    if labels is None or len(labels)==0 or not config.use_label:
                        labels = uncond_labels
                    idx = None
                    noise_recon = torch.randn_like(x) if config.use_noise_query else None

                    if config.train_ae:
                        if config.use_noise_query and config.diffusion_decoder:
                            recon_images, idx = euler_recon_loop(eval_enc, eval_dec, x, labels, config=config, diti_scheduler=diti_scheduler)
                        else:
                            decoded_recon, idx = reconstruction(eval_enc, eval_dec, x, labels, query_tokens=noise_recon,t = noise_t)
                            recon_images = decoded_recon if not config.l2_v_target else noise_recon-decoded_recon

                        recon_images = img_norm_to_uint8(unpatchify(recon_images, config.image_size, config.patch_size))
                        if fabric.world_size>1:
                            recon_images = fabric.all_gather(recon_images) 
                            recon_images = recon_images.flatten(0,1)
                        recons_buf.append(recon_images.permute(0,2,3,1).cpu().numpy())
                    if config.train_ar:
                        if config.use_noise_query and config.diffusion_decoder:
                            sample_images = euler_sample_loop(eval_enc, eval_dec, eval_ar, bz=x.shape[0], 
                                                class_label=labels,temperature=config.temperature, 
                                                topK=config.topK, topP=config.topP, 
                                                cfg=config.cfg, cfg_schedule=config.cfg_schedule, cfg_power=config.cfg_power, 
                                                cache_kv=config.cache_kv, config=config, diti_scheduler=diti_scheduler)
                        else:
                            noise_sample = torch.randn_like(x) if config.use_noise_query else None
                            decoded_sample = sampling(eval_enc, eval_dec, eval_ar, 
                                            bz=x.shape[0], class_label=labels,temperature=config.temperature, 
                                            topK=config.topK, topP=config.topP, cfg=config.cfg, cfg_schedule=config.cfg_schedule, 
                                            cfg_power=config.cfg_power, cache_kv=config.cache_kv, config=config,t=noise_t,query_tokens=noise_sample)
                            sample_images = decoded_sample if not config.l2_v_target else noise_sample-decoded_sample

                        sample_images = img_norm_to_uint8(unpatchify(sample_images, config.image_size, config.patch_size))
                        if fabric.world_size>1:
                            sample_images = fabric.all_gather(sample_images) 
                            sample_images = sample_images.flatten(0,1)
                        samples_buf.append(sample_images.permute(0,2,3,1).cpu().numpy())

        
                fabric.barrier()
                if fabric.is_global_zero:
                    if config.train_ar:
                        np.savez_compressed(sample_cached_path, np.concatenate(samples_buf,axis=0))
                    if config.train_ae:
                        np.savez_compressed(recon_cached_path, np.concatenate(recons_buf,axis=0))
                    
                    if config.train_ar:
                        gFID = adm_fid_evaluator(sample_cached_path, gt_cache_path, config, fabric)
                    if config.train_ae:
                        rFID = adm_fid_evaluator(recon_cached_path, gt_cache_path, config, fabric)
                fabric.barrier()


            if not config.ema_reconstruction and config.train_ae:
                toogle_train_eval(enc_raw,train=True) 
                toogle_train_eval(dec_raw,train=True) 
            if not config.ema_sampling and config.train_ar:
                toogle_train_eval(ar_model_raw,train=True) 

        # metrics
        if (global_step +1) % config.log_freq == 0:
            # current learning rates
            lr_enc = opt_enc.param_groups[0]['lr'] if opt_enc is not None else 0.
            lr_dec = opt_dec.param_groups[0]['lr'] if opt_dec is not None else 0.
            lr_ar = opt_ar.param_groups[0]['lr'] if opt_ar is not None else 0.
            lr_gan_loss = opt_gan_loss.param_groups[0]['lr'] if opt_gan_loss is not None else 0.
            metrics = {
                        'Phase':phase_idx,
                        'ae_loss/l1_loss':l1_loss,
                        'ae_loss/l2_loss':l2_loss,
                        'ae_loss/lpips_loss':lpips_loss,
                        'ae_loss/loss':ae_loss, 
                        'ae_loss/vqloss': vqloss,
                        'ae_loss/gan_G_loss':gan_G_loss,
                        'ae_loss/gan_D_loss':gan_D_loss,
                        'prior_loss/logq-logp':log_diff.mean() if target.DO_PRIOR_ENC else 0., 
                        'prior_loss/logq':ar_log_probs.mean() if target.DO_PRIOR_ENC else 0., 
                        'prior_loss/logp':ema_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0., 
                        'prior_loss/entropy_q':entropy_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0.,
                        'prior_loss/entropy_p':entropy_ema_ar_log_probs.mean() if target.DO_PRIOR_ENC else 0.,
                        'prior_loss/ar_loss':prior_ar_loss, 
                        'prior_loss/ema_ar_loss':prior_ema_ar_loss  if target.DO_PRIOR_ENC else 0.,
                        'prior_loss/correct_token_rate': correct_token_rate,
                        'FID/rFID':rFID,
                        'FID/gFID':gFID,
                        'lr/lr_enc': lr_enc,
                        'lr/lr_dec': lr_dec,
                        'lr/lr_ar': lr_ar,
                        'lr/lr_gan_loss': lr_gan_loss,
                      }
            gan_loss_dict = {**gan_G_loss_dict, **gan_D_loss_dict, 'gan_G_loss_weight':gan_G_loss_weight}
            metrics.update(gan_loss_dict)
            metrics = {k: fabric.all_reduce(v).item() if isinstance(v, torch.Tensor) else v for k, v in metrics.items() }
            metrics_logger = {k: v for k, v in metrics.items() if v!=0. }
            metrics_4f = {k: f"{v:.4f}" if k!='Phase' else int(v) for k, v in metrics.items() }
            fabric.log_dict(metrics_logger,step=global_step+1)
            pbar.set_postfix(metrics_4f)
            pbar.update(config.log_freq)

        if config.save_ckpt and (global_step==0 or (global_step+1) % config.ckpt_freq == 0):
            ckpt_name = f'Phase={phase_idx}-Step={global_step+1}-rFID={rFID:.4f}-gFID={gFID:.4f}' if rFID is not None and gFID is not None else "_".join([f'{k}={v:.4f}' for k, v in metrics.items()]) 
            ckpt_path = f'{save_dir}/ckpts/{ckpt_name}.ckpt'
            fabric.save(ckpt_path, {'enc':enc,
                        'ar_model':ar_model,
                        'freeze_ar_model': freeze_ar_model,
                        'dec':dec,
                        'gan_loss':gan_loss,
                        'enc_ema':enc_ema,
                        'dec_ema':dec_ema,
                        'ar_model_ema':ar_model_ema,
                        'opt_enc':opt_enc,
                        'opt_ar':opt_ar,
                        'opt_dec':opt_dec,
                        'opt_gan_loss':opt_gan_loss})
        global_step += 1
        prev_phase_idx = phase_idx
    if config.save_ckpt:
        ckpt_path = f'{save_dir}/ckpts/last.ckpt'
        fabric.save(ckpt_path, {'enc':enc,
                    'ar_model':ar_model,
                    'dec':dec,
                    'gan_loss':gan_loss,
                    'freeze_ar_model':freeze_ar_model,
                    'enc_ema':enc_ema,
                    'dec_ema':dec_ema,
                    'ar_model_ema':ar_model_ema,
                    'opt_enc':opt_enc,
                    'opt_ar':opt_ar,
                    'opt_dec':opt_dec,
                    'opt_gan_loss':opt_gan_loss})

    logger.finalize("Training completed")

if __name__ == "__main__":
    OmegaConf.register_new_resolver("eval", eval)
    cli_conf = OmegaConf.from_cli()

    config_path = cli_conf.get("--config")
    conf = OmegaConf.load(config_path)
    if cli_conf is not None:
        for k, v in cli_conf.items():
            OmegaConf.update(conf, k, v)

    train(conf)