RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Resize validation images with the bucket sizing strategy.

This commit is contained in:
Jaret Burkett 2026-07-20 11:22:33 -06:00
parent df354da23e
commit 088084e2c2
1 changed files with 49 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
jobs/process/BaseSDTrainProcess.py
View File

@ -25,6 +25,7 @@ from huggingface_hub import HfApi, interpreter_login
from toolkit.memory_management import MemoryManager from toolkit.memory_management import MemoryManager
from toolkit.basic import value_map from toolkit.basic import value_map
from toolkit.buckets import get_bucket_for_image_size
from toolkit.clip_vision_adapter import ClipVisionAdapter from toolkit.clip_vision_adapter import ClipVisionAdapter
from toolkit.custom_adapter import CustomAdapter from toolkit.custom_adapter import CustomAdapter
from toolkit.data_loader import get_dataloader_from_datasets, trigger_dataloader_setup_epoch from toolkit.data_loader import get_dataloader_from_datasets, trigger_dataloader_setup_epoch
@ -1602,14 +1603,21 @@ class BaseSDTrainProcess(BaseTrainProcess):
dtype = get_torch_dtype(self.train_config.dtype) dtype = get_torch_dtype(self.train_config.dtype)
resolution = val_config.resolution resolution = val_config.resolution
divisibility = self.sd.get_bucket_divisibility()
image_list = [] image_list = []
prompt_list = [] prompt_list = []
for item in validation_items: for item in validation_items:
img = Image.open(item.image_path) img = Image.open(item.image_path)
img = ImageOps.exif_transpose(img).convert('RGB') img = ImageOps.exif_transpose(img).convert('RGB')
# deterministic resize, shortest side to resolution then center crop # deterministic resize that keeps the aspect ratio, matches the pixel budget
img = transforms.Resize(resolution)(img) # of the resolution and the bucket divisibility of the model
img = transforms.CenterCrop(resolution)(img) bucket = get_bucket_for_image_size(
img.width, img.height,
resolution=resolution,
divisibility=divisibility,
)
img = img.resize((bucket['width'], bucket['height']), Image.BICUBIC)
tensor = transforms.ToTensor()(img) * 2.0 - 1.0 tensor = transforms.ToTensor()(img) * 2.0 - 1.0
image_list.append(tensor) image_list.append(tensor)
prompt = item.prompt prompt = item.prompt
@ -1638,21 +1646,24 @@ class BaseSDTrainProcess(BaseTrainProcess):
# seed so the vae latent dist sampling is always identical # seed so the vae latent dist sampling is always identical
torch.manual_seed(42) torch.manual_seed(42)
orig_vae_device = self.sd.vae.device orig_vae_device = self.sd.vae.device
latents = self.sd.encode_images(image_list, device=device, dtype=dtype) # images can have different aspect ratios so they are encoded one at a time
latent_list = [
self.sd.encode_images([image], device=device, dtype=dtype).to('cpu', dtype=torch.float32)
for image in image_list
]
self.sd.vae.to(orig_vae_device) self.sd.vae.to(orig_vae_device)
# fixed noise per image, seeds start at 42 and increment for each image # fixed noise per image, seeds start at 42 and increment for each image
noise_list = [] noise_list = []
for i in range(latents.shape[0]): for i, latent in enumerate(latent_list):
generator = torch.Generator(device='cpu').manual_seed(42 + i) generator = torch.Generator(device='cpu').manual_seed(42 + i)
noise_list.append( noise_list.append(
torch.randn(latents[i:i + 1].shape, generator=generator, dtype=torch.float32) torch.randn(latent.shape, generator=generator, dtype=torch.float32)
) )
noise = torch.cat(noise_list, dim=0)
self._validation_cache = { self._validation_cache = {
'latents': latents.to('cpu', dtype=torch.float32), 'latents': latent_list,
'noise': noise, 'noise': noise_list,
'embeds': embeds_list, 'embeds': embeds_list,
} }
flush() flush()
@ -1675,17 +1686,6 @@ class BaseSDTrainProcess(BaseTrainProcess):
start_multiplier = network.multiplier start_multiplier = network.multiplier
network.multiplier = 1.0 network.multiplier = 1.0
with torch.no_grad(), network: 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: if self.sd.is_flow_matching:
timestep_values = [sigma * 1000.0 for sigma in sigmas] timestep_values = [sigma * 1000.0 for sigma in sigmas]
else: else:
@ -1694,30 +1694,39 @@ class BaseSDTrainProcess(BaseTrainProcess):
min(int(round(sigma * num_train_timesteps)), num_train_timesteps - 1) min(int(round(sigma * num_train_timesteps)), num_train_timesteps - 1)
for sigma in sigmas for sigma in sigmas
] ]
timesteps = torch.tensor( timesteps = torch.tensor(timestep_values, device=device)
[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() # images can have different aspect ratios, so each image is predicted as its
# own batch with all sigmas at once
losses = []
for latents_cpu, noise_cpu, embeds_cpu in zip(cache['latents'], cache['noise'], cache['embeds']):
latents = latents_cpu.to(device, dtype=dtype)
noise = noise_cpu.to(device, dtype=dtype)
batch_latents = torch.cat([latents] * len(sigmas), dim=0)
batch_noise = torch.cat([noise] * len(sigmas), dim=0)
batch_embeds = concat_prompt_embeds([embeds_cpu.clone().to(device, dtype=dtype)] * len(sigmas))
noise_pred = self.sd.predict_noise( noisy_latents = self.sd.add_noise(batch_latents, batch_noise, timesteps).detach()
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: noise_pred = self.sd.predict_noise(
target = batch_noise - batch_latents latents=noisy_latents.to(device, dtype=dtype),
elif self.sd.prediction_type == 'v_prediction': conditional_embeddings=batch_embeds,
target = self.sd.noise_scheduler.get_velocity(batch_latents, batch_noise, timesteps) timestep=timesteps,
else: guidance_scale=1.0,
target = batch_noise guidance_embedding_scale=self.train_config.cfg_scale,
bypass_guidance_embedding=self.train_config.bypass_guidance_embedding,
)
val_loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float()) 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
losses.append(torch.nn.functional.mse_loss(noise_pred.float(), target.float()))
val_loss = torch.stack(losses).mean()
self.additional_logs['val/loss'] = val_loss.item() self.additional_logs['val/loss'] = val_loss.item()
network.multiplier = start_multiplier network.multiplier = start_multiplier
if was_unet_training: if was_unet_training: