semanticist/engine/gpt_trainer.py

import os, torch
import os.path as osp
import shutil
import numpy as np
import copy
import torch.nn as nn
from tqdm.auto import tqdm
from accelerate import Accelerator
from torchvision.utils import make_grid, save_image
from torch.utils.data import DataLoader, DistributedSampler
from semanticist.utils.logger import SmoothedValue, MetricLogger, empty_cache
from accelerate.utils import DistributedDataParallelKwargs
from semanticist.stage2.gpt import GPT_models
from semanticist.stage2.generate import generate
from pathlib import Path
import time

from semanticist.engine.trainer_utils import (
    instantiate_from_config, concat_all_gather,
    save_img_batch, get_fid_stats,
    EMAModel, create_scheduler, load_state_dict, load_safetensors,
    setup_result_folders, create_optimizer,
    CacheDataLoader
)

class GPTTrainer(nn.Module):
    def __init__(
        self,
        ae_model,
        gpt_model,
        dataset,
        test_only=False,
        num_test_images=50000,
        num_epoch=400,
        eval_classes=[1, 7, 282, 604, 724, 207, 250, 751, 404, 850], # goldfish, cock, tiger cat, hourglass, ship, golden retriever, husky, race car, airliner, teddy bear
        blr=1e-4,
        cosine_lr=False,
        lr_min=0,
        warmup_epochs=100,
        warmup_steps=None,
        warmup_lr_init=0,
        decay_steps=None,
        batch_size=32,
        cache_bs=8,
        test_bs=100,
        num_workers=8,
        pin_memory=False,
        max_grad_norm=None,
        grad_accum_steps=1,
        precision='bf16',
        save_every=10000,
        sample_every=1000,
        fid_every=50000,
        result_folder=None,
        log_dir="./log",
        ae_cfg=1.0,
        cfg=6.0,
        cfg_schedule="linear",
        temperature=1.0,
        train_num_slots=None,
        test_num_slots=None,
        eval_fid=False,
        fid_stats=None,
        enable_ema=False,
        compile=False,
        enable_cache_latents=True,
        cache_dir='/dev/shm/slot_cache'
    ):
        super().__init__()
        kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
        self.accelerator = Accelerator(
            kwargs_handlers=[kwargs],
            mixed_precision=precision,
            gradient_accumulation_steps=grad_accum_steps,
            log_with="tensorboard",
            project_dir=log_dir,
        )

        self.ae_model = instantiate_from_config(ae_model)
        ae_model_path = ae_model.params.ckpt_path
        assert ae_model_path.endswith(".safetensors") or ae_model_path.endswith(".pt") or ae_model_path.endswith(".pth") or ae_model_path.endswith(".pkl")
        assert osp.exists(ae_model_path), f"AE model checkpoint {ae_model_path} does not exist"
        self._load_checkpoint(ae_model_path, self.ae_model)

        self.ae_model.to(self.device)
        for param in self.ae_model.parameters():
            param.requires_grad = False
        self.ae_model.eval()

        self.model_name = gpt_model.target
        if 'GPT' in gpt_model.target:
            self.gpt_model = GPT_models[gpt_model.target](**gpt_model.params)
        else:
            raise ValueError(f"Unknown model type: {gpt_model.target}")
        self.num_slots = ae_model.params.num_slots
        self.slot_dim = ae_model.params.slot_dim

        self.test_only = test_only
        self.test_bs = test_bs
        self.num_test_images = num_test_images
        self.num_classes = gpt_model.params.num_classes
        self.batch_size = batch_size
        if not test_only:
            self.train_ds = instantiate_from_config(dataset)
            train_size = len(self.train_ds)
            if self.accelerator.is_main_process:
                print(f"train dataset size: {train_size}")

            sampler = DistributedSampler(
                self.train_ds,
                num_replicas=self.accelerator.num_processes,
                rank=self.accelerator.process_index,
                shuffle=True,
            )
            self.train_dl = DataLoader(
                self.train_ds,
                batch_size=batch_size if not enable_cache_latents else cache_bs,
                sampler=sampler,
                num_workers=num_workers,
                pin_memory=pin_memory,
                drop_last=True,
            )

            effective_bs = batch_size * grad_accum_steps * self.accelerator.num_processes
            lr = blr * effective_bs / 256
            if self.accelerator.is_main_process:
                print(f"Effective batch size is {effective_bs}")

            self.g_optim = create_optimizer(self.gpt_model, weight_decay=0.05, learning_rate=lr)
            
            if warmup_epochs is not None:
                warmup_steps = warmup_epochs * len(self.train_dl)
                
            self.g_sched = create_scheduler(
                self.g_optim,
                num_epoch,
                len(self.train_dl),
                lr_min,
                warmup_steps,
                warmup_lr_init,
                decay_steps,
                cosine_lr
            )
            self.accelerator.register_for_checkpointing(self.g_sched)
            self.gpt_model, self.g_optim, self.g_sched = self.accelerator.prepare(self.gpt_model, self.g_optim, self.g_sched)
        else:
            self.gpt_model = self.accelerator.prepare(self.gpt_model)

        self.steps = 0
        self.loaded_steps = -1

        if compile:
            self.ae_model = torch.compile(self.ae_model, mode="reduce-overhead")
            _model = self.accelerator.unwrap_model(self.gpt_model)
            _model = torch.compile(_model, mode="reduce-overhead")

        self.enable_ema = enable_ema
        if self.enable_ema and not self.test_only: # when testing, we directly load the ema dict and skip here
            self.ema_model = EMAModel(self.accelerator.unwrap_model(self.gpt_model), self.device)
            self.accelerator.register_for_checkpointing(self.ema_model)

        self._load_checkpoint(gpt_model.params.ckpt_path)
        if self.test_only:
            self.steps = self.loaded_steps

        self.num_epoch = num_epoch
        self.save_every = save_every
        self.sample_every = sample_every
        self.fid_every = fid_every
        self.max_grad_norm = max_grad_norm
        self.eval_classes = eval_classes
        self.cfg = cfg
        self.ae_cfg = ae_cfg
        self.cfg_schedule = cfg_schedule
        self.temperature = temperature
        self.train_num_slots = train_num_slots
        self.test_num_slots = test_num_slots
        if self.train_num_slots is not None:
            self.train_num_slots = min(self.train_num_slots, self.num_slots)
        else:
            self.train_num_slots = self.num_slots
        if self.test_num_slots is not None:
            self.num_slots_to_gen = min(self.test_num_slots, self.train_num_slots)
        else:
            self.num_slots_to_gen = self.train_num_slots
        self.eval_fid = eval_fid
        if eval_fid:
            assert fid_stats is not None
        self.fid_stats = fid_stats

        # Setup result folders
        self.result_folder = result_folder
        self.model_saved_dir, self.image_saved_dir = setup_result_folders(result_folder)

        # Setup cache
        self.cache_dir = Path(cache_dir)
        self.enable_cache_latents = enable_cache_latents
        self.cache_loader = None

    @property
    def device(self):
        return self.accelerator.device

    def _load_checkpoint(self, ckpt_path=None, model=None):
        if ckpt_path is None or not osp.exists(ckpt_path):
            return

        if model is None:
            model = self.accelerator.unwrap_model(self.gpt_model)

        if osp.isdir(ckpt_path):
            self.loaded_steps = int(
                ckpt_path.split("step")[-1].split("/")[0]
            )
            if not self.test_only:
                self.accelerator.load_state(ckpt_path)
            else:
                if self.enable_ema:
                    model_path = osp.join(ckpt_path, "custom_checkpoint_1.pkl")
                    if osp.exists(model_path):
                        state_dict = torch.load(model_path, map_location="cpu")
                        load_state_dict(state_dict, model)
                        if self.accelerator.is_main_process:
                            print(f"Loaded ema model from {model_path}")
                else:
                    model_path = osp.join(ckpt_path, "model.safetensors")
                    if osp.exists(model_path):
                        load_safetensors(model_path, model)
        else:
            if ckpt_path.endswith(".safetensors"):
                load_safetensors(ckpt_path, model)
            else:
                state_dict = torch.load(ckpt_path, map_location="cpu")
                load_state_dict(state_dict, model)
        if self.accelerator.is_main_process:
            print(f"Loaded checkpoint from {ckpt_path}")

    def _build_cache(self):
        """Build cache for slots and targets."""
        rank = self.accelerator.process_index
        world_size = self.accelerator.num_processes
        
        # Clean up any existing cache files first
        slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
        targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
        
        if slots_file.exists():
            os.remove(slots_file)
        if targets_file.exists():
            os.remove(targets_file)
        
        dataset_size = len(self.train_dl.dataset)
        shard_size = dataset_size // world_size
        
        # Detect number of augmentations from first batch
        with torch.no_grad():
            sample_batch = next(iter(self.train_dl))
            img, _ = sample_batch
            num_augs = img.shape[1] if len(img.shape) == 5 else 1
        
        print(f"Rank {rank}: Creating new cache with {num_augs} augmentations per image...")
        os.makedirs(self.cache_dir, exist_ok=True)
        slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
        targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
        
        # Create memory-mapped files
        slots_mmap = np.memmap(
            slots_file,
            dtype='float32',
            mode='w+',
            shape=(shard_size * num_augs, self.train_num_slots, self.slot_dim)
        )
        
        targets_mmap = np.memmap(
            targets_file,
            dtype='int64',
            mode='w+',
            shape=(shard_size * num_augs,)
        )
        
        # Cache data
        with torch.no_grad():
            for i, batch in enumerate(tqdm(
                self.train_dl, 
                desc=f"Rank {rank}: Caching data",
                disable=not self.accelerator.is_local_main_process
            )):
                imgs, targets = batch
                if len(imgs.shape) == 5:  # [B, num_augs, C, H, W]
                    B, A, C, H, W = imgs.shape
                    imgs = imgs.view(-1, C, H, W)  # [B*num_augs, C, H, W]
                    targets = targets.unsqueeze(1).expand(-1, A).reshape(-1)  # [B*num_augs]
                
                # Split imgs into n chunks
                num_splits = num_augs
                split_size = imgs.shape[0] // num_splits
                imgs_splits = torch.split(imgs, split_size)
                targets_splits = torch.split(targets, split_size)
                
                start_idx = i * self.train_dl.batch_size * num_augs
                
                for split_idx, (img_split, targets_split) in enumerate(zip(imgs_splits, targets_splits)):
                    img_split = img_split.to(self.device, non_blocking=True)
                    slots_split = self.ae_model.encode_slots(img_split)[:, :self.train_num_slots, :]
                    
                    split_start = start_idx + (split_idx * split_size)
                    split_end = split_start + img_split.shape[0]
                    
                    # Write directly to mmap files
                    slots_mmap[split_start:split_end] = slots_split.cpu().numpy()
                    targets_mmap[split_start:split_end] = targets_split.numpy()
        
        # Close the mmap files
        del slots_mmap
        del targets_mmap
        
        # Reopen in read mode
        self.cached_latents = np.memmap(
            slots_file,
            dtype='float32',
            mode='r',
            shape=(shard_size * num_augs, self.train_num_slots, self.slot_dim)
        )
        
        self.cached_targets = np.memmap(
            targets_file,
            dtype='int64',
            mode='r',
            shape=(shard_size * num_augs,)
        )
        
        # Store the number of augmentations for the cache loader
        self.num_augs = num_augs

    def _setup_cache(self):
        """Setup cache if enabled."""
        self._build_cache()
        self.accelerator.wait_for_everyone()

        # Initialize cache loader if cache exists
        if self.cached_latents is not None:
            self.cache_loader = CacheDataLoader(
                slots=self.cached_latents,
                targets=self.cached_targets,
                batch_size=self.batch_size,
                num_augs=self.num_augs,
                seed=42 + self.accelerator.process_index
            )

    def __del__(self):
        """Cleanup cache files."""
        if self.enable_cache_latents:
            rank = self.accelerator.process_index
            world_size = self.accelerator.num_processes
            
            # Clean up slots cache
            slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
            if slots_file.exists():
                os.remove(slots_file)
            
            # Clean up targets cache
            targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
            if targets_file.exists():
                os.remove(targets_file)

    def _train_step(self, slots, targets=None):
        """Execute single training step."""
        
        with self.accelerator.accumulate(self.gpt_model):
            with self.accelerator.autocast():
                loss = self.gpt_model(slots, targets)
            
            self.accelerator.backward(loss)
            if self.accelerator.sync_gradients and self.max_grad_norm is not None:
                self.accelerator.clip_grad_norm_(self.gpt_model.parameters(), self.max_grad_norm)
            self.g_optim.step()
            if self.g_sched is not None:
                self.g_sched.step_update(self.steps)
            self.g_optim.zero_grad()

        # Update EMA model if enabled
        if self.enable_ema:
            self.ema_model.update(self.accelerator.unwrap_model(self.gpt_model))
        
        return loss

    def _train_epoch_cached(self, epoch, logger):
        """Train one epoch using cached data."""
        self.cache_loader.set_epoch(epoch)
        header = f'Epoch: [{epoch}/{self.num_epoch}]'
        
        for batch in logger.log_every(self.cache_loader, 20, header):
            slots, targets = (b.to(self.device, non_blocking=True) for b in batch)
            
            self.steps += 1

            if self.steps == 1:
                print(f"Training batch size: {len(slots)}")
                print(f"Hello from index {self.accelerator.local_process_index}")
            
            loss = self._train_step(slots, targets)
            self._handle_periodic_ops(loss, logger)

    def _train_epoch_uncached(self, epoch, logger):
        """Train one epoch using raw data."""
        header = f'Epoch: [{epoch}/{self.num_epoch}]'
        
        for batch in logger.log_every(self.train_dl, 20, header):
            img, targets = (b.to(self.device, non_blocking=True) for b in batch)
            
            self.steps += 1
            
            if self.steps == 1:
                print(f"Training batch size: {img.size(0)}")
                print(f"Hello from index {self.accelerator.local_process_index}")

            slots = self.ae_model.encode_slots(img)[:, :self.train_num_slots, :]
            loss = self._train_step(slots, targets)
            self._handle_periodic_ops(loss, logger)

    def _handle_periodic_ops(self, loss, logger):
        """Handle periodic operations and logging."""
        logger.update(loss=loss.item())
        logger.update(lr=self.g_optim.param_groups[0]["lr"])
        
        if self.steps % self.save_every == 0:
            self.save()
        
        if (self.steps % self.sample_every == 0) or (self.eval_fid and self.steps % self.fid_every == 0):
            empty_cache()
            self.evaluate()
            self.accelerator.wait_for_everyone()
            empty_cache()

    def _save_config(self, config):
        """Save configuration file."""
        if config is not None and self.accelerator.is_main_process:
            import shutil
            from omegaconf import OmegaConf

            if isinstance(config, str) and osp.exists(config):
                shutil.copy(config, osp.join(self.result_folder, "config.yaml"))
            else:
                config_save_path = osp.join(self.result_folder, "config.yaml")
                OmegaConf.save(config, config_save_path)

    def _should_skip_epoch(self, epoch):
        """Check if epoch should be skipped due to loaded checkpoint."""
        loader = self.train_dl if not self.enable_cache_latents else self.cache_loader
        if ((epoch + 1) * len(loader)) <= self.loaded_steps:
            if self.accelerator.is_main_process:
                print(f"Epoch {epoch} is skipped because it is loaded from ckpt")
            self.steps += len(loader)
            return True
        
        if self.steps < self.loaded_steps:
            for _ in loader:
                self.steps += 1
                if self.steps >= self.loaded_steps:
                    break
        return False

    def train(self, config=None):
        """Main training loop."""
        # Initial setup
        n_parameters = sum(p.numel() for p in self.parameters() if p.requires_grad)
        if self.accelerator.is_main_process:
            print(f"number of learnable parameters: {n_parameters//1e6}M")
        
        self._save_config(config)
        self.accelerator.init_trackers("gpt")

        # Handle test-only mode
        if self.test_only:
            empty_cache()
            self.evaluate()
            self.accelerator.wait_for_everyone()
            empty_cache()
            return

        # Setup cache if enabled
        if self.enable_cache_latents:
            self._setup_cache()

        # Training loop
        for epoch in range(self.num_epoch):
            if self._should_skip_epoch(epoch):
                continue

            self.gpt_model.train()
            logger = MetricLogger(delimiter="  ")
            logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))

            # Choose training path based on cache availability
            if self.enable_cache_latents:
                self._train_epoch_cached(epoch, logger)
            else:
                self._train_epoch_uncached(epoch, logger)

            # Synchronize and log epoch stats
            logger.synchronize_between_processes()
            if self.accelerator.is_main_process:
                print("Averaged stats:", logger)

        # Finish training
        self.accelerator.end_training()
        self.save()
        if self.accelerator.is_main_process:
            print("Train finished!")

    def save(self):
        self.accelerator.wait_for_everyone()
        self.accelerator.save_state(
            os.path.join(self.model_saved_dir, f"step{self.steps}")
        )

    @torch.no_grad()
    def evaluate(self, use_ema=True):
        self.gpt_model.eval()
        unwraped_gpt_model = self.accelerator.unwrap_model(self.gpt_model)
        # switch to ema params, only when eval_fid is True
        # if test_only, we directly load the ema dict and skip here
        use_ema = use_ema and self.enable_ema and self.eval_fid and not self.test_only
        if use_ema:
            if hasattr(self, "ema_model"):
                model_without_ddp = self.accelerator.unwrap_model(self.gpt_model)
                model_state_dict = copy.deepcopy(model_without_ddp.state_dict())
                ema_state_dict = copy.deepcopy(model_without_ddp.state_dict())
                for i, (name, _value) in enumerate(model_without_ddp.named_parameters()):
                    if "nested_sampler" in name:
                        continue
                    ema_state_dict[name] = self.ema_model.state_dict()[name]
                if self.accelerator.is_main_process:
                    print("Switch to ema")
                model_without_ddp.load_state_dict(ema_state_dict)
            else:
                print("EMA model not found, using original model")
                use_ema = False

        if not self.test_only:
            classes = torch.tensor(self.eval_classes, device=self.device)
            with self.accelerator.autocast():
                slots = generate(unwraped_gpt_model, classes, self.num_slots_to_gen, cfg_scale=self.cfg, cfg_schedule=self.cfg_schedule, temperature=self.temperature)
                if self.num_slots_to_gen < self.num_slots:
                    null_slots = self.ae_model.dit.null_cond.expand(slots.shape[0], -1, -1)
                    null_slots = null_slots[:, self.num_slots_to_gen:, :]
                    slots = torch.cat([slots, null_slots], dim=1)
                imgs = self.ae_model.sample(slots, targets=classes, cfg=self.ae_cfg) # targets are not used for now

            imgs = concat_all_gather(imgs)
            if self.accelerator.num_processes > 16:
                imgs = imgs[:16*len(self.eval_classes)]
            imgs = imgs.detach().cpu()
            grid = make_grid(
                imgs, nrow=len(self.eval_classes), normalize=True, value_range=(0, 1)
            )
            if self.accelerator.is_main_process:
                save_image(
                    grid,
                    os.path.join(
                        self.image_saved_dir, f"step{self.steps}_aecfg-{self.ae_cfg}_cfg-{self.cfg_schedule}-{self.cfg}_slots{self.num_slots_to_gen}_temp{self.temperature}.jpg"
                    ),
                )
        if self.eval_fid and (self.test_only or (self.steps % self.fid_every == 0)):
            # Create output directory (only on main process)
            save_folder = os.path.join(self.image_saved_dir, f"gen_step{self.steps}_aecfg-{self.ae_cfg}_cfg-{self.cfg_schedule}-{self.cfg}_slots{self.num_slots_to_gen}_temp{self.temperature}")
            if self.accelerator.is_main_process:
                os.makedirs(save_folder, exist_ok=True)

            # Setup for distributed generation
            world_size = self.accelerator.num_processes
            local_rank = self.accelerator.process_index
            batch_size = self.test_bs
            
            # Create balanced class distribution
            num_classes = self.num_classes
            images_per_class = self.num_test_images // num_classes
            class_labels = np.repeat(np.arange(num_classes), images_per_class)
            
            # Shuffle the class labels to ensure random ordering
            np.random.shuffle(class_labels)
            
            total_images = len(class_labels)

            padding_size = world_size * batch_size - (total_images % (world_size * batch_size))
            class_labels = np.pad(class_labels, (0, padding_size), 'constant')
            padded_total_images = len(class_labels)

            # Distribute workload across GPUs
            images_per_gpu = padded_total_images // world_size
            start_idx = local_rank * images_per_gpu
            end_idx = min(start_idx + images_per_gpu, padded_total_images)
            local_class_labels = class_labels[start_idx:end_idx]
            local_num_steps = len(local_class_labels) // batch_size
            
            if self.accelerator.is_main_process:
                print(f"Generating {total_images} images ({images_per_class} per class) across {world_size} GPUs")
            
            used_time = 0
            gen_img_cnt = 0
            
            for i in range(local_num_steps):
                if self.accelerator.is_main_process and i % 10 == 0:
                    print(f"Generation step {i}/{local_num_steps}")

                # Get and pad labels for current batch
                batch_start = i * batch_size
                batch_end = batch_start + batch_size
                labels = local_class_labels[batch_start:batch_end]
                
                # Convert to tensors and track real vs padding
                labels = torch.tensor(labels, device=self.device)
                
                # Generate images
                self.accelerator.wait_for_everyone()
                start_time = time.time()
                with torch.no_grad():
                    with self.accelerator.autocast():
                        slots = generate(unwraped_gpt_model, labels, self.num_slots_to_gen, 
                                            cfg_scale=self.cfg, 
                                            cfg_schedule=self.cfg_schedule,
                                            temperature=self.temperature)
                        if self.num_slots_to_gen < self.num_slots:
                            null_slots = self.ae_model.dit.null_cond.expand(slots.shape[0], -1, -1)
                            null_slots = null_slots[:, self.num_slots_to_gen:, :]
                            slots = torch.cat([slots, null_slots], dim=1)
                        imgs = self.ae_model.sample(slots, targets=labels, cfg=self.ae_cfg)

                samples_in_batch = min(batch_size * world_size, total_images - gen_img_cnt)

                # Update timing stats
                used_time += time.time() - start_time
                gen_img_cnt += samples_in_batch
                if self.accelerator.is_main_process and i % 10 == 0:
                    print(f"Avg generation time: {used_time/gen_img_cnt:.5f} sec/image")

                gathered_imgs = concat_all_gather(imgs)
                gathered_imgs = gathered_imgs[:samples_in_batch]
                
                # Save images (only on main process)
                if self.accelerator.is_main_process:
                    real_imgs = gathered_imgs.detach().cpu()
                    
                    save_paths = [
                        os.path.join(save_folder, f"{str(idx).zfill(5)}.png")
                        for idx in range(gen_img_cnt - samples_in_batch, gen_img_cnt)
                    ]
                    save_img_batch(real_imgs, save_paths)

            # Calculate metrics (only on main process)
            self.accelerator.wait_for_everyone()
            if self.accelerator.is_main_process:
                generated_files = len(os.listdir(save_folder))
                print(f"Generated {generated_files} images out of {total_images} expected")
                
                metrics_dict = get_fid_stats(save_folder, None, self.fid_stats)
                fid = metrics_dict["frechet_inception_distance"]
                inception_score = metrics_dict["inception_score_mean"]
                
                metric_prefix = "fid_ema" if use_ema else "fid"
                isc_prefix = "isc_ema" if use_ema else "isc"
                
                self.accelerator.log({
                    metric_prefix: fid,
                    isc_prefix: inception_score,
                    "gpt_cfg": self.cfg,
                    "ae_cfg": self.ae_cfg,
                    "cfg_schedule": self.cfg_schedule,
                    "temperature": self.temperature,
                    "num_slots": self.test_num_slots if self.test_num_slots is not None else self.train_num_slots
                }, step=self.steps)
                
                # Print comprehensive CFG information
                cfg_info = (
                    f"{'EMA ' if use_ema else ''}CFG params: "
                    f"gpt_cfg={self.cfg}, ae_cfg={self.ae_cfg}, "
                    f"cfg_schedule={self.cfg_schedule}, "
                    f"num_slots={self.test_num_slots if self.test_num_slots is not None else self.train_num_slots}, "
                    f"temperature={self.temperature}"
                )
                print(cfg_info)
                print(f"FID: {fid:.2f}, ISC: {inception_score:.2f}")

                # Cleanup
                shutil.rmtree(save_folder)

        # back to no ema
        if use_ema:
            if self.accelerator.is_main_process:
                print("Switch back from ema")
            model_without_ddp.load_state_dict(model_state_dict)

        self.gpt_model.train()