kohya-ss / kohya-ss/sd-scripts
Trying to train with Masked Loss on a compound dataset containing both masked and normal datasets results in KeyError: 'conditioning_images'
- Dominant language
- Python
- Stars
- 7.2k
- Forks
- 1.2k
- Avg merge
- 11m
- Merged PRs (30d)
- 2
Description
When trying to train with a dataset config toml file containing a multi dataset structure, for example two datasets one containing subsets with conditioning data for masked loss, and a second dataset that does not have masked loss but rather has regularization images.
```
[[datasets]]
[[datasets.subsets]]
image_dir = 'path'
conditioning_data_dir = 'path'
[[datasets]]
[[datasets.subsets]]
image_dir = 'path'
[[datasets.subsets]]
is_reg = true
image_dir = 'path'
class_tokens = 'word'
```
attempting to train on this results in the training starting, getting 1 to 2 maybe 3 steps in before throwing the following error, copied the entire set of lines out of the console:
```
steps: 0%| | 0/279800 [00:00
trainer.train(args)
File "kohya_ss\sd-scripts\train_network.py", line 885, in train
loss = apply_masked_loss(loss, batch)
File "kohya_ss\sd-scripts\library\custom_train_functions.py", line 484, in apply_masked_loss
mask_image = batch["conditioning_images"].to(dtype=loss.dtype)[:, 0].unsqueeze(1) # use R channel
KeyError: 'conditioning_images'
steps: 0%| | 2/279800 [00:03<145:19:04, 1.87s/it, Average key norm=0.000256, Keys Scaled=0, avr_loss=0.0575]
Traceback (most recent call last):
File "Python\Python310\lib\runpy.py", line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File "Python\Python310\lib\runpy.py", line 86, in _run_code
exec(code, run_globals)
File "kohya_ss\venv\Scripts\accelerate.exe\__main__.py", line 7, in
File "kohya_ss\venv\lib\site-packages\accelerate\commands\accelerate_cli.py", line 47, in main
args.func(args)
File "kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 1017, in launch_command
simple_launcher(args)
File "kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 637, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['kohya_ss\\venv\\Scripts\\python.exe', 'kohya_ss/sd-scripts/sdxl_train_network.py', '--config_file', './outputs/tmpfilelora.toml']' returned non-zero exit status 1.
```
I am guessing that this is happening because the masked loss is trying to be applied at all times no matter what instead of only applying to the datasets that contain valid conditioning data?
I have verified that the toml config I am using will work and train with only the masked dataset and masked loss enabled, or only the normal dataset with regularization and masked loss disabled, but attempting to train with both datasets in the config file at the same time and masked loss enabled results in the error and training failing.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in sd-scripts/train_network.py at the call to apply_masked_loss, then inspect library/custom_train_functions.py at apply_masked_loss and the batch data produced for the two dataset types. Reproduce with the mixed TOML configuration and confirm training no longer raises KeyError when a batch lacks conditioning_images, while masked-loss batches still use their conditioning data.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100