jobs/process/TrainVAEProcess.py

import copy
import glob
import os
import shutil
import time
from collections import OrderedDict

from PIL import Image
from PIL.ImageOps import exif_transpose
from einops import rearrange
from safetensors.torch import save_file, load_file
from torch.utils.data import DataLoader, ConcatDataset
import torch
from torch import nn
from torchvision.transforms import transforms

from jobs.process import BaseTrainProcess
from toolkit.image_utils import show_tensors
from toolkit.kohya_model_util import load_vae, convert_diffusers_back_to_ldm
from toolkit.data_loader import ImageDataset
from toolkit.losses import ComparativeTotalVariation, get_gradient_penalty, PatternLoss, total_variation, total_variation_deltas
from toolkit.metadata import get_meta_for_safetensors
from toolkit.optimizer import get_optimizer
from toolkit.style import get_style_model_and_losses
from toolkit.train_tools import get_torch_dtype
from diffusers import AutoencoderKL, AutoencoderTiny
from toolkit.models.autoencoder_tiny_with_pooled_exits import AutoencoderTinyWithPooledExits
from tqdm import tqdm
import math
import torchvision.utils
import time
import numpy as np
from .models.critic import Critic
from torchvision.transforms import Resize
import lpips
import random
import traceback
from transformers import SiglipImageProcessor, SiglipVisionModel
import torch.nn.functional as F

IMAGE_TRANSFORMS = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)


def unnormalize(tensor):
    return (tensor / 2 + 0.5).clamp(0, 1)


def channel_dropout(x, p=0.5):
    keep_prob = 1 - p
    mask = torch.rand(x.size(0), x.size(1), 1, 1, device=x.device, dtype=x.dtype) < keep_prob
    mask = mask / keep_prob  # scale
    return x * mask


def sharpen_image(images: torch.Tensor) -> torch.Tensor:
    # Define sharpening kernel
    kernel = torch.tensor([
        [ 0, -1,  0],
        [-1,  5, -1],
        [ 0, -1,  0]
    ], dtype=images.dtype, device=images.device).view(1, 1, 3, 3)

    # Repeat kernel for each channel
    kernel = kernel.repeat(3, 1, 1, 1)  # (out_channels, in_channels/groups, kH, kW)

    # Apply the filter
    sharpened = F.conv2d(images, kernel, padding=1, groups=3)

    return sharpened

class TrainVAEProcess(BaseTrainProcess):
    def __init__(self, process_id: int, job, config: OrderedDict):
        super().__init__(process_id, job, config)
        self.data_loader = None
        self.vae = None
        self.target_latent_vae = None
        self.device = self.get_conf('device', self.job.device)
        self.vae_path = self.get_conf('vae_path', None)
        self.target_latent_vae_path = self.get_conf('target_latent_vae_path', None)
        self.eq_vae = self.get_conf('eq_vae', False)
        self.datasets_objects = self.get_conf('datasets', required=True)
        self.batch_size = self.get_conf('batch_size', 1, as_type=int)
        self.resolution = self.get_conf('resolution', 256, as_type=int)
        self.sample_resolution = self.get_conf('sample_resolution', self.resolution, as_type=int)
        self.learning_rate = self.get_conf('learning_rate', 1e-6, as_type=float)
        self.sample_every = self.get_conf('sample_every', None)
        self.optimizer_type = self.get_conf('optimizer', 'adam')
        self.epochs = self.get_conf('epochs', None, as_type=int)
        self.max_steps = self.get_conf('max_steps', None, as_type=int)
        self.save_every = self.get_conf('save_every', None)
        self.dtype = self.get_conf('dtype', 'float32')
        self.sample_sources = self.get_conf('sample_sources', None)
        self.log_every = self.get_conf('log_every', 100, as_type=int)
        self.style_weight = self.get_conf('style_weight', 0, as_type=float)
        self.content_weight = self.get_conf('content_weight', 0, as_type=float)
        self.kld_weight = self.get_conf('kld_weight', 0, as_type=float)
        self.clip_weight = self.get_conf('clip_weight', 0, as_type=float)
        self.mse_weight = self.get_conf('mse_weight', 1e0, as_type=float)
        self.mae_weight = self.get_conf('mae_weight', 0, as_type=float)
        self.mv_loss_weight = self.get_conf('mv_loss_weight', 0, as_type=float)
        self.tv_weight = self.get_conf('tv_weight', 0, as_type=float)
        self.ltv_weight = self.get_conf('ltv_weight', 0, as_type=float)
        self.lpm_weight = self.get_conf('lpm_weight', 0, as_type=float) # latent pixel matching
        self.lpips_weight = self.get_conf('lpips_weight', 0, as_type=float)
        self.critic_weight = self.get_conf('critic_weight', 1, as_type=float)
        self.pattern_weight = self.get_conf('pattern_weight', 0, as_type=float)
        self.optimizer_params = self.get_conf('optimizer_params', {})
        self.vae_config = self.get_conf('vae_config', None)
        self.dropout = self.get_conf('dropout', 0.0, as_type=float)
        self.train_encoder = self.get_conf('train_encoder', False, as_type=bool)
        self.random_scaling = self.get_conf('random_scaling', False, as_type=bool)
        self.vae_type = self.get_conf('vae_type', 'AutoencoderKL', as_type=str) # AutoencoderKL or AutoencoderTiny
        self.only_if_contains = self.get_conf('only_if_contains', None)

        self.do_pooled_exits = False
        self.VaeClass = AutoencoderKL
        if self.vae_type == 'AutoencoderTiny':
            self.VaeClass = AutoencoderTiny
        if self.vae_type == 'AutoencoderTinyWithPooledExits':
            self.VaeClass = AutoencoderTinyWithPooledExits
            self.do_pooled_exits = True
        
        if not self.train_encoder:
            # remove losses that only target encoder
            self.kld_weight = 0
            self.mv_loss_weight = 0
            self.ltv_weight = 0
            self.lpm_weight = 0

        self.blocks_to_train = self.get_conf('blocks_to_train', ['all'])
        self.torch_dtype = get_torch_dtype(self.dtype)
        self.vgg_19 = None
        self.clip = None
        self.clip_image_processor = None
        self.clip_image_size = 256
        self.style_weight_scalers = []
        self.content_weight_scalers = []
        self.lpips_loss:lpips.LPIPS = None

        self.vae_scale_factor = 8
        self.target_vae_scale_factor = 8

        self.step_num = 0
        self.epoch_num = 0

        self.use_critic = self.get_conf('use_critic', False, as_type=bool)
        self.critic = None

        if self.use_critic:
            self.critic = Critic(
                device=self.device,
                dtype=self.dtype,
                process=self,
                **self.get_conf('critic', {})  # pass any other params
            )

        if self.sample_every is not None and self.sample_sources is None:
            raise ValueError('sample_every is specified but sample_sources is not')

        if self.epochs is None and self.max_steps is None:
            raise ValueError('epochs or max_steps must be specified')

        self.data_loaders = []
        # check datasets
        assert isinstance(self.datasets_objects, list)
        for dataset in self.datasets_objects:
            if 'path' not in dataset:
                raise ValueError('dataset must have a path')
            # check if is dir
            if not os.path.isdir(dataset['path']):
                raise ValueError(f"dataset path does is not a directory: {dataset['path']}")

        # make training folder
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root, exist_ok=True)

        self._pattern_loss = None

    def update_training_metadata(self):
        self.add_meta(OrderedDict({"training_info": self.get_training_info()}))

    def get_training_info(self):
        info = OrderedDict({
            'step': self.step_num,
            'epoch': self.epoch_num,
        })
        return info

    def load_datasets(self):
        if self.data_loader is None:
            print(f"Loading datasets")
            datasets = []
            for dataset in self.datasets_objects:
                print(f" - Dataset: {dataset['path']}")
                ds = copy.copy(dataset)
                dataset_res = self.resolution
                if self.random_scaling:
                    # scale 2x to allow for random scaling
                    dataset_res = int(dataset_res * 2)
                ds['resolution'] = dataset_res
                image_dataset = ImageDataset(ds)
                datasets.append(image_dataset)

            concatenated_dataset = ConcatDataset(datasets)
            self.data_loader = DataLoader(
                concatenated_dataset,
                batch_size=self.batch_size,
                shuffle=True,
                num_workers=16
            )

    def remove_oldest_checkpoint(self):
        max_to_keep = 4
        folders = glob.glob(os.path.join(self.save_root, f"{self.job.name}*_diffusers"))
        if len(folders) > max_to_keep:
            folders.sort(key=os.path.getmtime)
            for folder in folders[:-max_to_keep]:
                print(f"Removing {folder}")
                shutil.rmtree(folder)
        # also handle CRITIC_vae_42_000000500.safetensors format for critic
        critic_files = glob.glob(os.path.join(self.save_root, f"CRITIC_{self.job.name}*.safetensors"))
        if len(critic_files) > max_to_keep:
            critic_files.sort(key=os.path.getmtime)
            for file in critic_files[:-max_to_keep]:
                print(f"Removing {file}")
                os.remove(file)

    def setup_vgg19(self):
        if self.vgg_19 is None:
            self.vgg_19, self.style_losses, self.content_losses, self.vgg19_pool_4 = get_style_model_and_losses(
                single_target=True,
                device=self.device,
                output_layer_name='pool_4',
                dtype=self.torch_dtype
            )
            self.vgg_19.to(self.device, dtype=self.torch_dtype)
            self.vgg_19.requires_grad_(False)

            # we run random noise through first to get layer scalers to normalize the loss per layer
            # bs of 2 because we run pred and target through stacked
            noise = torch.randn((2, 3, self.resolution, self.resolution), device=self.device, dtype=self.torch_dtype)
            self.vgg_19(noise)
            for style_loss in self.style_losses:
                # get a scaler  to normalize to 1
                scaler = 1 / torch.mean(style_loss.loss).item()
                self.style_weight_scalers.append(scaler)
            for content_loss in self.content_losses:
                # get a scaler  to normalize to 1
                scaler = 1 / torch.mean(content_loss.loss).item()
                self.content_weight_scalers.append(scaler)

            self.print(f"Style weight scalers: {self.style_weight_scalers}")
            self.print(f"Content weight scalers: {self.content_weight_scalers}")
    
    def setup_clip(self):
        ckpt = 'google/siglip2-base-patch16-256'
        if self.resolution == 512:
            ckpt = 'google/siglip2-so400m-patch16-512'
            # ckpt = 'google/siglip2-base-patch16-512'
            self.clip_image_size = 512
        self.print(f"Loading CLIP model from {ckpt}")
        vision_encoder = SiglipVisionModel.from_pretrained(ckpt, device_map="auto", torch_dtype=torch.bfloat16).eval()
        processor = SiglipImageProcessor.from_pretrained(ckpt)
        self.clip = vision_encoder
        self.clip_image_processor = processor
        
    def get_clip_embeddings(self, image_n1p1):
        tensors_0_1 = (image_n1p1 + 1) / 2
        # sharpen images
        tensors_0_1 = sharpen_image(tensors_0_1)
        
        tensors_0_1 = tensors_0_1.clamp(0, 1)

        # resize if needed
        if tensors_0_1.shape[-2:] != (self.clip_image_size, self.clip_image_size):
            tensors_0_1 = torch.nn.functional.interpolate(tensors_0_1, size=(self.clip_image_size, self.clip_image_size), mode='bilinear', align_corners=False)

        mean = torch.tensor([0.5, 0.5, 0.5]).to(
            tensors_0_1.device, dtype=tensors_0_1.dtype
        ).view([1, 3, 1, 1]).detach()
        std = torch.tensor([0.5, 0.5, 0.5]).to(
            tensors_0_1.device, dtype=tensors_0_1.dtype
        ).view([1, 3, 1, 1]).detach()
        
        # tensors_0_1 = torch.clip((255. * tensors_0_1), 0, 255).round() / 255.0
        clip_image = (tensors_0_1 - mean) / std
        
        id_embeds = self.clip(
            clip_image.to(self.clip.device, dtype=torch.bfloat16),
            output_hidden_states=True,
        )  
        last_hidden_state = id_embeds['last_hidden_state']
        return last_hidden_state

    def get_clip_loss(self, pred, target):
        # pred and target come in as -1 to 1.
        with torch.no_grad():
            target_embeddings = self.get_clip_embeddings(target).float()
        pred_embeddings = self.get_clip_embeddings(pred).float()
        return torch.nn.functional.mse_loss(pred_embeddings, target_embeddings)
    
    def get_pooled_output_loss(self, pooled_outputs, target):
        if pooled_outputs is None:
            return torch.tensor(0.0, device=self.device)

        # pooled_outputs is a list of tensors, each with shape (batch_size, 3, h, w)
        # target is a tensor with shape (batch_size, 3, h, w)
        loss = 0.0
        for pooled_output in pooled_outputs:
            with torch.no_grad():
                # resize target to match pooled_output size
                target_resized = torch.nn.functional.interpolate(target, size=pooled_output.shape[2:], mode='bilinear', align_corners=False)
            loss += torch.nn.functional.mse_loss(pooled_output.float(), target_resized.float())
        return loss / len(pooled_outputs) if len(pooled_outputs) > 0 else torch.tensor(0.0, device=self.device)

    def get_style_loss(self):
        if self.style_weight > 0:
            # scale all losses with loss scalers
            loss = torch.sum(
                torch.stack([loss.loss * scaler for loss, scaler in zip(self.style_losses, self.style_weight_scalers)]))
            return loss
        else:
            return torch.tensor(0.0, device=self.device)

    def get_content_loss(self):
        if self.content_weight > 0:
            # scale all losses with loss scalers
            loss = torch.sum(torch.stack(
                [loss.loss * scaler for loss, scaler in zip(self.content_losses, self.content_weight_scalers)]))
            return loss
        else:
            return torch.tensor(0.0, device=self.device)

    def get_mse_loss(self, pred, target):
        if self.mse_weight > 0:
            loss_fn = nn.MSELoss()
            loss_normal = loss_fn(pred, target)

            pred_sharp = sharpen_image(pred)
            with torch.no_grad():
                target_sharp = sharpen_image(target)
            
            loss_sharp = loss_fn(pred_sharp, target_sharp)

            return (loss_sharp + loss_normal) / 2
        else:
            return torch.tensor(0.0, device=self.device)
    
    def get_mae_loss(self, pred, target):
        if self.mae_weight > 0:
            loss_fn = nn.L1Loss()
            loss_normal = loss_fn(pred, target)
            
            pred_sharp = sharpen_image(pred)
            with torch.no_grad():
                target_sharp = sharpen_image(target)
            loss_sharp = loss_fn(pred_sharp, target_sharp)
            return (loss_sharp + loss_normal) / 2
        else:
            return torch.tensor(0.0, device=self.device)

    def get_kld_loss(self, mu, log_var):
        if self.kld_weight > 0:
            # Kullback-Leibler divergence
            # added here for full training (not implemented). Not needed for only decoder
            # as we are not changing the distribution of the latent space
            # normally it would help keep a normal distribution for latents
            KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())  # KL divergence
            return KLD
        else:
            return torch.tensor(0.0, device=self.device)

    def get_mean_variance_loss(self, latents: torch.Tensor):
        if self.mv_loss_weight > 0:
            # collapse rows into channels
            latents_col = rearrange(latents, 'b c h (gw w) -> b (c gw) h w', gw=latents.shape[-1])
            mean_col = latents_col.mean(dim=(2, 3), keepdim=True)
            std_col  = latents_col.std(dim=(2, 3), keepdim=True, unbiased=False)
            mean_loss_col = (mean_col ** 2).mean()
            std_loss_col  = ((std_col - 1) ** 2).mean()

            # collapse columns into channels
            latents_row = rearrange(latents, 'b c (gh h) w -> b (c gh) h w', gh=latents.shape[-2])
            mean_row = latents_row.mean(dim=(2, 3), keepdim=True)
            std_row  = latents_row.std(dim=(2, 3), keepdim=True, unbiased=False)
            mean_loss_row = (mean_row ** 2).mean()
            std_loss_row  = ((std_row - 1) ** 2).mean()

            # do a global one
            
            mean = latents.mean(dim=(2, 3), keepdim=True)
            std  = latents.std(dim=(2, 3), keepdim=True, unbiased=False)
            mean_loss_global = (mean ** 2).mean()
            std_loss_global  = ((std - 1) ** 2).mean()

            return (mean_loss_col + std_loss_col + mean_loss_row + std_loss_row + mean_loss_global + std_loss_global) / 3
        else:
            return torch.tensor(0.0, device=self.device)
        
    def get_ltv_loss(self, latent, images):
        # loss to reduce the latent space variance
        if self.ltv_weight > 0:
            with torch.no_grad():
                images = images.to(latent.device, dtype=latent.dtype)
                # resize down to latent size
                images = torch.nn.functional.interpolate(images, size=(latent.shape[2], latent.shape[3]), mode='bilinear', align_corners=False)
                
                # mean the color channel and then expand to latent size
                images = images.mean(dim=1, keepdim=True)
                images = images.repeat(1, latent.shape[1], 1, 1)
                
                # normalize to a mean of 0 and std of 1
                images_mean = images.mean(dim=(2, 3), keepdim=True)
                images_std = images.std(dim=(2, 3), keepdim=True)
                images = (images - images_mean) / (images_std + 1e-6)
                
                # now we target the same std of the image for the latent space as to not reduce to 0
            
            latent_tv = torch.abs(total_variation_deltas(latent))
            images_tv = torch.abs(total_variation_deltas(images))
            loss = torch.abs(latent_tv - images_tv) # keep it spatially aware
            loss = loss.mean(dim=2, keepdim=True)
            loss = loss.mean(dim=3, keepdim=True)  # mean over height and width
            loss = loss.mean(dim=1, keepdim=True)  # mean over channels
            loss = loss.mean()
            return loss
        else:
            return torch.tensor(0.0, device=self.device)
        
    def get_latent_pixel_matching_loss(self, latent, pixels):
        if self.lpm_weight > 0:
            with torch.no_grad():
                pixels = pixels.to(latent.device, dtype=latent.dtype)
                # resize down to latent size
                pixels = torch.nn.functional.interpolate(pixels, size=(latent.shape[2], latent.shape[3]), mode='bilinear', align_corners=False)
                
                # mean the color channel and then expand to latent size
                pixels = pixels.mean(dim=1, keepdim=True)
                pixels = pixels.repeat(1, latent.shape[1], 1, 1)
                # match the mean std of latent
                latent_mean = latent.mean(dim=(2, 3), keepdim=True)
                latent_std = latent.std(dim=(2, 3), keepdim=True)
                pixels_mean = pixels.mean(dim=(2, 3), keepdim=True)
                pixels_std = pixels.std(dim=(2, 3), keepdim=True)
                pixels = (pixels - pixels_mean) / (pixels_std + 1e-6) * latent_std + latent_mean
                
            return torch.nn.functional.mse_loss(latent.float(), pixels.float())
            
        else:
            return torch.tensor(0.0, device=self.device)

    def get_tv_loss(self, pred, target):
        if self.tv_weight > 0:
            get_tv_loss = ComparativeTotalVariation()
            loss = get_tv_loss(pred, target)
            return loss
        else:
            return torch.tensor(0.0, device=self.device)

    def get_pattern_loss(self, pred, target):
        if self._pattern_loss is None:
            self._pattern_loss = PatternLoss(pattern_size=16, dtype=self.torch_dtype).to(self.device,
                                                                                        dtype=self.torch_dtype)
        loss = torch.mean(self._pattern_loss(pred, target))
        return loss

    def save(self, step=None):
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root, exist_ok=True)

        step_num = ''
        if step is not None:
            # zeropad 9 digits
            step_num = f"_{str(step).zfill(9)}"

        self.update_training_metadata()
        filename = f'{self.job.name}{step_num}_diffusers'

        self.vae = self.vae.to("cpu", dtype=torch.float16)
        self.vae.save_pretrained(
            save_directory=os.path.join(self.save_root, filename)
        )
        self.vae = self.vae.to(self.device, dtype=self.torch_dtype)

        self.print(f"Saved to {os.path.join(self.save_root, filename)}")

        if self.use_critic:
            self.critic.save(step)

        self.remove_oldest_checkpoint()

    def sample(self, step=None):
        sample_folder = os.path.join(self.save_root, 'samples')
        if not os.path.exists(sample_folder):
            os.makedirs(sample_folder, exist_ok=True)

        with torch.no_grad():
            for i, img_url in enumerate(self.sample_sources):
                img = exif_transpose(Image.open(img_url))
                img = img.convert('RGB')
                # crop if not square
                if img.width != img.height:
                    min_dim = min(img.width, img.height)
                    img = img.crop((0, 0, min_dim, min_dim))
                # resize
                img = img.resize((self.sample_resolution, self.sample_resolution))

                input_img = img
                img = IMAGE_TRANSFORMS(img).unsqueeze(0).to(self.device, dtype=self.torch_dtype)
                img = img
                # latent = self.vae.encode(img).latent_dist.sample()

                target_latent = None
                if self.target_latent_vae is not None:
                    target_input_scale = self.target_vae_scale_factor / self.vae_scale_factor 
                    target_input_size = (int(img.shape[2] * target_input_scale), int(img.shape[3] * target_input_scale))
                    # resize to target input size
                    target_input_batch = Resize(target_input_size)(img).to(self.device, dtype=torch.float32)
                    target_latent = self.target_latent_vae.encode(target_input_batch).latent_dist.sample().detach()
                    shift = self.target_latent_vae.config['shift_factor'] if self.target_latent_vae.config['shift_factor'] is not None else 0
                    target_latent = self.target_latent_vae.config['scaling_factor'] * (target_latent - shift)
                    target_latent = target_latent.to(self.device, dtype=self.torch_dtype)
                latent = self.vae.encode(img, return_dict=False)[0]
                
                if hasattr(latent, 'sample'):
                    latent = latent.sample()
                
                shift = self.vae.config['shift_factor'] if self.vae.config['shift_factor'] is not None else 0
                latent = self.vae.config['scaling_factor'] * (latent - shift)
                
                latent_img = latent.clone()
                bs, ch, h, w = latent_img.shape
                grid_size = math.ceil(math.sqrt(ch))
                pad = grid_size * grid_size - ch

                # take first item in batch
                latent_img = latent_img[0]  # shape: (ch, h, w)

                if pad > 0:
                    padding = torch.zeros((pad, h, w), dtype=latent_img.dtype, device=latent_img.device)
                    latent_img = torch.cat([latent_img, padding], dim=0)

                # make grid
                new_img = torch.zeros((1, grid_size * h, grid_size * w), dtype=latent_img.dtype, device=latent_img.device)
                for x in range(grid_size):
                    for y in range(grid_size):
                        if x * grid_size + y < ch:
                            new_img[0, x * h:(x + 1) * h, y * w:(y + 1) * w] = latent_img[x * grid_size + y]
                latent_img = new_img
                # make rgb
                latent_img = latent_img.repeat(3, 1, 1).unsqueeze(0)
                latent_img = (latent_img / 2 + 0.5).clamp(0, 1)
                
                # resize to 256x256
                latent_img = torch.nn.functional.interpolate(latent_img, size=(self.sample_resolution, self.sample_resolution), mode='nearest')
                latent_img = latent_img.squeeze(0).cpu().permute(1, 2, 0).float().numpy()
                latent_img = (latent_img * 255).astype(np.uint8)
                # convert to pillow image
                latent_img = Image.fromarray(latent_img)
                
                if target_latent is not None:
                    latent = target_latent.to(latent.device, dtype=latent.dtype)
                    
                shift = self.vae.config['shift_factor'] if self.vae.config['shift_factor'] is not None else 0
                latent = latent / self.vae.config['scaling_factor'] + shift
                
                decoded = self.vae.decode(latent).sample
                decoded = (decoded / 2 + 0.5).clamp(0, 1)
                # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
                decoded = decoded.cpu().permute(0, 2, 3, 1).squeeze(0).float().numpy()

                # convert to pillow image
                decoded = Image.fromarray((decoded * 255).astype(np.uint8))

                # stack input image and decoded image
                input_img = input_img.resize((self.sample_resolution, self.sample_resolution))
                decoded = decoded.resize((self.sample_resolution, self.sample_resolution))

                output_img = Image.new('RGB', (self.sample_resolution * 3, self.sample_resolution))
                output_img.paste(input_img, (0, 0))
                output_img.paste(decoded, (self.sample_resolution, 0))
                output_img.paste(latent_img, (self.sample_resolution * 2, 0))

                scale_up = 2
                if output_img.height <= 300:
                    scale_up = 4
                if output_img.height >= 1000:
                    scale_up = 1

                # scale up using nearest neighbor
                output_img = output_img.resize((output_img.width * scale_up, output_img.height * scale_up), Image.NEAREST)

                step_num = ''
                if step is not None:
                    # zero-pad 9 digits
                    step_num = f"_{str(step).zfill(9)}"
                seconds_since_epoch = int(time.time())
                # zero-pad 2 digits
                i_str = str(i).zfill(2)
                filename = f"{seconds_since_epoch}{step_num}_{i_str}.png"
                output_img.save(os.path.join(sample_folder, filename))

    def load_vae(self):
        path_to_load = self.vae_path
        # see if we have a checkpoint in out output to resume from
        self.print(f"Looking for latest checkpoint in {self.save_root}")
        files = glob.glob(os.path.join(self.save_root, f"{self.job.name}*_diffusers"))
        if files and len(files) > 0:
            latest_file = max(files, key=os.path.getmtime)
            print(f" - Latest checkpoint is: {latest_file}")
            path_to_load = latest_file
            # todo update step and epoch count
        else:
            self.print(f" - No checkpoint found, starting from scratch")
        # load vae
        self.print(f"Loading VAE")
        self.print(f" - Loading VAE: {path_to_load}")
        if self.vae is None:
            if path_to_load is not None:
                self.vae = self.VaeClass.from_pretrained(path_to_load)
            elif self.vae_config is not None:
                self.vae = self.VaeClass(**self.vae_config)
            else:
                raise ValueError('vae_path or ae_config must be specified')

        # set decoder to train
        self.vae.to(self.device, dtype=self.torch_dtype)
        if self.eq_vae:
            self.vae.encoder.train()
        else:
            self.vae.requires_grad_(False)
            self.vae.eval()
        self.vae.decoder.train()
        self.vae_scale_factor = 2 ** (len(self.vae.config['block_out_channels']) - 1)
        
        if self.target_latent_vae_path is not None:
            self.print(f"Loading target latent VAE from {self.target_latent_vae_path}")
            self.target_latent_vae = AutoencoderKL.from_pretrained(self.target_latent_vae_path)
            self.target_latent_vae.to(self.device, dtype=torch.float32)
            self.target_latent_vae.eval()
            self.target_vae_scale_factor = 2 ** (len(self.target_latent_vae.config['block_out_channels']) - 1)
        else:
            self.target_latent_vae = None
            self.target_vae_scale_factor = self.vae_scale_factor

    def run(self):
        super().run()
        self.load_datasets()

        max_step_epochs = self.max_steps // len(self.data_loader)
        num_epochs = self.epochs
        if num_epochs is None or num_epochs > max_step_epochs:
            num_epochs = max_step_epochs

        max_epoch_steps = len(self.data_loader) * num_epochs
        num_steps = self.max_steps
        if num_steps is None or num_steps > max_epoch_steps:
            num_steps = max_epoch_steps
        self.max_steps = num_steps
        self.epochs = num_epochs
        start_step = self.step_num
        self.first_step = start_step

        self.print(f"Training VAE")
        self.print(f" - Training folder: {self.training_folder}")
        self.print(f" - Batch size: {self.batch_size}")
        self.print(f" - Learning rate: {self.learning_rate}")
        self.print(f" - Epochs: {num_epochs}")
        self.print(f" - Max steps: {self.max_steps}")

        # load vae
        self.load_vae()

        params = []

        # only set last 2 layers to trainable
        for param in self.vae.decoder.parameters():
            param.requires_grad = False

        train_all = 'all' in self.blocks_to_train

        if train_all:
            params = list(self.vae.decoder.named_parameters())
            self.vae.decoder.requires_grad_(True)
            if self.train_encoder:
                # encoder
                params += list(self.vae.encoder.named_parameters())
                self.vae.encoder.requires_grad_(True)
        else:
            # mid_block
            if train_all or 'mid_block' in self.blocks_to_train:
                params += list(self.vae.decoder.mid_block.named_parameters())
                self.vae.decoder.mid_block.requires_grad_(True)
            # up_blocks
            if train_all or 'up_blocks' in self.blocks_to_train:
                params += list(self.vae.decoder.up_blocks.named_parameters())
                self.vae.decoder.up_blocks.requires_grad_(True)
            # conv_out (single conv layer output)
            if train_all or 'conv_out' in self.blocks_to_train:
                params += list(self.vae.decoder.conv_out.named_parameters())
                self.vae.decoder.conv_out.requires_grad_(True)

        if self.style_weight > 0 or self.content_weight > 0:
            self.setup_vgg19()
            # self.vgg_19.requires_grad_(False)
            self.vgg_19.eval()
        
        if self.use_critic:
            self.critic.setup()

        if self.clip_weight > 0:
            self.setup_clip()

        if self.lpips_weight > 0 and self.lpips_loss is None:
            # self.lpips_loss = lpips.LPIPS(net='vgg')
            self.lpips_loss = lpips.LPIPS(net='vgg').to(self.device, dtype=torch.bfloat16)
            
        if self.only_if_contains is not None:
            orig_params = params
            params = []
            for name, param in orig_params:
                for contains in self.only_if_contains:
                    if contains in name:
                        params.append(param)
                        break

        optimizer = get_optimizer(params, self.optimizer_type, self.learning_rate,
                                  optimizer_params=self.optimizer_params)

        # setup scheduler
        # todo allow other schedulers
        scheduler = torch.optim.lr_scheduler.ConstantLR(
            optimizer,
            total_iters=num_steps,
            factor=1,
            # verbose=False
        )

        # setup tqdm progress bar
        self.progress_bar = tqdm(
            total=num_steps,
            desc='Training VAE',
            leave=True
        )

        # sample first
        self.sample()
        blank_losses = OrderedDict({
            "total": [],
            "lpips": [],
            "style": [],
            "content": [],
            "mse": [],
            "mae": [],
            "lat_mse": [],
            "mvl": [],
            "ltv": [],
            "lpm": [],
            "kl": [],
            "tv": [],
            "clip": [],
            "pool": [],
            "ptn": [],
            "crD": [],
            "crG": [],
        })
        epoch_losses = copy.deepcopy(blank_losses)
        log_losses = copy.deepcopy(blank_losses)
        # range start at self.epoch_num go to self.epochs
        
        latent_size = self.resolution // self.vae_scale_factor
        
        for epoch in range(self.epoch_num, self.epochs, 1):
            if self.step_num >= self.max_steps:
                break
            for batch in self.data_loader:
                if self.step_num >= self.max_steps:
                    break
                with torch.no_grad():
                    batch = batch.to(self.device, dtype=self.torch_dtype)
                    
                    if self.random_scaling:
                        # only random scale 0.5 of the time
                        if random.random() < 0.5:
                            # random scale the batch
                            scale_factor = 0.25
                        else:
                            scale_factor = 0.5
                        new_size = (int(batch.shape[2] * scale_factor), int(batch.shape[3] * scale_factor))
                        # make sure it is vae divisible
                        new_size = (new_size[0] // self.vae_scale_factor * self.vae_scale_factor,
                                    new_size[1] // self.vae_scale_factor * self.vae_scale_factor)
                            

                    # resize so it matches size of vae evenly
                    if batch.shape[2] % self.vae_scale_factor != 0 or batch.shape[3] % self.vae_scale_factor != 0:
                        batch = Resize((batch.shape[2] // self.vae_scale_factor * self.vae_scale_factor,
                                                batch.shape[3] // self.vae_scale_factor * self.vae_scale_factor))(batch)
                    
                    target_latent = None
                    lat_mse_loss = torch.tensor(0.0, device=self.device)
                    if self.target_latent_vae is not None:
                        target_input_scale = self.target_vae_scale_factor / self.vae_scale_factor 
                        target_input_size = (int(batch.shape[2] * target_input_scale), int(batch.shape[3] * target_input_scale))
                        # resize to target input size
                        target_input_batch = Resize(target_input_size)(batch).to(self.device, dtype=torch.float32)
                        target_latent = self.target_latent_vae.encode(target_input_batch).latent_dist.sample().detach()
                        # shift scale it
                        shift = self.target_latent_vae.config['shift_factor'] if self.target_latent_vae.config['shift_factor'] is not None else 0
                        target_latent = self.target_latent_vae.config['scaling_factor'] * (target_latent - shift)
                        target_latent = target_latent.to(self.device, dtype=self.torch_dtype)
                        

                    # forward pass
                # grad only if eq_vae
                with torch.set_grad_enabled(self.train_encoder):
                    if self.vae_type != 'AutoencoderKL':
                        # AutoencoderTiny cannot do latent distribution sampling
                        latents = self.vae.encode(batch, return_dict=False)[0]
                        mu, logvar = None, None
                    else:
                        dgd = self.vae.encode(batch).latent_dist
                        mu, logvar = dgd.mean, dgd.logvar
                        latents = dgd.sample()
                    
                    # scale shift latent to config
                    shift = self.vae.config['shift_factor'] if self.vae.config['shift_factor'] is not None else 0
                    latents = self.vae.config['scaling_factor'] * (latents - shift)
                    
                    if target_latent is not None and self.train_encoder:
                        # forward_latents = target_latent.detach()
                        lat_mse_loss = torch.nn.MSELoss()(target_latent.float(), latents.float())
                        latents = target_latent.detach()
                        forward_latents = target_latent.detach()
                        
                    elif self.eq_vae:
                        # process flips, rotate, scale
                        latent_chunks = list(latents.chunk(latents.shape[0], dim=0))
                        batch_chunks = list(batch.chunk(batch.shape[0], dim=0))
                        out_chunks = []
                        for i in range(len(latent_chunks)):
                            try:
                                do_rotate = random.randint(0, 3)
                                do_flip_x = random.randint(0, 1)
                                do_flip_y = random.randint(0, 1)
                                do_scale = random.randint(0, 1)
                                if do_rotate > 0:
                                    latent_chunks[i] = torch.rot90(latent_chunks[i], do_rotate, (2, 3))
                                    batch_chunks[i] = torch.rot90(batch_chunks[i], do_rotate, (2, 3))
                                if do_flip_x > 0:
                                    latent_chunks[i] = torch.flip(latent_chunks[i], [2])
                                    batch_chunks[i] = torch.flip(batch_chunks[i], [2])
                                if do_flip_y > 0:
                                    latent_chunks[i] = torch.flip(latent_chunks[i], [3])
                                    batch_chunks[i] = torch.flip(batch_chunks[i], [3])
                                
                                # resize latent to fit
                                if latent_chunks[i].shape[2] != latent_size or latent_chunks[i].shape[3] != latent_size:
                                    latent_chunks[i] = torch.nn.functional.interpolate(latent_chunks[i], size=(latent_size, latent_size), mode='bilinear', align_corners=False)
                                
                                # if do_scale > 0:
                                #     scale = 2
                                #     start_latent_h = latent_chunks[i].shape[2]
                                #     start_latent_w = latent_chunks[i].shape[3]
                                #     start_batch_h = batch_chunks[i].shape[2]
                                #     start_batch_w = batch_chunks[i].shape[3]
                                #     latent_chunks[i] = torch.nn.functional.interpolate(latent_chunks[i], scale_factor=scale, mode='bilinear', align_corners=False)
                                #     batch_chunks[i] = torch.nn.functional.interpolate(batch_chunks[i], scale_factor=scale, mode='bilinear', align_corners=False)
                                #     # random crop. latent is smaller than match but crops need to match
                                #     latent_x = random.randint(0, latent_chunks[i].shape[2] - start_latent_h)
                                #     latent_y = random.randint(0, latent_chunks[i].shape[3] - start_latent_w)
                                #     batch_x = latent_x * self.vae_scale_factor
                                #     batch_y = latent_y * self.vae_scale_factor
                                    
                                #     # crop
                                #     latent_chunks[i] = latent_chunks[i][:, :, latent_x:latent_x + start_latent_h, latent_y:latent_y + start_latent_w]
                                #     batch_chunks[i] = batch_chunks[i][:, :, batch_x:batch_x + start_batch_h, batch_y:batch_y + start_batch_w]
                            except Exception as e:
                                print(f"Error processing image {i}: {e}")
                                traceback.print_exc()
                                raise e
                            out_chunks.append(latent_chunks[i])
                        latents = torch.cat(out_chunks, dim=0)
                        # do dropout
                        if self.dropout > 0:
                            forward_latents = channel_dropout(latents, self.dropout)
                        else:
                            forward_latents = latents
                            
                        # resize batch to resolution if needed
                        if batch_chunks[0].shape[2] != self.resolution or batch_chunks[0].shape[3] != self.resolution:
                            batch_chunks = [torch.nn.functional.interpolate(b, size=(self.resolution, self.resolution), mode='bilinear', align_corners=False) for b in batch_chunks]
                        batch = torch.cat(batch_chunks, dim=0)
                                
                    else:
                        # latents.detach().requires_grad_(True)
                        forward_latents = latents
                    
                forward_latents = forward_latents.to(self.device, dtype=self.torch_dtype)
                
                if not self.train_encoder:
                    # detach latents if not training encoder
                    forward_latents = forward_latents.detach()
                
                # shift latents to match vae config
                shift = self.vae.config['shift_factor'] if self.vae.config['shift_factor'] is not None else 0
                forward_latents = forward_latents / self.vae.config['scaling_factor'] + shift

                pooled_outputs = None
                if self.do_pooled_exits:
                    pred, pooled_outputs = self.vae.decode_with_pooled_exits(forward_latents)
                else:
                    pred = self.vae.decode(forward_latents).sample

                # Run through VGG19
                if self.style_weight > 0 or self.content_weight > 0:
                    stacked = torch.cat([pred, batch], dim=0)
                    stacked = (stacked / 2 + 0.5).clamp(0, 1)
                    self.vgg_19(stacked)

                if self.use_critic:
                    stacked = torch.cat([pred, batch], dim=0)
                    critic_d_loss = self.critic.step(stacked.detach())
                else:
                    critic_d_loss = 0.0

                style_loss = self.get_style_loss() * self.style_weight
                content_loss = self.get_content_loss() * self.content_weight
                kld_loss = self.get_kld_loss(mu, logvar) * self.kld_weight
                mse_loss = self.get_mse_loss(pred, batch) * self.mse_weight
                mae_loss = self.get_mae_loss(pred, batch) * self.mae_weight
                pool_loss = self.get_pooled_output_loss(pooled_outputs, batch)
                if self.clip_weight > 0:
                    clip_loss = self.get_clip_loss(pred, batch) * self.clip_weight
                else:
                    clip_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)
                if self.lpips_weight > 0:
                    lpips_loss = self.lpips_loss(
                        pred.clamp(-1, 1).to(self.device, dtype=torch.bfloat16),
                        batch.clamp(-1, 1).to(self.device, dtype=torch.bfloat16)
                    ).float().mean() * self.lpips_weight
                else:
                    lpips_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)
                tv_loss = self.get_tv_loss(pred, batch) * self.tv_weight
                pattern_loss = self.get_pattern_loss(pred, batch) * self.pattern_weight
                if self.use_critic:
                    stacked = torch.cat([pred, batch], dim=0)
                    critic_gen_loss = self.critic.get_critic_loss(stacked) * self.critic_weight

                    # do not let abs critic gen loss be higher than abs lpips * 0.1 if using it
                    if self.lpips_weight > 0:
                        max_target = lpips_loss.abs() * 0.1
                        with torch.no_grad():
                            crit_g_scaler = 1.0
                            if critic_gen_loss.abs() > max_target:
                                crit_g_scaler = max_target / critic_gen_loss.abs()

                        critic_gen_loss *= crit_g_scaler
                else:
                    critic_gen_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)
                
                if self.mv_loss_weight > 0:
                    mv_loss = self.get_mean_variance_loss(latents) * self.mv_loss_weight
                else:
                    mv_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)
                
                if self.ltv_weight > 0:
                    ltv_loss = self.get_ltv_loss(latents, batch) * self.ltv_weight
                else:
                    ltv_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)
                    
                if self.lpm_weight > 0:
                    lpm_loss = self.get_latent_pixel_matching_loss(latents, batch) * self.lpm_weight
                else:
                    lpm_loss = torch.tensor(0.0, device=self.device, dtype=self.torch_dtype)

                loss = style_loss + content_loss + kld_loss + mse_loss + tv_loss + critic_gen_loss + pattern_loss + lpips_loss + mv_loss + ltv_loss + mae_loss + lat_mse_loss + clip_loss + pool_loss
                
                # check if loss is NaN or Inf
                if torch.isnan(loss) or torch.isinf(loss):
                    self.print(f"Loss is NaN or Inf, stopping at step {self.step_num}")
                    self.print(f" - Style loss: {style_loss.item()}")
                    self.print(f" - Content loss: {content_loss.item()}")
                    self.print(f" - KLD loss: {kld_loss.item()}")
                    self.print(f" - MSE loss: {mse_loss.item()}")
                    self.print(f" - MAE loss: {mae_loss.item()}")
                    self.print(f" - Latent MSE loss: {lat_mse_loss.item()}")
                    self.print(f" - LPIPS loss: {lpips_loss.item()}")
                    self.print(f" - TV loss: {tv_loss.item()}")
                    self.print(f" - Pattern loss: {pattern_loss.item()}")
                    self.print(f" - CLIP loss: {clip_loss.item()}")
                    self.print(f" - Pooled output loss: {pool_loss.item()}")
                    self.print(f" - Critic gen loss: {critic_gen_loss.item()}")
                    self.print(f" - Critic D loss: {critic_d_loss}")
                    self.print(f" - Mean variance loss: {mv_loss.item()}")
                    self.print(f" - Latent TV loss: {ltv_loss.item()}")
                    self.print(f" - Latent pixel matching loss: {lpm_loss.item()}")
                    self.print(f" - Total loss: {loss.item()}")
                    self.print(f" - Stopping training")
                    exit(1)

                # Backward pass and optimization
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                scheduler.step()

                # update progress bar
                loss_value = loss.item()
                # get exponent like 3.54e-4
                loss_string = f"loss: {loss_value:.2e}"
                if self.lpips_weight > 0:
                    loss_string += f" lpips: {lpips_loss.item():.2e}"
                if self.content_weight > 0:
                    loss_string += f" cnt: {content_loss.item():.2e}"
                if self.style_weight > 0:
                    loss_string += f" sty: {style_loss.item():.2e}"
                if self.kld_weight > 0:
                    loss_string += f" kld: {kld_loss.item():.2e}"
                if self.mse_weight > 0:
                    loss_string += f" mse: {mse_loss.item():.2e}"
                if self.mae_weight > 0:
                    loss_string += f" mae: {mae_loss.item():.2e}"
                if self.target_latent_vae:
                    loss_string += f" lat_mse: {lat_mse_loss.item():.2e}"
                if self.tv_weight > 0:
                    loss_string += f" tv: {tv_loss.item():.2e}"
                if self.pattern_weight > 0:
                    loss_string += f" ptn: {pattern_loss.item():.2e}"
                if self.clip_weight > 0:
                    loss_string += f" clip: {clip_loss.item():.2e}"
                if self.do_pooled_exits:
                    loss_string += f" pool: {pool_loss.item():.2e}"
                if self.use_critic and self.critic_weight > 0:
                    loss_string += f" crG: {critic_gen_loss.item():.2e}"
                if self.use_critic:
                    loss_string += f" crD: {critic_d_loss:.2e}"
                if self.mv_loss_weight > 0:
                    loss_string += f" mvl: {mv_loss:.2e}"
                if self.ltv_weight > 0:
                    loss_string += f" ltv: {ltv_loss:.2e}"
                if self.lpm_weight > 0:
                    loss_string += f" lpm: {lpm_loss:.2e}"
                

                if hasattr(optimizer, 'get_avg_learning_rate'):
                    learning_rate = optimizer.get_avg_learning_rate()
                elif self.optimizer_type.startswith('dadaptation') or \
                        self.optimizer_type.lower().startswith('prodigy'):
                    learning_rate = (
                            optimizer.param_groups[0]["d"] *
                            optimizer.param_groups[0]["lr"]
                    )
                else:
                    learning_rate = optimizer.param_groups[0]['lr']

                lr_critic_string = ''
                if self.use_critic:
                    lr_critic = self.critic.get_lr()
                    lr_critic_string = f" lrC: {lr_critic:.1e}"

                self.progress_bar.set_postfix_str(f"lr: {learning_rate:.1e}{lr_critic_string} {loss_string}")
                self.progress_bar.set_description(f"E: {epoch}")
                self.progress_bar.update(1)

                epoch_losses["total"].append(loss_value)
                epoch_losses["lpips"].append(lpips_loss.item())
                epoch_losses["style"].append(style_loss.item())
                epoch_losses["content"].append(content_loss.item())
                epoch_losses["mse"].append(mse_loss.item())
                epoch_losses["mae"].append(mae_loss.item())
                epoch_losses["lat_mse"].append(lat_mse_loss.item())
                epoch_losses["kl"].append(kld_loss.item())
                epoch_losses["tv"].append(tv_loss.item())
                epoch_losses["ptn"].append(pattern_loss.item())
                epoch_losses["clip"].append(clip_loss.item())
                epoch_losses["pool"].append(pool_loss.item())
                epoch_losses["crG"].append(critic_gen_loss.item())
                epoch_losses["crD"].append(critic_d_loss)
                epoch_losses["mvl"].append(mv_loss.item())
                epoch_losses["ltv"].append(ltv_loss.item())
                epoch_losses["lpm"].append(lpm_loss.item())

                log_losses["total"].append(loss_value)
                log_losses["lpips"].append(lpips_loss.item())
                log_losses["style"].append(style_loss.item())
                log_losses["content"].append(content_loss.item())
                log_losses["mse"].append(mse_loss.item())
                log_losses["mae"].append(mae_loss.item())
                log_losses["lat_mse"].append(lat_mse_loss.item())
                log_losses["kl"].append(kld_loss.item())
                log_losses["tv"].append(tv_loss.item())
                log_losses["ptn"].append(pattern_loss.item())
                log_losses["clip"].append(clip_loss.item())
                log_losses["pool"].append(pool_loss.item())
                log_losses["crG"].append(critic_gen_loss.item())
                log_losses["crD"].append(critic_d_loss)
                log_losses["mvl"].append(mv_loss.item())
                log_losses["ltv"].append(ltv_loss.item())
                log_losses["lpm"].append(lpm_loss.item())

                # don't do on first step
                if self.step_num != start_step:
                    if self.sample_every and self.step_num % self.sample_every == 0:
                        # print above the progress bar
                        self.print(f"Sampling at step {self.step_num}")
                        self.sample(self.step_num)

                    if self.save_every and self.step_num % self.save_every == 0:
                        # print above the progress bar
                        self.print(f"Saving at step {self.step_num}")
                        self.save(self.step_num)

                    if self.log_every and self.step_num % self.log_every == 0:
                        # log to tensorboard
                        if self.writer is not None:
                            # get avg loss
                            for key in log_losses:
                                log_losses[key] = sum(log_losses[key]) / (len(log_losses[key]) + 1e-6)
                                # if log_losses[key] > 0:
                                self.writer.add_scalar(f"loss/{key}", log_losses[key], self.step_num)
                        # reset log losses
                        log_losses = copy.deepcopy(blank_losses)

                self.step_num += 1
            # end epoch
            if self.writer is not None:
                eps = 1e-6
                # get avg loss
                for key in epoch_losses:
                    epoch_losses[key] = sum(log_losses[key]) / (len(log_losses[key]) + eps)
                    if epoch_losses[key] > 0:
                        self.writer.add_scalar(f"epoch loss/{key}", epoch_losses[key], epoch)
            # reset epoch losses
            epoch_losses = copy.deepcopy(blank_losses)

        self.save()