Added validation loss

jobs/process/BaseSDTrainProcess.py

@@ -12,6 +12,8 @@ from typing import Union, List, Optional
 
 import numpy as np
 import yaml
+from PIL import Image, ImageOps
+from torchvision import transforms
 from diffusers import T2IAdapter, ControlNetModel
 from diffusers.training_utils import compute_density_for_timestep_sampling
 from safetensors.torch import save_file, load_file
@@ -40,6 +42,7 @@ from toolkit.network_mixins import Network
 from toolkit.optimizer import get_optimizer
 from toolkit.paths import CONFIG_ROOT
 from toolkit.progress_bar import ToolkitProgressBar
+from toolkit.prompt_utils import concat_prompt_embeds
 from toolkit.reference_adapter import ReferenceAdapter
 from toolkit.sampler import get_sampler
 from toolkit.saving import save_t2i_from_diffusers, load_t2i_model, save_ip_adapter_from_diffusers, \
@@ -47,7 +50,7 @@ from toolkit.saving import save_t2i_from_diffusers, load_t2i_model, save_ip_adap
 
 from toolkit.scheduler import get_lr_scheduler
 from toolkit.sd_device_states_presets import get_train_sd_device_state_preset
-from toolkit.stable_diffusion_model import StableDiffusion
+from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork
 
 from jobs.process import BaseTrainProcess
 from toolkit.metadata import get_meta_for_safetensors, load_metadata_from_safetensors, add_base_model_info_to_meta, \
@@ -261,6 +264,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
         self.steps_this_boundary = 0
         self.num_consecutive_oom = 0
         self.additional_logs = {}
+        # cached latents, prompt embeds, and fixed noise for validation
+        self._validation_cache = None
 
     def post_process_generate_image_config_list(self, generate_image_config_list: List[GenerateImageConfig]):
         # override in subclass
@@ -1570,6 +1575,146 @@ class BaseSDTrainProcess(BaseTrainProcess):
             self.load_training_state_from_metadata(latest_save_path)
         # set trainable params
         self.sd.adapter = self.adapter
+    
+    def setup_validation(self):
+        # caches everything needed for validation (latents, prompt embeds, fixed noise)
+        # must be called while the vae and text encoder are still loaded, they may be
+        # dumped later to save memory
+        val_config = self.train_config.validation_config
+        if val_config is None:
+            return
+        if not self.accelerator.is_main_process:
+            return
+        validation_items = []
+        for item in val_config.validation_items:
+            if not item.image_path:
+                print_acc("Skipping validation item with no image")
+                continue
+            if not os.path.exists(item.image_path):
+                print_acc(f"Skipping validation item, image not found: {item.image_path}")
+                continue
+            validation_items.append(item)
+        if len(validation_items) == 0:
+            print_acc("Validation config has no valid validation_items, skipping validation")
+            return
+        print_acc(f"Caching validation latents and embeddings for {len(validation_items)} images")
+        device = self.device_torch
+        dtype = get_torch_dtype(self.train_config.dtype)
+        resolution = val_config.resolution
+
+        image_list = []
+        prompt_list = []
+        for item in validation_items:
+            img = Image.open(item.image_path)
+            img = ImageOps.exif_transpose(img).convert('RGB')
+            # deterministic resize, shortest side to resolution then center crop
+            img = transforms.Resize(resolution)(img)
+            img = transforms.CenterCrop(resolution)(img)
+            tensor = transforms.ToTensor()(img) * 2.0 - 1.0
+            image_list.append(tensor)
+            prompt_list.append(item.prompt)
+
+        fork_devices = [device] if device.type == 'cuda' else []
+        with torch.no_grad(), torch.random.fork_rng(devices=fork_devices):
+            # encode the prompts one at a time so they can be reassembled per sigma later
+            te = self.sd.text_encoder
+            te_list = te if isinstance(te, list) else ([te] if te is not None else [])
+            orig_te_devices = [next(t.parameters()).device for t in te_list]
+            self.sd.text_encoder_to(device)
+            embeds_list = [
+                self.sd.encode_prompt([prompt]).to('cpu', dtype=torch.float32).detach()
+                for prompt in prompt_list
+            ]
+            for t, te_device in zip(te_list, orig_te_devices):
+                t.to(te_device)
+
+            # seed so the vae latent dist sampling is always identical
+            torch.manual_seed(42)
+            orig_vae_device = self.sd.vae.device
+            latents = self.sd.encode_images(image_list, device=device, dtype=dtype)
+            self.sd.vae.to(orig_vae_device)
+
+            # fixed noise per image, seeds start at 42 and increment for each image
+            noise_list = []
+            for i in range(latents.shape[0]):
+                generator = torch.Generator(device='cpu').manual_seed(42 + i)
+                noise_list.append(
+                    torch.randn(latents[i:i + 1].shape, generator=generator, dtype=torch.float32)
+                )
+            noise = torch.cat(noise_list, dim=0)
+
+        self._validation_cache = {
+            'latents': latents.to('cpu', dtype=torch.float32),
+            'noise': noise,
+            'embeds': embeds_list,
+        }
+        flush()
+
+    def validate(self):
+        val_config = self.train_config.validation_config
+        if val_config is None or self._validation_cache is None:
+            return
+        if not self.accelerator.is_main_process:
+            return
+        device = self.device_torch
+        dtype = get_torch_dtype(self.train_config.dtype)
+        sigmas = val_config.validation_sigmas
+        cache = self._validation_cache
+
+        was_unet_training = self.sd.unet.training
+        self.sd.unet.eval()
+        # the network is only active inside its context, without this the base model is validated
+        network = self.network if self.network is not None else BlankNetwork()
+        start_multiplier = network.multiplier
+        network.multiplier = 1.0
+        with torch.no_grad(), network:
+            latents = cache['latents'].to(device, dtype=dtype)
+            noise = cache['noise'].to(device, dtype=dtype)
+            num_images = latents.shape[0]
+
+            # single batch of every (image, sigma) pair, ordered sigma major
+            batch_latents = torch.cat([latents] * len(sigmas), dim=0)
+            batch_noise = torch.cat([noise] * len(sigmas), dim=0)
+            batch_embeds = concat_prompt_embeds(
+                [e.clone().to(device, dtype=dtype) for e in cache['embeds']] * len(sigmas)
+            )
+
+            if self.sd.is_flow_matching:
+                timestep_values = [sigma * 1000.0 for sigma in sigmas]
+            else:
+                num_train_timesteps = self.sd.noise_scheduler.config.num_train_timesteps
+                timestep_values = [
+                    min(int(round(sigma * num_train_timesteps)), num_train_timesteps - 1)
+                    for sigma in sigmas
+                ]
+            timesteps = torch.tensor(
+                [t for t in timestep_values for _ in range(num_images)],
+                device=device,
+            )
+
+            noisy_latents = self.sd.add_noise(batch_latents, batch_noise, timesteps).detach()
+
+            noise_pred = self.sd.predict_noise(
+                latents=noisy_latents.to(device, dtype=dtype),
+                conditional_embeddings=batch_embeds,
+                timestep=timesteps,
+                guidance_scale=1.0,
+                guidance_embedding_scale=self.train_config.cfg_scale,
+                bypass_guidance_embedding=self.train_config.bypass_guidance_embedding,
+            )
+
+            if self.sd.is_flow_matching:
+                target = batch_noise - batch_latents
+            elif self.sd.prediction_type == 'v_prediction':
+                target = self.sd.noise_scheduler.get_velocity(batch_latents, batch_noise, timesteps)
+            else:
+                target = batch_noise
+
+            val_loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float())
+            self.additional_logs['val/loss'] = val_loss.item()
+        network.multiplier = start_multiplier
+        if was_unet_training:
+            self.sd.unet.train()
 
     def run(self):
         # torch.autograd.set_detect_anomaly(True)
@@ -2065,6 +2210,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
         )
         self.lr_scheduler = lr_scheduler
 
+        # cache validation latents and embeddings now, the vae and text encoder
+        # may be dumped before the train loop starts
+        self.setup_validation()
+
         ### HOOk ###
         self.before_dataset_load()
         # load datasets if passed in the root process
@@ -2480,6 +2629,18 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 print("\n==== Profile Results ====")
                 print(self.torch_profiler.key_averages().table(sort_by="cpu_time_total", row_limit=1000))
             self.timer.stop('train_loop')
+
+            # run validation before any possible sampling/logging for this step
+            if not did_oom and self.train_config.validation_config is not None:
+                val_config = self.train_config.validation_config
+                is_validate_step = (
+                    self.step_num == self.start_step
+                    or (val_config.validate_every_n_steps and self.step_num % val_config.validate_every_n_steps == 0)
+                )
+                if is_validate_step:
+                    with self.timer('validate'):
+                        self.validate()
+
             if not did_first_flush:
                 flush()
                 did_first_flush = True