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

Add features for models that may need a non masked loss such as inpainiting.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-22 10:59:17 -06:00
parent c133c55cf5
commit 820d534d6e
3 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -1382,7 +1382,7 @@ class SDTrainer(BaseSDTrainProcess):
clip_images = batch.clip_image_tensor.to(self.device_torch, dtype=dtype).detach() clip_images = batch.clip_image_tensor.to(self.device_torch, dtype=dtype).detach()
mask_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype) mask_multiplier = torch.ones((noisy_latents.shape[0], 1, 1, 1), device=self.device_torch, dtype=dtype)
if batch.mask_tensor is not None: if batch.mask_tensor is not None and self.sd.do_masked_loss:
with self.timer('get_mask_multiplier'): with self.timer('get_mask_multiplier'):
# upsampling no supported for bfloat16 # upsampling no supported for bfloat16
mask_multiplier = batch.mask_tensor.to(self.device_torch, dtype=torch.float16).detach() mask_multiplier = batch.mask_tensor.to(self.device_torch, dtype=torch.float16).detach()

3
toolkit/models/base_model.py
View File

@@ -194,6 +194,9 @@ class BaseModel:
# can be used on models to invalidate cache if things change. # can be used on models to invalidate cache if things change.
self.latent_space_version = None self.latent_space_version = None
# if a mask is passed, do the loss with the mask. May be set false for models that use a mask for other reasons.
self.do_masked_loss = True
# properties for old arch for backwards compatibility # properties for old arch for backwards compatibility
@property @property
def unet(self): def unet(self):

3
toolkit/stable_diffusion_model.py
View File

@@ -232,6 +232,9 @@ class StableDiffusion:
# can be used on models to invalidate cache if things change. # can be used on models to invalidate cache if things change.
self.latent_space_version = None self.latent_space_version = None
# if a mask is passed, do the loss with the mask. May be set false for models that use a mask for other reasons.
self.do_masked_loss = True
# properties for old arch for backwards compatibility # properties for old arch for backwards compatibility
@property @property
def is_xl(self): def is_xl(self):