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

Bug fixes. Added some functionality to help with private extensions

Browse Source
This commit is contained in:
Jaret Burkett
2023-10-05 07:09:34 -06:00
parent 579650eaf8
commit f73402473b
8 changed files with 99 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/data_loader.py
View File

@@ -478,6 +478,7 @@ def get_dataloader_from_datasets(
datasets = []
has_buckets = False
is_caching_latents = False
dataset_config_list = []
# preprocess them all
@@ -497,6 +498,8 @@ def get_dataloader_from_datasets(
datasets.append(dataset)
if config.buckets:
has_buckets = True
if config.cache_latents or config.cache_latents_to_disk:
is_caching_latents = True
else:
raise ValueError(f"invalid dataset type: {config.type}")
@@ -512,6 +515,9 @@ def get_dataloader_from_datasets(
)
return batch
# check if is caching latents
if has_buckets:
# make sure they all have buckets
for dataset in datasets:

18
toolkit/data_transfer_object/data_loader.py
View File

@@ -84,8 +84,22 @@ class DataLoaderBatchDTO:
if is_latents_cached:
self.latents = torch.cat([x.get_latent().unsqueeze(0) for x in self.file_items])
self.control_tensor: Union[torch.Tensor, None] = None
if self.file_items[0].control_tensor is not None:
self.control_tensor = torch.cat([x.control_tensor.unsqueeze(0) for x in self.file_items])
# if self.file_items[0].control_tensor is not None:
# if any have a control tensor, we concatenate them
if any([x.control_tensor is not None for x in self.file_items]):
# find one to use as a base
base_control_tensor = None
for x in self.file_items:
if x.control_tensor is not None:
base_control_tensor = x.control_tensor
break
control_tensors = []
for x in self.file_items:
if x.control_tensor is None:
control_tensors.append(torch.zeros_like(base_control_tensor))
else:
control_tensors.append(x.control_tensor)
self.control_tensor = torch.cat([x.unsqueeze(0) for x in control_tensors])
except Exception as e:
print(e)
raise e

1
toolkit/dataloader_mixins.py
View File

@@ -317,6 +317,7 @@ class ImageProcessingDTOMixin:
img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC)
# crop to x_crop, y_crop, x_crop + crop_width, y_crop + crop_height
if img.width < self.crop_x + self.crop_width or img.height < self.crop_y + self.crop_height:
# todo look into this. This still happens sometimes
print('size mismatch')
img = img.crop((
self.crop_x,

22
toolkit/llvae.py
View File

@@ -5,14 +5,18 @@ import itertools
class LosslessLatentDecoder(nn.Module):
def __init__(self, in_channels, latent_depth, dtype=torch.float32):
def __init__(self, in_channels, latent_depth, dtype=torch.float32, trainable=False):
super(LosslessLatentDecoder, self).__init__()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.latent_depth = latent_depth
self.in_channels = in_channels
self.out_channels = int(in_channels // (latent_depth * latent_depth))
numpy_kernel = self.build_kernel(in_channels, latent_depth)
self.kernel = torch.from_numpy(numpy_kernel).to(device=device, dtype=dtype)
numpy_kernel = torch.from_numpy(numpy_kernel).to(device=device, dtype=dtype)
if trainable:
self.kernel = nn.Parameter(numpy_kernel)
else:
self.kernel = numpy_kernel
def build_kernel(self, in_channels, latent_depth):
# my old code from tensorflow.
@@ -44,14 +48,18 @@ class LosslessLatentDecoder(nn.Module):
class LosslessLatentEncoder(nn.Module):
def __init__(self, in_channels, latent_depth, dtype=torch.float32):
def __init__(self, in_channels, latent_depth, dtype=torch.float32, trainable=False):
super(LosslessLatentEncoder, self).__init__()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.latent_depth = latent_depth
self.in_channels = in_channels
self.out_channels = int(in_channels * (latent_depth * latent_depth))
numpy_kernel = self.build_kernel(in_channels, latent_depth)
self.kernel = torch.from_numpy(numpy_kernel).to(device=device, dtype=dtype)
numpy_kernel = torch.from_numpy(numpy_kernel).to(device=device, dtype=dtype)
if trainable:
self.kernel = nn.Parameter(numpy_kernel)
else:
self.kernel = numpy_kernel
def build_kernel(self, in_channels, latent_depth):
@@ -82,13 +90,13 @@ class LosslessLatentEncoder(nn.Module):
class LosslessLatentVAE(nn.Module):
def __init__(self, in_channels, latent_depth, dtype=torch.float32):
def __init__(self, in_channels, latent_depth, dtype=torch.float32, trainable=False):
super(LosslessLatentVAE, self).__init__()
self.latent_depth = latent_depth
self.in_channels = in_channels
self.encoder = LosslessLatentEncoder(in_channels, latent_depth, dtype=dtype)
self.encoder = LosslessLatentEncoder(in_channels, latent_depth, dtype=dtype, trainable=trainable)
encoder_out_channels = self.encoder.out_channels
self.decoder = LosslessLatentDecoder(encoder_out_channels, latent_depth, dtype=dtype)
self.decoder = LosslessLatentDecoder(encoder_out_channels, latent_depth, dtype=dtype, trainable=trainable)
def forward(self, x):
latent = self.latent_encoder(x)

4
toolkit/stable_diffusion_model.py
View File

@@ -78,7 +78,7 @@ def flush():
UNET_IN_CHANNELS = 4 # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
VAE_SCALE_FACTOR = 8 # 2 ** (len(vae.config.block_out_channels) - 1) = 8
# VAE_SCALE_FACTOR = 8 # 2 ** (len(vae.config.block_out_channels) - 1) = 8
# if is type checking
if typing.TYPE_CHECKING:
@@ -471,6 +471,7 @@ class StableDiffusion:
batch_size=1,
noise_offset=0.0,
):
VAE_SCALE_FACTOR = 2 ** (len(self.vae.config['block_out_channels']) - 1)
if height is None and pixel_height is None:
raise ValueError("height or pixel_height must be specified")
if width is None and pixel_width is None:
@@ -493,6 +494,7 @@ class StableDiffusion:
return noise
def get_time_ids_from_latents(self, latents: torch.Tensor):
VAE_SCALE_FACTOR = 2 ** (len(self.vae.config['block_out_channels']) - 1)
if self.is_xl:
bs, ch, h, w = list(latents.shape)