diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index a9233d6..dc813b0 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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 diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index ef59002..3b0a0d6 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -337,6 +337,23 @@ class AdapterConfig: self.i2v_do_start_frame: bool = kwargs.get('i2v_do_start_frame', False) +class ValidationItem: + def __init__(self, **kwargs): + self.image_path: str = kwargs.get('image_path', '') + self.prompt: str = kwargs.get('prompt', '') + + +class ValidationConfig: + def __init__(self, **kwargs): + self.validation_items: List[ValidationItem] = [ + item if isinstance(item, ValidationItem) else ValidationItem(**item) + for item in kwargs.get('validation_items', []) + ] + self.resolution: int = kwargs.get('resolution', 512) + self.validate_every_n_steps: int = kwargs.get('validate_every_n_steps', 10) + self.validation_sigmas: List[float] = kwargs.get('validation_sigmas', [1.0, 0.75, 0.5, 0.25]) + + class EmbeddingConfig: def __init__(self, **kwargs): self.trigger = kwargs.get('trigger', 'custom_embedding') @@ -592,6 +609,10 @@ class TrainConfig: self.max_loss_debug: bool = kwargs.get("max_loss_debug", False) # will clip the loss to this amount to prevent wild outliers self.max_loss: Optional[float] = kwargs.get("max_loss", None) + self.validation_config: Optional[ValidationConfig] = None + validation = kwargs.get('validation_config', None) + if validation is not None: + self.validation_config: ValidationConfig = ValidationConfig(**validation) ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21', 'anima'] @@ -1417,5 +1438,3 @@ def validate_configs( if train_config.batch_size > 1 and any(dataset_config.auto_frame_count for dataset_config in dataset_configs): raise ValueError("Cannot use batch size greater than 1 with auto_frame_count. Please set batch_size to 1 or auto_frame_count to False.") - - diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index e5d51b7..5f9a22a 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -199,6 +199,8 @@ export default function SimpleJob({ const showGPUSelect = !isMac(); + const validationConfig = jobConfig.config.process[0].train.validation_config; + let numDatasetCols = 4; let numSampleTopCols = 4; let datasetStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6'; @@ -834,6 +836,136 @@ export default function SimpleJob({ +
+ { + if (value) { + setJobConfig( + { + validation_items: [{ image_path: '', prompt: '' }], + resolution: 512, + validate_every_n_steps: 10, + validation_sigmas: [1.0, 0.75, 0.5, 0.25], + }, + 'config.process[0].train.validation_config', + ); + } else { + setJobConfig(undefined, 'config.process[0].train.validation_config'); + } + }} + > + {validationConfig && ( + <> +

+ Validation runs a stable loss check on a fixed set of images. Each image is encoded once at startup + and predicted at the selected sigmas with fixed seeds, so the result is always deterministic and + comparable across the run. The average loss is logged as val/loss every time validation runs. The + images need to match the concept of your dataset, but{' '} + do not include the validation images in the dataset. + They must be images containing the concept you want to train, but not an image trained on. +

+
+ + setJobConfig(value, 'config.process[0].train.validation_config.validate_every_n_steps') + } + placeholder="eg. 10" + min={1} + required + /> + setJobConfig(value, 'config.process[0].train.validation_config.resolution')} + placeholder="eg. 512" + min={64} + required + /> + + setJobConfig( + value.split(',').map((v: string) => parseFloat(v)), + 'config.process[0].train.validation_config.validation_sigmas', + ) + } + options={[ + { value: '0.5', label: '0.5' }, + { value: '1, 0.5', label: '1.0, 0.5' }, + { value: '1, 0.66, 0.33', label: '1.0, 0.66, 0.33' }, + { value: '1, 0.75, 0.5, 0.25', label: '1.0, 0.75, 0.5, 0.25' }, + ]} + /> +
+
+ + {validationConfig.validation_items.map((item, i) => ( +
+
+ { + setJobConfig( + imagePath ?? '', + `config.process[0].train.validation_config.validation_items[${i}].image_path`, + ); + }} + /> +
+ + setJobConfig( + value, + `config.process[0].train.validation_config.validation_items[${i}].prompt`, + ) + } + placeholder="Enter prompt" + /> +
+
+ +
+
+
+ ))} + +
+ + )} + +
diff --git a/ui/src/components/Card.tsx b/ui/src/components/Card.tsx index 1487518..bdfc820 100644 --- a/ui/src/components/Card.tsx +++ b/ui/src/components/Card.tsx @@ -7,9 +7,52 @@ interface CardProps { children?: React.ReactNode; collapsible?: boolean; defaultOpen?: boolean; + // when provided, the card is opened/closed by a toggle switch on the right instead of a chevron + toggled?: boolean; + onToggle?: (value: boolean) => void; } -const Card: React.FC = ({ title, children, collapsible, defaultOpen }) => { +const Card: React.FC = ({ title, children, collapsible, defaultOpen, toggled, onToggle }) => { + if (onToggle) { + return ( +
+
+
+ {title && ( +

+ {title} +

+ )} +
+ +
+ {toggled && ( + <> + {children ?? null} +
+ + )} +
+ ); + } if (collapsible) { return ( diff --git a/ui/src/components/formInputs.tsx b/ui/src/components/formInputs.tsx index cd97e95..77ed7e1 100644 --- a/ui/src/components/formInputs.tsx +++ b/ui/src/components/formInputs.tsx @@ -198,18 +198,26 @@ export const NumberInput = (props: NumberInputProps) => { const numValue = Number(rawValue); - // Only apply constraints and call onChange when we have a valid number + // don't clamp to min/max while typing, it mangles partial input (typing 1024 with + // min 64 becomes 64024). Clamping happens on blur. if (!isNaN(numValue)) { - let constrainedValue = numValue; - - // Apply min/max constraints if they exist - if (min !== undefined && constrainedValue < min) { - constrainedValue = min; - } - if (max !== undefined && constrainedValue > max) { - constrainedValue = max; - } - + onChange(numValue); + } + }} + onBlur={() => { + const numValue = Number(inputValue); + if (inputValue === '' || isNaN(numValue)) { + return; + } + let constrainedValue = numValue; + if (min !== undefined && constrainedValue < min) { + constrainedValue = min; + } + if (max !== undefined && constrainedValue > max) { + constrainedValue = max; + } + if (constrainedValue !== numValue) { + setInputValue(constrainedValue); onChange(constrainedValue); } }} diff --git a/ui/src/types.ts b/ui/src/types.ts index 07f070c..8f24f08 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -117,6 +117,18 @@ export interface EMAConfig { ema_decay: number; } +export interface ValidationItem { + image_path: string; + prompt: string; +} + +export interface ValidationConfig { + validation_items: ValidationItem[]; + resolution: number; + validate_every_n_steps: number; + validation_sigmas?: number[]; +} + export interface TrainConfig { batch_size: number; bypass_guidance_embedding?: boolean; @@ -151,6 +163,7 @@ export interface TrainConfig { differential_guidance_scale?: number; audio_loss_multiplier?: number; max_loss?: number | null; + validation_config?: ValidationConfig; } export interface QuantizeKwargsConfig { diff --git a/version.py b/version.py index 9105263..b8f92a1 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.28" +VERSION = "0.10.29"