github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Bug fixes on vae trainer. Allow to target params for vae training.

Browse Source
This commit is contained in:
Jaret Burkett
2025-07-26 09:20:22 -06:00
parent 3e14a674ac
commit 0d89c44624
1 changed files with 39 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
jobs/process/TrainVAEProcess.py
View File

@@ -36,6 +36,7 @@ import lpips
import random import random
import traceback import traceback
from transformers import SiglipImageProcessor, SiglipVisionModel from transformers import SiglipImageProcessor, SiglipVisionModel
import torch.nn.functional as F
IMAGE_TRANSFORMS = transforms.Compose( IMAGE_TRANSFORMS = transforms.Compose(
[ [
@@ -56,6 +57,22 @@ def channel_dropout(x, p=0.5):
return x * mask 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): class TrainVAEProcess(BaseTrainProcess):
def __init__(self, process_id: int, job, config: OrderedDict): def __init__(self, process_id: int, job, config: OrderedDict):
super().__init__(process_id, job, config) super().__init__(process_id, job, config)
@@ -98,7 +115,8 @@ class TrainVAEProcess(BaseTrainProcess):
self.train_encoder = self.get_conf('train_encoder', False, as_type=bool) 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.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.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.do_pooled_exits = False
self.VaeClass = AutoencoderKL self.VaeClass = AutoencoderKL
if self.vae_type == 'AutoencoderTiny': if self.vae_type == 'AutoencoderTiny':
@@ -253,6 +271,9 @@ class TrainVAEProcess(BaseTrainProcess):
def get_clip_embeddings(self, image_n1p1): def get_clip_embeddings(self, image_n1p1):
tensors_0_1 = (image_n1p1 + 1) / 2 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) tensors_0_1 = tensors_0_1.clamp(0, 1)
# resize if needed # resize if needed
@@ -494,6 +515,9 @@ class TrainVAEProcess(BaseTrainProcess):
target_latent = target_latent.to(self.device, dtype=self.torch_dtype) target_latent = target_latent.to(self.device, dtype=self.torch_dtype)
latent = self.vae.encode(img, return_dict=False)[0] 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 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 = self.vae.config['scaling_factor'] * (latent - shift)
@@ -649,24 +673,24 @@ class TrainVAEProcess(BaseTrainProcess):
train_all = 'all' in self.blocks_to_train train_all = 'all' in self.blocks_to_train
if train_all: if train_all:
params = list(self.vae.decoder.parameters()) params = list(self.vae.decoder.named_parameters())
self.vae.decoder.requires_grad_(True) self.vae.decoder.requires_grad_(True)
if self.train_encoder: if self.train_encoder:
# encoder # encoder
params += list(self.vae.encoder.parameters()) params += list(self.vae.encoder.named_parameters())
self.vae.encoder.requires_grad_(True) self.vae.encoder.requires_grad_(True)
else: else:
# mid_block # mid_block
if train_all or 'mid_block' in self.blocks_to_train: if train_all or 'mid_block' in self.blocks_to_train:
params += list(self.vae.decoder.mid_block.parameters()) params += list(self.vae.decoder.mid_block.named_parameters())
self.vae.decoder.mid_block.requires_grad_(True) self.vae.decoder.mid_block.requires_grad_(True)
# up_blocks # up_blocks
if train_all or 'up_blocks' in self.blocks_to_train: if train_all or 'up_blocks' in self.blocks_to_train:
params += list(self.vae.decoder.up_blocks.parameters()) params += list(self.vae.decoder.up_blocks.named_parameters())
self.vae.decoder.up_blocks.requires_grad_(True) self.vae.decoder.up_blocks.requires_grad_(True)
# conv_out (single conv layer output) # conv_out (single conv layer output)
if train_all or 'conv_out' in self.blocks_to_train: if train_all or 'conv_out' in self.blocks_to_train:
params += list(self.vae.decoder.conv_out.parameters()) params += list(self.vae.decoder.conv_out.named_parameters())
self.vae.decoder.conv_out.requires_grad_(True) self.vae.decoder.conv_out.requires_grad_(True)
if self.style_weight > 0 or self.content_weight > 0: if self.style_weight > 0 or self.content_weight > 0:
@@ -683,6 +707,15 @@ class TrainVAEProcess(BaseTrainProcess):
if self.lpips_weight > 0 and self.lpips_loss is None: 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')
self.lpips_loss = lpips.LPIPS(net='vgg').to(self.device, dtype=torch.bfloat16) 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 = get_optimizer(params, self.optimizer_type, self.learning_rate,
optimizer_params=self.optimizer_params) optimizer_params=self.optimizer_params)