Alpha masked loss

[kohya-ss]

original PR #1223

[recris]

Is there a way we can support more than one "channel" in the mask tensor, like we have currently? Even though we're only using the R channel.

Context: I've been experimenting with dynamic masks where the actual mask values/strength varies based on the timestep. The idea is to make the training process to focus on certain things in high timesteps, such as the overall shapes and composition, then in low timesteps change the focus to specific details like faces. To achieve this I use a "high" mask and a "low" mask, and interpolate between the two. High mask is stored as Red, low mask as Green.

Here is my experimental apply_masked_loss version:

def apply_masked_loss(loss, batch, timesteps):
    # mask image is -1 to 1. we need to convert it to 0 to 1
    mask_image_hi = batch["conditioning_images"].to(dtype=loss.dtype)[:, 0].unsqueeze(1)  # use R channel
    mask_image_lo = batch["conditioning_images"].to(dtype=loss.dtype)[:, 1].unsqueeze(1)  # use G channel
    mask_image_hi = mask_image_hi / 2 + 0.5
    mask_image_lo = mask_image_lo / 2 + 0.5

    timesteps = (timesteps / 1000.0).reshape(timesteps.shape + (1,) * 3)
    mask = torch.lerp(mask_image_lo, mask_image_hi, timesteps)

    # resize to the same size as the loss
    mask = torch.nn.functional.interpolate(mask, size=loss.shape[2:], mode="area")
    mask = mask.broadcast_to(loss.shape)
    loss = loss * mask
    return loss

I am getting some decent results with this technique. For this to work it would be necessary to have access to RGB channels in this function (or some structure compatible with multiple masks).

[kohya-ss]

Is there a way we can support more than one "channel" in the mask tensor, like we have currently? Even though we're only using the R channel.

Thank you for your suggestion. The code is very interesting. I think it may possibly be extended in the future (e.g., using the B channel for something), so I will add it as an experimental feature.

[araleza]

Hello - I fetched the latest dev branch and tried passing --alpha_mask in as a command line parameter. (I do not have any .toml files). I got this error message:

Traceback (most recent call last):
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/./sdxl_train.py", line 944, in <module>
    train(args)
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/./sdxl_train.py", line 607, in train
    for step, batch in enumerate(train_dataloader):
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/accelerate/data_loader.py", line 448, in __iter__
    current_batch = next(dataloader_iter)
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/venv/lib/python3.10/site-packages/torch/utils/data/dataset.py", line 335, in __getitem__
    return self.datasets[dataset_idx][sample_idx]
  File "/home/ara/m.2/Dev/sdxl/sd-scripts/library/train_util.py", line 1207, in __getitem__
    alpha_mask = None if image_info.alpha_mask is None else torch.flip(image_info.alpha_mask, [1])
TypeError: flip(): argument 'input' (position 1) must be Tensor, not numpy.ndarray

I am doing fp32 sdxl_train.py training with huber loss, a batch size of 4, and I'm training te1 and the unet. I have a single directory of just png images, all with alpha channels, and I'm using multi-line captions.

Did I do something wrong? Do you need more information?

Command line (Linux):

accelerate launch --num_cpu_threads_per_process=2 "./sdxl_train.py" --pretrained_model_name_or_path="/home/ara/Documents/Dev/sdxl/training/earthscape/kohya/dreambooth/earthscape-step00002600.safetensors" --sdpa --enable_bucket --min_bucket_reso=64 --max_bucket_reso=1024 --train_data_dir="/home/ara/Documents/Dev/sdxl/training/earthscape/kohya/img" --resolution="1024,1024" --output_dir="/home/ara/Documents/Dev/sdxl/training/earthscape/kohya/dreambooth" --logging_dir="/home/ara/Documents/Dev/sdxl/training/earthscape/kohya/log" --save_model_as=safetensors --vae="/home/ara/Documents/Dev/sdxl/sdxl_vae.safetensors" --output_name="earthscape" --lr_scheduler_num_cycles="20000" --max_token_length=150 --max_data_loader_n_workers="0" --lr_scheduler="constant_with_warmup" --lr_warmup_steps="100" --max_train_steps="16000" --caption_extension=".txt" --optimizer_type="Adafactor" --optimizer_args scale_parameter=False relative_step=False warmup_init=False --max_data_loader_n_workers="0" --max_token_length=150 --bucket_reso_steps=32 --v_pred_like_loss="0.5" --save_every_n_steps="200" --save_last_n_steps="600" --min_snr_gamma=5 --gradient_checkpointing --xformers --bucket_no_upscale --noise_offset=0.0357 --adaptive_noise_scale=0.00357 --sample_sampler=k_dpm_2 --sample_prompts="/home/ara/Documents/Dev/sdxl/training/earthscape/kohya/dreambooth/sample/prompt.txt" --sample_every_n_steps="50" --fused_backward_pass --cache_latents --loss_type=huber --train_batch_size="4" --train_text_encoder --learning_rate_te1 1e-9 --learning_rate_te2 0 --learning_rate="4e-7" --flip_aug --enable_wildcard --shuffle_caption --alpha_mask

[araleza]

I think I may have found the fix. The code path which runs where latents are not being cached to disk does not seem to be complete. If I find these lines (around line 1200 of train_util.py):

                original_size = image_info.latents_original_size
                crop_ltrb = image_info.latents_crop_ltrb  # calc values later if flipped

and add

                if image_info.alpha_mask is not None:
                    image_info.alpha_mask = torch.FloatTensor(image_info.alpha_mask)

immediately after them, then it starts training again.

I haven't checked that alpha masking is working for the flipped versions of the images, because I don't know if I need to "copy to avoid negative stride problem" as the comment says in the .npz section below the non-npz.

If anyone reading this doesn't want to edit the code before this is fixed, maybe try using --cache_latents_to_disk too, as that goes through a different code path, and I think that code is already correct.

[kohya-ss]

Thank you for reporting. I will check it sooner.

[kohya-ss]

ndarray/Tensor issue is fixed. Sorry for waiting.