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

toolkit/models/base_model.py

@@ -193,6 +193,9 @@ class BaseModel:
         
         # can be used on models to invalidate cache if things change.
         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
     @property

toolkit/stable_diffusion_model.py

@@ -232,6 +232,9 @@ class StableDiffusion:
         # can be used on models to invalidate cache if things change.
         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
     @property
     def is_xl(self):