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

Added masking to slider training. Something is still weird though

Browse Source
This commit is contained in:
Jaret Burkett
2023-11-01 14:51:29 -06:00
parent a899ec91c8
commit 7d707b2fe6
6 changed files with 97 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
toolkit/config_modules.py
View File

@@ -283,7 +283,7 @@ class SliderConfig:
self.batch_full_slide: bool = kwargs.get('batch_full_slide', True)
self.use_adapter: bool = kwargs.get('use_adapter', None) # depth
self.adapter_img_dir = kwargs.get('adapter_img_dir', None)
self.high_ram = kwargs.get('high_ram', False)
self.low_ram = kwargs.get('low_ram', False)
# expand targets if shuffling
from toolkit.prompt_utils import get_slider_target_permutations
@@ -334,6 +334,7 @@ class DatasetConfig:
self.alpha_mask: bool = kwargs.get('alpha_mask', False) # if true, will use alpha channel as mask
self.mask_path: str = kwargs.get('mask_path',
None) # focus mask (black and white. White has higher loss than black)
self.invert_mask: bool = kwargs.get('invert_mask', False) # invert mask
self.mask_min_value: float = kwargs.get('mask_min_value', 0.01) # min value for . 0 - 1
self.poi: Union[str, None] = kwargs.get('poi',
None) # if one is set and in json data, will be used as auto crop scale point of interes

4
toolkit/dataloader_mixins.py
View File

@@ -18,7 +18,7 @@ from toolkit.buckets import get_bucket_for_image_size
from toolkit.metadata import get_meta_for_safetensors
from toolkit.prompt_utils import inject_trigger_into_prompt
from torchvision import transforms
from PIL import Image, ImageFilter
from PIL import Image, ImageFilter, ImageOps
from PIL.ImageOps import exif_transpose
import albumentations as A
@@ -612,6 +612,8 @@ class MaskFileItemDTOMixin:
img = Image.fromarray(np_img)
img = img.convert('RGB')
if self.dataset_config.invert_mask:
img = ImageOps.invert(img)
w, h = img.size
if w > h and self.scale_to_width < self.scale_to_height:
# throw error, they should match

4
toolkit/stable_diffusion_model.py
View File

@@ -535,7 +535,7 @@ class StableDiffusion:
text_embeddings: Union[PromptEmbeds, None] = None,
timestep: Union[int, torch.Tensor] = 1,
guidance_scale=7.5,
guidance_rescale=0, # 0.7 sdxl
guidance_rescale=0,
add_time_ids=None,
conditional_embeddings: Union[PromptEmbeds, None] = None,
unconditional_embeddings: Union[PromptEmbeds, None] = None,
@@ -674,7 +674,7 @@ class StableDiffusion:
add_time_ids=add_time_ids,
**kwargs,
)
latents = self.noise_scheduler.step(noise_pred, timestep, latents).prev_sample
latents = self.noise_scheduler.step(noise_pred, timestep, latents, return_dict=False)[0]
# return latents_steps
return latents

10
toolkit/train_tools.py
View File

@@ -691,12 +691,12 @@ class LearnableSNRGamma:
def __init__(self, noise_scheduler: Union['DDPMScheduler'], device='cuda'):
self.device = device
self.noise_scheduler: Union['DDPMScheduler'] = noise_scheduler
self.offset = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float32, device=device))
self.scale = torch.nn.Parameter(torch.tensor(0.001, dtype=torch.float32, device=device))
self.gamma = torch.nn.Parameter(torch.tensor(1.0, dtype=torch.float32, device=device))
self.optimizer = torch.optim.AdamW([self.offset, self.gamma, self.scale], lr=0.1)
self.offset = torch.nn.Parameter(torch.tensor(0.777, dtype=torch.float32, device=device))
self.scale = torch.nn.Parameter(torch.tensor(4.14, dtype=torch.float32, device=device))
self.gamma = torch.nn.Parameter(torch.tensor(2.03, dtype=torch.float32, device=device))
self.optimizer = torch.optim.AdamW([self.offset, self.gamma, self.scale], lr=0.01)
self.buffer = []
self.max_buffer_size = 100
self.max_buffer_size = 20
def forward(self, loss, timesteps):
# do a our train loop for lsnr here and return our values detached