From 8b013b0da8ea48ff29c7779b4f1623b7676222af Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 30 Dec 2025 17:00:40 +0300 Subject: [PATCH 01/67] initial CEP --- modules/modelSetup/BaseChromaSetup.py | 5 ++++ modules/modelSetup/BaseFluxSetup.py | 8 +++++ modules/modelSetup/BaseHiDreamSetup.py | 11 +++++++ modules/modelSetup/BaseHunyuanVideoSetup.py | 8 +++++ modules/modelSetup/BasePixArtAlphaSetup.py | 5 ++++ modules/modelSetup/BaseQwenSetup.py | 5 ++++ modules/modelSetup/BaseSanaSetup.py | 5 ++++ .../modelSetup/BaseStableDiffusion3Setup.py | 8 +++++ .../modelSetup/BaseStableDiffusionSetup.py | 5 ++++ .../modelSetup/BaseStableDiffusionXLSetup.py | 8 +++++ modules/modelSetup/BaseWuerstchenSetup.py | 8 +++++ modules/modelSetup/BaseZImageSetup.py | 6 ++++ .../modelSetup/mixin/ModelSetupNoiseMixin.py | 30 +++++++++++++++++++ modules/ui/TrainingTab.py | 10 ++++++- modules/util/config/TrainConfig.py | 4 +++ 15 files changed, 125 insertions(+), 1 deletion(-) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 7a7847df7..ed07280ba 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -190,6 +190,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = (latent_image - vae_shift_factor) * vae_scaling_factor diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 1865f382e..e72d67a29 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -235,6 +235,14 @@ def predict( apply_attention_mask=config.transformer.attention_mask, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + pooled_text_encoder_output = self._apply_conditional_embedding_perturbation( + pooled_text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = (latent_image - vae_shift_factor) * vae_scaling_factor diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 48e691ecd..4484883ae 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -330,6 +330,17 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) + if config.cep_enabled: + text_encoder_3_output = self._apply_conditional_embedding_perturbation( + text_encoder_3_output, config.cep_gamma, generator + ) + text_encoder_4_output = self._apply_conditional_embedding_perturbation( + text_encoder_4_output, config.cep_gamma, generator + ) + pooled_text_encoder_output = self._apply_conditional_embedding_perturbation( + pooled_text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = (latent_image - vae_shift_factor) * vae_scaling_factor diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index bbb90b71a..95ab1c766 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -234,6 +234,14 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + pooled_text_encoder_output = self._apply_conditional_embedding_perturbation( + pooled_text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = latent_image * vae_scaling_factor diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 8240fb5f4..bf7227656 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -182,6 +182,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = latent_image * vae_scaling_factor diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index dc7115274..0ab2fd887 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -107,6 +107,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = model.scale_latents(latent_image) latent_noise = self._create_noise(scaled_latent_image, config, generator) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 84078ff6f..be94d7f7a 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -193,6 +193,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = latent_image * vae_scaling_factor diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index 5015b21af..00e86c652 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -288,6 +288,14 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + pooled_text_encoder_output = self._apply_conditional_embedding_perturbation( + pooled_text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = (latent_image - vae_shift_factor) * vae_scaling_factor diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 0fc6ed0df..ff04b57f5 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -166,6 +166,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = latent_image * vae_scaling_factor diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 37121951a..42eac8c47 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -217,6 +217,14 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.dropout_probability, )) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + pooled_text_encoder_2_output = self._apply_conditional_embedding_perturbation( + pooled_text_encoder_2_output, config.cep_gamma, generator + ) + latent_image = batch['latent_image'] scaled_latent_image = latent_image * vae_scaling_factor diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 23b3440a5..98516a7de 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -247,6 +247,14 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_embedding = self._apply_conditional_embedding_perturbation( + text_embedding, config.cep_gamma, generator + ) + pooled_text_text_embedding = self._apply_conditional_embedding_perturbation( + pooled_text_text_embedding, config.cep_gamma, generator + ) + latent_input = scaled_noisy_latent_image if model.model_type.is_wuerstchen_v2(): diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index 727122f8b..3229e0f5d 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -106,6 +106,12 @@ def predict( text_encoder_output=batch.get('text_encoder_hidden_state'), text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + scaled_latent_image = model.scale_latents(batch['latent_image']) latent_noise = self._create_noise(scaled_latent_image, config, generator) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 1285caba8..f8b07cbbb 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -118,6 +118,36 @@ def _create_noise( return noise + def _apply_conditional_embedding_perturbation( + self, + embedding: Tensor, + gamma: float, + generator: Generator + ) -> Tensor: + """ + Applies Conditional Embedding Perturbation (CEP) as per Equation (8). + Paper: "Slight Corruption in Pre-training Data Makes Better Diffusion Models" + + delta ~ U(-sqrt(gamma/d), sqrt(gamma/d)) or N(0, sqrt(gamma/d)) + """ + # d denotes the dimension of c_theta(y) + d = embedding.shape[-1] + + # gamma controls perturbation magnitude (Paper uses gamma=1.0 as default baseline) + # Calculate scaling factor: sqrt(gamma / d) + scale = math.sqrt(gamma / d) + + # CEP-U (Uniform) scheme + noise = torch.rand( + embedding.shape, + generator=generator, + device=embedding.device, + dtype=embedding.dtype + ) + perturbation = (noise * 2.0 - 1.0) * scale + + return embedding + perturbation + def _get_timestep_discrete( self, num_train_timesteps: int, diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 9cc2bcec8..76e965b69 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -668,7 +668,15 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b tooltip="Dynamically shift the timestep distribution based on resolution.") components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - + # Conditional Embedding Perturbation (CEP) + cep_label = components.label(frame, 9, 0, "Conditional Embedding Perturbation (CEP)", + tooltip="Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images.") + cep_label.configure(wraplength=130, justify="left") + components.switch(frame, 9, 1, self.ui_state, "cep_enabled") + + components.label(frame, 10, 0, "CEP Gamma", + tooltip="Gamma controls perturbation noise magnitude, paper's default is 1. Only has an effect if CEP is enabled") + components.entry(frame, 10, 1, self.ui_state, "cep_gamma") def __create_masked_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index ddaee4b89..327cf89dc 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -450,6 +450,8 @@ class TrainConfig(BaseConfig): timestep_distribution: TimestepDistribution min_noising_strength: float max_noising_strength: float + cep_enabled: bool + cep_gamma: float noising_weight: float noising_bias: float @@ -1032,6 +1034,8 @@ def default_values() -> 'TrainConfig': data.append(("noising_bias", 0.0, float, False)) data.append(("timestep_shift", 1.0, float, False)) data.append(("dynamic_timestep_shifting", False, bool, False)) + data.append(("cep_enabled", False, bool, False)) + data.append(("cep_gamma", 0.0, float, False)) # unet From e4c2521878b8ffc96883d80d04c33aa857e90c97 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 30 Dec 2025 17:11:32 +0300 Subject: [PATCH 02/67] gamma 1 --- modules/util/config/TrainConfig.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 327cf89dc..deed4fa59 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -1035,7 +1035,7 @@ def default_values() -> 'TrainConfig': data.append(("timestep_shift", 1.0, float, False)) data.append(("dynamic_timestep_shifting", False, bool, False)) data.append(("cep_enabled", False, bool, False)) - data.append(("cep_gamma", 0.0, float, False)) + data.append(("cep_gamma", 1.0, float, False)) # unet From 7cd25b85f00b02fba2a6569891808ee0236c2153 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 30 Dec 2025 17:19:40 +0300 Subject: [PATCH 03/67] adjust position --- modules/ui/TrainingTab.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 76e965b69..73cf70195 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -669,14 +669,14 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") # Conditional Embedding Perturbation (CEP) - cep_label = components.label(frame, 9, 0, "Conditional Embedding Perturbation (CEP)", + cep_label = components.label(frame, 10, 0, "Conditional Embedding Perturbation (CEP)", tooltip="Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images.") cep_label.configure(wraplength=130, justify="left") - components.switch(frame, 9, 1, self.ui_state, "cep_enabled") + components.switch(frame, 10, 1, self.ui_state, "cep_enabled") - components.label(frame, 10, 0, "CEP Gamma", + components.label(frame, 11, 0, "CEP Gamma", tooltip="Gamma controls perturbation noise magnitude, paper's default is 1. Only has an effect if CEP is enabled") - components.entry(frame, 10, 1, self.ui_state, "cep_gamma") + components.entry(frame, 11, 1, self.ui_state, "cep_gamma") def __create_masked_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) From 39af0d3301809acd38aeb6c95bbd9d89e50a8208 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Thu, 29 Jan 2026 21:31:42 +0300 Subject: [PATCH 04/67] fix transformer (CEP) --- .../modelSetup/mixin/ModelSetupNoiseMixin.py | 37 +++++++++++-------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index f8b07cbbb..18961a12a 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -120,33 +120,38 @@ def _create_noise( def _apply_conditional_embedding_perturbation( self, - embedding: Tensor, + embedding: Tensor | list, gamma: float, generator: Generator - ) -> Tensor: + ) -> Tensor | list: """ Applies Conditional Embedding Perturbation (CEP) as per Equation (8). Paper: "Slight Corruption in Pre-training Data Makes Better Diffusion Models" delta ~ U(-sqrt(gamma/d), sqrt(gamma/d)) or N(0, sqrt(gamma/d)) """ - # d denotes the dimension of c_theta(y) - d = embedding.shape[-1] + def _perturb_cep(tensor: Tensor) -> Tensor: + # d denotes the dimension of c_theta(y) + d = tensor.shape[-1] - # gamma controls perturbation magnitude (Paper uses gamma=1.0 as default baseline) - # Calculate scaling factor: sqrt(gamma / d) - scale = math.sqrt(gamma / d) + # gamma controls perturbation magnitude (Paper uses gamma=1.0 as default baseline) + # Calculate scaling factor: sqrt(gamma / d) + scale = math.sqrt(gamma / d) - # CEP-U (Uniform) scheme - noise = torch.rand( - embedding.shape, - generator=generator, - device=embedding.device, - dtype=embedding.dtype - ) - perturbation = (noise * 2.0 - 1.0) * scale + # CEP-U (Uniform) scheme + noise = torch.rand( + tensor.shape, + generator=generator, + device=tensor.device, + dtype=tensor.dtype + ) + perturbation = (noise * 2.0 - 1.0) * scale + return tensor + perturbation - return embedding + perturbation + if isinstance(embedding, list): + return [_perturb_cep(emb) for emb in embedding] + else: + return _perturb_cep(embedding) def _get_timestep_discrete( self, From a860bb51a45175ea9ea93241dc8ed416ececed9a Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 15 Feb 2026 20:51:14 +0100 Subject: [PATCH 05/67] add flux2 --- modules/modelSetup/BaseFlux2Setup.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index 8c02b2873..437c8fbfc 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -111,6 +111,11 @@ def predict( text_encoder_output=batch.get('text_encoder_hidden_state'), text_encoder_dropout_probability=config.text_encoder.dropout_probability, ) + if config.cep_enabled: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + latent_image = model.patchify_latents(batch['latent_image'].float()) latent_height = latent_image.shape[-2] latent_width = latent_image.shape[-1] From 67dfd2d83d2596e472248d3b95a4d0b29d871ea7 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 22 Mar 2026 09:40:06 +0100 Subject: [PATCH 06/67] update --- modules/module/quantized/LinearW16A8.py | 121 ++++++++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 modules/module/quantized/LinearW16A8.py diff --git a/modules/module/quantized/LinearW16A8.py b/modules/module/quantized/LinearW16A8.py new file mode 100644 index 000000000..429a04377 --- /dev/null +++ b/modules/module/quantized/LinearW16A8.py @@ -0,0 +1,121 @@ +from modules.util.mm_8bit import mm_8bit as mm_8bit +from modules.util.quantization_util import ( + quantize_fp8_axiswise, + quantize_int8_axiswise, +) + +import torch +from torch import Tensor + +#TODO share code + +@torch.no_grad() +def int8_forward_axiswise(x: Tensor, weight: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: + x_8, x_scale = quantize_int8_axiswise(x, dim=-1) + w_8, w_scale = quantize_int8_axiswise(weight, dim=-1) + res = torch._int_mm(x_8, w_8.T) + res_scaled = res.float().mul_(w_scale.T).mul_(x_scale).to(compute_dtype) + if bias is not None: + res_scaled.add_(bias) + return res_scaled + +@torch.no_grad() +def fp8_forward_axiswise(x: Tensor, weight: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: + x_8, x_scale = quantize_fp8_axiswise(x, dim=-1) + w_8, w_scale = quantize_fp8_axiswise(weight, dim=-1) + one = torch.ones(1, device=x.device) + res = torch._scaled_mm(x_8, w_8.T, scale_a=one, scale_b=one, out_dtype=torch.float) + res_scaled = res.mul_(w_scale.T).mul_(x_scale).to(compute_dtype) #much faster than scaled by _scaled_mm + if bias is not None: + res_scaled.add_(bias) + return res_scaled + +def int8_backward_act_axiswise(output: Tensor, weight: Tensor) -> Tensor: + output_8, output_scale = quantize_int8_axiswise(output, dim=-1) + w_8, w_scale = quantize_int8_axiswise(weight, dim=0) + #almost always, grad outputs are already contiguous and this is a no-op. But there are some grad outputs from SDXL that are non-contiguous: + output_8 = output_8.contiguous() + mm_res = mm_8bit(output_8, w_8) + return mm_res.to(output.dtype).mul_(w_scale).mul_(output_scale) + +def fp8_backward_act_axiswise(output: Tensor, weight: Tensor) -> Tensor: + output_8, output_scale = quantize_fp8_axiswise(output, dim=-1) + w_8, w_scale = quantize_fp8_axiswise(weight, dim=0) + mm_res = mm_8bit(output_8.contiguous(), w_8) + return mm_res.to(output.dtype).mul_(w_scale).mul_(output_scale) + +def int8_backward_weight_axiswise(output: Tensor, x: Tensor) -> Tensor: + output_8, output_scale = quantize_int8_axiswise(output, dim=0) + x_8, x_scale = quantize_int8_axiswise(x, dim=0) + #TODO could be more efficient using a kernel that accepts a non-contiguous lhs matrix + mm_res = mm_8bit(output_8.T.contiguous(), x_8) + return mm_res.to(x.dtype).mul_(output_scale.T).mul_(x_scale) + +def fp8_backward_weight_axiswise(output: Tensor, x: Tensor) -> Tensor: + output_8, output_scale = quantize_fp8_axiswise(output, dim=0) + x_8, x_scale = quantize_fp8_axiswise(x, dim=0) + mm_res = mm_8bit(output_8.T.contiguous(), x_8) + return mm_res.to(x.dtype).mul_(output_scale.T).mul_(x_scale) + + +class LinearW16IntA8Function(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: + ctx.save_for_backward(x, weight) + return int8_forward_axiswise(x, weight, bias) + + @staticmethod + def backward(ctx, grad_output: Tensor): + x, weight = ctx.saved_tensors + + grad_x, grad_weight, grad_bias = None, None, None + if ctx.needs_input_grad[0]: + # grad_output @ weight.T + grad_x = int8_backward_act_axiswise(grad_output, weight) + if ctx.needs_input_grad[1]: + # grad_output.T @ x + grad_weight = int8_backward_weight_axiswise(grad_output, x) + if ctx.needs_input_grad[2]: + grad_bias = grad_output.sum(0) + + return grad_x, grad_weight, grad_bias + +class LinearW16FpA8Function(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: + ctx.save_for_backward(x, weight) + return fp8_forward_axiswise(x, weight, bias) + + @staticmethod + def backward(ctx, grad_output: Tensor): + x, weight = ctx.saved_tensors + + grad_x, grad_weight, grad_bias = None, None, None + if ctx.needs_input_grad[0]: + # grad_output @ weight.T + grad_x = fp8_backward_act_axiswise(grad_output, weight) + if ctx.needs_input_grad[1]: + # grad_output.T @ x + grad_weight = fp8_backward_weight_axiswise(grad_output, x) + if ctx.needs_input_grad[2]: + grad_bias = grad_output.sum(0) + + return grad_x, grad_weight, grad_bias + +class LinearW16A8(torch.nn.Linear): + def __init__(self, dtype, *args, **kwargs): + super().__init__(*args, **kwargs) + + assert dtype in [torch.int8, torch.float8_e4m3fn] + self._dtype = dtype + + def forward(self, x_orig: torch.Tensor) -> torch.Tensor: + x = x_orig.to(self.weight.dtype).reshape(-1, x_orig.shape[-1]) + if x.shape[0] > 16: + if self._dtype == torch.int8: + y = LinearW16IntA8Function.apply(x, self.weight, self.bias) + else: + y = LinearW16FpA8Function.apply(x, self.weight, self.bias) + return y.reshape(x_orig.shape[:-1] + (y.shape[-1], )) + else: + return super().forward(x_orig) From 230cae87fbde2a1d893357a846dd2d17d6c0cc20 Mon Sep 17 00:00:00 2001 From: dxqb Date: Tue, 24 Mar 2026 21:57:01 +0100 Subject: [PATCH 07/67] update --- .../quantized/{LinearW16A8.py => LinearA8.py} | 23 ++++++++++--------- modules/ui/ModelTab.py | 5 ++-- modules/util/enum/DataType.py | 6 +++++ modules/util/quantization_util.py | 11 +++++++-- 4 files changed, 30 insertions(+), 15 deletions(-) rename modules/module/quantized/{LinearW16A8.py => LinearA8.py} (86%) diff --git a/modules/module/quantized/LinearW16A8.py b/modules/module/quantized/LinearA8.py similarity index 86% rename from modules/module/quantized/LinearW16A8.py rename to modules/module/quantized/LinearA8.py index 429a04377..17a76ee38 100644 --- a/modules/module/quantized/LinearW16A8.py +++ b/modules/module/quantized/LinearA8.py @@ -58,11 +58,11 @@ def fp8_backward_weight_axiswise(output: Tensor, x: Tensor) -> Tensor: return mm_res.to(x.dtype).mul_(output_scale.T).mul_(x_scale) -class LinearW16IntA8Function(torch.autograd.Function): +class LinearIntA8Function(torch.autograd.Function): @staticmethod - def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: + def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: ctx.save_for_backward(x, weight) - return int8_forward_axiswise(x, weight, bias) + return int8_forward_axiswise(x, weight, bias, compute_dtype) @staticmethod def backward(ctx, grad_output: Tensor): @@ -78,13 +78,13 @@ def backward(ctx, grad_output: Tensor): if ctx.needs_input_grad[2]: grad_bias = grad_output.sum(0) - return grad_x, grad_weight, grad_bias + return grad_x, grad_weight, grad_bias, None -class LinearW16FpA8Function(torch.autograd.Function): +class LinearFpA8Function(torch.autograd.Function): @staticmethod - def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: + def forward(ctx, x: Tensor, weight: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: ctx.save_for_backward(x, weight) - return fp8_forward_axiswise(x, weight, bias) + return fp8_forward_axiswise(x, weight, bias, compute_dtype) @staticmethod def backward(ctx, grad_output: Tensor): @@ -100,22 +100,23 @@ def backward(ctx, grad_output: Tensor): if ctx.needs_input_grad[2]: grad_bias = grad_output.sum(0) - return grad_x, grad_weight, grad_bias + return grad_x, grad_weight, grad_bias, None -class LinearW16A8(torch.nn.Linear): +class LinearA8(torch.nn.Linear): def __init__(self, dtype, *args, **kwargs): super().__init__(*args, **kwargs) assert dtype in [torch.int8, torch.float8_e4m3fn] self._dtype = dtype + self.compute_dtype = None def forward(self, x_orig: torch.Tensor) -> torch.Tensor: x = x_orig.to(self.weight.dtype).reshape(-1, x_orig.shape[-1]) if x.shape[0] > 16: if self._dtype == torch.int8: - y = LinearW16IntA8Function.apply(x, self.weight, self.bias) + y = LinearIntA8Function.apply(x, self.weight, self.bias, self.compute_dtype) else: - y = LinearW16FpA8Function.apply(x, self.weight, self.bias) + y = LinearFpA8Function.apply(x, self.weight, self.bias, self.compute_dtype) return y.reshape(x_orig.shape[:-1] + (y.shape[-1], )) else: return super().forward(x_orig) diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index 6ff98d086..5c4e0facc 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -340,14 +340,15 @@ def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=Fals options += [ ("float W8A8", DataType.FLOAT_W8A8), ("int W8A8", DataType.INT_W8A8), - ] + ("bfloat16 A8 int", DataType.BFLOAT_16_A8_INT), + ("bfloat16 A8 float", DataType.BFLOAT_16_A8_FLOAT), ] if include_gguf: options.append(("GGUF", DataType.GGUF)) if include_a8: options += [ - ("GGUF A8 float", DataType.GGUF_A8_FLOAT), ("GGUF A8 int", DataType.GGUF_A8_INT), + ("GGUF A8 float", DataType.GGUF_A8_FLOAT), ] return options diff --git a/modules/util/enum/DataType.py b/modules/util/enum/DataType.py index 7e9bcb4c8..dfe836ef7 100644 --- a/modules/util/enum/DataType.py +++ b/modules/util/enum/DataType.py @@ -17,6 +17,8 @@ class DataType(Enum): GGUF = 'GGUF' GGUF_A8_FLOAT = 'GGUF_A8_FLOAT' GGUF_A8_INT = 'GGUF_A8_INT' + BFLOAT_16_A8_INT = 'BFLOAT_16_A8_INT' + BFLOAT_16_A8_FLOAT = 'BFLOAT_16_A8_FLOAT' def __str__(self): return self.value @@ -37,6 +39,10 @@ def torch_dtype( return torch.bfloat16 case DataType.TFLOAT_32: return torch.float32 + case DataType.BFLOAT_16_A8_FLOAT: + return torch.bfloat16 + case DataType.BFLOAT_16_A8_INT: + return torch.bfloat16 case _: return None diff --git a/modules/util/quantization_util.py b/modules/util/quantization_util.py index 9eae3030d..68570b30a 100644 --- a/modules/util/quantization_util.py +++ b/modules/util/quantization_util.py @@ -76,6 +76,7 @@ def dequantize(q: Tensor, scale: float | Tensor) -> Tensor: return q.float() * scale +from modules.module.quantized.LinearA8 import LinearA8 from modules.module.quantized.LinearFp8 import LinearFp8 from modules.module.quantized.LinearGGUFA8 import LinearGGUFA8 from modules.module.quantized.LinearSVD import BaseLinearSVD, make_svd_linear @@ -193,6 +194,12 @@ def replace_linear_with_quantized_layers( elif dtype == DataType.GGUF_A8_FLOAT: linear_class=LinearGGUFA8 kwargs = {'dtype': torch.float8_e4m3fn} + elif dtype == DataType.BFLOAT_16_A8_INT: + linear_class=LinearA8 + kwargs = {'dtype': torch.int8} + elif dtype == DataType.BFLOAT_16_A8_FLOAT: + linear_class=LinearA8 + kwargs = {'dtype': torch.float8_e4m3fn} else: return @@ -226,7 +233,7 @@ def replace_linear_with_quantized_layers( #https://github.com/Nerogar/OneTrainer/issues/1050 for name, module in parent_module.named_modules(): assert (not isinstance(module, convert_type) - or isinstance(module, (QuantizedLinearMixin, LinearGGUFA8)) + or isinstance(module, (QuantizedLinearMixin, LinearGGUFA8, LinearA8)) or any(s in name.split('.') for s in keep_in_fp32_modules) or (quant_filters is not None and len(quant_filters) > 0 and not any(f.matches(name) for f in quant_filters)) ), f"Linear layer {name} was not found in model for quantization" @@ -263,7 +270,7 @@ def quantize_layers(module: nn.Module, device: torch.device, train_dtype: DataTy child_modules = list(module.modules()) for _ in multi.master_first(): #avoid cache writing conflicts for child_module in tqdm(child_modules, desc="Quantizing model weights", total=len(child_modules), delay=5, smoothing=0.1): - if isinstance(child_module, (QuantizedModuleMixin, GGUFLinear)): + if isinstance(child_module, (QuantizedModuleMixin, GGUFLinear, LinearA8)): child_module.compute_dtype = train_dtype.torch_dtype() if isinstance(child_module, QuantizedModuleMixin): child_module.quantize(device=device) From 0c869b86f5976eaab139a15c562358548c0d5d90 Mon Sep 17 00:00:00 2001 From: dxqb Date: Wed, 25 Mar 2026 00:39:42 +0100 Subject: [PATCH 08/67] torch 2.11 --- modules/modelSampler/BaseModelSampler.py | 4 ++-- requirements-cuda.txt | 4 ++-- requirements-default.txt | 4 ++-- requirements-rocm.txt | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/modules/modelSampler/BaseModelSampler.py b/modules/modelSampler/BaseModelSampler.py index af68a5015..74724dbef 100644 --- a/modules/modelSampler/BaseModelSampler.py +++ b/modules/modelSampler/BaseModelSampler.py @@ -11,8 +11,8 @@ from modules.util.enum.VideoFormat import VideoFormat import torch -from torchvision.io import write_video +#from torchvision.io import write_video from PIL import Image @@ -95,6 +95,6 @@ def save_sampler_output( elif sampler_output.file_type == FileType.VIDEO: if video_format is None: raise ValueError("Video format required for sampling a video") - write_video(destination + video_format.extension(), options={"crf": "17"}, video_array=sampler_output.data, fps=24) + #write_video(destination + video_format.extension(), options={"crf": "17"}, video_array=sampler_output.data, fps=24) #FIXME elif sampler_output.file_type == FileType.AUDIO: pass # TODO diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 35c3059d7..e53a76215 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -1,7 +1,7 @@ # pytorch --extra-index-url https://download.pytorch.org/whl/cu128 -torch==2.9.1+cu128 -torchvision==0.24.1+cu128 +torch==2.11.0+cu128 +torchvision==0.26.0+cu128 onnxruntime-gpu==1.23.2 nvidia-nccl-cu12==2.27.5; sys_platform == "linux" triton-windows==3.5.1.post24; sys_platform == "win32" diff --git a/requirements-default.txt b/requirements-default.txt index 06590561c..7581e2ac4 100644 --- a/requirements-default.txt +++ b/requirements-default.txt @@ -1,6 +1,6 @@ # pytorch -torch==2.9.1 -torchvision==0.24.1 +torch==2.11.0 +torchvision==0.26.0 onnxruntime==1.23.2 # optimizers diff --git a/requirements-rocm.txt b/requirements-rocm.txt index 41a96d70a..4f92044cb 100644 --- a/requirements-rocm.txt +++ b/requirements-rocm.txt @@ -3,8 +3,8 @@ # pytorch --extra-index-url https://download.pytorch.org/whl/rocm6.3 -torch==2.9.1+rocm6.3 -torchvision==0.24.1+rocm6.3 +torch==2.11.0+rocm6.3 +torchvision==0.26.0+rocm6.3 onnxruntime==1.23.2 # optimizers From fe50e8d1df075abec8a265d5b164b085de792bb3 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:06:47 +0200 Subject: [PATCH 09/67] feat: split CTK UI into controller/base-view/ctk-view pattern Co-Authored-By: Claude Sonnet 4.6 --- ....py => BaseAdditionalEmbeddingsTabView.py} | 0 .../ui/{CaptionUI.py => BaseCaptionUIView.py} | 0 .../ui/{CloudTab.py => BaseCloudTabView.py} | 0 .../{ConceptTab.py => BaseConceptTabView.py} | 0 ...ceptWindow.py => BaseConceptWindowView.py} | 0 .../{ConfigList.py => BaseConfigListView.py} | 0 ...rtModelUI.py => BaseConvertModelUIView.py} | 0 ...w.py => BaseGenerateCaptionsWindowView.py} | 0 ...ndow.py => BaseGenerateMasksWindowView.py} | 0 modules/ui/{LoraTab.py => BaseLoraTabView.py} | 0 .../ui/{ModelTab.py => BaseModelTabView.py} | 0 ...damWindow.py => BaseMuonAdamWindowView.py} | 0 ...gWindow.py => BaseOffloadingWindowView.py} | 0 ...ow.py => BaseOptimizerParamsWindowView.py} | 0 ...ngWindow.py => BaseProfilingWindowView.py} | 0 ...{SampleFrame.py => BaseSampleFrameView.py} | 0 ...indow.py => BaseSampleParamsWindowView.py} | 0 ...ampleWindow.py => BaseSampleWindowView.py} | 0 ...{SamplingTab.py => BaseSamplingTabView.py} | 0 ...ow.py => BaseSchedulerParamsWindowView.py} | 0 ... => BaseTimestepDistributionWindowView.py} | 0 modules/ui/{TopBar.py => BaseTopBarView.py} | 0 modules/ui/{TrainUI.py => BaseTrainUIView.py} | 0 ...{TrainingTab.py => BaseTrainingTabView.py} | 0 ...{VideoToolUI.py => BaseVideoToolUIView.py} | 0 modules/ui/CaptionUIController.py | 572 +++++++++++ modules/ui/ConceptWindowController.py | 934 ++++++++++++++++++ modules/ui/ConvertModelUIController.py | 170 ++++ modules/ui/CtkAdditionalEmbeddingsTabView.py | 136 +++ modules/ui/CtkCaptionUIView.py | 572 +++++++++++ modules/ui/CtkCloudTabView.py | 221 +++++ modules/ui/CtkConceptTabView.py | 286 ++++++ modules/ui/CtkConceptWindowView.py | 934 ++++++++++++++++++ modules/ui/CtkConfigListView.py | 354 +++++++ modules/ui/CtkConvertModelUIView.py | 170 ++++ modules/ui/CtkGenerateCaptionsWindowView.py | 133 +++ modules/ui/CtkGenerateMasksWindowView.py | 151 +++ modules/ui/CtkLoraTabView.py | 154 +++ modules/ui/CtkModelTabView.py | 688 +++++++++++++ modules/ui/CtkMuonAdamWindowView.py | 107 ++ modules/ui/CtkOffloadingWindowView.py | 75 ++ modules/ui/CtkOptimizerParamsWindowView.py | 288 ++++++ modules/ui/CtkProfilingWindowView.py | 57 ++ modules/ui/CtkSampleFrameView.py | 134 +++ modules/ui/CtkSampleParamsWindowView.py | 39 + modules/ui/CtkSampleWindowView.py | 227 +++++ modules/ui/CtkSamplingTabView.py | 124 +++ modules/ui/CtkSchedulerParamsWindowView.py | 119 +++ .../ui/CtkTimestepDistributionWindowView.py | 186 ++++ modules/ui/CtkTopBarView.py | 260 +++++ modules/ui/CtkTrainUIView.py | 889 +++++++++++++++++ modules/ui/CtkTrainingTabView.py | 856 ++++++++++++++++ modules/ui/CtkVideoToolUIView.py | 877 ++++++++++++++++ .../ui/GenerateCaptionsWindowController.py | 133 +++ modules/ui/GenerateMasksWindowController.py | 151 +++ modules/ui/OptimizerParamsWindowController.py | 288 ++++++ modules/ui/SampleWindowController.py | 227 +++++ .../TimestepDistributionWindowController.py | 186 ++++ modules/ui/TopBarController.py | 260 +++++ modules/ui/TrainUIController.py | 889 +++++++++++++++++ modules/ui/VideoToolUIController.py | 877 ++++++++++++++++ modules/util/ui/ctk_validation.py | 501 ++++++++++ 62 files changed, 13225 insertions(+) rename modules/ui/{AdditionalEmbeddingsTab.py => BaseAdditionalEmbeddingsTabView.py} (100%) rename modules/ui/{CaptionUI.py => BaseCaptionUIView.py} (100%) rename modules/ui/{CloudTab.py => BaseCloudTabView.py} (100%) rename modules/ui/{ConceptTab.py => BaseConceptTabView.py} (100%) rename modules/ui/{ConceptWindow.py => BaseConceptWindowView.py} (100%) rename modules/ui/{ConfigList.py => BaseConfigListView.py} (100%) rename modules/ui/{ConvertModelUI.py => BaseConvertModelUIView.py} (100%) rename modules/ui/{GenerateCaptionsWindow.py => BaseGenerateCaptionsWindowView.py} (100%) rename modules/ui/{GenerateMasksWindow.py => BaseGenerateMasksWindowView.py} (100%) rename modules/ui/{LoraTab.py => BaseLoraTabView.py} (100%) rename modules/ui/{ModelTab.py => BaseModelTabView.py} (100%) rename modules/ui/{MuonAdamWindow.py => BaseMuonAdamWindowView.py} (100%) rename modules/ui/{OffloadingWindow.py => BaseOffloadingWindowView.py} (100%) rename modules/ui/{OptimizerParamsWindow.py => BaseOptimizerParamsWindowView.py} (100%) rename modules/ui/{ProfilingWindow.py => BaseProfilingWindowView.py} (100%) rename modules/ui/{SampleFrame.py => BaseSampleFrameView.py} (100%) rename modules/ui/{SampleParamsWindow.py => BaseSampleParamsWindowView.py} (100%) rename modules/ui/{SampleWindow.py => BaseSampleWindowView.py} (100%) rename modules/ui/{SamplingTab.py => BaseSamplingTabView.py} (100%) rename modules/ui/{SchedulerParamsWindow.py => BaseSchedulerParamsWindowView.py} (100%) rename modules/ui/{TimestepDistributionWindow.py => BaseTimestepDistributionWindowView.py} (100%) rename modules/ui/{TopBar.py => BaseTopBarView.py} (100%) rename modules/ui/{TrainUI.py => BaseTrainUIView.py} (100%) rename modules/ui/{TrainingTab.py => BaseTrainingTabView.py} (100%) rename modules/ui/{VideoToolUI.py => BaseVideoToolUIView.py} (100%) create mode 100644 modules/ui/CaptionUIController.py create mode 100644 modules/ui/ConceptWindowController.py create mode 100644 modules/ui/ConvertModelUIController.py create mode 100644 modules/ui/CtkAdditionalEmbeddingsTabView.py create mode 100644 modules/ui/CtkCaptionUIView.py create mode 100644 modules/ui/CtkCloudTabView.py create mode 100644 modules/ui/CtkConceptTabView.py create mode 100644 modules/ui/CtkConceptWindowView.py create mode 100644 modules/ui/CtkConfigListView.py create mode 100644 modules/ui/CtkConvertModelUIView.py create mode 100644 modules/ui/CtkGenerateCaptionsWindowView.py create mode 100644 modules/ui/CtkGenerateMasksWindowView.py create mode 100644 modules/ui/CtkLoraTabView.py create mode 100644 modules/ui/CtkModelTabView.py create mode 100644 modules/ui/CtkMuonAdamWindowView.py create mode 100644 modules/ui/CtkOffloadingWindowView.py create mode 100644 modules/ui/CtkOptimizerParamsWindowView.py create mode 100644 modules/ui/CtkProfilingWindowView.py create mode 100644 modules/ui/CtkSampleFrameView.py create mode 100644 modules/ui/CtkSampleParamsWindowView.py create mode 100644 modules/ui/CtkSampleWindowView.py create mode 100644 modules/ui/CtkSamplingTabView.py create mode 100644 modules/ui/CtkSchedulerParamsWindowView.py create mode 100644 modules/ui/CtkTimestepDistributionWindowView.py create mode 100644 modules/ui/CtkTopBarView.py create mode 100644 modules/ui/CtkTrainUIView.py create mode 100644 modules/ui/CtkTrainingTabView.py create mode 100644 modules/ui/CtkVideoToolUIView.py create mode 100644 modules/ui/GenerateCaptionsWindowController.py create mode 100644 modules/ui/GenerateMasksWindowController.py create mode 100644 modules/ui/OptimizerParamsWindowController.py create mode 100644 modules/ui/SampleWindowController.py create mode 100644 modules/ui/TimestepDistributionWindowController.py create mode 100644 modules/ui/TopBarController.py create mode 100644 modules/ui/TrainUIController.py create mode 100644 modules/ui/VideoToolUIController.py create mode 100644 modules/util/ui/ctk_validation.py diff --git a/modules/ui/AdditionalEmbeddingsTab.py b/modules/ui/BaseAdditionalEmbeddingsTabView.py similarity index 100% rename from modules/ui/AdditionalEmbeddingsTab.py rename to modules/ui/BaseAdditionalEmbeddingsTabView.py diff --git a/modules/ui/CaptionUI.py b/modules/ui/BaseCaptionUIView.py similarity index 100% rename from modules/ui/CaptionUI.py rename to modules/ui/BaseCaptionUIView.py diff --git a/modules/ui/CloudTab.py b/modules/ui/BaseCloudTabView.py similarity index 100% rename from modules/ui/CloudTab.py rename to modules/ui/BaseCloudTabView.py diff --git a/modules/ui/ConceptTab.py b/modules/ui/BaseConceptTabView.py similarity index 100% rename from modules/ui/ConceptTab.py rename to modules/ui/BaseConceptTabView.py diff --git a/modules/ui/ConceptWindow.py b/modules/ui/BaseConceptWindowView.py similarity index 100% rename from modules/ui/ConceptWindow.py rename to modules/ui/BaseConceptWindowView.py diff --git a/modules/ui/ConfigList.py b/modules/ui/BaseConfigListView.py similarity index 100% rename from modules/ui/ConfigList.py rename to modules/ui/BaseConfigListView.py diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/BaseConvertModelUIView.py similarity index 100% rename from modules/ui/ConvertModelUI.py rename to modules/ui/BaseConvertModelUIView.py diff --git a/modules/ui/GenerateCaptionsWindow.py b/modules/ui/BaseGenerateCaptionsWindowView.py similarity index 100% rename from modules/ui/GenerateCaptionsWindow.py rename to modules/ui/BaseGenerateCaptionsWindowView.py diff --git a/modules/ui/GenerateMasksWindow.py b/modules/ui/BaseGenerateMasksWindowView.py similarity index 100% rename from modules/ui/GenerateMasksWindow.py rename to modules/ui/BaseGenerateMasksWindowView.py diff --git a/modules/ui/LoraTab.py b/modules/ui/BaseLoraTabView.py similarity index 100% rename from modules/ui/LoraTab.py rename to modules/ui/BaseLoraTabView.py diff --git a/modules/ui/ModelTab.py b/modules/ui/BaseModelTabView.py similarity index 100% rename from modules/ui/ModelTab.py rename to modules/ui/BaseModelTabView.py diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/BaseMuonAdamWindowView.py similarity index 100% rename from modules/ui/MuonAdamWindow.py rename to modules/ui/BaseMuonAdamWindowView.py diff --git a/modules/ui/OffloadingWindow.py b/modules/ui/BaseOffloadingWindowView.py similarity index 100% rename from modules/ui/OffloadingWindow.py rename to modules/ui/BaseOffloadingWindowView.py diff --git a/modules/ui/OptimizerParamsWindow.py b/modules/ui/BaseOptimizerParamsWindowView.py similarity index 100% rename from modules/ui/OptimizerParamsWindow.py rename to modules/ui/BaseOptimizerParamsWindowView.py diff --git a/modules/ui/ProfilingWindow.py b/modules/ui/BaseProfilingWindowView.py similarity index 100% rename from modules/ui/ProfilingWindow.py rename to modules/ui/BaseProfilingWindowView.py diff --git a/modules/ui/SampleFrame.py b/modules/ui/BaseSampleFrameView.py similarity index 100% rename from modules/ui/SampleFrame.py rename to modules/ui/BaseSampleFrameView.py diff --git a/modules/ui/SampleParamsWindow.py b/modules/ui/BaseSampleParamsWindowView.py similarity index 100% rename from modules/ui/SampleParamsWindow.py rename to modules/ui/BaseSampleParamsWindowView.py diff --git a/modules/ui/SampleWindow.py b/modules/ui/BaseSampleWindowView.py similarity index 100% rename from modules/ui/SampleWindow.py rename to modules/ui/BaseSampleWindowView.py diff --git a/modules/ui/SamplingTab.py b/modules/ui/BaseSamplingTabView.py similarity index 100% rename from modules/ui/SamplingTab.py rename to modules/ui/BaseSamplingTabView.py diff --git a/modules/ui/SchedulerParamsWindow.py b/modules/ui/BaseSchedulerParamsWindowView.py similarity index 100% rename from modules/ui/SchedulerParamsWindow.py rename to modules/ui/BaseSchedulerParamsWindowView.py diff --git a/modules/ui/TimestepDistributionWindow.py b/modules/ui/BaseTimestepDistributionWindowView.py similarity index 100% rename from modules/ui/TimestepDistributionWindow.py rename to modules/ui/BaseTimestepDistributionWindowView.py diff --git a/modules/ui/TopBar.py b/modules/ui/BaseTopBarView.py similarity index 100% rename from modules/ui/TopBar.py rename to modules/ui/BaseTopBarView.py diff --git a/modules/ui/TrainUI.py b/modules/ui/BaseTrainUIView.py similarity index 100% rename from modules/ui/TrainUI.py rename to modules/ui/BaseTrainUIView.py diff --git a/modules/ui/TrainingTab.py b/modules/ui/BaseTrainingTabView.py similarity index 100% rename from modules/ui/TrainingTab.py rename to modules/ui/BaseTrainingTabView.py diff --git a/modules/ui/VideoToolUI.py b/modules/ui/BaseVideoToolUIView.py similarity index 100% rename from modules/ui/VideoToolUI.py rename to modules/ui/BaseVideoToolUIView.py diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py new file mode 100644 index 000000000..e6cc0551e --- /dev/null +++ b/modules/ui/CaptionUIController.py @@ -0,0 +1,572 @@ +import os +import platform +import subprocess +import traceback +from tkinter import filedialog + +from modules.module.Blip2Model import Blip2Model +from modules.module.BlipModel import BlipModel +from modules.module.ClipSegModel import ClipSegModel +from modules.module.MaskByColor import MaskByColor +from modules.module.RembgHumanModel import RembgHumanModel +from modules.module.RembgModel import RembgModel +from modules.module.WDModel import WDModel +from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow +from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.util import path_util +from modules.util.image_util import load_image +from modules.util.torch_util import default_device, torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +import cv2 +import numpy as np +from customtkinter import ScalingTracker, ThemeManager +from PIL import Image, ImageDraw + + +class CaptionUI(ctk.CTkToplevel): + def __init__( + self, + parent, + initial_dir: str | None, + initial_include_subdirectories: bool, + *args, + **kwargs, + ) -> None: + super().__init__(parent, *args, **kwargs) + self.protocol("WM_DELETE_WINDOW", self._on_close) + + self.dir = initial_dir + self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} + self.config_ui_state = UIState(self, self.config_ui_data) + self.image_size = 850 + self.help_text = """ + Keyboard shortcuts when focusing on the prompt input field: + Up arrow: previous image + Down arrow: next image + Return: save + Ctrl+M: only show the mask + Ctrl+D: draw mask editing mode + Ctrl+F: fill mask editing mode + + When editing masks: + Left click: add mask + Right click: remove mask + Mouse wheel: increase or decrease brush size""" + self.masking_model = None + self.captioning_model = None + self.image_rel_paths = [] + self.current_image_index = -1 + self.file_list = None + self.image_labels = [] + self.pil_image = None + self.image_width = 0 + self.image_height = 0 + self.pil_mask = None + self.mask_draw_x = 0 + self.mask_draw_y = 0 + self.mask_draw_radius = 0.01 + self.display_only_mask = False + self.image = None + self.image_label = None + self.mask_editing_mode = 'draw' + self.enable_mask_editing_var = ctk.BooleanVar() + self.mask_editing_alpha = None + self.prompt_var = None + self.prompt_component = None + + + self.title("OneTrainer") + self.geometry("1280x980") + self.resizable(False, False) + + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + + self.top_bar(self) + + self.bottom_frame = ctk.CTkFrame(self) + self.bottom_frame.grid(row=1, column=0, sticky="nsew") + self.bottom_frame.grid_rowconfigure(0, weight=1) + self.bottom_frame.grid_columnconfigure(0, weight=0) + self.bottom_frame.grid_columnconfigure(1, weight=1) + + self.file_list_column(self.bottom_frame) + self.content_column(self.bottom_frame) + self.load_directory() + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def top_bar(self, master): + top_frame = ctk.CTkFrame(master) + top_frame.grid(row=0, column=0, sticky="nsew") + + components.button(top_frame, 0, 0, "Open", self.open_directory, + tooltip="open a new directory") + components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + tooltip="open a dialog to automatically generate masks") + components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + tooltip="open a dialog to automatically generate captions") + + if platform.system() == "Windows": + components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + tooltip="open the current image in Explorer") + + components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + text="include subdirectories") + + top_frame.grid_columnconfigure(5, weight=1) + + components.button(top_frame, 0, 6, "Help", self.print_help, + tooltip=self.help_text) + + def file_list_column(self, master): + if self.file_list is not None: + self.image_labels = [] + self.file_list.destroy() + + self.file_list = ctk.CTkScrollableFrame(master, width=300) + self.file_list.grid(row=0, column=0, sticky="nsew") + + for i, filename in enumerate(self.image_rel_paths): + def __create_switch_image(index): + def __switch_image(event): + self.switch_image(index) + + return __switch_image + + label = ctk.CTkLabel(self.file_list, text=filename) + label.bind("", __create_switch_image(i)) + + self.image_labels.append(label) + label.grid(row=i, column=0, padx=5, sticky="nsw") + + def content_column(self, master): + image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) + + right_frame = ctk.CTkFrame(master, fg_color="transparent") + right_frame.grid(row=0, column=1, sticky="nsew") + + right_frame.grid_columnconfigure(4, weight=1) + right_frame.grid_rowconfigure(1, weight=1) + + components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + tooltip="draw a mask using a brush") + components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + tooltip="draw a mask using a fill tool") + + # checkbox to enable mask editing + self.enable_mask_editing_var = ctk.BooleanVar() + self.enable_mask_editing_var.set(False) + enable_mask_editing_checkbox = ctk.CTkCheckBox( + right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) + enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") + + # mask alpha textbox + self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") + self.mask_editing_alpha.insert(0, "1.0") + self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) + self.bind_key_events(self.mask_editing_alpha) + + mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) + mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + + # image + self.image = ctk.CTkImage( + light_image=image, + size=(self.image_size, self.image_size) + ) + self.image_label = ctk.CTkLabel( + master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + ) + self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") + + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) + + # prompt + self.prompt_var = ctk.StringVar() + self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) + self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") + self.bind_key_events(self.prompt_component) + self.prompt_component.focus_set() + + def bind_key_events(self, component): + component.bind("", self.next_image) + component.bind("", self.previous_image) + component.bind("", self.save) + component.bind("", self.toggle_mask) + component.bind("", self.draw_mask_editing_mode) + component.bind("", self.fill_mask_editing_mode) + + def load_directory(self, include_subdirectories: bool = False): + self.scan_directory(include_subdirectories) + self.file_list_column(self.bottom_frame) + + if len(self.image_rel_paths) > 0: + self.switch_image(0) + else: + self.switch_image(-1) + + self.prompt_component.focus_set() + + def scan_directory(self, include_subdirectories: bool = False): + def __is_supported_image_extension(filename): + name, ext = os.path.splitext(filename) + return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") + + self.image_rel_paths = [] + + if not self.dir or not os.path.isdir(self.dir): + return + + if include_subdirectories: + for root, _, files in os.walk(self.dir): + for filename in files: + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(root, filename), self.dir) + ) + else: + for _, filename in enumerate(os.listdir(self.dir)): + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(self.dir, filename), self.dir) + ) + + def load_image(self): + image_name = "resources/icons/icon.png" + + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.join(self.dir, image_name) + + try: + return load_image(image_name, convert_mode="RGB") + except Exception: + print(f'Could not open image {image_name}') + + def load_mask(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + return load_image(mask_name, convert_mode='RGB') + except Exception: + return None + else: + return None + + def load_prompt(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + try: + with open(prompt_name, "r", encoding='utf-8') as f: + return f.readlines()[0].strip() + except Exception: + return "" + else: + return "" + + def previous_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.switch_image(self.current_image_index - 1) + + def next_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.switch_image(self.current_image_index + 1) + + def switch_image(self, index): + if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): + self.image_labels[self.current_image_index].configure( + text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + + self.current_image_index = index + if index >= 0: + self.image_labels[index].configure(text_color="#FF0000") + + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() + + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) + + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + + self.refresh_image() + self.prompt_var.set(prompt) + else: + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) + + def refresh_image(self): + if self.pil_mask: + resized_pil_mask = self.pil_mask.resize( + (self.pil_image.width, self.pil_image.height), + Image.Resampling.NEAREST + ) + + if self.display_only_mask: + self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + else: + np_image = np.array(self.pil_image).astype(np.float32) / 255.0 + np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 + + # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha + norm_min = 0.3 + np_mask_min = np_mask.min() + if np_mask_min == 0: + # optimize for common case + np_mask = np_mask * (1.0 - norm_min) + norm_min + elif np_mask_min < 1: + # note: min of 1 means we get divide by 0 + np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min + + np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) + masked_image = Image.fromarray(np_masked_image, mode='RGB') + + self.image.configure(light_image=masked_image, size=masked_image.size) + else: + self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + + def draw_mask_radius(self, delta, raw_event): + # Wheel up = Increase radius. Wheel down = Decrease radius. + multiplier = 1.0 + (delta * 0.05) + self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + + def edit_mask(self, event): + if not self.enable_mask_editing_var.get(): + return + + if event.widget != self.image_label.children["!label"]: + return + + if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): + return + + display_scaling = ScalingTracker.get_window_scaling(self) + + event_x = event.x / display_scaling + event_y = event.y / display_scaling + + start_x = int(event_x / self.pil_image.width * self.image_width) + start_y = int(event_y / self.pil_image.height * self.image_height) + end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) + end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) + + self.mask_draw_x = event_x + self.mask_draw_y = event_y + + is_right = False + is_left = False + if event.state & 0x0100 or event.num == 1: # left mouse button + is_left = True + elif event.state & 0x0400 or event.num == 3: # right mouse button + is_right = True + + if self.mask_editing_mode == 'draw': + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + if self.mask_editing_mode == 'fill': + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) + + draw = ImageDraw.Draw(self.pil_mask) + draw.line((start_x, start_y, end_x, end_y), fill=color, + width=radius + radius + 1) + draw.ellipse((start_x - radius, start_y - radius, + start_x + radius, start_y + radius), fill=color, outline=None) + draw.ellipse((end_x - radius, end_y - radius, end_x + radius, + end_y + radius), fill=color, outline=None) + + self.refresh_image() + + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + np_mask = np.array(self.pil_mask).astype(np.uint8) + cv2.floodFill(np_mask, None, (start_x, start_y), color) + self.pil_mask = Image.fromarray(np_mask, 'RGB') + + self.refresh_image() + + def save(self, event): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(self.prompt_var.get()) + except Exception: + return + + if self.pil_mask: + self.pil_mask.save(mask_name) + + def draw_mask_editing_mode(self, *args): + self.mask_editing_mode = 'draw' + + if args: + # disable default event + return "break" + return None + + def fill_mask_editing_mode(self, *args): + self.mask_editing_mode = 'fill' + + def toggle_mask(self, *args): + self.display_only_mask = not self.display_only_mask + self.refresh_image() + + def open_directory(self): + new_dir = filedialog.askdirectory() + + if new_dir: + self.dir = new_dir + self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + + def open_mask_window(self): + dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_caption_window(self): + dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() + + def load_masking_model(self, model): + model_type = type(self.masking_model).__name__ if self.masking_model else None + + if model == "ClipSeg" and model_type != "ClipSegModel": + self._release_models() + print("loading ClipSeg model, this may take a while") + self.masking_model = ClipSegModel(default_device, torch.float32) + elif model == "Rembg" and model_type != "RembgModel": + self._release_models() + print("loading Rembg model, this may take a while") + self.masking_model = RembgModel(default_device, torch.float32) + elif model == "Rembg-Human" and model_type != "RembgHumanModel": + self._release_models() + print("loading Rembg-Human model, this may take a while") + self.masking_model = RembgHumanModel(default_device, torch.float32) + elif model == "Hex Color" and model_type != "MaskByColor": + self._release_models() + self.masking_model = MaskByColor(default_device, torch.float32) + + def load_captioning_model(self, model): + model_type = type(self.captioning_model).__name__ if self.captioning_model else None + + if model == "Blip" and model_type != "BlipModel": + self._release_models() + print("loading Blip model, this may take a while") + self.captioning_model = BlipModel(default_device, torch.float16) + elif model == "Blip2" and model_type != "Blip2Model": + self._release_models() + print("loading Blip2 model, this may take a while") + self.captioning_model = Blip2Model(default_device, torch.float16) + elif model == "WD14 VIT v2" and model_type != "WDModel": + self._release_models() + print("loading WD14_VIT_v2 model, this may take a while") + self.captioning_model = WDModel(default_device, torch.float16) + + def print_help(self): + print(self.help_text) + + def _release_models(self): + """Release all models from VRAM""" + freed = False + if self.captioning_model is not None: + self.captioning_model = None + freed = True + if self.masking_model is not None: + self.masking_model = None + freed = True + if freed: + torch_gc() + + def _on_close(self): + self._release_models() + self.destroy() + + def destroy(self): + self._release_models() + super().destroy() diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py new file mode 100644 index 000000000..f58879d5f --- /dev/null +++ b/modules/ui/ConceptWindowController.py @@ -0,0 +1,934 @@ +import fractions +import math +import os +import pathlib +import platform +import random +import threading +import time +import traceback + +from modules.util import concept_stats, path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.BalancingStrategy import BalancingStrategy +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +from mgds.LoadingPipeline import LoadingPipeline +from mgds.OutputPipelineModule import OutputPipelineModule +from mgds.PipelineModule import PipelineModule +from mgds.pipelineModules.CapitalizeTags import CapitalizeTags +from mgds.pipelineModules.DropTags import DropTags +from mgds.pipelineModules.RandomBrightness import RandomBrightness +from mgds.pipelineModules.RandomCircularMaskShrink import ( + RandomCircularMaskShrink, +) +from mgds.pipelineModules.RandomContrast import RandomContrast +from mgds.pipelineModules.RandomFlip import RandomFlip +from mgds.pipelineModules.RandomHue import RandomHue +from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop +from mgds.pipelineModules.RandomRotate import RandomRotate +from mgds.pipelineModules.RandomSaturation import RandomSaturation +from mgds.pipelineModules.ShuffleTags import ShuffleTags +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( + RandomAccessPipelineModule, +) + +import torch +from torchvision.transforms import functional + +import customtkinter as ctk +import huggingface_hub +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg +from PIL import Image + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data + + def length(self) -> int: + return 1 + + def get_inputs(self) -> list[str]: + return [] + + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data + + +class ConceptWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + concept: ConceptConfig, + ui_state: UIState, + image_ui_state: UIState, + text_ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.train_config = train_config + + self.concept = concept + self.ui_state = ui_state + self.image_ui_state = image_ui_state + self.text_ui_state = text_ui_state + self.image_preview_file_index = 0 + self.preview_augmentations = ctk.BooleanVar(self, True) + self.bucket_fig = None + + self.title("Concept") + self.geometry("800x700") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.__general_tab(tabview.add("general"), concept) + self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) + self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) + self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) + + #automatic concept scan + self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) + self.scan_thread.start() + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __general_tab(self, master, concept: ConceptConfig): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=1) + + # name + components.label(frame, 0, 0, "Name", + tooltip="Name of the concept") + components.entry(frame, 0, 1, self.ui_state, "name") + + # enabled + components.label(frame, 1, 0, "Enabled", + tooltip="Enable or disable this concept") + components.switch(frame, 1, 1, self.ui_state, "enabled") + + # concept type + components.label(frame, 2, 0, "Concept Type", + tooltip="STANDARD: Standard finetuning with the sample as training target\n" + "VALIDATION: Use concept for validation instead of training\n" + "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " + "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " + "Only implemented for LoRA.", + wide_tooltip=True) + components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + + # path + components.label(frame, 3, 0, "Path", + tooltip="Path where the training data is located") + components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") + components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, + tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.") + + # prompt source + components.label(frame, 4, 0, "Prompt Source", + tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts") + prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + + def set_prompt_path_entry_enabled(option: str): + if option == 'concept': + for child in prompt_path_entry.children.values(): + child.configure(state="normal") + else: + for child in prompt_path_entry.children.values(): + child.configure(state="disabled") + + components.options_kv(frame, 4, 1, [ + ("From text file per sample", 'sample'), + ("From single text file", 'concept'), + ("From image file name", 'filename'), + ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(concept.text.prompt_source) + + # include subdirectories + components.label(frame, 5, 0, "Include Subdirectories", + tooltip="Includes images from subdirectories into the dataset") + components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + + # image variations + components.label(frame, 6, 0, "Image Variations", + tooltip="The number of different image versions to cache if latent caching is enabled.") + components.entry(frame, 6, 1, self.ui_state, "image_variations") + + # text variations + components.label(frame, 7, 0, "Text Variations", + tooltip="The number of different text versions to cache if latent caching is enabled.") + components.entry(frame, 7, 1, self.ui_state, "text_variations") + + # balancing + components.label(frame, 8, 0, "Balancing", + tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") + components.entry(frame, 8, 1, self.ui_state, "balancing") + components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + + # loss weight + components.label(frame, 9, 0, "Loss Weight", + tooltip="The loss multiplyer for this concept.") + components.entry(frame, 9, 1, self.ui_state, "loss_weight") + + frame.pack(fill="both", expand=1) + return frame + + def __image_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # header + components.label(frame, 0, 1, "Random", + tooltip="Enable this augmentation with random values") + components.label(frame, 0, 2, "Fixed", + tooltip="Enable this augmentation with fixed values") + + # crop jitter + components.label(frame, 1, 0, "Crop Jitter", + tooltip="Enables random cropping of samples") + components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + + # random flip + components.label(frame, 2, 0, "Random Flip", + tooltip="Randomly flip the sample during training") + components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") + components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + + # random rotation + components.label(frame, 3, 0, "Random Rotation", + tooltip="Randomly rotates the sample during training") + components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") + components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") + components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + + # random brightness + components.label(frame, 4, 0, "Random Brightness", + tooltip="Randomly adjusts the brightness of the sample during training") + components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") + components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") + components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + + # random contrast + components.label(frame, 5, 0, "Random Contrast", + tooltip="Randomly adjusts the contrast of the sample during training") + components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") + components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") + components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + + # random saturation + components.label(frame, 6, 0, "Random Saturation", + tooltip="Randomly adjusts the saturation of the sample during training") + components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") + components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") + components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + + # random hue + components.label(frame, 7, 0, "Random Hue", + tooltip="Randomly adjusts the hue of the sample during training") + components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") + components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") + components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + + # random circular mask shrink + components.label(frame, 8, 0, "Circular Mask Generation", + tooltip="Automatically create circular masks for masked training") + components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + + # random rotate and crop + components.label(frame, 9, 0, "Random Rotate and Crop", + tooltip="Randomly rotate the training samples and crop to the masked region") + components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + + # circular mask generation + components.label(frame, 10, 0, "Resolution Override", + tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") + components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + + # image + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image = ctk.CTkImage( + light_image=image_preview, + size=image_preview.size, + ) + image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label.grid(row=0, column=4, rowspan=6) + + # refresh preview + update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") + update_button_frame.grid_columnconfigure(1, weight=1) + + prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) + components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) + next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) + + prev_preview_button.configure(width=40) + next_preview_button.configure(width=40) + + #caption and filename preview + self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) + self.filename_preview.grid(row=2, column=0, columnspan=3) + self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) + + frame.pack(fill="both", expand=1) + return frame + + def __text_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # tag shuffling + components.label(frame, 0, 0, "Tag Shuffling", + tooltip="Enables tag shuffling") + components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + + # keep tag count + components.label(frame, 1, 0, "Tag Delimiter", + tooltip="The delimiter between tags") + components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + + # keep tag count + components.label(frame, 2, 0, "Keep Tag Count", + tooltip="The number of tags at the start of the caption that are not shuffled or dropped") + components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + + # tag dropout + components.label(frame, 3, 0, "Tag Dropout", + tooltip="Enables random dropout for tags in the captions.") + components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") + components.label(frame, 4, 0, "Dropout Mode", + tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") + components.options_kv(frame, 4, 1, [ + ("Full", 'FULL'), + ("Random", 'RANDOM'), + ("Random Weighted", 'RANDOM WEIGHTED'), + ], self.text_ui_state, "tag_dropout_mode", None) + components.label(frame, 4, 2, "Probability", + tooltip="Probability to drop tags, from 0 to 1.") + components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + + components.label(frame, 5, 0, "Special Dropout Tags", + tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") + components.options_kv(frame, 5, 1, [ + ("None", 'NONE'), + ("Blacklist", 'BLACKLIST'), + ("Whitelist", 'WHITELIST'), + ], self.text_ui_state, "tag_dropout_special_tags_mode", None) + components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") + components.label(frame, 6, 0, "Special Tags Regex", + tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") + components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + + #capitalization randomization + components.label(frame, 7, 0, "Randomize Capitalization", + tooltip="Enables randomization of capitalization for tags in the caption.") + components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") + components.label(frame, 7, 2, "Force Lowercase", + tooltip="If enabled, converts the caption to lowercase before any further processing.") + components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + + components.label(frame, 8, 0, "Captialization Mode", + tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") + components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") + components.label(frame, 8, 2, "Probability", + tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") + components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") + + frame.pack(fill="both", expand=1) + return frame + + def __concept_stats_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=150) + frame.grid_columnconfigure(1, weight=0, minsize=150) + frame.grid_columnconfigure(2, weight=0, minsize=150) + frame.grid_columnconfigure(3, weight=0, minsize=150) + + self.cancel_scan_flag = threading.Event() + + #file size + self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB") + self.file_size_label.configure(font=ctk.CTkFont(underline=True)) + self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + + #subdirectory count + self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") + self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) + self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + + #basic img/vid stats - count of each type in the concept + #the \n at the start of the label gives it better vertical spacing with other rows + self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") + self.image_count_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) + self.video_count_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'") + self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") + self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + + #advanced img/vid stats - how many img/vid files have a mask or caption of the same name + self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask") + self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") + self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + #currently no masks for videos? + + self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption") + self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption") + self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") + self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + + #resolution info + self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)") + self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)") + self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)") + self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + + #video length info + self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames") + self.length_max_label.configure(font=ctk.CTkFont(underline=True)) + self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames") + self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames") + self.length_min_label.configure(font=ctk.CTkFont(underline=True)) + self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + + #video fps info + self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps") + self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept") + self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps") + self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + + #caption info + self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + + #aspect bucket info + self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ + Images which don't match a bucket exactly are cropped to the nearest one.") + self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") + self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + + #aspect bucketing plot, mostly copied from timestep preview graph + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + plt.set_loglevel('WARNING') #suppress errors about data type in bar chart + + assert self.bucket_fig is None + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + #refresh stats - must be after all labels are defined or will give error + self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), + tooltip="Reload basic statistics for the concept directory") + self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster + self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") + self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + + frame.pack(fill="both", expand=1) + return frame + + def __prev_image_preview(self): + self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) + self.__update_image_preview() + + def __next_image_preview(self): + self.image_preview_file_index += 1 + self.__update_image_preview() + + def __update_image_preview(self): + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image.configure(light_image=image_preview, size=image_preview.size) + self.filename_preview.configure(text=filename_preview) + self.caption_preview.configure(state="normal") + self.caption_preview.delete(index1="1.0", index2="end") + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + + @staticmethod + def get_concept_path(path: str) -> str | None: + if os.path.isdir(path): + return path + try: + #don't download, only check if available locally: + return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) + except Exception: + return None + + def __download_dataset(self): + try: + huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") + except Exception: + traceback.print_exc() + + def __download_dataset_threaded(self): + download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + download_thread.start() + + def _read_text_file_for_preview(self, file_path: str) -> str: + empty_msg = "[Empty prompt]" + try: + with open(file_path, "r") as f: + if self.preview_augmentations.get(): + lines = [line.strip() for line in f if line.strip()] + return random.choice(lines) if lines else empty_msg + content = f.read().strip() + return content if content else empty_msg + except FileNotFoundError: + return "File not found, please check the path" + except IsADirectoryError: + return "[Provided path is a directory, please correct the caption path]" + except PermissionError: + if platform.system() == "Windows": + return "[Permission denied, please check the file permissions or Windows Defender settings]" + else: + return "[Permission denied, please check the file permissions]" + except UnicodeDecodeError: + return "[Invalid file encoding. This should not happen, please report this issue]" + + def __get_preview_image(self): + preview_image_path = "resources/icons/icon.png" + file_index = -1 + glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" + + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if path.is_file() and path_util.is_supported_image_extension(extension) \ + and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): + preview_image_path = path_util.canonical_join(concept_path, path) + file_index += 1 + if file_index == self.image_preview_file_index: + break + + image = load_image(preview_image_path, 'RGB') + image_tensor = functional.to_tensor(image) + + splitext = os.path.splitext(preview_image_path) + preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") + if not os.path.isfile(preview_mask_path): + preview_mask_path = None + + if preview_mask_path: + mask = Image.open(preview_mask_path).convert("L") + mask_tensor = functional.to_tensor(mask) + else: + mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + + source = self.concept.text.prompt_source + preview_p = pathlib.Path(preview_image_path) + if source == "filename": + prompt_output = preview_p.stem or "[Empty prompt]" + else: + file_map = { + "sample": preview_p.with_suffix(".txt"), + "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, + } + file_path = file_map.get(source) + prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + + modules = [] + if self.preview_augmentations.get(): + input_module = InputPipelineModule({ + 'true': True, + 'image': image_tensor, + 'mask': mask_tensor, + 'enable_random_flip': self.concept.image.enable_random_flip, + 'enable_fixed_flip': self.concept.image.enable_fixed_flip, + 'enable_random_rotate': self.concept.image.enable_random_rotate, + 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, + 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, + 'enable_random_brightness': self.concept.image.enable_random_brightness, + 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, + 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, + 'enable_random_contrast': self.concept.image.enable_random_contrast, + 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, + 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, + 'enable_random_saturation': self.concept.image.enable_random_saturation, + 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, + 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, + 'enable_random_hue': self.concept.image.enable_random_hue, + 'enable_fixed_hue': self.concept.image.enable_fixed_hue, + 'random_hue_max_strength': self.concept.image.random_hue_max_strength, + 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, + 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, + + 'prompt' : prompt_output, + 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, + 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, + 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, + 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, + 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, + 'tag_delimiter' : self.concept.text.tag_delimiter, + 'keep_tags_count' : self.concept.text.keep_tags_count, + 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, + 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, + 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, + 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, + 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, + 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, + }) + + circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') + random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') + random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') + random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') + random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') + random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') + random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') + random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') + drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', + special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', + keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') + caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', + capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') + shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') + output_module = OutputPipelineModule(['image', 'mask', 'prompt']) + + modules = [ + input_module, + circular_mask_shrink, + random_mask_rotate_crop, + random_flip, + random_rotate, + random_brightness, + random_contrast, + random_saturation, + random_hue, + drop_tags, + caps_randomize, + shuffle_tags, + output_module, + ] + + pipeline = LoadingPipeline( + device=torch.device('cpu'), + modules=modules, + batch_size=1, + seed=random.randint(0, 2**30), + state=None, + initial_epoch=0, + initial_index=0, + ) + + data = pipeline.__next__() + image_tensor = data['image'] + mask_tensor = data['mask'] + prompt_output = data['prompt'] + + filename_output = os.path.basename(preview_image_path) + + mask_tensor = torch.clamp(mask_tensor, 0.3, 1) + image_tensor = image_tensor * mask_tensor + + image = functional.to_pil_image(image_tensor) + + image.thumbnail((300, 300)) + + return image, filename_output, prompt_output + + def __update_concept_stats(self): + #file size + self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") + self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") + + #directory count + self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + + #image count + self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) + self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) + self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + + #video count + self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) + #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) + self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + + #mask count + self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) + self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + + #caption count + if self.concept.concept_stats["subcaption_count"] > 0: + self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + else: + self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) + self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + + #resolution info + max_pixels = self.concept.concept_stats["max_pixels"] + avg_pixels = self.concept.concept_stats["avg_pixels"] + min_pixels = self.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.pixel_max_preview.configure(text="-") + self.pixel_avg_preview.configure(text="-") + self.pixel_min_preview.configure(text="-") + else: + #formatted as (#pixels/1000000) MP, width x height, \n filename + self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + + #video length and fps info + max_length = self.concept.concept_stats["max_length"] + avg_length = self.concept.concept_stats["avg_length"] + min_length = self.concept.concept_stats["min_length"] + max_fps = self.concept.concept_stats["max_fps"] + avg_fps = self.concept.concept_stats["avg_fps"] + min_fps = self.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.length_max_preview.configure(text="-") + self.length_avg_preview.configure(text="-") + self.length_min_preview.configure(text="-") + self.fps_max_preview.configure(text="-") + self.fps_avg_preview.configure(text="-") + self.fps_min_preview.configure(text="-") + else: + #formatted as (#frames) frames \n filename + self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') + self.length_avg_preview.configure(text=f'{int(avg_length)} frames') + self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + #formatted as (#fps) fps \n filename + self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') + self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + + #caption info + max_caption_length = self.concept.concept_stats["max_caption_length"] + avg_caption_length = self.concept.concept_stats["avg_caption_length"] + min_caption_length = self.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.caption_max_preview.configure(text="-") + self.caption_avg_preview.configure(text="-") + self.caption_min_preview.configure(text="-") + else: + #formatted as (#chars) chars, (#words) words, \n filename + self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + + #aspect bucketing + aspect_buckets = self.concept.concept_stats["aspect_buckets"] + if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero + min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values + if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be + min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket + else: + min_val2 = min_val #if no second smallest bucket exists set to min_val + min_aspect_buckets = {key: val for key,val in aspect_buckets.items() if val in (min_val, min_val2)} + min_bucket_str = "" + for key, val in min_aspect_buckets.items(): + min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' + min_bucket_str.strip() + self.small_bucket_preview.configure(text=min_bucket_str) + + self.bucket_ax.cla() + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = self.bucket_ax.bar(aspect_ratios, counts) + self.bucket_ax.bar_label(b, color=self.text_color) + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + self.canvas.draw() + + def decimal_to_aspect_ratio(self, value : float): + #find closest fraction to decimal aspect value and convert to a:b format + aspect_fraction = fractions.Fraction(value).limit_denominator(16) + aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' + return aspect_string + + def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + start_time = time.perf_counter() + last_update = time.perf_counter() + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__disable_scan_buttons) + concept_path = self.get_concept_path(self.concept.path) + + if not concept_path: + print(f"Unable to get statistics for concept path: {self.concept.path}") + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + return + subfolders = [concept_path] + + stats_dict = concept_stats.init_concept_stats(advanced_checks) + for path in subfolders: + if self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: + break + stats_dict = concept_stats.folder_scan(path, stats_dict, advanced_checks, self.concept, start_time, wait_time, self.cancel_scan_flag) + if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set(): #add all subfolders of current directory to for loop + subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) + self.concept.concept_stats = stats_dict + #update GUI approx every half second + if time.perf_counter() > (last_update + 0.5): + last_update = time.perf_counter() + self.concept_stats_tab.after(0, self.__update_concept_stats) + + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + self.concept_stats_tab.after(0, self.__update_concept_stats) + + def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): + self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) + self.scan_thread.start() + + def __disable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="disabled") + self.refresh_advanced_stats_button.configure(state="disabled") + + def __enable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="normal") + self.refresh_advanced_stats_button.configure(state="normal") + + def __cancel_concept_stats(self): + self.cancel_scan_flag.set() + + def __auto_update_concept_stats(self): + try: + self.__update_concept_stats() #load stats from config if available, else raises KeyError + if self.concept.concept_stats["file_size"] == 0: #force rescan if empty + raise KeyError + except KeyError: + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + if self.concept.concept_stats["processing_time"] < 0.1: + self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + + def destroy(self): + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + + super().destroy() + + def __ok(self): + self.destroy() diff --git a/modules/ui/ConvertModelUIController.py b/modules/ui/ConvertModelUIController.py new file mode 100644 index 000000000..6cb1b507a --- /dev/null +++ b/modules/ui/ConvertModelUIController.py @@ -0,0 +1,170 @@ +import traceback +from uuid import uuid4 + +from modules.util import create +from modules.util.args.ConvertModelArgs import ConvertModelArgs +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModelNames import EmbeddingName, ModelNames +from modules.util.torch_util import torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ConvertModelUI(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.parent = parent + self.convert_model_args = ConvertModelArgs.default_values() + self.ui_state = UIState(self, self.convert_model_args) + self.button = None + + + self.title("Convert models") + self.geometry("550x350") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.main_frame(self.frame) + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # model type + components.label(master, 0, 0, "Model Type", + tooltip="Type of the model") + components.options_kv(master, 0, 1, [ #TODO simplify + ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), + ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), + ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), + ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), + ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), + ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2", ModelType.FLUX_2), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here + ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here + ("ZImage", ModelType.Z_IMAGE), + ], self.ui_state, "model_type") + + # training method + components.label(master, 1, 0, "Model Type", + tooltip="The type of model to convert") + components.options_kv(master, 1, 1, [ + ("Base Model", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ], self.ui_state, "training_method") + + # input name + components.label(master, 2, 0, "Input name", + tooltip="Filename, directory or hugging face repository of the base model") + components.path_entry( + master, 2, 1, self.ui_state, "input_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # output data type + components.label(master, 3, 0, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "output_dtype") + + # output format + components.label(master, 4, 0, "Output Format", + tooltip="Format to use when saving the output model") + components.options_kv(master, 4, 1, [ + ("Safetensors", ModelFormat.SAFETENSORS), + ("Diffusers", ModelFormat.DIFFUSERS), + ], self.ui_state, "output_model_format") + + # output model destination + components.label(master, 5, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + master, 5, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + self.button = components.button(master, 6, 1, "Convert", self.convert_model) + + def convert_model(self): + try: + self.button.configure(state="disabled") + model_loader = create.create_model_loader( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + model_saver = create.create_model_saver( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + + print("Loading model " + self.convert_model_args.input_name) + if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=self.convert_model_args.input_name, + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=None, + lora=self.convert_model_args.input_name, + embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + else: + raise Exception("could not load model: " + self.convert_model_args.input_name) + + print("Saving model " + self.convert_model_args.output_model_destination) + model_saver.save( + model=model, + model_type=self.convert_model_args.model_type, + output_model_format=self.convert_model_args.output_model_format, + output_model_destination=self.convert_model_args.output_model_destination, + dtype=self.convert_model_args.output_dtype.torch_dtype(), + ) + print("Model converted") + except Exception: + traceback.print_exc() + + torch_gc() + self.button.configure(state="normal") diff --git a/modules/ui/CtkAdditionalEmbeddingsTabView.py b/modules/ui/CtkAdditionalEmbeddingsTabView.py new file mode 100644 index 000000000..6a5e3fbe7 --- /dev/null +++ b/modules/ui/CtkAdditionalEmbeddingsTabView.py @@ -0,0 +1,136 @@ + +from modules.ui.ConfigList import ConfigList +from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class AdditionalEmbeddingsTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + attr_name="additional_embeddings", + enable_key="train", + from_external_file=False, + add_button_text="add embedding", + is_full_width=True, + show_toggle_button=True + ) + + def refresh_ui(self): + if self.element_list is not None: + self.element_list.destroy() + self.element_list = None + self.widgets_initialized = False + self._create_element_list() + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return TrainEmbeddingConfig.default_values() + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + +class EmbeddingWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, corner_radius=10, bg_color="transparent" + ) + + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(0, weight=1) + + top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + top_frame.grid(row=0, column=0, sticky="nsew") + top_frame.grid_columnconfigure(3, weight=1) + top_frame.grid_columnconfigure(5, weight=1) + + bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + bottom_frame.grid(row=1, column=0, sticky="nsew") + bottom_frame.grid_columnconfigure(7, weight=1) + + # close button + close_button = ctk.CTkButton( + master=top_frame, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.grid(row=0, column=0) + + # clone button + clone_button = ctk.CTkButton( + master=top_frame, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i, self.__randomize_uuid), + ) + clone_button.grid(row=0, column=1, padx=5) + + # embedding model names + components.label(top_frame, 0, 2, "base embedding:", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + components.path_entry( + top_frame, 0, 3, self.ui_state, "model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # placeholder + components.label(top_frame, 0, 4, "placeholder:", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(top_frame, 0, 5, self.ui_state, "placeholder") + + # token count + components.label(top_frame, 0, 6, "token count:", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") + token_count_entry.configure(width=40) + + # trainable + components.label(bottom_frame, 0, 0, "train:") + trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) + trainable_switch.configure(width=40) + + # output embedding + components.label(bottom_frame, 0, 2, "output embedding:", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") + output_embedding_switch.configure(width=40) + + # stop training after + components.label(bottom_frame, 0, 4, "stop training after:", + tooltip="When to stop training the embedding") + components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") + + # initial embedding text + components.label(bottom_frame, 0, 6, "initial embedding text:", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") + + def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): + embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid + return embedding_config + + def configure_element(self): + pass + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkCaptionUIView.py b/modules/ui/CtkCaptionUIView.py new file mode 100644 index 000000000..e6cc0551e --- /dev/null +++ b/modules/ui/CtkCaptionUIView.py @@ -0,0 +1,572 @@ +import os +import platform +import subprocess +import traceback +from tkinter import filedialog + +from modules.module.Blip2Model import Blip2Model +from modules.module.BlipModel import BlipModel +from modules.module.ClipSegModel import ClipSegModel +from modules.module.MaskByColor import MaskByColor +from modules.module.RembgHumanModel import RembgHumanModel +from modules.module.RembgModel import RembgModel +from modules.module.WDModel import WDModel +from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow +from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.util import path_util +from modules.util.image_util import load_image +from modules.util.torch_util import default_device, torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +import cv2 +import numpy as np +from customtkinter import ScalingTracker, ThemeManager +from PIL import Image, ImageDraw + + +class CaptionUI(ctk.CTkToplevel): + def __init__( + self, + parent, + initial_dir: str | None, + initial_include_subdirectories: bool, + *args, + **kwargs, + ) -> None: + super().__init__(parent, *args, **kwargs) + self.protocol("WM_DELETE_WINDOW", self._on_close) + + self.dir = initial_dir + self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} + self.config_ui_state = UIState(self, self.config_ui_data) + self.image_size = 850 + self.help_text = """ + Keyboard shortcuts when focusing on the prompt input field: + Up arrow: previous image + Down arrow: next image + Return: save + Ctrl+M: only show the mask + Ctrl+D: draw mask editing mode + Ctrl+F: fill mask editing mode + + When editing masks: + Left click: add mask + Right click: remove mask + Mouse wheel: increase or decrease brush size""" + self.masking_model = None + self.captioning_model = None + self.image_rel_paths = [] + self.current_image_index = -1 + self.file_list = None + self.image_labels = [] + self.pil_image = None + self.image_width = 0 + self.image_height = 0 + self.pil_mask = None + self.mask_draw_x = 0 + self.mask_draw_y = 0 + self.mask_draw_radius = 0.01 + self.display_only_mask = False + self.image = None + self.image_label = None + self.mask_editing_mode = 'draw' + self.enable_mask_editing_var = ctk.BooleanVar() + self.mask_editing_alpha = None + self.prompt_var = None + self.prompt_component = None + + + self.title("OneTrainer") + self.geometry("1280x980") + self.resizable(False, False) + + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + + self.top_bar(self) + + self.bottom_frame = ctk.CTkFrame(self) + self.bottom_frame.grid(row=1, column=0, sticky="nsew") + self.bottom_frame.grid_rowconfigure(0, weight=1) + self.bottom_frame.grid_columnconfigure(0, weight=0) + self.bottom_frame.grid_columnconfigure(1, weight=1) + + self.file_list_column(self.bottom_frame) + self.content_column(self.bottom_frame) + self.load_directory() + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def top_bar(self, master): + top_frame = ctk.CTkFrame(master) + top_frame.grid(row=0, column=0, sticky="nsew") + + components.button(top_frame, 0, 0, "Open", self.open_directory, + tooltip="open a new directory") + components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + tooltip="open a dialog to automatically generate masks") + components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + tooltip="open a dialog to automatically generate captions") + + if platform.system() == "Windows": + components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + tooltip="open the current image in Explorer") + + components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + text="include subdirectories") + + top_frame.grid_columnconfigure(5, weight=1) + + components.button(top_frame, 0, 6, "Help", self.print_help, + tooltip=self.help_text) + + def file_list_column(self, master): + if self.file_list is not None: + self.image_labels = [] + self.file_list.destroy() + + self.file_list = ctk.CTkScrollableFrame(master, width=300) + self.file_list.grid(row=0, column=0, sticky="nsew") + + for i, filename in enumerate(self.image_rel_paths): + def __create_switch_image(index): + def __switch_image(event): + self.switch_image(index) + + return __switch_image + + label = ctk.CTkLabel(self.file_list, text=filename) + label.bind("", __create_switch_image(i)) + + self.image_labels.append(label) + label.grid(row=i, column=0, padx=5, sticky="nsw") + + def content_column(self, master): + image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) + + right_frame = ctk.CTkFrame(master, fg_color="transparent") + right_frame.grid(row=0, column=1, sticky="nsew") + + right_frame.grid_columnconfigure(4, weight=1) + right_frame.grid_rowconfigure(1, weight=1) + + components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + tooltip="draw a mask using a brush") + components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + tooltip="draw a mask using a fill tool") + + # checkbox to enable mask editing + self.enable_mask_editing_var = ctk.BooleanVar() + self.enable_mask_editing_var.set(False) + enable_mask_editing_checkbox = ctk.CTkCheckBox( + right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) + enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") + + # mask alpha textbox + self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") + self.mask_editing_alpha.insert(0, "1.0") + self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) + self.bind_key_events(self.mask_editing_alpha) + + mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) + mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + + # image + self.image = ctk.CTkImage( + light_image=image, + size=(self.image_size, self.image_size) + ) + self.image_label = ctk.CTkLabel( + master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + ) + self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") + + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) + + # prompt + self.prompt_var = ctk.StringVar() + self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) + self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") + self.bind_key_events(self.prompt_component) + self.prompt_component.focus_set() + + def bind_key_events(self, component): + component.bind("", self.next_image) + component.bind("", self.previous_image) + component.bind("", self.save) + component.bind("", self.toggle_mask) + component.bind("", self.draw_mask_editing_mode) + component.bind("", self.fill_mask_editing_mode) + + def load_directory(self, include_subdirectories: bool = False): + self.scan_directory(include_subdirectories) + self.file_list_column(self.bottom_frame) + + if len(self.image_rel_paths) > 0: + self.switch_image(0) + else: + self.switch_image(-1) + + self.prompt_component.focus_set() + + def scan_directory(self, include_subdirectories: bool = False): + def __is_supported_image_extension(filename): + name, ext = os.path.splitext(filename) + return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") + + self.image_rel_paths = [] + + if not self.dir or not os.path.isdir(self.dir): + return + + if include_subdirectories: + for root, _, files in os.walk(self.dir): + for filename in files: + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(root, filename), self.dir) + ) + else: + for _, filename in enumerate(os.listdir(self.dir)): + if __is_supported_image_extension(filename): + self.image_rel_paths.append( + os.path.relpath(os.path.join(self.dir, filename), self.dir) + ) + + def load_image(self): + image_name = "resources/icons/icon.png" + + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.join(self.dir, image_name) + + try: + return load_image(image_name, convert_mode="RGB") + except Exception: + print(f'Could not open image {image_name}') + + def load_mask(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + return load_image(mask_name, convert_mode='RGB') + except Exception: + return None + else: + return None + + def load_prompt(self): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + try: + with open(prompt_name, "r", encoding='utf-8') as f: + return f.readlines()[0].strip() + except Exception: + return "" + else: + return "" + + def previous_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.switch_image(self.current_image_index - 1) + + def next_image(self, event): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.switch_image(self.current_image_index + 1) + + def switch_image(self, index): + if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): + self.image_labels[self.current_image_index].configure( + text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + + self.current_image_index = index + if index >= 0: + self.image_labels[index].configure(text_color="#FF0000") + + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() + + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) + + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + + self.refresh_image() + self.prompt_var.set(prompt) + else: + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) + + def refresh_image(self): + if self.pil_mask: + resized_pil_mask = self.pil_mask.resize( + (self.pil_image.width, self.pil_image.height), + Image.Resampling.NEAREST + ) + + if self.display_only_mask: + self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + else: + np_image = np.array(self.pil_image).astype(np.float32) / 255.0 + np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 + + # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha + norm_min = 0.3 + np_mask_min = np_mask.min() + if np_mask_min == 0: + # optimize for common case + np_mask = np_mask * (1.0 - norm_min) + norm_min + elif np_mask_min < 1: + # note: min of 1 means we get divide by 0 + np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min + + np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) + masked_image = Image.fromarray(np_masked_image, mode='RGB') + + self.image.configure(light_image=masked_image, size=masked_image.size) + else: + self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + + def draw_mask_radius(self, delta, raw_event): + # Wheel up = Increase radius. Wheel down = Decrease radius. + multiplier = 1.0 + (delta * 0.05) + self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + + def edit_mask(self, event): + if not self.enable_mask_editing_var.get(): + return + + if event.widget != self.image_label.children["!label"]: + return + + if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): + return + + display_scaling = ScalingTracker.get_window_scaling(self) + + event_x = event.x / display_scaling + event_y = event.y / display_scaling + + start_x = int(event_x / self.pil_image.width * self.image_width) + start_y = int(event_y / self.pil_image.height * self.image_height) + end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) + end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) + + self.mask_draw_x = event_x + self.mask_draw_y = event_y + + is_right = False + is_left = False + if event.state & 0x0100 or event.num == 1: # left mouse button + is_left = True + elif event.state & 0x0400 or event.num == 3: # right mouse button + is_right = True + + if self.mask_editing_mode == 'draw': + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + if self.mask_editing_mode == 'fill': + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) + + draw = ImageDraw.Draw(self.pil_mask) + draw.line((start_x, start_y, end_x, end_y), fill=color, + width=radius + radius + 1) + draw.ellipse((start_x - radius, start_y - radius, + start_x + radius, start_y + radius), fill=color, outline=None) + draw.ellipse((end_x - radius, end_y - radius, end_x + radius, + end_y + radius), fill=color, outline=None) + + self.refresh_image() + + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + color = None + + adding_to_mask = True + if is_left: + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range + color = (rgb_value, rgb_value, rgb_value) + + elif is_right: + color = (0, 0, 0) + adding_to_mask = False + + if color is not None: + if self.pil_mask is None: + if adding_to_mask: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) + else: + self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) + + np_mask = np.array(self.pil_mask).astype(np.uint8) + cv2.floodFill(np_mask, None, (start_x, start_y), color) + self.pil_mask = Image.fromarray(np_mask, 'RGB') + + self.refresh_image() + + def save(self, event): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] + + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) + + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) + + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(self.prompt_var.get()) + except Exception: + return + + if self.pil_mask: + self.pil_mask.save(mask_name) + + def draw_mask_editing_mode(self, *args): + self.mask_editing_mode = 'draw' + + if args: + # disable default event + return "break" + return None + + def fill_mask_editing_mode(self, *args): + self.mask_editing_mode = 'fill' + + def toggle_mask(self, *args): + self.display_only_mask = not self.display_only_mask + self.refresh_image() + + def open_directory(self): + new_dir = filedialog.askdirectory() + + if new_dir: + self.dir = new_dir + self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + + def open_mask_window(self): + dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_caption_window(self): + dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) + self.wait_window(dialog) + self.switch_image(self.current_image_index) + + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() + + def load_masking_model(self, model): + model_type = type(self.masking_model).__name__ if self.masking_model else None + + if model == "ClipSeg" and model_type != "ClipSegModel": + self._release_models() + print("loading ClipSeg model, this may take a while") + self.masking_model = ClipSegModel(default_device, torch.float32) + elif model == "Rembg" and model_type != "RembgModel": + self._release_models() + print("loading Rembg model, this may take a while") + self.masking_model = RembgModel(default_device, torch.float32) + elif model == "Rembg-Human" and model_type != "RembgHumanModel": + self._release_models() + print("loading Rembg-Human model, this may take a while") + self.masking_model = RembgHumanModel(default_device, torch.float32) + elif model == "Hex Color" and model_type != "MaskByColor": + self._release_models() + self.masking_model = MaskByColor(default_device, torch.float32) + + def load_captioning_model(self, model): + model_type = type(self.captioning_model).__name__ if self.captioning_model else None + + if model == "Blip" and model_type != "BlipModel": + self._release_models() + print("loading Blip model, this may take a while") + self.captioning_model = BlipModel(default_device, torch.float16) + elif model == "Blip2" and model_type != "Blip2Model": + self._release_models() + print("loading Blip2 model, this may take a while") + self.captioning_model = Blip2Model(default_device, torch.float16) + elif model == "WD14 VIT v2" and model_type != "WDModel": + self._release_models() + print("loading WD14_VIT_v2 model, this may take a while") + self.captioning_model = WDModel(default_device, torch.float16) + + def print_help(self): + print(self.help_text) + + def _release_models(self): + """Release all models from VRAM""" + freed = False + if self.captioning_model is not None: + self.captioning_model = None + freed = True + if self.masking_model is not None: + self.masking_model = None + freed = True + if freed: + torch_gc() + + def _on_close(self): + self._release_models() + self.destroy() + + def destroy(self): + self._release_models() + super().destroy() diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py new file mode 100644 index 000000000..99057e428 --- /dev/null +++ b/modules/ui/CtkCloudTabView.py @@ -0,0 +1,221 @@ + +import webbrowser + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.CloudAction import CloudAction +from modules.util.enum.CloudFileSync import CloudFileSync +from modules.util.enum.CloudType import CloudType +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class CloudTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.parent = parent + self.reattach = False + + self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, weight=0) + self.frame.grid_columnconfigure(3, weight=1) + self.frame.grid_columnconfigure(4, weight=0) + self.frame.grid_columnconfigure(5, weight=1) + + components.label(self.frame, 0, 0, "Enabled", + tooltip="Enable cloud training") + components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") + + components.label(self.frame, 1, 0, "Type", + tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") + components.options_kv(self.frame, 1, 1, [ + ("RUNPOD", CloudType.RUNPOD), + ("LINUX", CloudType.LINUX), + ], self.ui_state, "cloud.type") + + components.label(self.frame, 2, 0, "File sync method", + tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") + components.options_kv(self.frame, 2, 1, [ + ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), + ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), + ], self.ui_state, "cloud.file_sync") + + components.label(self.frame, 3, 0, "API key", + tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") + components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") + + components.label(self.frame, 4, 0, "Hostname", + tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") + components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") + + components.label(self.frame, 5, 0, "Port", + tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") + components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") + + components.label(self.frame, 6, 0, "User", + tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') + components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") + + components.label(self.frame, 7, 0, "SSH keyfile path", + tooltip="Absolute path to the private key file used for SSH connections. Leave empty to rely on your system SSH configuration.") + components.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="file") + + components.label(self.frame, 8, 0, "SSH password", + tooltip="SSH password for password-based authentication. If you try to use native SCP requires sshpass to be installed. Leave empty to use key-based authentication.") + components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") + + components.label(self.frame, 9, 0, "Cloud id", + tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") + components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") + + components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", + tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") + components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + + + + components.label(self.frame, 1, 2, "Remote Directory", + tooltip="The directory on the cloud where files will be uploaded and downloaded.") + components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") + components.label(self.frame, 2, 2, "OneTrainer Directory", + tooltip="The directory for OneTrainer on the cloud.") + components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") + components.label(self.frame, 3, 2, "Huggingface cache Directory", + tooltip="Huggingface models are downloaded to this remote directory.") + components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") + components.label(self.frame, 4, 2, "Install OneTrainer", + tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") + components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") + components.label(self.frame, 5, 2, "Install command", + tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") + components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") + components.label(self.frame, 6, 2, "Update OneTrainer", + tooltip="Update OneTrainer if it already exists on the cloud.") + components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + + components.label(self.frame, 8, 2, "Detach remote trainer", + tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") + components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") + components.label(self.frame, 9, 2, "Reattach id", + tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") + reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") + reattach_frame.grid_columnconfigure(0, weight=1) + reattach_frame.grid_columnconfigure(1, weight=1) + components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) + components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) + + components.label(self.frame, 11, 2, "Download samples", + tooltip="Download samples from the remote workspace directory to your local machine.") + components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") + components.label(self.frame, 12, 2, "Download output model", + tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") + components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") + components.label(self.frame, 13, 2, "Download saved checkpoints", + tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") + components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") + components.label(self.frame, 14, 2, "Download backups", + tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") + components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") + components.label(self.frame, 15, 2, "Download tensorboard logs", + tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") + components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") + components.label(self.frame, 16, 2, "Delete remote workspace", + tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") + components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + + components.label(self.frame, 1, 4, "Create cloud via API", + tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") + create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") + create_frame.grid_columnconfigure(0, weight=0) + create_frame.grid_columnconfigure(1, weight=1) + components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") + components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) + + components.label(self.frame, 2, 4, "Cloud name", + tooltip="The name of the new cloud instance.") + components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") + components.label(self.frame, 3, 4, "Type", + tooltip="Select the RunPod cloud type. See RunPod's website for details.") + components.options_kv(self.frame, 3, 5, [ + ("", ""), + ("Community", "COMMUNITY"), + ("Secure", "SECURE"), + ], self.ui_state, "cloud.sub_type") + + + components.label(self.frame, 4, 4, "GPU", + tooltip="Select the GPU type. Enter an API key before pressing the button.") + + _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) + self.gpu_types_menu=gpu_components['component'] + + components.label(self.frame, 5, 4, "Volume size", + tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") + components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") + + components.label(self.frame, 6, 4, "Min download", + tooltip="Set the minimum download speed of the cloud in Mbps.") + components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") + + components.label(self.frame, 8, 4, "Action on finish", + tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") + components.options_kv(self.frame, 8, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_finish") + + components.label(self.frame, 9, 4, "Action on error", + tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") + components.options_kv(self.frame, 9, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_error") + + components.label(self.frame, 10, 4, "Action on detached finish", + tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") + components.options_kv(self.frame, 10, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_detached_finish") + + components.label(self.frame, 11, 4, "Action on detached error", + tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") + components.options_kv(self.frame, 11, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], self.ui_state, "cloud.on_detached_error") + + self.frame.pack(fill="both", expand=1) + + def __set_gpu_types(self): + self.gpu_types_menu.configure(values=[]) + if self.train_config.cloud.type == CloudType.RUNPOD: + import runpod + runpod.api_key=self.train_config.secrets.cloud.api_key + gpus=runpod.get_gpus() + self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) + + def __reattach(self): + self.reattach=True + try: + self.parent.start_training() + finally: + self.reattach=False + + def __create_cloud(self): + if self.train_config.cloud.type == CloudType.RUNPOD: + webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py new file mode 100644 index 000000000..0b6505694 --- /dev/null +++ b/modules/ui/CtkConceptTabView.py @@ -0,0 +1,286 @@ +import os +import pathlib +from tkinter import BooleanVar, StringVar + +from modules.ui.ConceptWindow import ConceptWindow +from modules.ui.ConfigList import ConfigList +from modules.util import path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import DebounceTimer + +import customtkinter as ctk +from PIL import Image + + +class ConceptTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + self.search_var = StringVar() + self.filter_var = StringVar(value="ALL") + self.show_disabled_var = BooleanVar(value=True) + + super().__init__( + master, + train_config, + ui_state, + from_external_file=True, + attr_name="concept_file_name", + config_dir="training_concepts", + default_config_name="concepts.json", + add_button_text="Add Concept", + add_button_tooltip="Adds a new concept to the current config.", + is_full_width=False, + show_toggle_button=True + ) + self._toolbar = None + self._toolbar_is_wrapped = False + self._add_search_bar() + # wrap toolbar if too narrow + self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return ConceptConfig.default_values() + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) + + def _add_search_bar(self): + toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") + toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") + toolbar.grid_columnconfigure(2, weight=1) + self._toolbar = toolbar + + # Search + ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) + self.search_var = StringVar() + self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, + placeholder_text="Filter...", width=200) + self.search_entry.grid(row=0, column=1) + self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) + self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) + + # Spacer + ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) + + # Type filter + ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) + self.filter_var = StringVar(value="ALL") + ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], + variable=self.filter_var, command=lambda x: self._update_filters(), + width=150).grid(row=0, column=4) + + # Show disabled checkbox + self.show_disabled_var = BooleanVar(value=True) + self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, + command=self._update_filters, width=100) + self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) + self._refresh_show_disabled_text() + + # Clear button + ctk.CTkButton(toolbar, text="Clear", width=50, + command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) + + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() + + def _reset_filters(self): + self.search_var.set("") + self.filter_var.set("ALL") + self.show_disabled_var.set(True) + self._update_filters() + + def _element_matches_filters(self, element): + # Check enabled status + if not self.filters.get("show_disabled", True): + if hasattr(element, 'enabled') and not element.enabled: + return False + + # Search filter + search = self.filters.get("search", "").lower() + if search: + if not hasattr(element, '_search_cache'): + cache = [] + try: + if getattr(element, 'name', None): + cache.append(element.name.lower()) + p = getattr(element, 'path', None) + if p: + try: + cache.append(os.path.basename(p).lower()) + cache.append(p.lower()) + except (TypeError, AttributeError): + pass + except (AttributeError, TypeError): + pass + element._search_cache = cache + if not any(search in text for text in getattr(element, '_search_cache', [])): + return False + + # Type filter + type_filter = self.filters.get("type", "ALL") + if type_filter != "ALL": + if hasattr(element, 'type') and element.type: + try: + return ConceptType(element.type).value == type_filter + except (ValueError, AttributeError): + return False + return False + + return True + + def _maybe_reposition_toolbar(self, width): + if not self._toolbar: + return + threshold = 1070 + want_wrapped = width < threshold + if want_wrapped == self._toolbar_is_wrapped: + return + self._toolbar_is_wrapped = want_wrapped + if want_wrapped: + self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) + else: + self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + + def _refresh_show_disabled_text(self): + try: + disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) + except (AttributeError, TypeError): + disabled_count = 0 + text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" + try: + if getattr(self, 'show_disabled_checkbox', None): + self.show_disabled_checkbox.configure(text=text) + except (AttributeError, RuntimeError): + pass + + +class ConceptWidget(ctk.CTkFrame): + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, width=150, height=170, corner_radius=10, bg_color="transparent" + ) + + self.concept = concept + self.ui_state = UIState(self, concept) + self.image_ui_state = UIState(self, concept.image) + self.text_ui_state = UIState(self, concept.text) + self.i = i + + self.grid_rowconfigure(1, weight=1) + + # image + self.image = ctk.CTkImage( + light_image=self.__get_preview_image(), + size=(150, 150) + ) + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150) + image_label.grid(row=0, column=0) + + # name + self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) + + # close button + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.place(x=0, y=0) + + # clone button + clone_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i, self.__randomize_seed), + ) + clone_button.place(x=25, y=0) + + # enabled switch + enabled_switch = ctk.CTkSwitch( + master=self, + width=40, + variable=self.ui_state.get_var("enabled"), + text="", + command=save_command, + ) + enabled_switch.place(x=110, y=0) + + image_label.bind( + "", + lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) + ) + + def __randomize_seed(self, concept: ConceptConfig): + concept.seed = ConceptConfig.default_values().seed + return concept + + def __get_display_name(self): + if self.concept.name: + return self.concept.name + elif self.concept.path: + return os.path.basename(self.concept.path) + else: + return "" + + def configure_element(self): + self.name_label.configure(text=self.__get_display_name()) + self.image.configure(light_image=self.__get_preview_image()) + try: + if hasattr(self.concept, '_search_cache'): + delattr(self.concept, '_search_cache') + except AttributeError: + pass + + def __get_preview_image(self): + preview_path = "resources/icons/icon.png" + glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" + + concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if (path.is_file() + and path_util.is_supported_image_extension(extension) + and not path.name.endswith("-masklabel.png") + and not path.name.endswith("-condlabel.png")): + preview_path = path_util.canonical_join(concept_path, path) + break + try: + image = load_image(preview_path, convert_mode="RGBA") + except (OSError): + image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) + size = min(image.width, image.height) + image = image.crop(( + (image.width - size) // 2, + (image.height - size) // 2, + (image.width - size) // 2 + size, + (image.height - size) // 2 + size, + )) + return image.resize((150, 150), Image.Resampling.BILINEAR) + + def place_in_list(self): + index = getattr(self, 'visible_index', self.i) + x = index % 6 + y = index // 6 + self.grid(row=y, column=x, pady=5, padx=5) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py new file mode 100644 index 000000000..f58879d5f --- /dev/null +++ b/modules/ui/CtkConceptWindowView.py @@ -0,0 +1,934 @@ +import fractions +import math +import os +import pathlib +import platform +import random +import threading +import time +import traceback + +from modules.util import concept_stats, path_util +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.BalancingStrategy import BalancingStrategy +from modules.util.enum.ConceptType import ConceptType +from modules.util.image_util import load_image +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +from mgds.LoadingPipeline import LoadingPipeline +from mgds.OutputPipelineModule import OutputPipelineModule +from mgds.PipelineModule import PipelineModule +from mgds.pipelineModules.CapitalizeTags import CapitalizeTags +from mgds.pipelineModules.DropTags import DropTags +from mgds.pipelineModules.RandomBrightness import RandomBrightness +from mgds.pipelineModules.RandomCircularMaskShrink import ( + RandomCircularMaskShrink, +) +from mgds.pipelineModules.RandomContrast import RandomContrast +from mgds.pipelineModules.RandomFlip import RandomFlip +from mgds.pipelineModules.RandomHue import RandomHue +from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop +from mgds.pipelineModules.RandomRotate import RandomRotate +from mgds.pipelineModules.RandomSaturation import RandomSaturation +from mgds.pipelineModules.ShuffleTags import ShuffleTags +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( + RandomAccessPipelineModule, +) + +import torch +from torchvision.transforms import functional + +import customtkinter as ctk +import huggingface_hub +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg +from PIL import Image + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data + + def length(self) -> int: + return 1 + + def get_inputs(self) -> list[str]: + return [] + + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data + + +class ConceptWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + concept: ConceptConfig, + ui_state: UIState, + image_ui_state: UIState, + text_ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.train_config = train_config + + self.concept = concept + self.ui_state = ui_state + self.image_ui_state = image_ui_state + self.text_ui_state = text_ui_state + self.image_preview_file_index = 0 + self.preview_augmentations = ctk.BooleanVar(self, True) + self.bucket_fig = None + + self.title("Concept") + self.geometry("800x700") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.__general_tab(tabview.add("general"), concept) + self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) + self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) + self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) + + #automatic concept scan + self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) + self.scan_thread.start() + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __general_tab(self, master, concept: ConceptConfig): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=1) + + # name + components.label(frame, 0, 0, "Name", + tooltip="Name of the concept") + components.entry(frame, 0, 1, self.ui_state, "name") + + # enabled + components.label(frame, 1, 0, "Enabled", + tooltip="Enable or disable this concept") + components.switch(frame, 1, 1, self.ui_state, "enabled") + + # concept type + components.label(frame, 2, 0, "Concept Type", + tooltip="STANDARD: Standard finetuning with the sample as training target\n" + "VALIDATION: Use concept for validation instead of training\n" + "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " + "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " + "Only implemented for LoRA.", + wide_tooltip=True) + components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + + # path + components.label(frame, 3, 0, "Path", + tooltip="Path where the training data is located") + components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") + components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, + tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.") + + # prompt source + components.label(frame, 4, 0, "Prompt Source", + tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts") + prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + + def set_prompt_path_entry_enabled(option: str): + if option == 'concept': + for child in prompt_path_entry.children.values(): + child.configure(state="normal") + else: + for child in prompt_path_entry.children.values(): + child.configure(state="disabled") + + components.options_kv(frame, 4, 1, [ + ("From text file per sample", 'sample'), + ("From single text file", 'concept'), + ("From image file name", 'filename'), + ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(concept.text.prompt_source) + + # include subdirectories + components.label(frame, 5, 0, "Include Subdirectories", + tooltip="Includes images from subdirectories into the dataset") + components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + + # image variations + components.label(frame, 6, 0, "Image Variations", + tooltip="The number of different image versions to cache if latent caching is enabled.") + components.entry(frame, 6, 1, self.ui_state, "image_variations") + + # text variations + components.label(frame, 7, 0, "Text Variations", + tooltip="The number of different text versions to cache if latent caching is enabled.") + components.entry(frame, 7, 1, self.ui_state, "text_variations") + + # balancing + components.label(frame, 8, 0, "Balancing", + tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") + components.entry(frame, 8, 1, self.ui_state, "balancing") + components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + + # loss weight + components.label(frame, 9, 0, "Loss Weight", + tooltip="The loss multiplyer for this concept.") + components.entry(frame, 9, 1, self.ui_state, "loss_weight") + + frame.pack(fill="both", expand=1) + return frame + + def __image_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # header + components.label(frame, 0, 1, "Random", + tooltip="Enable this augmentation with random values") + components.label(frame, 0, 2, "Fixed", + tooltip="Enable this augmentation with fixed values") + + # crop jitter + components.label(frame, 1, 0, "Crop Jitter", + tooltip="Enables random cropping of samples") + components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + + # random flip + components.label(frame, 2, 0, "Random Flip", + tooltip="Randomly flip the sample during training") + components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") + components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + + # random rotation + components.label(frame, 3, 0, "Random Rotation", + tooltip="Randomly rotates the sample during training") + components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") + components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") + components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + + # random brightness + components.label(frame, 4, 0, "Random Brightness", + tooltip="Randomly adjusts the brightness of the sample during training") + components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") + components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") + components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + + # random contrast + components.label(frame, 5, 0, "Random Contrast", + tooltip="Randomly adjusts the contrast of the sample during training") + components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") + components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") + components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + + # random saturation + components.label(frame, 6, 0, "Random Saturation", + tooltip="Randomly adjusts the saturation of the sample during training") + components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") + components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") + components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + + # random hue + components.label(frame, 7, 0, "Random Hue", + tooltip="Randomly adjusts the hue of the sample during training") + components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") + components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") + components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + + # random circular mask shrink + components.label(frame, 8, 0, "Circular Mask Generation", + tooltip="Automatically create circular masks for masked training") + components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + + # random rotate and crop + components.label(frame, 9, 0, "Random Rotate and Crop", + tooltip="Randomly rotate the training samples and crop to the masked region") + components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + + # circular mask generation + components.label(frame, 10, 0, "Resolution Override", + tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") + components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + + # image + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image = ctk.CTkImage( + light_image=image_preview, + size=image_preview.size, + ) + image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label.grid(row=0, column=4, rowspan=6) + + # refresh preview + update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") + update_button_frame.grid_columnconfigure(1, weight=1) + + prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) + components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) + next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) + + prev_preview_button.configure(width=40) + next_preview_button.configure(width=40) + + #caption and filename preview + self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) + self.filename_preview.grid(row=2, column=0, columnspan=3) + self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) + + frame.pack(fill="both", expand=1) + return frame + + def __text_augmentation_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # tag shuffling + components.label(frame, 0, 0, "Tag Shuffling", + tooltip="Enables tag shuffling") + components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + + # keep tag count + components.label(frame, 1, 0, "Tag Delimiter", + tooltip="The delimiter between tags") + components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + + # keep tag count + components.label(frame, 2, 0, "Keep Tag Count", + tooltip="The number of tags at the start of the caption that are not shuffled or dropped") + components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + + # tag dropout + components.label(frame, 3, 0, "Tag Dropout", + tooltip="Enables random dropout for tags in the captions.") + components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") + components.label(frame, 4, 0, "Dropout Mode", + tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") + components.options_kv(frame, 4, 1, [ + ("Full", 'FULL'), + ("Random", 'RANDOM'), + ("Random Weighted", 'RANDOM WEIGHTED'), + ], self.text_ui_state, "tag_dropout_mode", None) + components.label(frame, 4, 2, "Probability", + tooltip="Probability to drop tags, from 0 to 1.") + components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + + components.label(frame, 5, 0, "Special Dropout Tags", + tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") + components.options_kv(frame, 5, 1, [ + ("None", 'NONE'), + ("Blacklist", 'BLACKLIST'), + ("Whitelist", 'WHITELIST'), + ], self.text_ui_state, "tag_dropout_special_tags_mode", None) + components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") + components.label(frame, 6, 0, "Special Tags Regex", + tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") + components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + + #capitalization randomization + components.label(frame, 7, 0, "Randomize Capitalization", + tooltip="Enables randomization of capitalization for tags in the caption.") + components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") + components.label(frame, 7, 2, "Force Lowercase", + tooltip="If enabled, converts the caption to lowercase before any further processing.") + components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + + components.label(frame, 8, 0, "Captialization Mode", + tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") + components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") + components.label(frame, 8, 2, "Probability", + tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") + components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") + + frame.pack(fill="both", expand=1) + return frame + + def __concept_stats_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=150) + frame.grid_columnconfigure(1, weight=0, minsize=150) + frame.grid_columnconfigure(2, weight=0, minsize=150) + frame.grid_columnconfigure(3, weight=0, minsize=150) + + self.cancel_scan_flag = threading.Event() + + #file size + self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB") + self.file_size_label.configure(font=ctk.CTkFont(underline=True)) + self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + + #subdirectory count + self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") + self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) + self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + + #basic img/vid stats - count of each type in the concept + #the \n at the start of the label gives it better vertical spacing with other rows + self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") + self.image_count_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) + self.video_count_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'") + self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") + self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + + #advanced img/vid stats - how many img/vid files have a mask or caption of the same name + self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask") + self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") + self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + #currently no masks for videos? + + self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption") + self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption") + self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) + self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") + self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) + self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + + #resolution info + self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)") + self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)") + self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)") + self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) + self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + + #video length info + self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames") + self.length_max_label.configure(font=ctk.CTkFont(underline=True)) + self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames") + self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames") + self.length_min_label.configure(font=ctk.CTkFont(underline=True)) + self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + + #video fps info + self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps") + self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept") + self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps") + self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) + self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + + #caption info + self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") + self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) + self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + + #aspect bucket info + self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ + Images which don't match a bucket exactly are cropped to the nearest one.") + self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") + self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) + self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + + #aspect bucketing plot, mostly copied from timestep preview graph + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + plt.set_loglevel('WARNING') #suppress errors about data type in bar chart + + assert self.bucket_fig is None + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + #refresh stats - must be after all labels are defined or will give error + self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), + tooltip="Reload basic statistics for the concept directory") + self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster + self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") + self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + + frame.pack(fill="both", expand=1) + return frame + + def __prev_image_preview(self): + self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) + self.__update_image_preview() + + def __next_image_preview(self): + self.image_preview_file_index += 1 + self.__update_image_preview() + + def __update_image_preview(self): + image_preview, filename_preview, caption_preview = self.__get_preview_image() + self.image.configure(light_image=image_preview, size=image_preview.size) + self.filename_preview.configure(text=filename_preview) + self.caption_preview.configure(state="normal") + self.caption_preview.delete(index1="1.0", index2="end") + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + + @staticmethod + def get_concept_path(path: str) -> str | None: + if os.path.isdir(path): + return path + try: + #don't download, only check if available locally: + return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) + except Exception: + return None + + def __download_dataset(self): + try: + huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") + except Exception: + traceback.print_exc() + + def __download_dataset_threaded(self): + download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + download_thread.start() + + def _read_text_file_for_preview(self, file_path: str) -> str: + empty_msg = "[Empty prompt]" + try: + with open(file_path, "r") as f: + if self.preview_augmentations.get(): + lines = [line.strip() for line in f if line.strip()] + return random.choice(lines) if lines else empty_msg + content = f.read().strip() + return content if content else empty_msg + except FileNotFoundError: + return "File not found, please check the path" + except IsADirectoryError: + return "[Provided path is a directory, please correct the caption path]" + except PermissionError: + if platform.system() == "Windows": + return "[Permission denied, please check the file permissions or Windows Defender settings]" + else: + return "[Permission denied, please check the file permissions]" + except UnicodeDecodeError: + return "[Invalid file encoding. This should not happen, please report this issue]" + + def __get_preview_image(self): + preview_image_path = "resources/icons/icon.png" + file_index = -1 + glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" + + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + for path in pathlib.Path(concept_path).glob(glob_pattern): + if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): + continue + extension = os.path.splitext(path)[1] + if path.is_file() and path_util.is_supported_image_extension(extension) \ + and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): + preview_image_path = path_util.canonical_join(concept_path, path) + file_index += 1 + if file_index == self.image_preview_file_index: + break + + image = load_image(preview_image_path, 'RGB') + image_tensor = functional.to_tensor(image) + + splitext = os.path.splitext(preview_image_path) + preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") + if not os.path.isfile(preview_mask_path): + preview_mask_path = None + + if preview_mask_path: + mask = Image.open(preview_mask_path).convert("L") + mask_tensor = functional.to_tensor(mask) + else: + mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + + source = self.concept.text.prompt_source + preview_p = pathlib.Path(preview_image_path) + if source == "filename": + prompt_output = preview_p.stem or "[Empty prompt]" + else: + file_map = { + "sample": preview_p.with_suffix(".txt"), + "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, + } + file_path = file_map.get(source) + prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + + modules = [] + if self.preview_augmentations.get(): + input_module = InputPipelineModule({ + 'true': True, + 'image': image_tensor, + 'mask': mask_tensor, + 'enable_random_flip': self.concept.image.enable_random_flip, + 'enable_fixed_flip': self.concept.image.enable_fixed_flip, + 'enable_random_rotate': self.concept.image.enable_random_rotate, + 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, + 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, + 'enable_random_brightness': self.concept.image.enable_random_brightness, + 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, + 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, + 'enable_random_contrast': self.concept.image.enable_random_contrast, + 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, + 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, + 'enable_random_saturation': self.concept.image.enable_random_saturation, + 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, + 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, + 'enable_random_hue': self.concept.image.enable_random_hue, + 'enable_fixed_hue': self.concept.image.enable_fixed_hue, + 'random_hue_max_strength': self.concept.image.random_hue_max_strength, + 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, + 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, + + 'prompt' : prompt_output, + 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, + 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, + 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, + 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, + 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, + 'tag_delimiter' : self.concept.text.tag_delimiter, + 'keep_tags_count' : self.concept.text.keep_tags_count, + 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, + 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, + 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, + 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, + 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, + 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, + }) + + circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') + random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') + random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') + random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') + random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') + random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') + random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') + random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') + drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', + special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', + keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') + caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', + capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') + shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') + output_module = OutputPipelineModule(['image', 'mask', 'prompt']) + + modules = [ + input_module, + circular_mask_shrink, + random_mask_rotate_crop, + random_flip, + random_rotate, + random_brightness, + random_contrast, + random_saturation, + random_hue, + drop_tags, + caps_randomize, + shuffle_tags, + output_module, + ] + + pipeline = LoadingPipeline( + device=torch.device('cpu'), + modules=modules, + batch_size=1, + seed=random.randint(0, 2**30), + state=None, + initial_epoch=0, + initial_index=0, + ) + + data = pipeline.__next__() + image_tensor = data['image'] + mask_tensor = data['mask'] + prompt_output = data['prompt'] + + filename_output = os.path.basename(preview_image_path) + + mask_tensor = torch.clamp(mask_tensor, 0.3, 1) + image_tensor = image_tensor * mask_tensor + + image = functional.to_pil_image(image_tensor) + + image.thumbnail((300, 300)) + + return image, filename_output, prompt_output + + def __update_concept_stats(self): + #file size + self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") + self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") + + #directory count + self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + + #image count + self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) + self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) + self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + + #video count + self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) + #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) + self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + + #mask count + self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) + self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + + #caption count + if self.concept.concept_stats["subcaption_count"] > 0: + self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + else: + self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) + self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + + #resolution info + max_pixels = self.concept.concept_stats["max_pixels"] + avg_pixels = self.concept.concept_stats["avg_pixels"] + min_pixels = self.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.pixel_max_preview.configure(text="-") + self.pixel_avg_preview.configure(text="-") + self.pixel_min_preview.configure(text="-") + else: + #formatted as (#pixels/1000000) MP, width x height, \n filename + self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + + #video length and fps info + max_length = self.concept.concept_stats["max_length"] + avg_length = self.concept.concept_stats["avg_length"] + min_length = self.concept.concept_stats["min_length"] + max_fps = self.concept.concept_stats["max_fps"] + avg_fps = self.concept.concept_stats["avg_fps"] + min_fps = self.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.length_max_preview.configure(text="-") + self.length_avg_preview.configure(text="-") + self.length_min_preview.configure(text="-") + self.fps_max_preview.configure(text="-") + self.fps_avg_preview.configure(text="-") + self.fps_min_preview.configure(text="-") + else: + #formatted as (#frames) frames \n filename + self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') + self.length_avg_preview.configure(text=f'{int(avg_length)} frames') + self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + #formatted as (#fps) fps \n filename + self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') + self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + + #caption info + max_caption_length = self.concept.concept_stats["max_caption_length"] + avg_caption_length = self.concept.concept_stats["avg_caption_length"] + min_caption_length = self.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.caption_max_preview.configure(text="-") + self.caption_avg_preview.configure(text="-") + self.caption_min_preview.configure(text="-") + else: + #formatted as (#chars) chars, (#words) words, \n filename + self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + + #aspect bucketing + aspect_buckets = self.concept.concept_stats["aspect_buckets"] + if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero + min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values + if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be + min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket + else: + min_val2 = min_val #if no second smallest bucket exists set to min_val + min_aspect_buckets = {key: val for key,val in aspect_buckets.items() if val in (min_val, min_val2)} + min_bucket_str = "" + for key, val in min_aspect_buckets.items(): + min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' + min_bucket_str.strip() + self.small_bucket_preview.configure(text=min_bucket_str) + + self.bucket_ax.cla() + aspects = [str(x) for x in list(aspect_buckets.keys())] + aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] + counts = list(aspect_buckets.values()) + b = self.bucket_ax.bar(aspect_ratios, counts) + self.bucket_ax.bar_label(b, color=self.text_color) + sec = self.bucket_ax.secondary_xaxis(location=-0.1) + sec.spines["bottom"].set_linewidth(0) + sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) + sec.tick_params('x', length=0) + self.canvas.draw() + + def decimal_to_aspect_ratio(self, value : float): + #find closest fraction to decimal aspect value and convert to a:b format + aspect_fraction = fractions.Fraction(value).limit_denominator(16) + aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' + return aspect_string + + def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + start_time = time.perf_counter() + last_update = time.perf_counter() + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__disable_scan_buttons) + concept_path = self.get_concept_path(self.concept.path) + + if not concept_path: + print(f"Unable to get statistics for concept path: {self.concept.path}") + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + return + subfolders = [concept_path] + + stats_dict = concept_stats.init_concept_stats(advanced_checks) + for path in subfolders: + if self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: + break + stats_dict = concept_stats.folder_scan(path, stats_dict, advanced_checks, self.concept, start_time, wait_time, self.cancel_scan_flag) + if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set(): #add all subfolders of current directory to for loop + subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) + self.concept.concept_stats = stats_dict + #update GUI approx every half second + if time.perf_counter() > (last_update + 0.5): + last_update = time.perf_counter() + self.concept_stats_tab.after(0, self.__update_concept_stats) + + self.cancel_scan_flag.clear() + self.concept_stats_tab.after(0, self.__enable_scan_buttons) + self.concept_stats_tab.after(0, self.__update_concept_stats) + + def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): + self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) + self.scan_thread.start() + + def __disable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="disabled") + self.refresh_advanced_stats_button.configure(state="disabled") + + def __enable_scan_buttons(self): + self.refresh_basic_stats_button.configure(state="normal") + self.refresh_advanced_stats_button.configure(state="normal") + + def __cancel_concept_stats(self): + self.cancel_scan_flag.set() + + def __auto_update_concept_stats(self): + try: + self.__update_concept_stats() #load stats from config if available, else raises KeyError + if self.concept.concept_stats["file_size"] == 0: #force rescan if empty + raise KeyError + except KeyError: + concept_path = self.get_concept_path(self.concept.path) + if concept_path: + self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + if self.concept.concept_stats["processing_time"] < 0.1: + self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + + def destroy(self): + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + + super().destroy() + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkConfigListView.py b/modules/ui/CtkConfigListView.py new file mode 100644 index 000000000..75d69252a --- /dev/null +++ b/modules/ui/CtkConfigListView.py @@ -0,0 +1,354 @@ +import contextlib +import copy +import json +import os +import tkinter as tk +from abc import ABCMeta, abstractmethod + +from modules.util import path_util +from modules.util.config.BaseConfig import BaseConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.path_util import write_json_atomic +from modules.util.ui import components, dialogs +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ConfigList(metaclass=ABCMeta): + + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + from_external_file: bool, + attr_name: str = "", + enable_key: str = "enabled", + config_dir: str = "", + default_config_name: str = "", + add_button_text: str = "", + add_button_tooltip: str = "", + is_full_width: bool = "", + show_toggle_button: bool = False, + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.from_external_file = from_external_file + self.attr_name = attr_name + self.enable_key = enable_key + + self.config_dir = config_dir + self.default_config_name = default_config_name + + self.is_full_width = is_full_width + + # From search-concepts + self.filters = {"search": "", "type": "ALL", "show_disabled": True} + self.widgets_initialized = False + + # From master + self.toggle_button = None + self.show_toggle_button = show_toggle_button + self.is_opening_window = False + self._is_current_item_enabled = False + + self.master.grid_rowconfigure(0, weight=0) + self.master.grid_rowconfigure(1, weight=1) + self.master.grid_columnconfigure(0, weight=1) + + if self.from_external_file: + self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.top_frame.grid(row=0, column=0, sticky="nsew") + + self.configs_dropdown = None + self.element_list = None + + self.configs = [] + self.__load_available_config_names() + + self.current_config = getattr(self.train_config, self.attr_name) + self.widgets = [] + self.__load_current_config(getattr(self.train_config, self.attr_name)) + + self.__create_configs_dropdown() + components.button(self.top_frame, 0, 1, "Add Config", self.__add_config, tooltip="Adds a new config, which are containers for concepts, which themselves contain your dataset", width=20, padx=5) + components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) + else: + self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.top_frame.grid(row=0, column=0, sticky="nsew") + components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, width=20, padx=5) + + self.current_config = getattr(self.train_config, self.attr_name) + + self.element_list = None + self._create_element_list() + + if show_toggle_button: + # tooltips break if you initialize with an empty string, default to a single space + self.toggle_button = components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) + self._update_toggle_button_text() + + + + @abstractmethod + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + pass + + @abstractmethod + def create_new_element(self) -> BaseConfig: + pass + + @abstractmethod + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + def _refresh_show_disabled_text(self): + return + + def _reset_filters(self): # pragma: no cover - default noop + search_var = getattr(self, 'search_var', None) + filter_var = getattr(self, 'filter_var', None) + show_disabled_var = getattr(self, 'show_disabled_var', None) + + if search_var: + search_var.set("") + if filter_var: + filter_var.set("ALL") + if show_disabled_var: + show_disabled_var.set(True) + if search_var and hasattr(self, '_update_filters'): + self._update_filters() + + def _update_item_enabled_state(self): + # Only count items that match current filters + self._is_current_item_enabled = any( + item.ui_state.get_var(self.enable_key).get() + for i, item in enumerate(self.widgets) + if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]) + ) + + def _update_toggle_button_text(self): + if not self.show_toggle_button: + return + self._update_item_enabled_state() + if self.toggle_button is not None: + self.toggle_button.configure(text="Disable" if self._is_current_item_enabled else "Enable") + + def _toggle(self): + self._toggle_items() + + def _toggle_items(self): + enable_state = not self._is_current_item_enabled + + # Only toggle items that match current filters + for i, widget in enumerate(self.widgets): + if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]): + widget.ui_state.get_var(self.enable_key).set(enable_state) + self.save_current_config() + + self._update_widget_visibility() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.destroy() + + self.configs_dropdown = components.options_kv( + self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config + ) + self._update_toggle_button_text() + + def _create_element_list(self, **filters): + if not self.from_external_file: + self.current_config = getattr(self.train_config, self.attr_name) + + self.filters.update(filters) + + if not self.widgets_initialized: + self._initialize_all_widgets() + self.widgets_initialized = True + + self._update_widget_visibility() + self._update_toggle_button_text() + + def _initialize_all_widgets(self): + self.widgets = [] + if self.element_list is not None: + self.element_list.destroy() + + self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.element_list.grid(row=1, column=0, sticky="nsew") + + if self.is_full_width: + self.element_list.grid_columnconfigure(0, weight=1) + + for i, element in enumerate(self.current_config): + widget = self.create_widget( + self.element_list, element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + + def _update_widget_visibility(self): + visible_index = 0 + + for i, widget in enumerate(self.widgets): + if i < len(self.current_config): + element = self.current_config[i] + + if self._element_matches_filters(element): + widget.visible_index = visible_index + widget.place_in_list() + visible_index += 1 + else: + widget.grid_remove() + + def __load_available_config_names(self): + if os.path.isdir(self.config_dir): + for path in os.listdir(self.config_dir): + path = path_util.canonical_join(self.config_dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + + if len(self.configs) == 0: + name = self.default_config_name.removesuffix(".json") + self.__create_config(name) + self.save_current_config() + + def __create_config(self, name: str): + name = path_util.safe_filename(name) + path = path_util.canonical_join(self.config_dir, f"{name}.json") + self.configs.append((name, path)) + self.__create_configs_dropdown() + + def __add_config(self): + dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) + + def __add_element(self): + new_element = self.create_new_element() + self.current_config.append(new_element) + # incremental insertion if widgets already initialized, else fall back to full rebuild + if self.widgets_initialized and self.element_list is not None: + i = len(self.current_config) - 1 + widget = self.create_widget( + self.element_list, new_element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + self._update_widget_visibility() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + def __clone_element(self, clone_i, modify_element_fun=None): + new_element = copy.deepcopy(self.current_config[clone_i]) + + if modify_element_fun is not None: + new_element = modify_element_fun(new_element) + self.current_config.append(new_element) + if self.widgets_initialized and self.element_list is not None: + i = len(self.current_config) - 1 + widget = self.create_widget( + self.element_list, new_element, i, + self.__open_element_window, + self.__remove_element, + self.__clone_element, + self.save_current_config + ) + self.widgets.append(widget) + self._update_widget_visibility() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + def __remove_element(self, remove_i): + self.current_config.pop(remove_i) + if self.widgets_initialized and 0 <= remove_i < len(self.widgets): + removed = self.widgets.pop(remove_i) + with contextlib.suppress(tk.TclError, AttributeError): + removed.destroy() + # Reindex remaining widgets + for idx, widget in enumerate(self.widgets): + widget.i = idx + self._update_widget_visibility() + else: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + + def __load_current_config(self, filename): + try: + with open(filename, "r") as f: + self.current_config = [] + + loaded_config_json = json.load(f) + for element_json in loaded_config_json: + element = self.create_new_element().from_dict(element_json) + self.current_config.append(element) + except (FileNotFoundError, json.JSONDecodeError) as e: + print(f"Failed to load config from {filename}: {e}") + self.current_config = [] + + # reset filters when switching configs + if hasattr(self, '_reset_filters') and self.widgets_initialized: + self._reset_filters() + + self.widgets_initialized = False + self._create_element_list() + self._update_toggle_button_text() + + def save_current_config(self): + if self.from_external_file: + try: + if not os.path.exists(self.config_dir): + os.makedirs(self.config_dir, exist_ok=True) + + write_json_atomic( + getattr(self.train_config, self.attr_name), + [element.to_dict() for element in self.current_config] + ) + except (OSError) as e: + print(f"Failed to save config: {e}") + + self._update_toggle_button_text() + + if self.widgets_initialized: + try: + self._update_widget_visibility() + except (tk.TclError, AttributeError) as e: + print.debug(f"Widget visibility update failed: {e}") + + # let subclass refresh any show-disabled UI + if hasattr(self, '_refresh_show_disabled_text'): + self._refresh_show_disabled_text() + + def _element_matches_filters(self, element): + return True # Show all by default + + def __open_element_window(self, i, ui_state): + if self.is_opening_window: + return + self.is_opening_window = True + try: + window = self.open_element_window(i, ui_state) + self.master.wait_window(window) + try: + if self.widgets is not None and 0 <= i < len(self.widgets): + self.widgets[i].configure_element() + except Exception: + self.widgets_initialized = False + self._create_element_list() + self.save_current_config() + finally: + self.is_opening_window = False diff --git a/modules/ui/CtkConvertModelUIView.py b/modules/ui/CtkConvertModelUIView.py new file mode 100644 index 000000000..6cb1b507a --- /dev/null +++ b/modules/ui/CtkConvertModelUIView.py @@ -0,0 +1,170 @@ +import traceback +from uuid import uuid4 + +from modules.util import create +from modules.util.args.ConvertModelArgs import ConvertModelArgs +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModelNames import EmbeddingName, ModelNames +from modules.util.torch_util import torch_gc +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ConvertModelUI(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.parent = parent + self.convert_model_args = ConvertModelArgs.default_values() + self.ui_state = UIState(self, self.convert_model_args) + self.button = None + + + self.title("Convert models") + self.geometry("550x350") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.main_frame(self.frame) + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # model type + components.label(master, 0, 0, "Model Type", + tooltip="Type of the model") + components.options_kv(master, 0, 1, [ #TODO simplify + ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), + ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), + ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), + ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), + ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), + ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2", ModelType.FLUX_2), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here + ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here + ("ZImage", ModelType.Z_IMAGE), + ], self.ui_state, "model_type") + + # training method + components.label(master, 1, 0, "Model Type", + tooltip="The type of model to convert") + components.options_kv(master, 1, 1, [ + ("Base Model", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ], self.ui_state, "training_method") + + # input name + components.label(master, 2, 0, "Input name", + tooltip="Filename, directory or hugging face repository of the base model") + components.path_entry( + master, 2, 1, self.ui_state, "input_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # output data type + components.label(master, 3, 0, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "output_dtype") + + # output format + components.label(master, 4, 0, "Output Format", + tooltip="Format to use when saving the output model") + components.options_kv(master, 4, 1, [ + ("Safetensors", ModelFormat.SAFETENSORS), + ("Diffusers", ModelFormat.DIFFUSERS), + ], self.ui_state, "output_model_format") + + # output model destination + components.label(master, 5, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + master, 5, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + self.button = components.button(master, 6, 1, "Convert", self.convert_model) + + def convert_model(self): + try: + self.button.configure(state="disabled") + model_loader = create.create_model_loader( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + model_saver = create.create_model_saver( + model_type=self.convert_model_args.model_type, + training_method=self.convert_model_args.training_method + ) + + print("Loading model " + self.convert_model_args.input_name) + if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=self.convert_model_args.input_name, + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: + model = model_loader.load( + model_type=self.convert_model_args.model_type, + model_names=ModelNames( + base_model=None, + lora=self.convert_model_args.input_name, + embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), + ), + weight_dtypes=self.convert_model_args.weight_dtypes(), + quantization=QuantizationConfig.default_values(), + ) + else: + raise Exception("could not load model: " + self.convert_model_args.input_name) + + print("Saving model " + self.convert_model_args.output_model_destination) + model_saver.save( + model=model, + model_type=self.convert_model_args.model_type, + output_model_format=self.convert_model_args.output_model_format, + output_model_destination=self.convert_model_args.output_model_destination, + dtype=self.convert_model_args.output_dtype.torch_dtype(), + ) + print("Model converted") + except Exception: + traceback.print_exc() + + torch_gc() + self.button.configure(state="normal") diff --git a/modules/ui/CtkGenerateCaptionsWindowView.py b/modules/ui/CtkGenerateCaptionsWindowView.py new file mode 100644 index 000000000..1690879f1 --- /dev/null +++ b/modules/ui/CtkGenerateCaptionsWindowView.py @@ -0,0 +1,133 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateCaptionsWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating captions for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + self.parent = parent + + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all captions", "Create if absent", "Add as new line"] + self.model_var = ctk.StringVar(self, "Blip") + self.models = ["Blip", "Blip2", "WD14 VIT v2"] + + self.title("Batch generate captions") + self.geometry("360x360") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) + self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) + self.caption_entry = ctk.CTkEntry(self.frame, width=200) + self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) + self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.prefix_entry = ctk.CTkEntry(self.frame, width=200) + self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) + self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.postfix_entry = ctk.CTkEntry(self.frame, width=200) + self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_captions(self): + self.parent.load_captioning_model(self.model_var.get()) + + mode = { + "Replace all captions": "replace", + "Create if absent": "fill", + "Add as new line": "add", + }[self.mode_var.get()] + + self.parent.captioning_model.caption_folder( + sample_dir=self.path_entry.get(), + initial_caption=self.caption_entry.get(), + caption_prefix=self.prefix_entry.get(), + caption_postfix=self.postfix_entry.get(), + mode=mode, + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/CtkGenerateMasksWindowView.py b/modules/ui/CtkGenerateMasksWindowView.py new file mode 100644 index 000000000..daff0d3d5 --- /dev/null +++ b/modules/ui/CtkGenerateMasksWindowView.py @@ -0,0 +1,151 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateMasksWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating masks for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + + self.parent = parent + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] + self.model_var = ctk.StringVar(self, "ClipSeg") + self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] + + self.title("Batch generate masks") + self.geometry("360x430") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) + self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) + self.prompt_entry = ctk.CTkEntry(self.frame, width=200) + self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) + self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") + self.threshold_entry.insert(0, "0.3") + self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) + self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") + self.smooth_entry.insert(0, 5) + self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) + self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") + self.expand_entry.insert(0, 10) + self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) + self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") + self.alpha_entry.insert(0, 1) + self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) + + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_masks(self): + self.parent.load_masking_model(self.model_var.get()) + + mode = { + "Replace all masks": "replace", + "Create if absent": "fill", + "Add to existing": "add", + "Subtract from existing": "subtract", + "Blend with existing": "blend", + }[self.mode_var.get()] + + self.parent.masking_model.mask_folder( + sample_dir=self.path_entry.get(), + prompts=[self.prompt_entry.get()], + mode=mode, + alpha=float(self.alpha_entry.get()), + threshold=float(self.threshold_entry.get()), + smooth_pixels=int(self.smooth_entry.get()), + expand_pixels=int(self.expand_entry.get()), + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py new file mode 100644 index 000000000..1c73d90ce --- /dev/null +++ b/modules/ui/CtkLoraTabView.py @@ -0,0 +1,154 @@ + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelType import PeftType +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation_helpers import check_range + +import customtkinter as ctk + + +class LoraTab: + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + self.scroll_frame = None + self.options_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + + self.scroll_frame.grid_columnconfigure(0, weight=0) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=2) + + components.label(self.scroll_frame, 0, 0, "Type", + tooltip="The type of low-parameter finetuning method.") + # This will instantly call self.setup_lora. + components.options_kv(self.scroll_frame, 0, 1, [ + ("LoRA", PeftType.LORA), + ("LoHa", PeftType.LOHA), + ("OFT v2", PeftType.OFT_2), + ], self.ui_state, "peft_type", command=self.setup_lora) + + def setup_lora(self, peft_type: PeftType): + if peft_type == PeftType.LOHA: + name = "LoHa" + elif peft_type == PeftType.OFT_2: + name = "OFT v2" + else: + name = "LoRA" + + if self.options_frame: + self.options_frame.destroy() + self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") + self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") + master = self.options_frame + + master.grid_columnconfigure(0, weight=0, uniform="a") + master.grid_columnconfigure(1, weight=1, uniform="a") + master.grid_columnconfigure(2, minsize=50, uniform="a") + master.grid_columnconfigure(3, weight=0, uniform="a") + master.grid_columnconfigure(4, weight=1, uniform="a") + + # lora model name + components.label(master, 0, 0, f"{name} base model", + tooltip=f"The base {name} to train on. Leave empty to create a new {name}") + entry = components.path_entry( + master, 0, 1, self.ui_state, "lora_model_name", + mode="file", path_modifier=components.json_path_modifier + ) + entry.grid(row=0, column=1, columnspan=4) + + + # LoRA decomposition + if peft_type == PeftType.LORA: + components.label(master, 1, 3, "Decompose Weights (DoRA)", + tooltip="Decompose LoRA Weights (aka, DoRA).") + components.switch(master, 1, 4, self.ui_state, "lora_decompose") + + components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", + tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") + components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") + components.label(master, 3, 3, "Apply on output axis (DoRA Only)", + tooltip="Apply the weight decomposition on the output axis instead of the input axis.") + components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") + + # LoRA and LoHA shared settings + if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: + # rank + components.label(master, 1, 0, f"{name} rank", + tooltip=f"The rank parameter used when creating a new {name}") + components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) + + # alpha + components.label(master, 2, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) + + # Dropout Percentage + components.label(master, 3, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 3, 1, self.ui_state, "dropout_probability") + + # weight dtype + components.label(master, 4, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 4, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # For use with additional embeddings. + components.label(master, 5, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") + + # OFTv2 + elif peft_type == PeftType.OFT_2: + # Block Size + components.label(master, 1, 0, f"{name} Block Size", + tooltip=f"The block size parameter used when creating a new {name}") + components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) + + # COFT + components.label(master, 1, 3, "Constrained OFT (COFT)", + tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") + components.switch(master, 1, 4, self.ui_state, "oft_coft") + + components.label(master, 2, 3, "COFT Epsilon", + tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") + components.entry(master, 2, 4, self.ui_state, "coft_eps") + + # Block Share + components.label(master, 3, 3, "Block Share", + tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") + components.switch(master, 3, 4, self.ui_state, "oft_block_share") + + # Dropout Percentage + components.label(master, 2, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 2, 1, self.ui_state, "dropout_probability") + + # OFT weight dtype + components.label(master, 3, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # For use with additional embeddings. + components.label(master, 4, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/CtkModelTabView.py b/modules/ui/CtkModelTabView.py new file mode 100644 index 000000000..ff17ea3ba --- /dev/null +++ b/modules/ui/CtkModelTabView.py @@ -0,0 +1,688 @@ + +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ConfigPart import ConfigPart +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class ModelTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.scroll_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + self.scroll_frame.grid_columnconfigure(0, weight=1) + + base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) + base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") + + base_frame.grid_columnconfigure(0, weight=0) + base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) + base_frame.grid_columnconfigure(2, minsize=50) + base_frame.grid_columnconfigure(3, weight=0) + base_frame.grid_columnconfigure(4, weight=1) + + if self.train_config.model_type.is_stable_diffusion(): #TODO simplify + self.__setup_stable_diffusion_ui(base_frame) + if self.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(base_frame) + elif self.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(base_frame) + elif self.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(base_frame) + elif self.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(base_frame) + elif self.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(base_frame) + elif self.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(base_frame) + elif self.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(base_frame) + elif self.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(base_frame) + elif self.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(base_frame) + elif self.train_config.model_type.is_sana(): + self.__setup_sana_ui(base_frame) + elif self.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(base_frame) + elif self.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(base_frame) + elif self.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(base_frame) + + def __setup_stable_diffusion_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_unet=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method in [ + TrainingMethod.FINE_TUNE, + TrainingMethod.FINE_TUNE_VAE, + ], + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_stable_diffusion_3_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_text_encoder_3=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_flux_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_flux_2_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_z_image_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_ernie_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_chroma_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_qwen_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_stable_diffusion_xl_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_unet=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_wuerstchen_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_prior=True, + allow_override_prior=self.train_config.model_type.is_stable_cascade(), + has_text_encoder=True, + ) + row = self.__create_effnet_encoder_components(frame, row) + row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) + row = self.__create_output_components( + frame, + row, + allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE + or self.train_config.model_type.is_stable_cascade(), + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_pixart_alpha_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_sana_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_hunyuan_video_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_hi_dream_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + has_text_encoder_1=True, + has_text_encoder_2=True, + has_text_encoder_3=True, + has_text_encoder_4=True, + allow_override_text_encoder_4=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + + def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: + options = [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ("float16", DataType.FLOAT_16), + ("float8 (W8)", DataType.FLOAT_8), + # ("int8", DataType.INT_8), # TODO: reactivate when the int8 implementation is fixed in bitsandbytes: https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1332 + ("nfloat4", DataType.NFLOAT_4), + ] + if include_a8: + options += [ + ("float W8A8", DataType.FLOAT_W8A8), + ("int W8A8", DataType.INT_W8A8), + ] + + if include_gguf: + options.append(("GGUF", DataType.GGUF)) + if include_a8: + options += [ + ("GGUF A8 float", DataType.GGUF_A8_FLOAT), + ("GGUF A8 int", DataType.GGUF_A8_INT), + ] + + return options + + def __create_base_dtype_components(self, frame, row: int) -> int: + # huggingface token + components.label(frame, row, 0, "Hugging Face Token", + tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " + "Go to https://huggingface.co/settings/tokens to create an access token.", + wide_tooltip=True) + components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") + + row += 1 + + # base model + components.label(frame, row, 0, "Base Model", + tooltip="Filename, directory or Hugging Face repository of the base model") + components.path_entry( + frame, row, 1, self.ui_state, "base_model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # compile + components.label(frame, row, 3, "Compile transformer blocks", + tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") + components.switch(frame, row, 4, self.ui_state, "compile") + + row += 1 + + return row + + def __create_base_components( + self, + frame, + row: int, + has_unet: bool = False, + has_prior: bool = False, + allow_override_prior: bool = False, + has_transformer: bool = False, + allow_override_transformer: bool = False, + allow_override_text_encoder_4: bool = False, + has_text_encoder: bool = False, + has_text_encoder_1: bool = False, + has_text_encoder_2: bool = False, + has_text_encoder_3: bool = False, + has_text_encoder_4: bool = False, + has_vae: bool = False, + ) -> int: + if has_unet: + # unet weight dtype + components.label(frame, row, 3, "UNet Data Type", + tooltip="The unet weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), + self.ui_state, "unet.weight_dtype") + + row += 1 + + if has_prior: + if allow_override_prior: + # prior model + components.label(frame, row, 0, "Prior Model", + tooltip="Filename, directory or Hugging Face repository of the prior model") + components.path_entry( + frame, row, 1, self.ui_state, "prior.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # prior weight dtype + components.label(frame, row, 3, "Prior Data Type", + tooltip="The prior weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "prior.weight_dtype") + + row += 1 + + if has_transformer: + if allow_override_transformer: + # transformer model + components.label(frame, row, 0, "Override Transformer / GGUF", + tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") + components.path_entry( + frame, row, 1, self.ui_state, "transformer.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # transformer weight dtype + components.label(frame, row, 3, "Transformer Data Type", + tooltip="The transformer weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), + self.ui_state, "transformer.weight_dtype") + + row += 1 + + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + + components.label(frame, row, 0, "Quantization") + components.layer_filter_entry(frame, row, 1, self.ui_state, + preset_var_name="quantization.layer_filter_preset", presets=presets, + preset_label="Quantization Layer Filter", + preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", + entry_var_name="quantization.layer_filter", + entry_tooltip="Comma-separated list of layers to quantize. Regular expressions (if toggled) are supported. Any model layer with a matching name will be quantized", + regex_var_name="quantization.layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + frame_color="transparent", + ) + + # SVDQuant - create vertical grids to match the size of layer_filter_entry + svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") + svd_label_frame.grid(row=row, column=3, sticky="nsew") + svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") + svd_entry_frame.grid(row=row, column=4, sticky="nsew") + components.label(svd_label_frame, 0, 0, "SVDQuant", + tooltip="What datatype to use for SVDQuant weights decomposition.") + components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], + self.ui_state, "quantization.svd_dtype") + components.label(svd_label_frame, 1, 0, "SVDQuant Rank", + tooltip="Rank for SVDQuant weights decomposition") + components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") + row += 1 + + + if has_text_encoder: + # text encoder weight dtype + components.label(frame, row, 3, "Text Encoder Data Type", + tooltip="The text encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder.weight_dtype") + + row += 1 + + if has_text_encoder_1: + # text encoder 1 weight dtype + components.label(frame, row, 3, "Text Encoder 1 Data Type", + tooltip="The text encoder 1 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder.weight_dtype") + + row += 1 + + if has_text_encoder_2: + # text encoder 2 weight dtype + components.label(frame, row, 3, "Text Encoder 2 Data Type", + tooltip="The text encoder 2 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_2.weight_dtype") + + row += 1 + + if has_text_encoder_3: + # text encoder 3 weight dtype + components.label(frame, row, 3, "Text Encoder 3 Data Type", + tooltip="The text encoder 3 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_3.weight_dtype") + + row += 1 + + if has_text_encoder_4: + if allow_override_text_encoder_4: + # text encoder 4 weight dtype + components.label(frame, row, 0, "Text Encoder 4 Override", + tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") + components.path_entry( + frame, row, 1, self.ui_state, "text_encoder_4.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # text encoder 4 weight dtype + components.label(frame, row, 3, "Text Encoder 4 Data Type", + tooltip="The text encoder 4 weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "text_encoder_4.weight_dtype") + + row += 1 + + if has_vae: + # base model + components.label(frame, row, 0, "VAE Override", + tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. Using a safetensor VAE file will cause an error that the model cannot be loaded.") + components.path_entry( + frame, row, 1, self.ui_state, "vae.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # vae weight dtype + components.label(frame, row, 3, "VAE Data Type", + tooltip="The vae weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "vae.weight_dtype") + + row += 1 + + return row + + def __create_effnet_encoder_components(self, frame, row: int): + # effnet encoder model + components.label(frame, row, 0, "Effnet Encoder Model", + tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") + components.path_entry( + frame, row, 1, self.ui_state, "effnet_encoder.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # effnet encoder weight dtype + components.label(frame, row, 3, "Effnet Encoder Data Type", + tooltip="The effnet encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "effnet_encoder.weight_dtype") + + row += 1 + + return row + + def __create_decoder_components( + self, + frame, + row: int, + has_text_encoder: bool, + ) -> int: + # decoder model + components.label(frame, row, 0, "Decoder Model", + tooltip="Filename, directory or Hugging Face repository of the decoder model") + components.path_entry( + frame, row, 1, self.ui_state, "decoder.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # decoder weight dtype + components.label(frame, row, 3, "Decoder Data Type", + tooltip="The decoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder.weight_dtype") + + row += 1 + + if has_text_encoder: + # decoder text encoder weight dtype + components.label(frame, row, 3, "Decoder Text Encoder Data Type", + tooltip="The decoder text encoder weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder_text_encoder.weight_dtype") + + row += 1 + + # decoder vqgan weight dtype + components.label(frame, row, 3, "Decoder VQGAN Data Type", + tooltip="The decoder vqgan weight data type") + components.options_kv(frame, row, 4, self.__create_dtype_options(), + self.ui_state, "decoder_vqgan.weight_dtype") + + row += 1 + + return row + + def __create_output_components( + self, + frame, + row: int, + allow_safetensors: bool = False, + allow_diffusers: bool = False, + allow_legacy_safetensors: bool = False, + allow_comfy: bool = False, + ) -> int: + # output model destination + components.label(frame, row, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + components.path_entry( + frame, row, 1, self.ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + # output data type + components.label(frame, row, 3, "Output Data Type", + tooltip="Precision to use when saving the output model") + components.options_kv(frame, row, 4, [ + ("float16", DataType.FLOAT_16), + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ("float8", DataType.FLOAT_8), + ("nfloat4", DataType.NFLOAT_4), + ], self.ui_state, "output_dtype") + + row += 1 + + # output format + formats = [] + if allow_safetensors: + formats.append(("Safetensors", ModelFormat.SAFETENSORS)) + if allow_diffusers: + formats.append(("Diffusers", ModelFormat.DIFFUSERS)) + # if allow_legacy_safetensors: + # formats.append(("Legacy Safetensors", ModelFormat.LEGACY_SAFETENSORS)) + if allow_comfy: + formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) + + components.label(frame, row, 0, "Output Format", + tooltip="Format to use when saving the output model") + components.options_kv(frame, row, 1, formats, self.ui_state, "output_model_format") + + # include config + components.label(frame, row, 3, "Include Config", + tooltip="Include the training configuration in the final model. Only supported for safetensors files. " + "None: No config is included. " + "Settings: All training settings are included. " + "All: All settings, including the samples and concepts are included.") + components.options_kv(frame, row, 4, [ + ("None", ConfigPart.NONE), + ("Settings", ConfigPart.SETTINGS), + ("All", ConfigPart.ALL), + ], self.ui_state, "include_train_config") + + row += 1 + + return row diff --git a/modules/ui/CtkMuonAdamWindowView.py b/modules/ui/CtkMuonAdamWindowView.py new file mode 100644 index 000000000..5879ab432 --- /dev/null +++ b/modules/ui/CtkMuonAdamWindowView.py @@ -0,0 +1,107 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + +MUON_AUX_ADAM_DEFAULTS = { + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.0, +} + +class MuonAdamWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state: UIState, + parent_optimizer_type: Optimizer, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.adam_ui_state = ui_state + self.parent_optimizer_type = parent_optimizer_type + + if self.parent_optimizer_type == Optimizer.MUON: + self.title("Muon's Auxiliary AdamW Settings") + self.adam_params_def = MUON_AUX_ADAM_DEFAULTS + else: + self.title("Muon_adv's Auxiliary AdamW_adv Settings") + self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.destroy) + self.create_adam_params_ui(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def create_adam_params_ui(self, master): + # This is a large map, copied from OptimizerParamsWindow for simplicity. + # @formatter:off + KEY_DETAIL_MAP = { + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + } + # @formatter:on + + adam_params = self.adam_params_def + + for index, key in enumerate(adam_params.keys()): + if key not in KEY_DETAIL_MAP: + continue + + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + param_type = arg_info['type'] + + row = index // 2 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if param_type != 'bool': + components.entry(master, row, col + 1, self.adam_ui_state, key) + else: + components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/CtkOffloadingWindowView.py b/modules/ui/CtkOffloadingWindowView.py new file mode 100644 index 000000000..54035e121 --- /dev/null +++ b/modules/ui/CtkOffloadingWindowView.py @@ -0,0 +1,75 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.GradientCheckpointingMethod import ( + GradientCheckpointingMethod, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OffloadingWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.title("Offloading") + self.geometry("800x400") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=1) + frame.grid_columnconfigure(1, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, + "gradient_checkpointing") + + # gradient checkpointing layer offloading + components.label(frame, 1, 0, "Async Offloading", + tooltip="Enables Asynchronous offloading.") + components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") + + # gradient checkpointing layer offloading + components.label(frame, 2, 0, "Offload Activations", + tooltip="Enables Activation Offloading") + components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") + + # gradient checkpointing layer offloading + components.label(frame, 3, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") + + frame.pack(fill="both", expand=1) + return frame + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkOptimizerParamsWindowView.py b/modules/ui/CtkOptimizerParamsWindowView.py new file mode 100644 index 000000000..16063c26c --- /dev/null +++ b/modules/ui/CtkOptimizerParamsWindowView.py @@ -0,0 +1,288 @@ +import contextlib +from tkinter import TclError + +from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import ( + OPTIMIZER_DEFAULT_PARAMETERS, + change_optimizer, + load_optimizer_defaults, + update_optimizer_config, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OptimizerParamsWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + self.optimizer_ui_state = ui_state.get_var("optimizer") + self.protocol("WM_DELETE_WINDOW", self.on_window_close) + self.muon_adam_button = None + + self.title("Optimizer Settings") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # Optimizer + components.label(master, 0, 0, "Optimizer", + tooltip="The type of optimizer") + + # Create the optimizer dropdown menu and set the command + components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", + command=self.on_optimizer_change) + + # Defaults Button + components.label(master, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, + tooltip="Load default settings for the selected optimizer") + + self.create_dynamic_ui(master) + + def clear_dynamic_ui(self, master): + with contextlib.suppress(TclError): + for widget in master.winfo_children(): + grid_info = widget.grid_info() + if int(grid_info["row"]) >= 1: + widget.destroy() + + def create_dynamic_ui( + self, + master, + ): + + # Lookup for the title and tooltip for a key + # @formatter:off + KEY_DETAIL_MAP = { + 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, + 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, + 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, + 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, + 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, + 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, + 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, + 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, + 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, + 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, + 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, + 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, + 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, + 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, + 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, + 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, + 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, + 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, + 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, + 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, + 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, + 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, + 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, + 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, + 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, + 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, + 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, + 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, + 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, + 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, + 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, + 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, + 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, + 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, + 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, + 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, + 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, + 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, + 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, + 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, + 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, + 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, + 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, + 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, + 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, + 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, + 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, + 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, + 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, + 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, + 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, + 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, + 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, + 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, + 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, + 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, + 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, + 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, + 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, + 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, + 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, + 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, + 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, + 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, + 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, + 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, + 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, + 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, + 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, + 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, + 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, + 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, + 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, + 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, + 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, + 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, + } + # @formatter:on + + if not self.winfo_exists(): # check if this window isn't open + return + + selected_optimizer = self.train_config.optimizer.optimizer + + # Extract the keys for the selected optimizer + for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): + if key not in KEY_DETAIL_MAP: + continue + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + type = arg_info['type'] + + row = (index // 2) + 1 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if key == 'MuonWithAuxAdam': + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + + components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + + self.muon_adam_button = components.button( + frame, 0, 1, "...", self.open_muon_adam_window, + tooltip="Configure the auxiliary AdamW_adv optimizer", + width=20, padx=5 ) + self.toggle_muon_adam_button() + elif type != 'bool': + components.entry(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + else: + components.switch(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + + def update_user_pref(self, *args): + update_optimizer_config(self.train_config) + self.toggle_muon_adam_button() + + def on_optimizer_change(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.clear_dynamic_ui(self.frame) + self.create_dynamic_ui(self.frame) + + def load_defaults(self, *args): + optimizer_config = load_optimizer_defaults(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def on_window_close(self): + self.destroy() + + def toggle_muon_adam_button(self): + if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + + def open_muon_adam_window(self): + current_optimizer = self.train_config.optimizer.optimizer + + adam_config = TrainOptimizerConfig.default_values() + current_state = self.train_config.optimizer.muon_adam_config + + if current_optimizer == Optimizer.MUON: + defaults = MUON_AUX_ADAM_DEFAULTS + else: + defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + if current_state is None: + adam_config.from_dict(defaults) + if current_optimizer != Optimizer.MUON: + adam_config.optimizer = Optimizer.ADAMW_ADV + elif isinstance(current_state, dict): + adam_config.from_dict(current_state) + else: + # Should not happen if TrainConfig defines it as dict, but for safety + adam_config = current_state + + temp_adam_ui_state = UIState(self, adam_config) + window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + self.wait_window(window) + + self.train_config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/CtkProfilingWindowView.py b/modules/ui/CtkProfilingWindowView.py new file mode 100644 index 000000000..8d298abe3 --- /dev/null +++ b/modules/ui/CtkProfilingWindowView.py @@ -0,0 +1,57 @@ +import faulthandler + +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from scalene import scalene_profiler + + +class ProfilingWindow(ctk.CTkToplevel): + def __init__(self, parent, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + self.parent = parent + + self.title("Profiling") + self.geometry("512x512") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=0) + self.grid_rowconfigure(2, weight=1) + self.grid_columnconfigure(0, weight=1) + + components.button(self, 0, 0, "Dump stack", self._dump_stack) + self._profile_button = components.button( + self, 1, 0, "Start Profiling", self._start_profiler, + tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") + + # Bottom bar + self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) + self._bottom_bar.grid(row=2, column=0, sticky="sew") + self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") + + self.protocol("WM_DELETE_WINDOW", self.withdraw) + self.withdraw() + self.after(200, lambda: set_window_icon(self)) + + def _dump_stack(self): + with open('stacks.txt', 'w') as f: + faulthandler.dump_traceback(f) + self._message_label.configure(text='Stack dumped to stacks.txt') + + def _end_profiler(self): + scalene_profiler.stop() + + self._message_label.configure(text='Inactive') + self._profile_button.configure(text='Start Profiling') + self._profile_button.configure(command=self._start_profiler) + + def _start_profiler(self): + scalene_profiler.start() + + self._message_label.configure(text='Profiling active...') + self._profile_button.configure(text='End Profiling') + self._profile_button.configure(command=self._end_profiler) diff --git a/modules/ui/CtkSampleFrameView.py b/modules/ui/CtkSampleFrameView.py new file mode 100644 index 000000000..297caac29 --- /dev/null +++ b/modules/ui/CtkSampleFrameView.py @@ -0,0 +1,134 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SampleFrame(ctk.CTkFrame): + def __init__( + self, + parent, + sample: SampleConfig, + ui_state: UIState, + model_type: ModelType, + include_prompt: bool = True, + include_settings: bool = True, + ): + ctk.CTkFrame.__init__(self, parent, fg_color="transparent") + + self.sample = sample + self.ui_state = ui_state + self.model_type = model_type + + is_flow_matching = model_type.is_flow_matching() + is_inpainting_model = model_type.has_conditioning_image_input() + is_video_model = model_type.is_video_model() + + if include_prompt and include_prompt: + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + if include_prompt: + top_frame = ctk.CTkFrame(self, fg_color="transparent") + top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + top_frame.grid_columnconfigure(0, weight=0) + top_frame.grid_columnconfigure(1, weight=1) + + if include_settings: + bottom_frame = ctk.CTkFrame(self, fg_color="transparent") + bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + bottom_frame.grid_columnconfigure(0, weight=0) + bottom_frame.grid_columnconfigure(1, weight=1) + bottom_frame.grid_columnconfigure(2, weight=0) + bottom_frame.grid_columnconfigure(3, weight=1) + + if include_prompt: + # prompt + components.label(top_frame, 0, 0, "prompt:") + components.entry(top_frame, 0, 1, self.ui_state, "prompt") + + # negative prompt + components.label(top_frame, 1, 0, "negative prompt:") + components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") + + if include_settings: + # width + components.label(bottom_frame, 0, 0, "width:") + components.entry(bottom_frame, 0, 1, self.ui_state, "width") + + # height + components.label(bottom_frame, 0, 2, "height:") + components.entry(bottom_frame, 0, 3, self.ui_state, "height") + + if is_video_model: + # frames + components.label(bottom_frame, 1, 0, "frames:", + tooltip="Number of frames to generate. Only used when generating videos.") + components.entry(bottom_frame, 1, 1, self.ui_state, "frames") + + # length + components.label(bottom_frame, 1, 2, "length:", + tooltip="Length in seconds of audio output.") + components.entry(bottom_frame, 1, 3, self.ui_state, "length") + + # seed + components.label(bottom_frame, 2, 0, "seed:") + components.entry(bottom_frame, 2, 1, self.ui_state, "seed") + + # random seed + components.label(bottom_frame, 2, 2, "random seed:") + components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") + + # cfg scale + components.label(bottom_frame, 3, 0, "cfg scale:") + components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") + + # sampler + if not is_flow_matching: + components.label(bottom_frame, 4, 2, "sampler:") + components.options_kv(bottom_frame, 4, 3, [ + ("DDIM", NoiseScheduler.DDIM), + ("Euler", NoiseScheduler.EULER), + ("Euler A", NoiseScheduler.EULER_A), + # ("DPM++", NoiseScheduler.DPMPP), # TODO: produces noisy samples + # ("DPM++ SDE", NoiseScheduler.DPMPP_SDE), # TODO: produces noisy samples + ("UniPC", NoiseScheduler.UNIPC), + ("Euler Karras", NoiseScheduler.EULER_KARRAS), + ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), + ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), + ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) + ], self.ui_state, "noise_scheduler") + + # steps + components.label(bottom_frame, 4, 0, "steps:") + components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") + + # inpainting + if is_inpainting_model: + components.label(bottom_frame, 5, 0, "inpainting:", + tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") + components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") + + # base image path + components.label(bottom_frame, 6, 0, "base image path:", + tooltip="The base image used when inpainting.") + components.file_entry(bottom_frame, 6, 1, self.ui_state, "base_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) + + # mask image path + components.label(bottom_frame, 6, 2, "mask image path:", + tooltip="The mask used when inpainting.") + components.file_entry(bottom_frame, 6, 3, self.ui_state, "mask_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) diff --git a/modules/ui/CtkSampleParamsWindowView.py b/modules/ui/CtkSampleParamsWindowView.py new file mode 100644 index 000000000..2b0b3f3f1 --- /dev/null +++ b/modules/ui/CtkSampleParamsWindowView.py @@ -0,0 +1,39 @@ +from modules.ui.SampleFrame import SampleFrame +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SampleParamsWindow(ctk.CTkToplevel): + def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + + self.sample = sample + self.ui_state = ui_state + self.model_type = model_type + + self.title("Sample") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) + frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkSampleWindowView.py b/modules/ui/CtkSampleWindowView.py new file mode 100644 index 000000000..0f91ad2fa --- /dev/null +++ b/modules/ui/CtkSampleWindowView.py @@ -0,0 +1,227 @@ +import contextlib +import copy +import os +import tkinter as tk +import traceback + +from modules.model.BaseModel import BaseModel +from modules.modelSampler.BaseModelSampler import ( + BaseModelSampler, + ModelSamplerOutput, +) +from modules.ui.SampleFrame import SampleFrame +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.FileType import FileType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.time_util import get_string_timestamp +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +from PIL import Image + + +class SampleWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + use_external_model: bool, + callbacks: TrainCallbacks | None = None, + commands: TrainCommands | None = None, + *args, **kwargs + ): + super().__init__(parent, *args, **kwargs) + + self.title("Sample") + self.geometry("1200x800") + self.resizable(True, True) + + if not use_external_model: + self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) + # remove some settings to speed up model loading for sampling + self.initial_train_config.optimizer.optimizer = None + self.initial_train_config.ema = EMAMode.OFF + else: + self.initial_train_config = None + + #TODO why is there a current_train_config and an initial_train_config? + #current_train_config doesn't seem to ever change + self.current_train_config = train_config + self.callbacks = callbacks + self.commands = commands + + # get model specific defaults + model_type = train_config.model_type + self.sample = SampleConfig.default_values(model_type) + self.ui_state = UIState(self, self.sample) + + if use_external_model: + self.callbacks.set_on_sample_custom(self.__update_preview) + self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) + else: + self.model = None + self.model_sampler = None + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_rowconfigure(3, weight=0) + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1) + + prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + + settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + # image + self.image = ctk.CTkImage( + light_image=self.__dummy_image(), + size=(512, 512) + ) + + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) + image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + + self.progress = components.progress(self, 2, 0) + components.button(self, 3, 0, "sample", self.__sample) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def __load_model(self) -> BaseModel: + model_loader = create.create_model_loader( + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + model_setup = create.create_model_setup( + model_type=self.initial_train_config.model_type, + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + training_method=self.initial_train_config.training_method, + ) + + model_names = self.initial_train_config.model_names() + if self.initial_train_config.continue_last_backup: + last_backup_path = self.initial_train_config.get_last_backup_path() + + if last_backup_path: + if self.initial_train_config.training_method == TrainingMethod.LORA: + model_names.lora = last_backup_path + elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: + model_names.embedding.model_name = last_backup_path + else: # fine-tunes + model_names.base_model = last_backup_path + + print(f"Loading from backup '{last_backup_path}'...") + else: + print("No backup found, loading without backup...") + + if self.initial_train_config.quantization.cache_dir is None: + self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" + os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) + + model = model_loader.load( + model_type=self.initial_train_config.model_type, + model_names=model_names, + weight_dtypes=self.initial_train_config.weight_dtypes(), + quantization=self.initial_train_config.quantization, + ) + model.train_config = self.initial_train_config + + model_setup.setup_optimizations(model, self.initial_train_config) + model_setup.setup_train_device(model, self.initial_train_config) + model_setup.setup_model(model, self.initial_train_config) + model.to(torch.device(self.initial_train_config.temp_device)) + + return model + + def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + return create.create_model_sampler( + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + model=model, + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + def __update_preview(self, sampler_output: ModelSamplerOutput): + if sampler_output.file_type == FileType.IMAGE: + image = sampler_output.data + self.image.configure( + light_image=image, + size=(image.width, image.height), + ) + + def __update_progress(self, progress: int, max_progress: int): + self.progress.set(progress / max_progress) + self.update() + + def __dummy_image(self) -> Image: + return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) + + def __sample(self): + sample = copy.copy(self.sample) + + if self.commands: + self.commands.sample_custom(sample) + else: + if self.model is None: + # lazy initialization + self.model = self.__load_model() + self.model_sampler = self.__create_sampler(self.model) + + sample.from_train_config(self.current_train_config) + + sample_dir = os.path.join( + self.initial_train_config.workspace_dir, + "samples", + "custom", + ) + + progress = self.model.train_progress + sample_path = os.path.join( + sample_dir, + f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" + ) + + self.model.eval() + + self.model_sampler.sample( + sample_config=sample, + destination=sample_path, + image_format=self.current_train_config.sample_image_format, + video_format=self.current_train_config.sample_video_format, + audio_format=self.current_train_config.sample_audio_format, + on_sample=self.__update_preview, + on_update_progress=self.__update_progress, + ) + + def destroy(self): + try: + if hasattr(self, "_icon_image_ref"): + del self._icon_image_ref + + # Remove any pending after callbacks + for after_id in self.tk.call('after', 'info'): + with contextlib.suppress(tk.TclError, RuntimeError): + self.after_cancel(after_id) + + super().destroy() + except (tk.TclError, RuntimeError) as e: + print(f"Error destroying window: {e}") + except Exception as e: + print(f"Unexpected error destroying window: {e}") + traceback.print_exc() diff --git a/modules/ui/CtkSamplingTabView.py b/modules/ui/CtkSamplingTabView.py new file mode 100644 index 000000000..5a3c44f08 --- /dev/null +++ b/modules/ui/CtkSamplingTabView.py @@ -0,0 +1,124 @@ +from modules.ui.ConfigList import ConfigList +from modules.ui.SampleParamsWindow import SampleParamsWindow +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.ui import components +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class SamplingTab(ConfigList): + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + from_external_file=True, + attr_name="sample_definition_file_name", + config_dir="training_samples", + default_config_name="samples.json", + add_button_text="Add Sample", + add_button_tooltip="Add a new sample configuration.", + is_full_width=True, + show_toggle_button=True + ) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict: + return SampleConfig.default_values(self.train_config.model_type) + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) + + +class SampleWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__( + master=master, corner_radius=10, bg_color="transparent" + ) + + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(10, weight=1) + + # close button + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.grid(row=0, column=0) + + # clone button + clone_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i), + ) + clone_button.grid(row=0, column=1, padx=5) + + # enabled + self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) + self.enabled_switch.configure(width=40) + + # width + components.label(self, 0, 3, "width:") + self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") + self.width_entry.bind('', lambda _: save_command()) + self.width_entry.configure(width=50) + + # height + components.label(self, 0, 5, "height:") + self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") + self.height_entry.bind('', lambda _: save_command()) + self.height_entry.configure(width=50) + + # seed + components.label(self, 0, 7, "seed:") + self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") + self.seed_entry.bind('', lambda _: save_command()) + self.seed_entry.configure(width=80) + + # prompt + components.label(self, 0, 9, "prompt:") + self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") + self.prompt_entry.bind('', lambda _: save_command()) + + # button + self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) + self.button.configure(width=40) + + self.__set_enabled() + + def __switch_enabled(self): + self.save_command() + self.__set_enabled() + + def __set_enabled(self): + enabled = self.element.enabled + self.width_entry.configure(state="normal" if enabled else "disabled") + self.height_entry.configure(state="normal" if enabled else "disabled") + self.prompt_entry.configure(state="normal" if enabled else "disabled") + self.seed_entry.configure(state="normal" if enabled else "disabled") + self.button.configure(state="normal" if enabled else "disabled") + + def configure_element(self): + pass + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkSchedulerParamsWindowView.py b/modules/ui/CtkSchedulerParamsWindowView.py new file mode 100644 index 000000000..f96ed4876 --- /dev/null +++ b/modules/ui/CtkSchedulerParamsWindowView.py @@ -0,0 +1,119 @@ +from modules.ui.ConfigList import ConfigList +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.LearningRateScheduler import LearningRateScheduler +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class KvParams(ConfigList): + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__( + master, + train_config, + ui_state, + attr_name="scheduler_params", + from_external_file=False, + add_button_text="add parameter", + is_full_width=True + ) + + def refresh_ui(self): + self._create_element_list() + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + def create_new_element(self) -> dict[str, str]: + return {"key": "", "value": ""} + + def open_element_window(self, i, ui_state): + pass + + +class KvWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__(master=master, bg_color="transparent") + self.element = element + self.ui_state = UIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1, uniform=1) + self.grid_columnconfigure(2, weight=1, uniform=1) + + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i)) + close_button.grid(row=0, column=0) + + # Key + tooltip_key = "Key name for an argument in your scheduler" + self.key = components.entry(self, 0, 1, self.ui_state, "key", + tooltip=tooltip_key, wide_tooltip=True) + self.key.bind("", lambda _: save_command()) + self.key.configure(width=50) + + # Value + tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." + self.value = components.entry(self, 0, 2, self.ui_state, "value", + tooltip=tooltip_val, wide_tooltip=True) + self.value.bind("", lambda _: save_command()) + self.value.configure(width=50) + + def place_in_list(self): + self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") + + +class SchedulerParamsWindow(ctk.CTkToplevel): + def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + + self.title("Learning Rate Scheduler Settings") + self.geometry("800x400") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkFrame(self) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") + self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: + components.label(master, 0, 0, "Class Name", + tooltip="Python class module and name for the custom scheduler class, in the form of ..") + components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") + + # Any additional parameters, in key-value form. + self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) + + def on_window_close(self): + self.destroy() diff --git a/modules/ui/CtkTimestepDistributionWindowView.py b/modules/ui/CtkTimestepDistributionWindowView.py new file mode 100644 index 000000000..21e41ce3e --- /dev/null +++ b/modules/ui/CtkTimestepDistributionWindowView.py @@ -0,0 +1,186 @@ + +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( + ModelSetupNoiseMixin, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch +from torch import Tensor + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class TimestepGenerator(ModelSetupNoiseMixin): + + def __init__( + self, + timestep_distribution: TimestepDistribution, + min_noising_strength: float, + max_noising_strength: float, + noising_weight: float, + noising_bias: float, + timestep_shift: float, + ): + super().__init__() + + self.timestep_distribution = timestep_distribution + self.min_noising_strength = min_noising_strength + self.max_noising_strength = max_noising_strength + self.noising_weight = noising_weight + self.noising_bias = noising_bias + self.timestep_shift = timestep_shift + + def generate(self) -> Tensor: + generator = torch.Generator() + generator.seed() + + config = TrainConfig.default_values() + config.timestep_distribution = self.timestep_distribution + config.min_noising_strength = self.min_noising_strength + config.max_noising_strength = self.max_noising_strength + config.noising_weight = self.noising_weight + config.noising_bias = self.noising_bias + config.timestep_shift = self.timestep_shift + + + return self._get_timestep_discrete( + num_train_timesteps=1000, + deterministic=False, + generator=generator, + batch_size=1000000, + config=config, + ) + + +class TimestepDistributionWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.title("Timestep Distribution") + self.geometry("900x600") + self.resizable(True, True) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + frame.grid_rowconfigure(7, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + "timestep_distribution") + + # min noising strength + components.label(frame, 1, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + + # max noising strength + components.label(frame, 2, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + + # noising weight + components.label(frame, 3, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 3, 1, self.ui_state, "noising_weight") + + # noising bias + components.label(frame, 4, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 4, 1, self.ui_state, "noising_bias") + + # timestep shift + components.label(frame, 5, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + + # dynamic timestep shifting + components.label(frame, 6, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") + + + # plot + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + fig, ax = plt.subplots() + self.ax = ax + self.canvas = FigureCanvasTkAgg(fig, master=frame) + self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) + + fig.set_facecolor(background_color) + ax.set_facecolor(background_color) + ax.spines['bottom'].set_color(text_color) + ax.spines['left'].set_color(text_color) + ax.spines['top'].set_color(text_color) + ax.spines['right'].set_color(text_color) + ax.tick_params(axis='x', colors=text_color, which="both") + ax.tick_params(axis='y', colors=text_color, which="both") + ax.xaxis.label.set_color(text_color) + ax.yaxis.label.set_color(text_color) + + self.__update_preview() + + # update button + components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + + frame.pack(fill="both", expand=1) + return frame + + def __update_preview(self): + generator = TimestepGenerator( + timestep_distribution=self.config.timestep_distribution, + min_noising_strength=self.config.min_noising_strength, + max_noising_strength=self.config.max_noising_strength, + noising_weight=self.config.noising_weight, + noising_bias=self.config.noising_bias, + timestep_shift=self.config.timestep_shift, + ) + + self.ax.cla() + self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.canvas.draw() + + def __ok(self): + self.destroy() diff --git a/modules/ui/CtkTopBarView.py b/modules/ui/CtkTopBarView.py new file mode 100644 index 000000000..820fdb71a --- /dev/null +++ b/modules/ui/CtkTopBarView.py @@ -0,0 +1,260 @@ +import json +import os +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress + +from modules.util import path_util +from modules.util.config.SecretsConfig import SecretsConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.optimizer_util import change_optimizer +from modules.util.path_util import write_json_atomic +from modules.util.ui import components, dialogs +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class TopBar: + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + change_model_type_callback: Callable[[ModelType], None], + change_training_method_callback: Callable[[TrainingMethod], None], + load_preset_callback: Callable[[], None], + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.change_model_type_callback = change_model_type_callback + self.change_training_method_callback = change_training_method_callback + self.load_preset_callback = load_preset_callback + + self.dir = "training_presets" + + self.config_ui_data = { + "config_name": path_util.canonical_join(self.dir, "#.json") + } + self.config_ui_state = UIState(master, self.config_ui_data) + + self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] + self.__load_available_config_names() + + self.current_config = [] + + self.frame = ctk.CTkFrame(master=master, corner_radius=0) + self.frame.grid(row=0, column=0, sticky="nsew") + + self.training_method = None + + # title + components.app_title(self.frame, 0, 0) + + # dropdown + self.configs_dropdown = None + self.__create_configs_dropdown() + + # remove button + # TODO + # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + + # Wiki button + components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + + # save button + components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) + + # padding + self.frame.grid_columnconfigure(5, weight=1) + + # model type + components.options_kv( + master=self.frame, + row=0, + column=6, + values=[ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ], + ui_state=self.ui_state, + var_name="model_type", + command=self.__change_model_type, + ) + + def __create_training_method(self): + if self.training_method: + self.training_method.destroy() + + values = [] + #TODO simplify + if self.train_config.model_type.is_stable_diffusion(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), + ] + elif self.train_config.model_type.is_stable_diffusion_3() \ + or self.train_config.model_type.is_stable_diffusion_xl() \ + or self.train_config.model_type.is_wuerstchen() \ + or self.train_config.model_type.is_pixart() \ + or self.train_config.model_type.is_flux_1() \ + or self.train_config.model_type.is_sana() \ + or self.train_config.model_type.is_hunyuan_video() \ + or self.train_config.model_type.is_hi_dream() \ + or self.train_config.model_type.is_chroma(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ] + elif self.train_config.model_type.is_qwen() \ + or self.train_config.model_type.is_z_image() \ + or self.train_config.model_type.is_flux_2() \ + or self.train_config.model_type.is_ernie(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ] + + # training method + self.training_method = components.options_kv( + master=self.frame, + row=0, + column=7, + values=values, + ui_state=self.ui_state, + var_name="training_method", + command=self.change_training_method_callback, + ) + + def __change_model_type(self, model_type: ModelType): + self.change_model_type_callback(model_type) + self.__create_training_method() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.grid_forget() + + self.configs_dropdown = components.options_kv( + self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config + ) + + def __load_available_config_names(self): + if os.path.isdir(self.dir): + for path in os.listdir(self.dir): + if path != "#.json": + path = path_util.canonical_join(self.dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + self.configs.sort() + + def __save_to_file(self, name) -> str: + name = path_util.safe_filename(name) + path = path_util.canonical_join("training_presets", f"{name}.json") + + write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) + + return path + + def __save_secrets(self, path) -> str: + write_json_atomic(path, self.train_config.secrets.to_dict()) + return path + + def open_wiki(self): + webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + + def __save_new_config(self, name): + path = self.__save_to_file(name) + + is_new_config = name not in [x[0] for x in self.configs] + + if is_new_config: + self.configs.append((name, path)) + self.configs.sort() + + if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): + self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) + + if is_new_config: + self.__create_configs_dropdown() + + def __save_config(self): + default_value = self.configs_dropdown.get() + while default_value.startswith('#'): + default_value = default_value[1:] + + dialogs.StringInputDialog( + parent=self.master, + title="name", + question="Config Name", + callback=self.__save_new_config, + default_value=default_value, + validate_callback=lambda x: not x.startswith("#") + ) + + def __load_current_config(self, filename): + try: + basename = os.path.basename(filename) + is_built_in_preset = basename.startswith("#") and basename != "#.json" + + with open(filename, "r") as f: + loaded_dict = json.load(f) + default_config = TrainConfig.default_values() + if is_built_in_preset: + # always assume built-in configs are saved in the most recent version + loaded_dict["__version"] = default_config.config_version + loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() + + with suppress(FileNotFoundError), open("secrets.json", "r") as f: + secrets_dict=json.load(f) + loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) + + self.train_config.from_dict(loaded_config.to_dict()) + self.ui_state.update(loaded_config) + + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.load_preset_callback() + except FileNotFoundError: + pass + except Exception: + print(traceback.format_exc()) + + def __remove_config(self): + # TODO + pass + + def save_default(self): + self.__save_to_file("#") + self.__save_secrets("secrets.json") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py new file mode 100644 index 000000000..ba90d2e64 --- /dev/null +++ b/modules/ui/CtkTrainUIView.py @@ -0,0 +1,889 @@ +import ctypes +import datetime +import json +import os +import platform +import subprocess +import sys +import threading +import time +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from tkinter import filedialog, messagebox + +import scripts.generate_debug_report +from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab +from modules.ui.CaptionUI import CaptionUI +from modules.ui.CloudTab import CloudTab +from modules.ui.ConceptTab import ConceptTab +from modules.ui.ConvertModelUI import ConvertModelUI +from modules.ui.LoraTab import LoraTab +from modules.ui.ModelTab import ModelTab +from modules.ui.ProfilingWindow import ProfilingWindow +from modules.ui.SampleWindow import SampleWindow +from modules.ui.SamplingTab import SamplingTab +from modules.ui.TopBar import TopBar +from modules.ui.TrainingTab import TrainingTab +from modules.ui.VideoToolUI import VideoToolUI +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.GradientReducePrecision import GradientReducePrecision +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import flush_and_validate_all + +import torch + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker + +# chunk for forcing Windows to ignore DPI scaling when moving between monitors +# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 +if platform.system() == "Windows": + with suppress(Exception): + # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically + ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE + +class TrainUI(ctk.CTk): + set_step_progress: Callable[[int, int], None] + set_epoch_progress: Callable[[int, int], None] + + status_label: ctk.CTkLabel | None + training_button: ctk.CTkButton | None + training_callbacks: TrainCallbacks | None + training_commands: TrainCommands | None + + _TRAIN_BUTTON_STYLES = { + "idle": { + "text": "Start Training", + "state": "normal", + "fg_color": "#198754", + "hover_color": "#146c43", + "text_color": "white", + "text_color_disabled": "white", + }, + "running": { + "text": "Stop Training", + "state": "normal", + "fg_color": "#dc3545", + "hover_color": "#bb2d3b", + "text_color": "white", + }, + "stopping": { + "text": "Stopping...", + "state": "disabled", + "fg_color": "#dc3545", + "hover_color": "#dc3545", + "text_color": "white", + "text_color_disabled": "white", + }, + } + + def __init__(self): + super().__init__() + + self.title("OneTrainer") + self.geometry("1100x740") + + self.after(100, lambda: self._set_icon()) + + # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration + ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") + ctk.set_default_color_theme("blue") + + self.train_config = TrainConfig.default_values() + self.ui_state = UIState(self, self.train_config) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = None + self.tabview = None + + self.model_tab = None + self.training_tab = None + self.lora_tab = None + self.cloud_tab = None + self.additional_embeddings_tab = None + + self.top_bar_component = self.top_bar(self) + self.content_frame(self) + self.bottom_bar(self) + + self.training_thread = None + self.training_callbacks = None + self.training_commands = None + + self.always_on_tensorboard_subprocess = None + self.current_workspace_dir = self.train_config.workspace_dir + self._check_start_always_on_tensorboard() + + self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + + # Persistent profiling window. + self.profiling_window = ProfilingWindow(self) + + self.protocol("WM_DELETE_WINDOW", self.__close) + + def __close(self): + self.top_bar_component.save_default() + self._stop_always_on_tensorboard() + if hasattr(self, 'workspace_dir_trace_id'): + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + self.quit() + + def top_bar(self, master): + return TopBar( + master, + self.train_config, + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + + def bottom_bar(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=2, column=0, sticky="nsew") + + self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + + # status + ETA container + self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + self.status_frame.grid(row=0, column=1, sticky="w") + self.status_frame.grid_rowconfigure(0, weight=0) + self.status_frame.grid_rowconfigure(1, weight=0) + self.status_frame.grid_columnconfigure(0, weight=1) + + self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + + # padding + frame.grid_columnconfigure(2, weight=1) + + + # export button + self.export_button = components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + # debug button + components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + # tensorboard button + components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + # training button + self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + self._set_training_button_style("idle") # centralized styling + + return frame + + def content_frame(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_columnconfigure(0, weight=1) + + self.tabview = ctk.CTkTabview(frame) + self.tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.create_general_tab(self.tabview.add("general")) + self.model_tab = self.create_model_tab(self.tabview.add("model")) + self.data_tab = self.create_data_tab(self.tabview.add("data")) + self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) + self.training_tab = self.create_training_tab(self.tabview.add("training")) + self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) + self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) + self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) + self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) + self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) + + self.change_training_method(self.train_config.training_method) + + return frame + + def create_general_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # workspace dir + components.label(frame, 0, 0, "Workspace Directory", + tooltip="The directory where all files of this training run are saved") + components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + + # cache dir + components.label(frame, 0, 2, "Cache Directory", + tooltip="The directory where cached data is saved") + components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + + # continue from previous backup + components.label(frame, 2, 0, "Continue from last backup", + tooltip="Automatically continues training from the last backup saved in /backup") + components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + + # only cache + components.label(frame, 2, 2, "Only Cache", + tooltip="Only populate the cache, without any training") + components.switch(frame, 2, 3, self.ui_state, "only_cache") + + # TODO: In Phase 4 rework the general tab. + # prevent overwrites + components.label(frame, 3, 0, "Prevent Overwrites", + tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") + components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + + # debug + components.label(frame, 4, 0, "Debug mode", + tooltip="Save debug information during the training into the debug directory") + components.switch(frame, 4, 1, self.ui_state, "debug_mode") + + components.label(frame, 4, 2, "Debug Directory", + tooltip="The directory where debug data is saved") + components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + + # tensorboard + components.label(frame, 6, 0, "Tensorboard", + tooltip="Starts the Tensorboard Web UI during training") + components.switch(frame, 6, 1, self.ui_state, "tensorboard") + + components.label(frame, 6, 2, "Always-On Tensorboard", + tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") + components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) + + components.label(frame, 7, 0, "Expose Tensorboard", + tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") + components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") + components.label(frame, 7, 2, "Tensorboard Port", + tooltip="Port to use for Tensorboard link") + components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") + + + # validation + components.label(frame, 8, 0, "Validation", + tooltip="Enable validation steps and add new graph in tensorboard") + components.switch(frame, 8, 1, self.ui_state, "validation") + + components.label(frame, 8, 2, "Validate after", + tooltip="The interval used when validate training") + components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + + # device + components.label(frame, 10, 0, "Dataloader Threads", + tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") + components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + + components.label(frame, 11, 0, "Train Device", + tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") + components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + + components.label(frame, 12, 0, "Multi-GPU", + tooltip="Enable multi-GPU training") + components.switch(frame, 12, 1, self.ui_state, "multi_gpu") + components.label(frame, 12, 2, "Device Indexes", + tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") + components.entry(frame, 12, 3, self.ui_state, "device_indexes") + + components.label(frame, 13, 0, "Gradient Reduce Precision", + tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" + "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" + "FLOAT_32: Reduce gradients in float32\n" + "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", + wide_tooltip=True) + components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + "gradient_reduce_precision") + + components.label(frame, 13, 2, "Fused Gradient Reduce", + tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") + components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") + + components.label(frame, 14, 0, "Async Gradient Reduce", + tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") + components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") + components.label(frame, 14, 2, "Buffer size (MB)", + tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") + components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") + + components.label(frame, 15, 0, "Temp Device", + tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") + components.entry(frame, 15, 1, self.ui_state, "temp_device") + + frame.pack(fill="both", expand=1) + return frame + + def create_model_tab(self, master): + return ModelTab(master, self.train_config, self.ui_state) + + def create_data_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # aspect ratio bucketing + components.label(frame, 0, 0, "Aspect Ratio Bucketing", + tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") + components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + + # latent caching + components.label(frame, 1, 0, "Latent Caching", + tooltip="Caching of intermediate training data that can be re-used between epochs") + components.switch(frame, 1, 1, self.ui_state, "latent_caching") + + # clear cache before training + components.label(frame, 2, 0, "Clear cache before training", + tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") + components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + + frame.pack(fill="both", expand=1) + return frame + + def create_concepts_tab(self, master): + return ConceptTab(master, self.train_config, self.ui_state) + + def create_training_tab(self, master) -> TrainingTab: + return TrainingTab(master, self.train_config, self.ui_state) + + def create_cloud_tab(self, master) -> CloudTab: + return CloudTab(master, self.train_config, self.ui_state,parent=self) + + def create_sampling_tab(self, master): + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + # sample after + top_frame = ctk.CTkFrame(master=master, corner_radius=0) + top_frame.grid(row=0, column=0, sticky="nsew") + sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") + sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + + components.label(top_frame, 0, 0, "Sample After", + tooltip="The interval used when automatically sampling from the model during training") + components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + + components.label(top_frame, 0, 2, "Skip First", + tooltip="Start sampling automatically after this interval has elapsed.") + components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + + components.label(top_frame, 0, 4, "Format", + tooltip="File Format used when saving samples") + components.options_kv(top_frame, 0, 5, [ + ("PNG", ImageFormat.PNG), + ("JPG", ImageFormat.JPG), + ], self.ui_state, "sample_image_format") + + components.button(top_frame, 0, 6, "sample now", self.sample_now) + + components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + + components.label(sub_frame, 0, 0, "Non-EMA Sampling", + tooltip="Whether to include non-ema sampling when using ema.") + components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + + components.label(sub_frame, 0, 2, "Samples to Tensorboard", + tooltip="Whether to include sample images in the Tensorboard output.") + components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + + # table + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + return SamplingTab(frame, self.train_config, self.ui_state) + + def create_backup_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # backup after + components.label(frame, 0, 0, "Backup After", + tooltip="The interval used when automatically creating model backups during training") + components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + + # backup now + components.button(frame, 0, 3, "backup now", self.backup_now) + + # rolling backup + components.label(frame, 1, 0, "Rolling Backup", + tooltip="If rolling backups are enabled, older backups are deleted automatically") + components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + + # rolling backup count + components.label(frame, 1, 3, "Rolling Backup Count", + tooltip="Defines the number of backups to keep if rolling backups are enabled") + components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + + # backup before save + components.label(frame, 2, 0, "Backup Before Save", + tooltip="Create a full backup before saving the final model") + components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + + # save after + components.label(frame, 3, 0, "Save Every", + tooltip="The interval used when automatically saving the model during training") + components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + + # save now + components.button(frame, 3, 3, "save now", self.save_now) + + # skip save + components.label(frame, 4, 0, "Skip First", + tooltip="Start saving automatically after this interval has elapsed") + components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + + # save filename prefix + components.label(frame, 5, 0, "Save Filename Prefix", + tooltip="The prefix for filenames used when saving the model during training") + components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + + frame.pack(fill="both", expand=1) + return frame + + def embedding_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # embedding model name + components.label(frame, 0, 0, "Base embedding", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + components.path_entry( + frame, 0, 1, self.ui_state, "embedding.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # token count + components.label(frame, 1, 0, "Token count", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + + # initial embedding text + components.label(frame, 2, 0, "Initial embedding text", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + + # embedding weight dtype + components.label(frame, 3, 0, "Embedding Weight Data Type", + tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "embedding_weight_dtype") + + # placeholder + components.label(frame, 4, 0, "Placeholder", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + + # output embedding + components.label(frame, 5, 0, "Output embedding", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") + + frame.pack(fill="both", expand=1) + return frame + + def create_additional_embeddings_tab(self, master): + return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + + def create_tools_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # dataset + components.label(frame, 0, 0, "Dataset Tools", + tooltip="Open the captioning tool") + components.button(frame, 0, 1, "Open", self.open_dataset_tool) + + # video tools + components.label(frame, 1, 0, "Video Tools", + tooltip="Open the video tools") + components.button(frame, 1, 1, "Open", self.open_video_tool) + + # convert model + components.label(frame, 2, 0, "Convert Model Tools", + tooltip="Open the model conversion tool") + components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + + # sample + components.label(frame, 3, 0, "Sampling Tool", + tooltip="Open the model sampling tool") + components.button(frame, 3, 1, "Open", self.open_sampling_tool) + + components.label(frame, 4, 0, "Profiling Tool", + tooltip="Open the profiling tools.") + components.button(frame, 4, 1, "Open", self.open_profiling_tool) + + frame.pack(fill="both", expand=1) + return frame + + def change_model_type(self, model_type: ModelType): + if self.model_tab: + self.model_tab.refresh_ui() + + if self.training_tab: + self.training_tab.refresh_ui() + + if self.lora_tab: + self.lora_tab.refresh_ui() + + def change_training_method(self, training_method: TrainingMethod): + if not self.tabview: + return + + if self.model_tab: + self.model_tab.refresh_ui() + + if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: + self.tabview.delete("LoRA") + self.lora_tab = None + if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: + self.tabview.delete("embedding") + + if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: + self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: + self.embedding_tab(self.tabview.add("embedding")) + + def load_preset(self): + if not self.tabview: + return + + if self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + spent_total = time.monotonic() - self.start_time + steps_done = train_progress.epoch * max_step + train_progress.epoch_step + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) + total_eta = spent_total / steps_done * remaining_steps + + if train_progress.global_step <= 30: + return "Estimating ..." + + td = datetime.timedelta(seconds=total_eta) + days = td.days + hours, remainder = divmod(td.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + if days > 0: + return f"{days}d {hours}h" + elif hours > 0: + return f"{hours}h {minutes}m" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + + def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def delete_eta_label(self): + self.eta_label.configure(text="") + + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + self.set_step_progress(train_progress.epoch_step, max_step) + self.set_epoch_progress(train_progress.epoch, max_epoch) + self.set_eta_label(train_progress, max_step, max_epoch) + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def open_dataset_tool(self): + window = CaptionUI(self, None, False) + self.wait_window(window) + + def open_video_tool(self): + window = VideoToolUI(self) + self.wait_window(window) + + def open_convert_model_tool(self): + window = ConvertModelUI(self) + self.wait_window(window) + + def open_sampling_tool(self): + if not self.training_callbacks and not self.training_commands: + window = SampleWindow( + self, + use_external_model=False, + train_config=self.train_config, + ) + self.wait_window(window) + torch_gc() + + def open_profiling_tool(self): + self.profiling_window.deiconify() + + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + + if not zip_path: + return + + zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + + self.on_update_status("Generating debug package...") + + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.on_update_status(f"Error generating debug package: {e}") + + + def open_manual_sample_window (self): + training_callbacks = self.training_callbacks + training_commands = self.training_commands + + if training_callbacks and training_commands: + window = SampleWindow( + self, + train_config=self.train_config, + use_external_model=True, + callbacks=training_callbacks, + commands=training_commands, + ) + self.wait_window(window) + training_callbacks.set_on_sample_custom() + + def __training_thread_function(self): + error_caught = False + + self.training_callbacks = TrainCallbacks( + on_update_train_progress=self.on_update_train_progress, + on_update_status=self.on_update_status, + ) + + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + try: + trainer.start() + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + + self.start_time = time.monotonic() + trainer.train() + except Exception: + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + error_caught = True + traceback.print_exc() + + trainer.end() + + # clear gpu memory + del trainer + + self.training_thread = None + self.training_commands = None + torch.clear_autocast_cache() + torch_gc() + + if error_caught: + self.on_update_status("Error: check the console for details") + else: + self.on_update_status("Stopped") + self.delete_eta_label() + + # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash + self.after(0, self._set_training_button_idle) + + if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + self.after(0, self._start_always_on_tensorboard) + + def start_training(self): + if self.training_thread is None: + self.save_default() + + # --- pre-training validation gate --- + errors = flush_and_validate_all() + + if errors: + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + return + + self._set_training_button_running() + + if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() + + self.training_commands = TrainCommands() + torch_gc() + + self.training_thread = threading.Thread(target=self.__training_thread_function) + self.training_thread.start() + else: + self._set_training_button_stopping() + self.on_update_status("Stopping ...") + self.training_commands.stop() + + def save_default(self): + self.top_bar_component.save_default() + self.concepts_tab.save_current_config() + self.sampling_tab.save_current_config() + self.additional_embeddings_tab.save_current_config() + + def export_training(self): + file_path = filedialog.asksaveasfilename(filetypes=[ + ("All Files", "*.*"), + ("json", "*.json"), + ], initialdir=".", initialfile="config.json") + + if file_path: + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.save() + + def _check_start_always_on_tensorboard(self): + if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _start_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() + + tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") + tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") + + os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) + + tensorboard_args = [ + tensorboard_executable, + "--logdir", + tensorboard_log_dir, + "--port", + str(self.train_config.tensorboard_port), + "--samples_per_plugin=images=100,scalars=10000", + ] + + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") + + try: + self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) + except Exception: + self.always_on_tensorboard_subprocess = None + + def _stop_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + try: + self.always_on_tensorboard_subprocess.terminate() + self.always_on_tensorboard_subprocess.wait(timeout=5) + except subprocess.TimeoutExpired: + self.always_on_tensorboard_subprocess.kill() + except Exception: + pass + finally: + self.always_on_tensorboard_subprocess = None + + def _on_workspace_dir_change(self, new_workspace_dir: str): + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir + + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _on_workspace_dir_change_trace(self, *args): + new_workspace_dir = self.train_config.workspace_dir + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir + + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _on_always_on_tensorboard_toggle(self): + if self.train_config.tensorboard_always_on: + if not (self.training_thread and self.train_config.tensorboard): + self._start_always_on_tensorboard() + else: + if not (self.training_thread and self.train_config.tensorboard): + self._stop_always_on_tensorboard() + + def _set_training_button_style(self, mode: str): + if not self.training_button: + return + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) + + def _set_training_button_idle(self): + self._set_training_button_style("idle") + + def _set_training_button_running(self): + self._set_training_button_style("running") + + def _set_training_button_stopping(self): + self._set_training_button_style("stopping") diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py new file mode 100644 index 000000000..bcca11ae9 --- /dev/null +++ b/modules/ui/CtkTrainingTabView.py @@ -0,0 +1,856 @@ +from modules.ui.OffloadingWindow import OffloadingWindow +from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow +from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow +from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod +from modules.util.enum.LearningRateScaler import LearningRateScaler +from modules.util.enum.LearningRateScheduler import LearningRateScheduler +from modules.util.enum.LossScaler import LossScaler +from modules.util.enum.LossWeight import LossWeight +from modules.util.enum.Optimizer import Optimizer +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.optimizer_util import change_optimizer +from modules.util.ui import components +from modules.util.ui.UIState import UIState +from modules.util.ui.validation_helpers import check_range, validate_resolution + +import customtkinter as ctk + + +class TrainingTab: + + def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + super().__init__() + + self.master = master + self.train_config = train_config + self.ui_state = ui_state + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.scroll_frame = None + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + + self.scroll_frame.grid_columnconfigure(0, weight=1) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=1) + + column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_0.grid(row=0, column=0, sticky="nsew") + column_0.grid_columnconfigure(0, weight=1) + + column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_1.grid(row=0, column=1, sticky="nsew") + column_1.grid_columnconfigure(0, weight=1) + + column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_2.grid(row=0, column=2, sticky="nsew") + column_2.grid_columnconfigure(0, weight=1) + + if self.train_config.model_type.is_stable_diffusion(): + self.__setup_stable_diffusion_ui(column_0, column_1, column_2) + if self.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_sana(): + self.__setup_sana_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(column_0, column_1, column_2) + + + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1) + self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1) + self.__create_text_encoder_n_frame(column_0, 2, i=2) + self.__create_embedding_frame(column_0, 3) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1) + self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_wuerstchen_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0, supports_circular_padding=True) + self.__create_prior_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 0) + self.__create_loss_frame(column_2, 1) + self.__create_layer_frame(column_2, 2) + + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2, supports_vb_loss=True) + self.__create_layer_frame(column_2, 3) + + def __setup_flux_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_flux_2_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_chroma_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_qwen_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_z_image_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_ernie_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_sana_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1) + self.__create_embedding_frame(column_0, 2) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_embedding_frame(column_0, 4) + + self.__create_base2_frame(column_1, 0, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __setup_hi_dream_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) + self.__create_embedding_frame(column_0, 5) + + self.__create_base2_frame(column_1, 0, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1) + self.__create_noise_frame(column_1, 2) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + + def __create_base_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # optimizer + components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") + components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", + command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) + + # learning rate scheduler + # Wackiness will ensue when reloading configs if we don't check and clear this first. + if hasattr(self, "lr_scheduler_comp"): + delattr(self, "lr_scheduler_comp") + delattr(self, "lr_scheduler_adv_comp") + components.label(frame, 1, 0, "Learning Rate Scheduler", + tooltip="Learning rate scheduler that automatically changes the learning rate during training") + _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, + "learning_rate_scheduler", command=self.__restore_scheduler_config, + adv_command=self.__open_scheduler_params_window) + self.lr_scheduler_comp = d['component'] + self.lr_scheduler_adv_comp = d['button_component'] + # Initial call requires the presence of self.lr_scheduler_adv_comp. + self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) + + # learning rate + components.label(frame, 2, 0, "Learning Rate", + tooltip="The base learning rate") + components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) + + # learning rate warmup steps + components.label(frame, 3, 0, "Learning Rate Warmup Steps", + tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") + components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") + + # learning rate min factor + components.label(frame, 4, 0, "Learning Rate Min Factor", + tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") + components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", + extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) + + # learning rate cycles + components.label(frame, 5, 0, "Learning Rate Cycles", + tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") + components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") + + # epochs + components.label(frame, 6, 0, "Epochs", + tooltip="The number of epochs for a full training run") + components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) + + # batch size + components.label(frame, 7, 0, "Local Batch Size", + tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") + components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) + + # accumulation steps + components.label(frame, 8, 0, "Accumulation Steps", + tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") + components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) + + # Learning Rate Scaler + components.label(frame, 9, 0, "Learning Rate Scaler", + tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") + components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, + "learning_rate_scaler") + + # clip grad norm + components.label(frame, 10, 0, "Clip Grad Norm", + tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") + components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") + + def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + # ema + components.label(frame, row, 0, "EMA", + tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") + components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") + row += 1 + + # ema decay + components.label(frame, row, 0, "EMA Decay", + tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") + components.entry(frame, row, 1, self.ui_state, "ema_decay", + extra_validate=check_range(lower=0.5, upper=1, + message="EMA decay must be between 0.5 and 1")) + row += 1 + + # ema update step interval + components.label(frame, row, 0, "EMA Update Step Interval", + tooltip="Number of steps between EMA update steps") + components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") + row += 1 + + # gradient checkpointing + components.label(frame, row, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, + "gradient_checkpointing", adv_command=self.__open_offloading_window) + row += 1 + + # gradient checkpointing layer offloading + components.label(frame, row, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") + row += 1 + + # train dtype + components.label(frame, row, 0, "Train Data Type", + tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") + components.options_kv(frame, row, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ("tfloat32", DataType.TFLOAT_32), + ], self.ui_state, "train_dtype") + row += 1 + + # fallback train dtype + components.label(frame, row, 0, "Fallback Train Data Type", + tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") + components.options_kv(frame, row, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "fallback_train_dtype") + row += 1 + + # autocast cache + components.label(frame, row, 0, "Autocast Cache", + tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") + components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") + row += 1 + + # resolution + components.label(frame, row, 0, "Resolution", + tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + components.entry(frame, row, 1, self.ui_state, "resolution", required=True, + extra_validate=validate_resolution()) + row += 1 + + # frames + if video_training_enabled: + components.label(frame, row, 0, "Frames", + tooltip="The number of frames used for training.") + components.entry(frame, row, 1, self.ui_state, "frames", required=True) + row += 1 + + # force circular padding + if supports_circular_padding: + components.label(frame, row, 0, "Force Circular Padding", + tooltip="Enables circular padding for all conv layers to better train seamless images") + components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + if supports_training: + components.label(frame, 0, 0, "Train Text Encoder", + tooltip="Enables training the text encoder model") + components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + + # dropout + components.label(frame, 1, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + + if supports_training: + # train text encoder epochs + components.label(frame, 2, 0, "Stop Training After", + tooltip="When to stop training the text encoder") + components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + "text_encoder.stop_training_after_unit", supports_time_units=False) + + # text encoder learning rate + components.label(frame, 3, 0, "Text Encoder Learning Rate", + tooltip="The learning rate of the text encoder. Overrides the base learning rate") + components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + + if supports_clip_skip: + # text encoder layer skip (clip skip) + components.label(frame, 4, 0, "Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + + if supports_sequence_length: + # text encoder sequence length + components.label(frame, row, 0, "Text Encoder Sequence Length", + tooltip="Number of tokens for captions") + components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") + row += 1 + + def __create_text_encoder_n_frame( + self, + master, + row: int, + i: int, + supports_include: bool = False, + supports_layer_skip: bool = True, + supports_sequence_length: bool = False, + ): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + suffix = f"_{i}" if i > 1 else "" + + if supports_include: + # include text encoder + components.label(frame, row, 0, f"Include Text Encoder {i}", + tooltip=f"Includes text encoder {i} in the training run") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") + row += 1 + + # train text encoder + components.label(frame, row, 0, f"Train Text Encoder {i}", + tooltip=f"Enables training the text encoder {i} model") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") + row += 1 + + # train text encoder embedding + components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", + tooltip=f"Enables training embeddings for the text encoder {i} model") + components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") + row += 1 + + # dropout + components.label(frame, row, 0, "Dropout Probability", + tooltip=f"The Probability for dropping the text encoder {i} conditioning") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") + row += 1 + + # train text encoder epochs + components.label(frame, row, 0, "Stop Training After", + tooltip=f"When to stop training the text encoder {i}") + components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", + f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) + row += 1 + + # text encoder learning rate + components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", + tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") + row += 1 + + if supports_layer_skip: + # text encoder layer skip (clip skip) + components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") + row += 1 + + if supports_sequence_length: + # text encoder sequence length + components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", + tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") + components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") + row += 1 + + def __create_embedding_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + + # embedding learning rate + components.label(frame, 0, 0, "Embeddings Learning Rate", + tooltip="The learning rate of embeddings. Overrides the base learning rate") + components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") + + # preserve embedding norm + components.label(frame, 1, 0, "Preserve Embedding Norm", + tooltip="Rescales each trained embedding to the median embedding norm") + components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") + + def __create_unet_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train unet + components.label(frame, 0, 0, "Train UNet", + tooltip="Enables training the UNet model") + components.switch(frame, 0, 1, self.ui_state, "unet.train") + + # train unet epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the UNet") + components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + supports_time_units=False) + + # unet learning rate + components.label(frame, 2, 0, "UNet Learning Rate", + tooltip="The learning rate of the UNet. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + + # rescale noise scheduler to zero terminal SNR + rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") + rescale_label.configure(wraplength=130, justify="left") + components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + + def __create_prior_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train prior + components.label(frame, 0, 0, "Train Prior", + tooltip="Enables training the Prior model") + components.switch(frame, 0, 1, self.ui_state, "prior.train") + + # train prior epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Prior") + components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", + supports_time_units=False) + + # prior learning rate + components.label(frame, 2, 0, "Prior Learning Rate", + tooltip="The learning rate of the Prior. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + + def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # train transformer + components.label(frame, 0, 0, "Train Transformer", + tooltip="Enables training the Transformer model") + components.switch(frame, 0, 1, self.ui_state, "transformer.train") + + # train transformer epochs + components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Transformer") + components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", + supports_time_units=False) + + # transformer learning rate + components.label(frame, 2, 0, "Transformer Learning Rate", + tooltip="The learning rate of the Transformer. Overrides the base learning rate") + components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + + if supports_force_attention_mask: + # transformer learning rate + components.label(frame, 3, 0, "Force Attention Mask", + tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") + components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + + if supports_guidance_scale: + # guidance scale + components.label(frame, 4, 0, "Guidance Scale", + tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") + components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + + def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # offset noise weight + components.label(frame, 0, 0, "Offset Noise Weight", + tooltip="The weight of offset noise added to each training step") + components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") + + if supports_generalized_offset_noise: + # generalized offset noise weight + generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", + tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") + generalised_offset_label.configure(wraplength=130, justify="left") + components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") + + # perturbation noise weight + components.label(frame, 2, 0, "Perturbation Noise Weight", + tooltip="The weight of perturbation noise added to each training step") + components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") + + # timestep distribution + components.label(frame, 3, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", + adv_command=self.__open_timestep_distribution_window) + + # min noising strength + components.label(frame, 4, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) + + # max noising strength + components.label(frame, 5, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) + + # noising weight + components.label(frame, 6, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) + + # noising bias + components.label(frame, 7, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) + + # timestep shift + components.label(frame, 8, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) + + if supports_dynamic_timestep_shifting: + # dynamic timestep shifting + components.label(frame, 9, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") + + + + def __create_masked_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # Masked Training + components.label(frame, 0, 0, "Masked Training", + tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") + components.switch(frame, 0, 1, self.ui_state, "masked_training") + + # unmasked probability + components.label(frame, 1, 0, "Unmasked Probability", + tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") + components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", + extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) + + # unmasked weight + components.label(frame, 2, 0, "Unmasked Weight", + tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") + components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", + extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) + + # normalize masked area loss + components.label(frame, 3, 0, "Normalize Masked Area Loss", + tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") + components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") + + # masked prior preservation + components.label(frame, 4, 0, "Masked Prior Preservation Weight", + tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") + components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", + extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) + + # use custom conditioning image + components.label(frame, 5, 0, "Custom Conditioning Image", + tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") + components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") + + def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + # MSE Strength + components.label(frame, 0, 0, "MSE Strength", + tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) + + # MAE Strength + components.label(frame, 1, 0, "MAE Strength", + tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) + + # log-cosh Strength + components.label(frame, 2, 0, "log-cosh Strength", + tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") + components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) + + # Huber Strength + components.label(frame, 3, 0, "Huber Strength", + tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") + components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) + + # Huber Delta + components.label(frame, 4, 0, "Huber Delta", + tooltip="Delta parameter for huber loss") + components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) + + if supports_vb_loss: + # VB Strength + components.label(frame, 5, 0, "VB Strength", + tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") + components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) + + # Loss Weight function + components.label(frame, 6, 0, "Loss Weight Function", + tooltip="Choice of loss weight function. Can help the model learn details more accurately.") + components.options(frame, 6, 1, [str(x) for x in list(LossWeight) + if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() + or x == LossWeight.CONSTANT + ], + self.ui_state, "loss_weight_fn") + + row = 7 + + # Loss weight strength + if not self.train_config.model_type.is_flow_matching(): + components.label(frame, row, 0, "Gamma", + tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") + components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", + extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) + row += 1 + + # Loss Scaler + components.label(frame, row, 0, "Loss Scaler", + tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") + components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") + row += 1 + + def __create_layer_frame(self, master, row): + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + components.layer_filter_entry(master, row, 0, self.ui_state, + preset_var_name="layer_filter_preset", presets=presets, + preset_label="Layer Filter", + preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", + entry_var_name="layer_filter", + entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", + regex_var_name="layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + ) + + + def __on_layer_filter_preset_change(self): + if not self.layer_selector: + return + selected = self.ui_state.get_var("layer_filter_preset").get() + self.__preset_set_layer_choice(selected) + + def __hide_layer_entry(self): + if self.layer_entry and self.layer_entry.winfo_manager(): + self.layer_entry.grid_remove() + + def __show_layer_entry(self): + if self.layer_entry and not self.layer_entry.winfo_manager(): + self.layer_entry.grid() + + def __open_optimizer_params_window(self): + window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_scheduler_params_window(self): + window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_timestep_distribution_window(self): + window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __open_offloading_window(self): + window = OffloadingWindow(self.master, self.train_config, self.ui_state) + self.master.wait_window(window) + + def __restore_optimizer_config(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def __restore_scheduler_config(self, variable): + if not hasattr(self, 'lr_scheduler_adv_comp'): + return + + if variable == "CUSTOM": + self.lr_scheduler_adv_comp.configure(state="normal") + else: + self.lr_scheduler_adv_comp.configure(state="disabled") diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py new file mode 100644 index 000000000..c3291e6ea --- /dev/null +++ b/modules/ui/CtkVideoToolUIView.py @@ -0,0 +1,877 @@ +import concurrent.futures +import math +import os +import pathlib +import random +import shlex +import subprocess +import threading +import webbrowser +from fractions import Fraction +from tkinter import filedialog + +from modules.util.image_util import load_image +from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS +from modules.util.ui import components + +import av +import customtkinter as ctk +import cv2 +import scenedetect +from PIL import Image + + +class VideoToolUI(ctk.CTkToplevel): + def __init__( + self, + parent, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + self.title("Video Tools") + self.geometry("600x720") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) + self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) + self.video_download_tab = self.__video_download_tab(tabview.add("download")) + self.status_bar(self) + + def status_bar(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=1, column=0) + frame.grid_columnconfigure(0, weight=0, minsize=160) + frame.grid_columnconfigure(1, weight=0, minsize=300) + frame.grid_columnconfigure(2, weight=1) + + #create preview image + preview_path = "resources/icons/icon.png" + preview = load_image(preview_path, 'RGB') + preview.thumbnail((150, 150)) + self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image_label = ctk.CTkLabel( + master=frame, text="Preview image", image=self.preview_image, height=150, width=150, + compound="top") + self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) + + #displays progress and messages that also go to terminal + self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) + self.status_label.insert(index="1.0", text="Current status") + self.status_label.configure(state="disabled") + self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + + def __clip_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.clip_single_entry = ctk.CTkEntry(frame, width=190) + self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.clip_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_clips_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.clip_time_start_entry.insert(0, "00:00:00") + self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.clip_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.clip_list_entry = ctk.CTkEntry(frame, width=190) + self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_list_entry)) + self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_clips_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted clips will be saved.") + self.clip_output_entry = ctk.CTkEntry(frame, width=190) + self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_output_entry)) + self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_clip = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") + self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # split at cuts + self.split_at_cuts = ctk.BooleanVar(self, False) + components.label(frame, 5, 0, "Split at Cuts", + tooltip="If enabled, detect cuts in the input video and split at those points. \ + Otherwise will split at any point, and clips may contain cuts.") + self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") + self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + # maximum length + components.label(frame, 6, 0, "Max Length (s)", + tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") + self.clip_length_entry = ctk.CTkEntry(frame, width=220) + self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.clip_length_entry.insert(0, "3") + + # Set FPS + components.label(frame, 7, 0, "Set FPS", + tooltip="FPS to convert output videos to, set to 0 to keep original rate.") + self.clip_fps_entry = ctk.CTkEntry(frame, width=220) + self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.clip_fps_entry.insert(0, "24.0") + + # Remove borders + self.clip_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 8, 0, "Remove Borders", + tooltip="Remove black borders from output clip") + self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") + self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 9, 0, "Crop Variation", + tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square videos. Set to 0 to use only base aspect.") + self.clip_crop_entry = ctk.CTkEntry(frame, width=220) + self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) + self.clip_crop_entry.insert(0, "0.2") + + # object filter - currently unused, may implement in future + # components.label(frame, 9, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __image_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.image_single_entry = ctk.CTkEntry(frame, width=190) + self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.image_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.image_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_images_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.image_time_start_entry = ctk.CTkEntry(frame, width=100) + self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.image_time_start_entry.insert(0, "00:00:00") + self.image_time_end_entry = ctk.CTkEntry(frame, width=100) + self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.image_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.image_list_entry = ctk.CTkEntry(frame, width=190) + self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.image_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_list_entry)) + self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_images_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted images will be saved.") + self.image_output_entry = ctk.CTkEntry(frame, width=190) + self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.image_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_output_entry)) + self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_img = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") + self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # image capture rate + components.label(frame, 5, 0, "Images/sec", + tooltip="Number of images to capture per second of video. \ + Images will be taken at semi-random frames around the specified frequency.") + self.capture_rate_entry = ctk.CTkEntry(frame, width=220) + self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.capture_rate_entry.insert(0, "0.5") + + # blur removal + components.label(frame, 6, 0, "Blur Removal", + tooltip="Threshold for removal of blurry images, relative to all others. \ + For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") + self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) + self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.blur_threshold_entry.insert(0, "0.2") + + # Remove borders + self.image_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 7, 0, "Remove Borders", + tooltip="Remove black borders from output image") + self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") + self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 8, 0, "Crop Variation", + tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square images. Set to 0 to use only base sapect.") + self.image_crop_entry = ctk.CTkEntry(frame, width=220) + self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.image_crop_entry.insert(0, "0.2") + + # # object filter - currently unused, may implement in future + # components.label(frame, 5, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __video_download_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # link + components.label(frame, 0, 0, "Single Link", + tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") + self.download_link_entry = ctk.CTkEntry(frame, width=220) + self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + + # link list + components.label(frame, 1, 0, "Link List", + tooltip="Path to txt file with list of links separated by newlines.") + self.download_list_entry = ctk.CTkEntry(frame, width=190) + self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.download_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) + self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + + # output directory + components.label(frame, 2, 0, "Output", + tooltip="Path to folder where downloaded videos will be saved.") + self.download_output_entry = ctk.CTkEntry(frame, width=190) + self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) + self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + + # additional args + components.label(frame, 3, 0, "Additional Args", + tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ + Default args will hide most terminal outputs.") + self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) + self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) + self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress --format mp4") + components.button(frame, 3, 2, "yt-dlp info", + command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) + + frame.pack(fill="both", expand=1) + return frame + + def __browse_for_dir(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __browse_for_file(self, entry_box, filetypes): + # get the path from the user + path = filedialog.askopenfilename(filetypes=filetypes) + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): + input_videos = [] + if not batch_mode: + path = pathlib.Path(input_path_single) + if path.is_file(): + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + return [path] + else: + self.__update_status("Invalid video file!") + return [] + else: + self.__update_status("No file specified, or invalid file path!") + return [] + else: + input_videos = [] + if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": + self.__update_status("Invalid input directory!") + return [] + # Only traverse supported extensions to avoid opening every file. + lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} + for path in pathlib.Path(input_path_dir).rglob("*"): + if path.is_file() and path.suffix.lower() in lower_exts: + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + input_videos.append(path) + self.__update_status(f'Found {len(input_videos)} videos to process') + return input_videos + + def __run_in_thread(self, target, *args): + """Clear status box and run target function in a daemon thread.""" + self.status_label.configure(state="normal") + self.status_label.delete(index1="1.0", index2="end") + self.status_label.configure(state="disabled") + t = threading.Thread(target=target, args=args) + t.daemon = True + t.start() + + @staticmethod + def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: + return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) + + def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: + fps = video.get(cv2.CAP_PROP_FPS) or 0.0 + if fps <= 0: + self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') + return 30.0 + return fps + + @staticmethod + def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, + video_path, input_dir: str) -> str: + if use_subdir and batch_mode: + return os.path.join(output_entry, + os.path.splitext(os.path.relpath(video_path, input_dir))[0]) + elif use_subdir: + return os.path.join(output_entry, + os.path.splitext(os.path.basename(video_path))[0]) + return output_entry + + def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: + # Return original dimensions and no offset if variation is zero + if variation == 0: + return 0, height, 0, width + + old_aspect = height/width + variation_scaled = old_aspect*variation + if old_aspect > 1.2: #tall image + new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) + elif old_aspect < 0.85: #wide image + new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) + else: #square image + new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) + + new_aspect = round(new_aspect, 2) + #keep the height the same if reducing width, and vice versa + if new_aspect > old_aspect: + new_height = int(height) + new_width = int(width*(old_aspect/new_aspect)) + elif new_aspect < old_aspect: + new_height = int(height*(new_aspect/old_aspect)) + new_width = int(width) + else: + new_height = int(height) + new_width = int(width) + + #random offset in dimension that was cropped + position_x = random.randint(0, width-new_width) + position_y = random.randint(0, height-new_height) + return position_y, new_height, position_x, new_width + + def find_main_contour(self, frame): + #outline image to find main content and exclude black bars often present on letterboxed videos + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) + frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if frame_contours: + #select largest contour by area + frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) + x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) + else: #fallback if no contours detected + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + #if bounding box did not detect the correct area, likely due to all-black frame + if not frame_contours or h1 < 10 or w1 < 10: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + return x1, y1, w1, h1 + + def __extract_clips_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_clips_multi, batch_mode) + + def __extract_clips_multi(self, batch_mode: bool): + if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + max_length = float(self.clip_length_entry.get()) + crop_variation = float(self.clip_crop_entry.get()) + target_fps = float(self.clip_fps_entry.get()) + input_single_entry = self.clip_single_entry.get() + input_multiple_entry = self.clip_list_entry.get() + output_entry = self.clip_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") + return + if max_length <= 0.25: + self.__update_status("Max Length of clips must be > 0.25 seconds.") + return + if target_fps < 0: + self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if len(input_videos) == 0: # exit if no paths found + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_clip_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + executor.submit(self.__extract_clips, + str(video_path), time_start, time_end, max_length, + self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + crop_variation, target_fps, output_directory) + + if batch_mode: + self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + else: + self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + + def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, + split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + max_length_frames = int(max_length * vid_fps) + min_length_frames = max(int(0.25 * vid_fps), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + + if split_at_cuts: + #use scenedetect to find cuts, based on start/end frame number + self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') + timecode_list = scenedetect.detect( + video_path=str(video_path), + detector=scenedetect.AdaptiveDetector(), + start_time=int(timestamp_min_frame), + end_time=int(timestamp_max_frame)) + scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] + if not scene_list: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + else: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + + scene_list_split = [] + for scene in scene_list: + length = scene[1]-scene[0] + if length > max_length_frames: #check for any scenes longer than max length + n = math.ceil(length/max_length_frames) #divide into n new scenes + new_length = int(length/n) + new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length + for i, _n in enumerate(new_splits[:-1]): + if new_splits[i + 1] - new_splits[i] > min_length_frames: + scene_list_split.append((new_splits[i], new_splits[i + 1])) + elif length > (min_length_frames + 2): + # Trim first/last frame to avoid transition artifacts + scene_list_split.append((scene[0] + 1, scene[1] - 1)) + + self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(self.__save_clip, scene, video_path, target_fps, + remove_borders, crop_variation, output_dir) + for scene in scene_list_split + ] + for future in concurrent.futures.as_completed(futures): + exc = future.exception() + if exc is not None: + self.__update_status(f'Error saving clip: {exc}') + + video.release() + + def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, + remove_borders: bool, crop_variation: float, output_dir: str): + basename, ext = os.path.splitext(os.path.basename(video_path)) + video = cv2.VideoCapture(str(video_path)) + fps = self.__get_safe_fps(video, video_path) + os.makedirs(output_dir, exist_ok=True) + output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' + output_ext = ".mp4" + + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene + frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) + success, frame = video.read() + if not success or frame is None: + self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') + video.release() + return + + # Blend random frames to detect borders, avoiding incorrect crop from black frames + if remove_borders: + frame_blend = frame + for i in range(5): + random_frame = random.randint(scene[0], scene[1]) + video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) + success, frame = video.read() + if not success or frame is None: + continue + a = 1/(i+1) + b = 1-a + frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) + x1, y1, w1, h1 = self.find_main_contour(frame_blend) + else: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + # Ensure dimensions are even, required + h2 -= h2 % 2 + w2 -= w2 % 2 + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) + success, frame = video.read() + if success: + try: + preview = Image.fromarray( + cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + self.preview_image.configure(light_image=preview, size=preview.size) + #truncate filename of long files so UI doesn't shift around + filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext + self.preview_image_label.configure( + text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + except Exception: + pass + video.release() + + if target_fps <= 0: + target_fps = fps + + output_path = f'{output_name}{output_ext}' + self.__write_clip_av(video_path, output_path, scene, fps, target_fps, + x1 + x2, y1 + y2, w2, h2) + + @staticmethod + def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], + src_fps: float, target_fps: float, + crop_x: int, crop_y: int, crop_w: int, crop_h: int): + start_sec = scene[0] / src_fps + end_sec = scene[1] / src_fps + rate_frac = Fraction(target_fps).limit_denominator(10000) + stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) + + with av.open(video_path) as input_container: + in_video = input_container.streams.video[0] + in_video.thread_type = 'AUTO' + in_audio = input_container.streams.audio[0] if input_container.streams.audio else None + + with av.open(output_path, mode='w') as output_container: + out_video = output_container.add_stream('libx264', rate=rate_frac) + out_video.width = crop_w + out_video.height = crop_h + out_video.pix_fmt = 'yuv420p' + out_video.time_base = stream_time_base + + out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None + + input_container.seek(int(start_sec * 1_000_000)) + + out_frame_idx = 0 + out_time_step = 1.0 / target_fps + video_done = False + decode_streams = [s for s in (in_video, in_audio) if s is not None] + + for packet in input_container.demux(decode_streams): + if packet.stream == in_video: + if video_done: + continue + for frame in packet.decode(): + if frame.time is None or frame.time < start_sec: + continue + if frame.time >= end_sec: + video_done = True + break + + # FPS conversion: skip frames when source fps > target fps + if frame.time < start_sec + out_frame_idx * out_time_step: + continue + + img = frame.to_ndarray(format='bgr24') + cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] + out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') + out_frame.pts = out_frame_idx + out_frame.time_base = stream_time_base + + for out_pkt in out_video.encode(out_frame): + output_container.mux(out_pkt) + out_frame_idx += 1 + + elif packet.stream == in_audio and out_audio is not None: + if packet.dts is None: + continue + pkt_time = float(packet.pts * packet.time_base) + if pkt_time < start_sec or pkt_time >= end_sec: + continue + # Re-timestamp audio relative to clip start + packet.pts = int((pkt_time - start_sec) / packet.time_base) + packet.dts = packet.pts + packet.stream = out_audio + output_container.mux(packet) + + # Flush video encoder + for pkt in out_video.encode(): + output_container.mux(pkt) + + def __extract_images_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_images_multi, batch_mode) + + def __extract_images_multi(self, batch_mode : bool): + if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + capture_rate = float(self.capture_rate_entry.get()) + blur_threshold = float(self.blur_threshold_entry.get()) + crop_variation = float(self.image_crop_entry.get()) + input_single_entry = self.image_single_entry.get() + input_multiple_entry = self.image_list_entry.get() + output_entry = self.image_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") + return + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + return + if not (0.0 <= blur_threshold < 1.0): + self.__update_status("Blur Removal must be between 0.0 and 1.0.") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if not input_videos: + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_img_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + executor.submit(self.__save_frames, + str(video_path), time_start, time_end, capture_rate, + blur_threshold, self.image_bordercrop.get(), + crop_variation, output_directory) + if batch_mode: + self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + else: + self.__update_status(f'Image extraction from "{input_single_entry}" complete') + + def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, + blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + video.release() + return + image_rate = max(int(vid_fps / capture_rate), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) + frame_list = [] + + for n in frame_range: + #pick frame from random triangular distribution around center of each "chunk" of the video + frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) + frame = max(0, min(frame, max(total_frames - 1, 0))) + frame_list.append(frame) + + self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') + + output_list = [] + for f in frame_list: + video.set(cv2.CAP_PROP_POS_FRAMES, f) + success, frame = video.read() + if success and frame is not None: + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() + output_list.append((f, frame_sharpness)) + + if not output_list: + self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') + video.release() + return + + output_list_sorted = sorted(output_list, key=lambda x: x[1]) + cutoff = int(blur_threshold * len(output_list_sorted)) + output_list_cut = output_list_sorted[cutoff:] + self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') + + basename, ext = os.path.splitext(os.path.basename(video_path)) + os.makedirs(output_dir, exist_ok=True) + + for f in output_list_cut: + filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' + video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) + success, frame = video.read() + + #crop out borders of frame + if remove_borders and success and frame is not None: + x1, y1, w1, h1 = self.find_main_contour(frame) + frame_cropped = frame[y1:y1+h1, x1:x1+w1] + else: + frame_cropped = frame if success and frame is not None else None + if frame_cropped is not None: + x1 = 0 + y1 = 0 + h1, w1, _ = frame_cropped.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + + if success and frame is not None and frame_cropped is not None: + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") + try: + preview = Image.fromarray( + cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext + self.preview_image.configure(light_image=preview, size=preview.size) + self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + except Exception: + pass # preview update is non-critical + + cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) + video.release() + + def __download_button(self, batch_mode: bool): + self.__run_in_thread(self.__download_multi, batch_mode) + + def __update_status(self, status_text: str): + print(status_text) + self.status_label.configure(state="normal") + self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.configure(state="disabled") + + def __download_multi(self, batch_mode: bool): + if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + if not batch_mode: + ydl_urls = [self.download_link_entry.get()] + elif batch_mode: + ydl_path = pathlib.Path(self.download_list_entry.get()) + if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": + with open(ydl_path) as file: + ydl_urls = file.readlines() + else: + self.__update_status("Invalid link list!") + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for url in ydl_urls: + executor.submit(self.__download_video, + url.strip(), self.download_output_entry.get(), + self.download_args_entry.get("0.0", ctk.END)) + + self.__update_status(f'Completed {len(ydl_urls)} downloads.') + + def __download_video(self, url: str, output_dir: str, output_args: str): + url = (url or "").strip() + if not url: + self.__update_status("Empty URL, skipping download.") + return + + #Respect quotes and split into list to run as yt-dlp command + additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] + cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] + + self.__update_status(f'Downloading {url}') + subprocess.run(cmd) + self.__update_status(f'Download {url} done!') diff --git a/modules/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py new file mode 100644 index 000000000..1690879f1 --- /dev/null +++ b/modules/ui/GenerateCaptionsWindowController.py @@ -0,0 +1,133 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateCaptionsWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating captions for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + self.parent = parent + + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all captions", "Create if absent", "Add as new line"] + self.model_var = ctk.StringVar(self, "Blip") + self.models = ["Blip", "Blip2", "WD14 VIT v2"] + + self.title("Batch generate captions") + self.geometry("360x360") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) + self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) + self.caption_entry = ctk.CTkEntry(self.frame, width=200) + self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) + self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.prefix_entry = ctk.CTkEntry(self.frame, width=200) + self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) + self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.postfix_entry = ctk.CTkEntry(self.frame, width=200) + self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_captions(self): + self.parent.load_captioning_model(self.model_var.get()) + + mode = { + "Replace all captions": "replace", + "Create if absent": "fill", + "Add as new line": "add", + }[self.mode_var.get()] + + self.parent.captioning_model.caption_folder( + sample_dir=self.path_entry.get(), + initial_caption=self.caption_entry.get(), + caption_prefix=self.prefix_entry.get(), + caption_postfix=self.postfix_entry.get(), + mode=mode, + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/GenerateMasksWindowController.py b/modules/ui/GenerateMasksWindowController.py new file mode 100644 index 000000000..daff0d3d5 --- /dev/null +++ b/modules/ui/GenerateMasksWindowController.py @@ -0,0 +1,151 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class GenerateMasksWindow(ctk.CTkToplevel): + def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating masks for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + super().__init__(parent, *args, **kwargs) + + self.parent = parent + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] + self.model_var = ctk.StringVar(self, "ClipSeg") + self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] + + self.title("Batch generate masks") + self.geometry("360x430") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) + self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) + self.prompt_entry = ctk.CTkEntry(self.frame, width=200) + self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) + self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") + self.threshold_entry.insert(0, "0.3") + self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) + self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") + self.smooth_entry.insert(0, 5) + self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) + self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") + self.expand_entry.insert(0, 10) + self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) + self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") + self.alpha_entry.insert(0, 1) + self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) + + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def create_masks(self): + self.parent.load_masking_model(self.model_var.get()) + + mode = { + "Replace all masks": "replace", + "Create if absent": "fill", + "Add to existing": "add", + "Subtract from existing": "subtract", + "Blend with existing": "blend", + }[self.mode_var.get()] + + self.parent.masking_model.mask_folder( + sample_dir=self.path_entry.get(), + prompts=[self.prompt_entry.get()], + mode=mode, + alpha=float(self.alpha_entry.get()), + threshold=float(self.threshold_entry.get()), + smooth_pixels=int(self.smooth_entry.get()), + expand_pixels=int(self.expand_entry.get()), + progress_callback=self.set_progress, + include_subdirectories=self.include_subdirectories_var.get(), + ) + self.parent.load_image() + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py new file mode 100644 index 000000000..16063c26c --- /dev/null +++ b/modules/ui/OptimizerParamsWindowController.py @@ -0,0 +1,288 @@ +import contextlib +from tkinter import TclError + +from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import ( + OPTIMIZER_DEFAULT_PARAMETERS, + change_optimizer, + load_optimizer_defaults, + update_optimizer_config, +) +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class OptimizerParamsWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + ui_state, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.parent = parent + self.train_config = train_config + self.ui_state = ui_state + self.optimizer_ui_state = ui_state.get_var("optimizer") + self.protocol("WM_DELETE_WINDOW", self.on_window_close) + self.muon_adam_button = None + + self.title("Optimizer Settings") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + components.button(self, 1, 0, "ok", command=self.on_window_close) + self.main_frame(self.frame) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + + def main_frame(self, master): + # Optimizer + components.label(master, 0, 0, "Optimizer", + tooltip="The type of optimizer") + + # Create the optimizer dropdown menu and set the command + components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", + command=self.on_optimizer_change) + + # Defaults Button + components.label(master, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, + tooltip="Load default settings for the selected optimizer") + + self.create_dynamic_ui(master) + + def clear_dynamic_ui(self, master): + with contextlib.suppress(TclError): + for widget in master.winfo_children(): + grid_info = widget.grid_info() + if int(grid_info["row"]) >= 1: + widget.destroy() + + def create_dynamic_ui( + self, + master, + ): + + # Lookup for the title and tooltip for a key + # @formatter:off + KEY_DETAIL_MAP = { + 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, + 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, + 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, + 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, + 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, + 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, + 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, + 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, + 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, + 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, + 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, + 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, + 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, + 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, + 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, + 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, + 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, + 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, + 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, + 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, + 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, + 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, + 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, + 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, + 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, + 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, + 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, + 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, + 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, + 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, + 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, + 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, + 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, + 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, + 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, + 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, + 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, + 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, + 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, + 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, + 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, + 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, + 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, + 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, + 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, + 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, + 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, + 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, + 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, + 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, + 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, + 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, + 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, + 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, + 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, + 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, + 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, + 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, + 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, + 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, + 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, + 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, + 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, + 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, + 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, + 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, + 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, + 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, + 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, + 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, + 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, + 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, + 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, + 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, + 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, + 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, + 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, + 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, + 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, + 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, + 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, + 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, + 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, + 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, + 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, + 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, + 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, + 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, + 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, + 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, + 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, + 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, + 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, + 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, + 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, + 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, + 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, + 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, + 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, + 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, + } + # @formatter:on + + if not self.winfo_exists(): # check if this window isn't open + return + + selected_optimizer = self.train_config.optimizer.optimizer + + # Extract the keys for the selected optimizer + for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): + if key not in KEY_DETAIL_MAP: + continue + arg_info = KEY_DETAIL_MAP[key] + + title = arg_info['title'] + tooltip = arg_info['tooltip'] + type = arg_info['type'] + + row = (index // 2) + 1 + col = 3 * (index % 2) + + components.label(master, row, col, title, tooltip=tooltip) + + if key == 'MuonWithAuxAdam': + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + + components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + + self.muon_adam_button = components.button( + frame, 0, 1, "...", self.open_muon_adam_window, + tooltip="Configure the auxiliary AdamW_adv optimizer", + width=20, padx=5 ) + self.toggle_muon_adam_button() + elif type != 'bool': + components.entry(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + else: + components.switch(master, row, col + 1, self.optimizer_ui_state, key, + command=self.update_user_pref) + + def update_user_pref(self, *args): + update_optimizer_config(self.train_config) + self.toggle_muon_adam_button() + + def on_optimizer_change(self, *args): + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.clear_dynamic_ui(self.frame) + self.create_dynamic_ui(self.frame) + + def load_defaults(self, *args): + optimizer_config = load_optimizer_defaults(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + def on_window_close(self): + self.destroy() + + def toggle_muon_adam_button(self): + if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + + def open_muon_adam_window(self): + current_optimizer = self.train_config.optimizer.optimizer + + adam_config = TrainOptimizerConfig.default_values() + current_state = self.train_config.optimizer.muon_adam_config + + if current_optimizer == Optimizer.MUON: + defaults = MUON_AUX_ADAM_DEFAULTS + else: + defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] + + if current_state is None: + adam_config.from_dict(defaults) + if current_optimizer != Optimizer.MUON: + adam_config.optimizer = Optimizer.ADAMW_ADV + elif isinstance(current_state, dict): + adam_config.from_dict(current_state) + else: + # Should not happen if TrainConfig defines it as dict, but for safety + adam_config = current_state + + temp_adam_ui_state = UIState(self, adam_config) + window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + self.wait_window(window) + + self.train_config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py new file mode 100644 index 000000000..0f91ad2fa --- /dev/null +++ b/modules/ui/SampleWindowController.py @@ -0,0 +1,227 @@ +import contextlib +import copy +import os +import tkinter as tk +import traceback + +from modules.model.BaseModel import BaseModel +from modules.modelSampler.BaseModelSampler import ( + BaseModelSampler, + ModelSamplerOutput, +) +from modules.ui.SampleFrame import SampleFrame +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.EMAMode import EMAMode +from modules.util.enum.FileType import FileType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.time_util import get_string_timestamp +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch + +import customtkinter as ctk +from PIL import Image + + +class SampleWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + train_config: TrainConfig, + use_external_model: bool, + callbacks: TrainCallbacks | None = None, + commands: TrainCommands | None = None, + *args, **kwargs + ): + super().__init__(parent, *args, **kwargs) + + self.title("Sample") + self.geometry("1200x800") + self.resizable(True, True) + + if not use_external_model: + self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) + # remove some settings to speed up model loading for sampling + self.initial_train_config.optimizer.optimizer = None + self.initial_train_config.ema = EMAMode.OFF + else: + self.initial_train_config = None + + #TODO why is there a current_train_config and an initial_train_config? + #current_train_config doesn't seem to ever change + self.current_train_config = train_config + self.callbacks = callbacks + self.commands = commands + + # get model specific defaults + model_type = train_config.model_type + self.sample = SampleConfig.default_values(model_type) + self.ui_state = UIState(self, self.sample) + + if use_external_model: + self.callbacks.set_on_sample_custom(self.__update_preview) + self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) + else: + self.model = None + self.model_sampler = None + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_rowconfigure(3, weight=0) + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1) + + prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + + settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + # image + self.image = ctk.CTkImage( + light_image=self.__dummy_image(), + size=(512, 512) + ) + + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) + image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + + self.progress = components.progress(self, 2, 0) + components.button(self, 3, 0, "sample", self.__sample) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def __load_model(self) -> BaseModel: + model_loader = create.create_model_loader( + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + model_setup = create.create_model_setup( + model_type=self.initial_train_config.model_type, + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + training_method=self.initial_train_config.training_method, + ) + + model_names = self.initial_train_config.model_names() + if self.initial_train_config.continue_last_backup: + last_backup_path = self.initial_train_config.get_last_backup_path() + + if last_backup_path: + if self.initial_train_config.training_method == TrainingMethod.LORA: + model_names.lora = last_backup_path + elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: + model_names.embedding.model_name = last_backup_path + else: # fine-tunes + model_names.base_model = last_backup_path + + print(f"Loading from backup '{last_backup_path}'...") + else: + print("No backup found, loading without backup...") + + if self.initial_train_config.quantization.cache_dir is None: + self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" + os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) + + model = model_loader.load( + model_type=self.initial_train_config.model_type, + model_names=model_names, + weight_dtypes=self.initial_train_config.weight_dtypes(), + quantization=self.initial_train_config.quantization, + ) + model.train_config = self.initial_train_config + + model_setup.setup_optimizations(model, self.initial_train_config) + model_setup.setup_train_device(model, self.initial_train_config) + model_setup.setup_model(model, self.initial_train_config) + model.to(torch.device(self.initial_train_config.temp_device)) + + return model + + def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + return create.create_model_sampler( + train_device=torch.device(self.initial_train_config.train_device), + temp_device=torch.device(self.initial_train_config.temp_device), + model=model, + model_type=self.initial_train_config.model_type, + training_method=self.initial_train_config.training_method, + ) + + def __update_preview(self, sampler_output: ModelSamplerOutput): + if sampler_output.file_type == FileType.IMAGE: + image = sampler_output.data + self.image.configure( + light_image=image, + size=(image.width, image.height), + ) + + def __update_progress(self, progress: int, max_progress: int): + self.progress.set(progress / max_progress) + self.update() + + def __dummy_image(self) -> Image: + return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) + + def __sample(self): + sample = copy.copy(self.sample) + + if self.commands: + self.commands.sample_custom(sample) + else: + if self.model is None: + # lazy initialization + self.model = self.__load_model() + self.model_sampler = self.__create_sampler(self.model) + + sample.from_train_config(self.current_train_config) + + sample_dir = os.path.join( + self.initial_train_config.workspace_dir, + "samples", + "custom", + ) + + progress = self.model.train_progress + sample_path = os.path.join( + sample_dir, + f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" + ) + + self.model.eval() + + self.model_sampler.sample( + sample_config=sample, + destination=sample_path, + image_format=self.current_train_config.sample_image_format, + video_format=self.current_train_config.sample_video_format, + audio_format=self.current_train_config.sample_audio_format, + on_sample=self.__update_preview, + on_update_progress=self.__update_progress, + ) + + def destroy(self): + try: + if hasattr(self, "_icon_image_ref"): + del self._icon_image_ref + + # Remove any pending after callbacks + for after_id in self.tk.call('after', 'info'): + with contextlib.suppress(tk.TclError, RuntimeError): + self.after_cancel(after_id) + + super().destroy() + except (tk.TclError, RuntimeError) as e: + print(f"Error destroying window: {e}") + except Exception as e: + print(f"Unexpected error destroying window: {e}") + traceback.print_exc() diff --git a/modules/ui/TimestepDistributionWindowController.py b/modules/ui/TimestepDistributionWindowController.py new file mode 100644 index 000000000..21e41ce3e --- /dev/null +++ b/modules/ui/TimestepDistributionWindowController.py @@ -0,0 +1,186 @@ + +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( + ModelSetupNoiseMixin, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState + +import torch +from torch import Tensor + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class TimestepGenerator(ModelSetupNoiseMixin): + + def __init__( + self, + timestep_distribution: TimestepDistribution, + min_noising_strength: float, + max_noising_strength: float, + noising_weight: float, + noising_bias: float, + timestep_shift: float, + ): + super().__init__() + + self.timestep_distribution = timestep_distribution + self.min_noising_strength = min_noising_strength + self.max_noising_strength = max_noising_strength + self.noising_weight = noising_weight + self.noising_bias = noising_bias + self.timestep_shift = timestep_shift + + def generate(self) -> Tensor: + generator = torch.Generator() + generator.seed() + + config = TrainConfig.default_values() + config.timestep_distribution = self.timestep_distribution + config.min_noising_strength = self.min_noising_strength + config.max_noising_strength = self.max_noising_strength + config.noising_weight = self.noising_weight + config.noising_bias = self.noising_bias + config.timestep_shift = self.timestep_shift + + + return self._get_timestep_discrete( + num_train_timesteps=1000, + deterministic=False, + generator=generator, + batch_size=1000000, + config=config, + ) + + +class TimestepDistributionWindow(ctk.CTkToplevel): + def __init__( + self, + parent, + config: TrainConfig, + ui_state: UIState, + *args, **kwargs, + ): + super().__init__(parent, *args, **kwargs) + + self.title("Timestep Distribution") + self.geometry("900x600") + self.resizable(True, True) + + self.config = config + self.ui_state = ui_state + self.image_preview_file_index = 0 + self.ax = None + self.canvas = None + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = self.__content_frame(self) + frame.grid(row=0, column=0, sticky='nsew') + components.button(self, 1, 0, "ok", self.__ok) + + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + + def __content_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + frame.grid_rowconfigure(7, weight=1) + + # timestep distribution + components.label(frame, 0, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + "timestep_distribution") + + # min noising strength + components.label(frame, 1, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + + # max noising strength + components.label(frame, 2, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + + # noising weight + components.label(frame, 3, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 3, 1, self.ui_state, "noising_weight") + + # noising bias + components.label(frame, 4, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + components.entry(frame, 4, 1, self.ui_state, "noising_bias") + + # timestep shift + components.label(frame, 5, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + + # dynamic timestep shifting + components.label(frame, 6, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") + + + # plot + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + fig, ax = plt.subplots() + self.ax = ax + self.canvas = FigureCanvasTkAgg(fig, master=frame) + self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) + + fig.set_facecolor(background_color) + ax.set_facecolor(background_color) + ax.spines['bottom'].set_color(text_color) + ax.spines['left'].set_color(text_color) + ax.spines['top'].set_color(text_color) + ax.spines['right'].set_color(text_color) + ax.tick_params(axis='x', colors=text_color, which="both") + ax.tick_params(axis='y', colors=text_color, which="both") + ax.xaxis.label.set_color(text_color) + ax.yaxis.label.set_color(text_color) + + self.__update_preview() + + # update button + components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + + frame.pack(fill="both", expand=1) + return frame + + def __update_preview(self): + generator = TimestepGenerator( + timestep_distribution=self.config.timestep_distribution, + min_noising_strength=self.config.min_noising_strength, + max_noising_strength=self.config.max_noising_strength, + noising_weight=self.config.noising_weight, + noising_bias=self.config.noising_bias, + timestep_shift=self.config.timestep_shift, + ) + + self.ax.cla() + self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.canvas.draw() + + def __ok(self): + self.destroy() diff --git a/modules/ui/TopBarController.py b/modules/ui/TopBarController.py new file mode 100644 index 000000000..820fdb71a --- /dev/null +++ b/modules/ui/TopBarController.py @@ -0,0 +1,260 @@ +import json +import os +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress + +from modules.util import path_util +from modules.util.config.SecretsConfig import SecretsConfig +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.optimizer_util import change_optimizer +from modules.util.path_util import write_json_atomic +from modules.util.ui import components, dialogs +from modules.util.ui.UIState import UIState + +import customtkinter as ctk + + +class TopBar: + def __init__( + self, + master, + train_config: TrainConfig, + ui_state: UIState, + change_model_type_callback: Callable[[ModelType], None], + change_training_method_callback: Callable[[TrainingMethod], None], + load_preset_callback: Callable[[], None], + ): + self.master = master + self.train_config = train_config + self.ui_state = ui_state + self.change_model_type_callback = change_model_type_callback + self.change_training_method_callback = change_training_method_callback + self.load_preset_callback = load_preset_callback + + self.dir = "training_presets" + + self.config_ui_data = { + "config_name": path_util.canonical_join(self.dir, "#.json") + } + self.config_ui_state = UIState(master, self.config_ui_data) + + self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] + self.__load_available_config_names() + + self.current_config = [] + + self.frame = ctk.CTkFrame(master=master, corner_radius=0) + self.frame.grid(row=0, column=0, sticky="nsew") + + self.training_method = None + + # title + components.app_title(self.frame, 0, 0) + + # dropdown + self.configs_dropdown = None + self.__create_configs_dropdown() + + # remove button + # TODO + # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + + # Wiki button + components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + + # save button + components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) + + # padding + self.frame.grid_columnconfigure(5, weight=1) + + # model type + components.options_kv( + master=self.frame, + row=0, + column=6, + values=[ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ], + ui_state=self.ui_state, + var_name="model_type", + command=self.__change_model_type, + ) + + def __create_training_method(self): + if self.training_method: + self.training_method.destroy() + + values = [] + #TODO simplify + if self.train_config.model_type.is_stable_diffusion(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), + ] + elif self.train_config.model_type.is_stable_diffusion_3() \ + or self.train_config.model_type.is_stable_diffusion_xl() \ + or self.train_config.model_type.is_wuerstchen() \ + or self.train_config.model_type.is_pixart() \ + or self.train_config.model_type.is_flux_1() \ + or self.train_config.model_type.is_sana() \ + or self.train_config.model_type.is_hunyuan_video() \ + or self.train_config.model_type.is_hi_dream() \ + or self.train_config.model_type.is_chroma(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ] + elif self.train_config.model_type.is_qwen() \ + or self.train_config.model_type.is_z_image() \ + or self.train_config.model_type.is_flux_2() \ + or self.train_config.model_type.is_ernie(): + values = [ + ("Fine Tune", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ] + + # training method + self.training_method = components.options_kv( + master=self.frame, + row=0, + column=7, + values=values, + ui_state=self.ui_state, + var_name="training_method", + command=self.change_training_method_callback, + ) + + def __change_model_type(self, model_type: ModelType): + self.change_model_type_callback(model_type) + self.__create_training_method() + + def __create_configs_dropdown(self): + if self.configs_dropdown is not None: + self.configs_dropdown.grid_forget() + + self.configs_dropdown = components.options_kv( + self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config + ) + + def __load_available_config_names(self): + if os.path.isdir(self.dir): + for path in os.listdir(self.dir): + if path != "#.json": + path = path_util.canonical_join(self.dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + self.configs.append((name, path)) + self.configs.sort() + + def __save_to_file(self, name) -> str: + name = path_util.safe_filename(name) + path = path_util.canonical_join("training_presets", f"{name}.json") + + write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) + + return path + + def __save_secrets(self, path) -> str: + write_json_atomic(path, self.train_config.secrets.to_dict()) + return path + + def open_wiki(self): + webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + + def __save_new_config(self, name): + path = self.__save_to_file(name) + + is_new_config = name not in [x[0] for x in self.configs] + + if is_new_config: + self.configs.append((name, path)) + self.configs.sort() + + if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): + self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) + + if is_new_config: + self.__create_configs_dropdown() + + def __save_config(self): + default_value = self.configs_dropdown.get() + while default_value.startswith('#'): + default_value = default_value[1:] + + dialogs.StringInputDialog( + parent=self.master, + title="name", + question="Config Name", + callback=self.__save_new_config, + default_value=default_value, + validate_callback=lambda x: not x.startswith("#") + ) + + def __load_current_config(self, filename): + try: + basename = os.path.basename(filename) + is_built_in_preset = basename.startswith("#") and basename != "#.json" + + with open(filename, "r") as f: + loaded_dict = json.load(f) + default_config = TrainConfig.default_values() + if is_built_in_preset: + # always assume built-in configs are saved in the most recent version + loaded_dict["__version"] = default_config.config_version + loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() + + with suppress(FileNotFoundError), open("secrets.json", "r") as f: + secrets_dict=json.load(f) + loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) + + self.train_config.from_dict(loaded_config.to_dict()) + self.ui_state.update(loaded_config) + + optimizer_config = change_optimizer(self.train_config) + self.ui_state.get_var("optimizer").update(optimizer_config) + + self.load_preset_callback() + except FileNotFoundError: + pass + except Exception: + print(traceback.format_exc()) + + def __remove_config(self): + # TODO + pass + + def save_default(self): + self.__save_to_file("#") + self.__save_secrets("secrets.json") diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py new file mode 100644 index 000000000..ba90d2e64 --- /dev/null +++ b/modules/ui/TrainUIController.py @@ -0,0 +1,889 @@ +import ctypes +import datetime +import json +import os +import platform +import subprocess +import sys +import threading +import time +import traceback +import webbrowser +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from tkinter import filedialog, messagebox + +import scripts.generate_debug_report +from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab +from modules.ui.CaptionUI import CaptionUI +from modules.ui.CloudTab import CloudTab +from modules.ui.ConceptTab import ConceptTab +from modules.ui.ConvertModelUI import ConvertModelUI +from modules.ui.LoraTab import LoraTab +from modules.ui.ModelTab import ModelTab +from modules.ui.ProfilingWindow import ProfilingWindow +from modules.ui.SampleWindow import SampleWindow +from modules.ui.SamplingTab import SamplingTab +from modules.ui.TopBar import TopBar +from modules.ui.TrainingTab import TrainingTab +from modules.ui.VideoToolUI import VideoToolUI +from modules.util import create +from modules.util.callbacks.TrainCallbacks import TrainCallbacks +from modules.util.commands.TrainCommands import TrainCommands +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.GradientReducePrecision import GradientReducePrecision +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress +from modules.util.ui import components +from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui.UIState import UIState +from modules.util.ui.validation import flush_and_validate_all + +import torch + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker + +# chunk for forcing Windows to ignore DPI scaling when moving between monitors +# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 +if platform.system() == "Windows": + with suppress(Exception): + # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically + ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE + +class TrainUI(ctk.CTk): + set_step_progress: Callable[[int, int], None] + set_epoch_progress: Callable[[int, int], None] + + status_label: ctk.CTkLabel | None + training_button: ctk.CTkButton | None + training_callbacks: TrainCallbacks | None + training_commands: TrainCommands | None + + _TRAIN_BUTTON_STYLES = { + "idle": { + "text": "Start Training", + "state": "normal", + "fg_color": "#198754", + "hover_color": "#146c43", + "text_color": "white", + "text_color_disabled": "white", + }, + "running": { + "text": "Stop Training", + "state": "normal", + "fg_color": "#dc3545", + "hover_color": "#bb2d3b", + "text_color": "white", + }, + "stopping": { + "text": "Stopping...", + "state": "disabled", + "fg_color": "#dc3545", + "hover_color": "#dc3545", + "text_color": "white", + "text_color_disabled": "white", + }, + } + + def __init__(self): + super().__init__() + + self.title("OneTrainer") + self.geometry("1100x740") + + self.after(100, lambda: self._set_icon()) + + # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration + ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") + ctk.set_default_color_theme("blue") + + self.train_config = TrainConfig.default_values() + self.ui_state = UIState(self, self.train_config) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = None + self.tabview = None + + self.model_tab = None + self.training_tab = None + self.lora_tab = None + self.cloud_tab = None + self.additional_embeddings_tab = None + + self.top_bar_component = self.top_bar(self) + self.content_frame(self) + self.bottom_bar(self) + + self.training_thread = None + self.training_callbacks = None + self.training_commands = None + + self.always_on_tensorboard_subprocess = None + self.current_workspace_dir = self.train_config.workspace_dir + self._check_start_always_on_tensorboard() + + self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + + # Persistent profiling window. + self.profiling_window = ProfilingWindow(self) + + self.protocol("WM_DELETE_WINDOW", self.__close) + + def __close(self): + self.top_bar_component.save_default() + self._stop_always_on_tensorboard() + if hasattr(self, 'workspace_dir_trace_id'): + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + self.quit() + + def top_bar(self, master): + return TopBar( + master, + self.train_config, + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + + def bottom_bar(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=2, column=0, sticky="nsew") + + self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + + # status + ETA container + self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + self.status_frame.grid(row=0, column=1, sticky="w") + self.status_frame.grid_rowconfigure(0, weight=0) + self.status_frame.grid_rowconfigure(1, weight=0) + self.status_frame.grid_columnconfigure(0, weight=1) + + self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + + # padding + frame.grid_columnconfigure(2, weight=1) + + + # export button + self.export_button = components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + # debug button + components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + # tensorboard button + components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + # training button + self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + self._set_training_button_style("idle") # centralized styling + + return frame + + def content_frame(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_columnconfigure(0, weight=1) + + self.tabview = ctk.CTkTabview(frame) + self.tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.create_general_tab(self.tabview.add("general")) + self.model_tab = self.create_model_tab(self.tabview.add("model")) + self.data_tab = self.create_data_tab(self.tabview.add("data")) + self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) + self.training_tab = self.create_training_tab(self.tabview.add("training")) + self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) + self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) + self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) + self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) + self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) + + self.change_training_method(self.train_config.training_method) + + return frame + + def create_general_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # workspace dir + components.label(frame, 0, 0, "Workspace Directory", + tooltip="The directory where all files of this training run are saved") + components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + + # cache dir + components.label(frame, 0, 2, "Cache Directory", + tooltip="The directory where cached data is saved") + components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + + # continue from previous backup + components.label(frame, 2, 0, "Continue from last backup", + tooltip="Automatically continues training from the last backup saved in /backup") + components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + + # only cache + components.label(frame, 2, 2, "Only Cache", + tooltip="Only populate the cache, without any training") + components.switch(frame, 2, 3, self.ui_state, "only_cache") + + # TODO: In Phase 4 rework the general tab. + # prevent overwrites + components.label(frame, 3, 0, "Prevent Overwrites", + tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") + components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + + # debug + components.label(frame, 4, 0, "Debug mode", + tooltip="Save debug information during the training into the debug directory") + components.switch(frame, 4, 1, self.ui_state, "debug_mode") + + components.label(frame, 4, 2, "Debug Directory", + tooltip="The directory where debug data is saved") + components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + + # tensorboard + components.label(frame, 6, 0, "Tensorboard", + tooltip="Starts the Tensorboard Web UI during training") + components.switch(frame, 6, 1, self.ui_state, "tensorboard") + + components.label(frame, 6, 2, "Always-On Tensorboard", + tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") + components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) + + components.label(frame, 7, 0, "Expose Tensorboard", + tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") + components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") + components.label(frame, 7, 2, "Tensorboard Port", + tooltip="Port to use for Tensorboard link") + components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") + + + # validation + components.label(frame, 8, 0, "Validation", + tooltip="Enable validation steps and add new graph in tensorboard") + components.switch(frame, 8, 1, self.ui_state, "validation") + + components.label(frame, 8, 2, "Validate after", + tooltip="The interval used when validate training") + components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + + # device + components.label(frame, 10, 0, "Dataloader Threads", + tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") + components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + + components.label(frame, 11, 0, "Train Device", + tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") + components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + + components.label(frame, 12, 0, "Multi-GPU", + tooltip="Enable multi-GPU training") + components.switch(frame, 12, 1, self.ui_state, "multi_gpu") + components.label(frame, 12, 2, "Device Indexes", + tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") + components.entry(frame, 12, 3, self.ui_state, "device_indexes") + + components.label(frame, 13, 0, "Gradient Reduce Precision", + tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" + "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" + "FLOAT_32: Reduce gradients in float32\n" + "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", + wide_tooltip=True) + components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + "gradient_reduce_precision") + + components.label(frame, 13, 2, "Fused Gradient Reduce", + tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") + components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") + + components.label(frame, 14, 0, "Async Gradient Reduce", + tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") + components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") + components.label(frame, 14, 2, "Buffer size (MB)", + tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") + components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") + + components.label(frame, 15, 0, "Temp Device", + tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") + components.entry(frame, 15, 1, self.ui_state, "temp_device") + + frame.pack(fill="both", expand=1) + return frame + + def create_model_tab(self, master): + return ModelTab(master, self.train_config, self.ui_state) + + def create_data_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # aspect ratio bucketing + components.label(frame, 0, 0, "Aspect Ratio Bucketing", + tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") + components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + + # latent caching + components.label(frame, 1, 0, "Latent Caching", + tooltip="Caching of intermediate training data that can be re-used between epochs") + components.switch(frame, 1, 1, self.ui_state, "latent_caching") + + # clear cache before training + components.label(frame, 2, 0, "Clear cache before training", + tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") + components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + + frame.pack(fill="both", expand=1) + return frame + + def create_concepts_tab(self, master): + return ConceptTab(master, self.train_config, self.ui_state) + + def create_training_tab(self, master) -> TrainingTab: + return TrainingTab(master, self.train_config, self.ui_state) + + def create_cloud_tab(self, master) -> CloudTab: + return CloudTab(master, self.train_config, self.ui_state,parent=self) + + def create_sampling_tab(self, master): + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + # sample after + top_frame = ctk.CTkFrame(master=master, corner_radius=0) + top_frame.grid(row=0, column=0, sticky="nsew") + sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") + sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + + components.label(top_frame, 0, 0, "Sample After", + tooltip="The interval used when automatically sampling from the model during training") + components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + + components.label(top_frame, 0, 2, "Skip First", + tooltip="Start sampling automatically after this interval has elapsed.") + components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + + components.label(top_frame, 0, 4, "Format", + tooltip="File Format used when saving samples") + components.options_kv(top_frame, 0, 5, [ + ("PNG", ImageFormat.PNG), + ("JPG", ImageFormat.JPG), + ], self.ui_state, "sample_image_format") + + components.button(top_frame, 0, 6, "sample now", self.sample_now) + + components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + + components.label(sub_frame, 0, 0, "Non-EMA Sampling", + tooltip="Whether to include non-ema sampling when using ema.") + components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + + components.label(sub_frame, 0, 2, "Samples to Tensorboard", + tooltip="Whether to include sample images in the Tensorboard output.") + components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + + # table + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + return SamplingTab(frame, self.train_config, self.ui_state) + + def create_backup_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # backup after + components.label(frame, 0, 0, "Backup After", + tooltip="The interval used when automatically creating model backups during training") + components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + + # backup now + components.button(frame, 0, 3, "backup now", self.backup_now) + + # rolling backup + components.label(frame, 1, 0, "Rolling Backup", + tooltip="If rolling backups are enabled, older backups are deleted automatically") + components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + + # rolling backup count + components.label(frame, 1, 3, "Rolling Backup Count", + tooltip="Defines the number of backups to keep if rolling backups are enabled") + components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + + # backup before save + components.label(frame, 2, 0, "Backup Before Save", + tooltip="Create a full backup before saving the final model") + components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + + # save after + components.label(frame, 3, 0, "Save Every", + tooltip="The interval used when automatically saving the model during training") + components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + + # save now + components.button(frame, 3, 3, "save now", self.save_now) + + # skip save + components.label(frame, 4, 0, "Skip First", + tooltip="Start saving automatically after this interval has elapsed") + components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + + # save filename prefix + components.label(frame, 5, 0, "Save Filename Prefix", + tooltip="The prefix for filenames used when saving the model during training") + components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + + frame.pack(fill="both", expand=1) + return frame + + def embedding_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # embedding model name + components.label(frame, 0, 0, "Base embedding", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + components.path_entry( + frame, 0, 1, self.ui_state, "embedding.model_name", + mode="file", path_modifier=components.json_path_modifier + ) + + # token count + components.label(frame, 1, 0, "Token count", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + + # initial embedding text + components.label(frame, 2, 0, "Initial embedding text", + tooltip="The initial embedding text used when creating a new embedding") + components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + + # embedding weight dtype + components.label(frame, 3, 0, "Embedding Weight Data Type", + tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "embedding_weight_dtype") + + # placeholder + components.label(frame, 4, 0, "Placeholder", + tooltip="The placeholder used when using the embedding in a prompt") + components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + + # output embedding + components.label(frame, 5, 0, "Output embedding", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") + + frame.pack(fill="both", expand=1) + return frame + + def create_additional_embeddings_tab(self, master): + return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + + def create_tools_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + # dataset + components.label(frame, 0, 0, "Dataset Tools", + tooltip="Open the captioning tool") + components.button(frame, 0, 1, "Open", self.open_dataset_tool) + + # video tools + components.label(frame, 1, 0, "Video Tools", + tooltip="Open the video tools") + components.button(frame, 1, 1, "Open", self.open_video_tool) + + # convert model + components.label(frame, 2, 0, "Convert Model Tools", + tooltip="Open the model conversion tool") + components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + + # sample + components.label(frame, 3, 0, "Sampling Tool", + tooltip="Open the model sampling tool") + components.button(frame, 3, 1, "Open", self.open_sampling_tool) + + components.label(frame, 4, 0, "Profiling Tool", + tooltip="Open the profiling tools.") + components.button(frame, 4, 1, "Open", self.open_profiling_tool) + + frame.pack(fill="both", expand=1) + return frame + + def change_model_type(self, model_type: ModelType): + if self.model_tab: + self.model_tab.refresh_ui() + + if self.training_tab: + self.training_tab.refresh_ui() + + if self.lora_tab: + self.lora_tab.refresh_ui() + + def change_training_method(self, training_method: TrainingMethod): + if not self.tabview: + return + + if self.model_tab: + self.model_tab.refresh_ui() + + if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: + self.tabview.delete("LoRA") + self.lora_tab = None + if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: + self.tabview.delete("embedding") + + if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: + self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: + self.embedding_tab(self.tabview.add("embedding")) + + def load_preset(self): + if not self.tabview: + return + + if self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + spent_total = time.monotonic() - self.start_time + steps_done = train_progress.epoch * max_step + train_progress.epoch_step + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) + total_eta = spent_total / steps_done * remaining_steps + + if train_progress.global_step <= 30: + return "Estimating ..." + + td = datetime.timedelta(seconds=total_eta) + days = td.days + hours, remainder = divmod(td.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + if days > 0: + return f"{days}d {hours}h" + elif hours > 0: + return f"{hours}h {minutes}m" + elif minutes > 0: + return f"{minutes}m {seconds}s" + else: + return f"{seconds}s" + + def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def delete_eta_label(self): + self.eta_label.configure(text="") + + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + self.set_step_progress(train_progress.epoch_step, max_step) + self.set_epoch_progress(train_progress.epoch, max_epoch) + self.set_eta_label(train_progress, max_step, max_epoch) + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def open_dataset_tool(self): + window = CaptionUI(self, None, False) + self.wait_window(window) + + def open_video_tool(self): + window = VideoToolUI(self) + self.wait_window(window) + + def open_convert_model_tool(self): + window = ConvertModelUI(self) + self.wait_window(window) + + def open_sampling_tool(self): + if not self.training_callbacks and not self.training_commands: + window = SampleWindow( + self, + use_external_model=False, + train_config=self.train_config, + ) + self.wait_window(window) + torch_gc() + + def open_profiling_tool(self): + self.profiling_window.deiconify() + + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + + if not zip_path: + return + + zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + + self.on_update_status("Generating debug package...") + + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.on_update_status(f"Error generating debug package: {e}") + + + def open_manual_sample_window (self): + training_callbacks = self.training_callbacks + training_commands = self.training_commands + + if training_callbacks and training_commands: + window = SampleWindow( + self, + train_config=self.train_config, + use_external_model=True, + callbacks=training_callbacks, + commands=training_commands, + ) + self.wait_window(window) + training_callbacks.set_on_sample_custom() + + def __training_thread_function(self): + error_caught = False + + self.training_callbacks = TrainCallbacks( + on_update_train_progress=self.on_update_train_progress, + on_update_status=self.on_update_status, + ) + + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + try: + trainer.start() + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + + self.start_time = time.monotonic() + trainer.train() + except Exception: + if self.train_config.cloud.enabled: + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + error_caught = True + traceback.print_exc() + + trainer.end() + + # clear gpu memory + del trainer + + self.training_thread = None + self.training_commands = None + torch.clear_autocast_cache() + torch_gc() + + if error_caught: + self.on_update_status("Error: check the console for details") + else: + self.on_update_status("Stopped") + self.delete_eta_label() + + # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash + self.after(0, self._set_training_button_idle) + + if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + self.after(0, self._start_always_on_tensorboard) + + def start_training(self): + if self.training_thread is None: + self.save_default() + + # --- pre-training validation gate --- + errors = flush_and_validate_all() + + if errors: + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + return + + self._set_training_button_running() + + if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() + + self.training_commands = TrainCommands() + torch_gc() + + self.training_thread = threading.Thread(target=self.__training_thread_function) + self.training_thread.start() + else: + self._set_training_button_stopping() + self.on_update_status("Stopping ...") + self.training_commands.stop() + + def save_default(self): + self.top_bar_component.save_default() + self.concepts_tab.save_current_config() + self.sampling_tab.save_current_config() + self.additional_embeddings_tab.save_current_config() + + def export_training(self): + file_path = filedialog.asksaveasfilename(filetypes=[ + ("All Files", "*.*"), + ("json", "*.json"), + ], initialdir=".", initialfile="config.json") + + if file_path: + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.save() + + def _check_start_always_on_tensorboard(self): + if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _start_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() + + tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") + tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") + + os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) + + tensorboard_args = [ + tensorboard_executable, + "--logdir", + tensorboard_log_dir, + "--port", + str(self.train_config.tensorboard_port), + "--samples_per_plugin=images=100,scalars=10000", + ] + + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") + + try: + self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) + except Exception: + self.always_on_tensorboard_subprocess = None + + def _stop_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + try: + self.always_on_tensorboard_subprocess.terminate() + self.always_on_tensorboard_subprocess.wait(timeout=5) + except subprocess.TimeoutExpired: + self.always_on_tensorboard_subprocess.kill() + except Exception: + pass + finally: + self.always_on_tensorboard_subprocess = None + + def _on_workspace_dir_change(self, new_workspace_dir: str): + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir + + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _on_workspace_dir_change_trace(self, *args): + new_workspace_dir = self.train_config.workspace_dir + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir + + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() + + def _on_always_on_tensorboard_toggle(self): + if self.train_config.tensorboard_always_on: + if not (self.training_thread and self.train_config.tensorboard): + self._start_always_on_tensorboard() + else: + if not (self.training_thread and self.train_config.tensorboard): + self._stop_always_on_tensorboard() + + def _set_training_button_style(self, mode: str): + if not self.training_button: + return + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) + + def _set_training_button_idle(self): + self._set_training_button_style("idle") + + def _set_training_button_running(self): + self._set_training_button_style("running") + + def _set_training_button_stopping(self): + self._set_training_button_style("stopping") diff --git a/modules/ui/VideoToolUIController.py b/modules/ui/VideoToolUIController.py new file mode 100644 index 000000000..c3291e6ea --- /dev/null +++ b/modules/ui/VideoToolUIController.py @@ -0,0 +1,877 @@ +import concurrent.futures +import math +import os +import pathlib +import random +import shlex +import subprocess +import threading +import webbrowser +from fractions import Fraction +from tkinter import filedialog + +from modules.util.image_util import load_image +from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS +from modules.util.ui import components + +import av +import customtkinter as ctk +import cv2 +import scenedetect +from PIL import Image + + +class VideoToolUI(ctk.CTkToplevel): + def __init__( + self, + parent, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + self.title("Video Tools") + self.geometry("600x720") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) + self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) + self.video_download_tab = self.__video_download_tab(tabview.add("download")) + self.status_bar(self) + + def status_bar(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=1, column=0) + frame.grid_columnconfigure(0, weight=0, minsize=160) + frame.grid_columnconfigure(1, weight=0, minsize=300) + frame.grid_columnconfigure(2, weight=1) + + #create preview image + preview_path = "resources/icons/icon.png" + preview = load_image(preview_path, 'RGB') + preview.thumbnail((150, 150)) + self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image_label = ctk.CTkLabel( + master=frame, text="Preview image", image=self.preview_image, height=150, width=150, + compound="top") + self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) + + #displays progress and messages that also go to terminal + self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) + self.status_label.insert(index="1.0", text="Current status") + self.status_label.configure(state="disabled") + self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + + def __clip_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.clip_single_entry = ctk.CTkEntry(frame, width=190) + self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.clip_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_clips_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.clip_time_start_entry.insert(0, "00:00:00") + self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) + self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.clip_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.clip_list_entry = ctk.CTkEntry(frame, width=190) + self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_list_entry)) + self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_clips_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted clips will be saved.") + self.clip_output_entry = ctk.CTkEntry(frame, width=190) + self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.clip_output_entry)) + self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_clip = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") + self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # split at cuts + self.split_at_cuts = ctk.BooleanVar(self, False) + components.label(frame, 5, 0, "Split at Cuts", + tooltip="If enabled, detect cuts in the input video and split at those points. \ + Otherwise will split at any point, and clips may contain cuts.") + self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") + self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + # maximum length + components.label(frame, 6, 0, "Max Length (s)", + tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") + self.clip_length_entry = ctk.CTkEntry(frame, width=220) + self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.clip_length_entry.insert(0, "3") + + # Set FPS + components.label(frame, 7, 0, "Set FPS", + tooltip="FPS to convert output videos to, set to 0 to keep original rate.") + self.clip_fps_entry = ctk.CTkEntry(frame, width=220) + self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.clip_fps_entry.insert(0, "24.0") + + # Remove borders + self.clip_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 8, 0, "Remove Borders", + tooltip="Remove black borders from output clip") + self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") + self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 9, 0, "Crop Variation", + tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square videos. Set to 0 to use only base aspect.") + self.clip_crop_entry = ctk.CTkEntry(frame, width=220) + self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) + self.clip_crop_entry.insert(0, "0.2") + + # object filter - currently unused, may implement in future + # components.label(frame, 9, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __image_extract_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # single video + components.label(frame, 0, 0, "Single Video", + tooltip="Link to single video file to process.") + self.image_single_entry = ctk.CTkEntry(frame, width=190) + self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + self.image_single_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.image_single_entry, + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] + )) + self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 0, 2, "Extract Single", + command=lambda: self.__extract_images_button(False)) + + # time range + components.label(frame, 1, 0, " Time Range", + tooltip="Time range to limit selection for single video, \ + format as hour:minute:second, minute:second, or seconds.") + self.image_time_start_entry = ctk.CTkEntry(frame, width=100) + self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.image_time_start_entry.insert(0, "00:00:00") + self.image_time_end_entry = ctk.CTkEntry(frame, width=100) + self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) + self.image_time_end_entry.insert(0, "99:99:99") + + # directory of videos + components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + self.image_list_entry = ctk.CTkEntry(frame, width=190) + self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.image_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_list_entry)) + self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 2, 2, "Extract Directory", + command=lambda: self.__extract_images_button(True)) + + # output directory + components.label(frame, 3, 0, "Output", + tooltip="Path to folder where extracted images will be saved.") + self.image_output_entry = ctk.CTkEntry(frame, width=190) + self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.image_output_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_dir(self.image_output_entry)) + self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + + # output to subdirectories + self.output_subdir_img = ctk.BooleanVar(self, False) + components.label(frame, 4, 0, "Output to\nSubdirectories", + tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ + Otherwise will all be saved to the top level of the output directory.") + self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") + self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + # image capture rate + components.label(frame, 5, 0, "Images/sec", + tooltip="Number of images to capture per second of video. \ + Images will be taken at semi-random frames around the specified frequency.") + self.capture_rate_entry = ctk.CTkEntry(frame, width=220) + self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.capture_rate_entry.insert(0, "0.5") + + # blur removal + components.label(frame, 6, 0, "Blur Removal", + tooltip="Threshold for removal of blurry images, relative to all others. \ + For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") + self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) + self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + self.blur_threshold_entry.insert(0, "0.2") + + # Remove borders + self.image_bordercrop = ctk.BooleanVar(self, False) + components.label(frame, 7, 0, "Remove Borders", + tooltip="Remove black borders from output image") + self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") + self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + # Crop Variation + components.label(frame, 8, 0, "Crop Variation", + tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ + somewhat biased towards making square images. Set to 0 to use only base sapect.") + self.image_crop_entry = ctk.CTkEntry(frame, width=220) + self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.image_crop_entry.insert(0, "0.2") + + # # object filter - currently unused, may implement in future + # components.label(frame, 5, 0, "Object Filter", + # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") + # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") + # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") + + frame.pack(fill="both", expand=1) + return frame + + def __video_download_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0, minsize=120) + frame.grid_columnconfigure(1, weight=0, minsize=200) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + + # link + components.label(frame, 0, 0, "Single Link", + tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") + self.download_link_entry = ctk.CTkEntry(frame, width=220) + self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) + components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + + # link list + components.label(frame, 1, 0, "Link List", + tooltip="Path to txt file with list of links separated by newlines.") + self.download_list_entry = ctk.CTkEntry(frame, width=190) + self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.download_list_button = ctk.CTkButton(frame, width=30, text="...", + command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) + self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + + # output directory + components.label(frame, 2, 0, "Output", + tooltip="Path to folder where downloaded videos will be saved.") + self.download_output_entry = ctk.CTkEntry(frame, width=190) + self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) + self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + + # additional args + components.label(frame, 3, 0, "Additional Args", + tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ + Default args will hide most terminal outputs.") + self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) + self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) + self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress --format mp4") + components.button(frame, 3, 2, "yt-dlp info", + command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) + + frame.pack(fill="both", expand=1) + return frame + + def __browse_for_dir(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __browse_for_file(self, entry_box, filetypes): + # get the path from the user + path = filedialog.askopenfilename(filetypes=filetypes) + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, ctk.END) + entry_box.insert(0, path) + self.focus_set() + + def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): + input_videos = [] + if not batch_mode: + path = pathlib.Path(input_path_single) + if path.is_file(): + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + return [path] + else: + self.__update_status("Invalid video file!") + return [] + else: + self.__update_status("No file specified, or invalid file path!") + return [] + else: + input_videos = [] + if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": + self.__update_status("Invalid input directory!") + return [] + # Only traverse supported extensions to avoid opening every file. + lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} + for path in pathlib.Path(input_path_dir).rglob("*"): + if path.is_file() and path.suffix.lower() in lower_exts: + vid = cv2.VideoCapture(str(path)) + ok = False + try: + if vid.isOpened(): + ok, _ = vid.read() + finally: + vid.release() + if ok: + input_videos.append(path) + self.__update_status(f'Found {len(input_videos)} videos to process') + return input_videos + + def __run_in_thread(self, target, *args): + """Clear status box and run target function in a daemon thread.""" + self.status_label.configure(state="normal") + self.status_label.delete(index1="1.0", index2="end") + self.status_label.configure(state="disabled") + t = threading.Thread(target=target, args=args) + t.daemon = True + t.start() + + @staticmethod + def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: + return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) + + def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: + fps = video.get(cv2.CAP_PROP_FPS) or 0.0 + if fps <= 0: + self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') + return 30.0 + return fps + + @staticmethod + def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, + video_path, input_dir: str) -> str: + if use_subdir and batch_mode: + return os.path.join(output_entry, + os.path.splitext(os.path.relpath(video_path, input_dir))[0]) + elif use_subdir: + return os.path.join(output_entry, + os.path.splitext(os.path.basename(video_path))[0]) + return output_entry + + def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: + # Return original dimensions and no offset if variation is zero + if variation == 0: + return 0, height, 0, width + + old_aspect = height/width + variation_scaled = old_aspect*variation + if old_aspect > 1.2: #tall image + new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) + elif old_aspect < 0.85: #wide image + new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) + else: #square image + new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) + + new_aspect = round(new_aspect, 2) + #keep the height the same if reducing width, and vice versa + if new_aspect > old_aspect: + new_height = int(height) + new_width = int(width*(old_aspect/new_aspect)) + elif new_aspect < old_aspect: + new_height = int(height*(new_aspect/old_aspect)) + new_width = int(width) + else: + new_height = int(height) + new_width = int(width) + + #random offset in dimension that was cropped + position_x = random.randint(0, width-new_width) + position_y = random.randint(0, height-new_height) + return position_y, new_height, position_x, new_width + + def find_main_contour(self, frame): + #outline image to find main content and exclude black bars often present on letterboxed videos + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) + frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if frame_contours: + #select largest contour by area + frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) + x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) + else: #fallback if no contours detected + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + #if bounding box did not detect the correct area, likely due to all-black frame + if not frame_contours or h1 < 10 or w1 < 10: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + return x1, y1, w1, h1 + + def __extract_clips_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_clips_multi, batch_mode) + + def __extract_clips_multi(self, batch_mode: bool): + if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + max_length = float(self.clip_length_entry.get()) + crop_variation = float(self.clip_crop_entry.get()) + target_fps = float(self.clip_fps_entry.get()) + input_single_entry = self.clip_single_entry.get() + input_multiple_entry = self.clip_list_entry.get() + output_entry = self.clip_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") + return + if max_length <= 0.25: + self.__update_status("Max Length of clips must be > 0.25 seconds.") + return + if target_fps < 0: + self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if len(input_videos) == 0: # exit if no paths found + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_clip_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + executor.submit(self.__extract_clips, + str(video_path), time_start, time_end, max_length, + self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + crop_variation, target_fps, output_directory) + + if batch_mode: + self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + else: + self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + + def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, + split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + max_length_frames = int(max_length * vid_fps) + min_length_frames = max(int(0.25 * vid_fps), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + + if split_at_cuts: + #use scenedetect to find cuts, based on start/end frame number + self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') + timecode_list = scenedetect.detect( + video_path=str(video_path), + detector=scenedetect.AdaptiveDetector(), + start_time=int(timestamp_min_frame), + end_time=int(timestamp_max_frame)) + scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] + if not scene_list: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + else: + scene_list = [(timestamp_min_frame, timestamp_max_frame)] + + scene_list_split = [] + for scene in scene_list: + length = scene[1]-scene[0] + if length > max_length_frames: #check for any scenes longer than max length + n = math.ceil(length/max_length_frames) #divide into n new scenes + new_length = int(length/n) + new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length + for i, _n in enumerate(new_splits[:-1]): + if new_splits[i + 1] - new_splits[i] > min_length_frames: + scene_list_split.append((new_splits[i], new_splits[i + 1])) + elif length > (min_length_frames + 2): + # Trim first/last frame to avoid transition artifacts + scene_list_split.append((scene[0] + 1, scene[1] - 1)) + + self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + futures = [ + executor.submit(self.__save_clip, scene, video_path, target_fps, + remove_borders, crop_variation, output_dir) + for scene in scene_list_split + ] + for future in concurrent.futures.as_completed(futures): + exc = future.exception() + if exc is not None: + self.__update_status(f'Error saving clip: {exc}') + + video.release() + + def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, + remove_borders: bool, crop_variation: float, output_dir: str): + basename, ext = os.path.splitext(os.path.basename(video_path)) + video = cv2.VideoCapture(str(video_path)) + fps = self.__get_safe_fps(video, video_path) + os.makedirs(output_dir, exist_ok=True) + output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' + output_ext = ".mp4" + + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene + frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) + success, frame = video.read() + if not success or frame is None: + self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') + video.release() + return + + # Blend random frames to detect borders, avoiding incorrect crop from black frames + if remove_borders: + frame_blend = frame + for i in range(5): + random_frame = random.randint(scene[0], scene[1]) + video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) + success, frame = video.read() + if not success or frame is None: + continue + a = 1/(i+1) + b = 1-a + frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) + x1, y1, w1, h1 = self.find_main_contour(frame_blend) + else: + x1 = 0 + y1 = 0 + h1, w1, _ = frame.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + # Ensure dimensions are even, required + h2 -= h2 % 2 + w2 -= w2 % 2 + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") + video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) + success, frame = video.read() + if success: + try: + preview = Image.fromarray( + cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + self.preview_image.configure(light_image=preview, size=preview.size) + #truncate filename of long files so UI doesn't shift around + filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext + self.preview_image_label.configure( + text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + except Exception: + pass + video.release() + + if target_fps <= 0: + target_fps = fps + + output_path = f'{output_name}{output_ext}' + self.__write_clip_av(video_path, output_path, scene, fps, target_fps, + x1 + x2, y1 + y2, w2, h2) + + @staticmethod + def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], + src_fps: float, target_fps: float, + crop_x: int, crop_y: int, crop_w: int, crop_h: int): + start_sec = scene[0] / src_fps + end_sec = scene[1] / src_fps + rate_frac = Fraction(target_fps).limit_denominator(10000) + stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) + + with av.open(video_path) as input_container: + in_video = input_container.streams.video[0] + in_video.thread_type = 'AUTO' + in_audio = input_container.streams.audio[0] if input_container.streams.audio else None + + with av.open(output_path, mode='w') as output_container: + out_video = output_container.add_stream('libx264', rate=rate_frac) + out_video.width = crop_w + out_video.height = crop_h + out_video.pix_fmt = 'yuv420p' + out_video.time_base = stream_time_base + + out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None + + input_container.seek(int(start_sec * 1_000_000)) + + out_frame_idx = 0 + out_time_step = 1.0 / target_fps + video_done = False + decode_streams = [s for s in (in_video, in_audio) if s is not None] + + for packet in input_container.demux(decode_streams): + if packet.stream == in_video: + if video_done: + continue + for frame in packet.decode(): + if frame.time is None or frame.time < start_sec: + continue + if frame.time >= end_sec: + video_done = True + break + + # FPS conversion: skip frames when source fps > target fps + if frame.time < start_sec + out_frame_idx * out_time_step: + continue + + img = frame.to_ndarray(format='bgr24') + cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] + out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') + out_frame.pts = out_frame_idx + out_frame.time_base = stream_time_base + + for out_pkt in out_video.encode(out_frame): + output_container.mux(out_pkt) + out_frame_idx += 1 + + elif packet.stream == in_audio and out_audio is not None: + if packet.dts is None: + continue + pkt_time = float(packet.pts * packet.time_base) + if pkt_time < start_sec or pkt_time >= end_sec: + continue + # Re-timestamp audio relative to clip start + packet.pts = int((pkt_time - start_sec) / packet.time_base) + packet.dts = packet.pts + packet.stream = out_audio + output_container.mux(packet) + + # Flush video encoder + for pkt in out_video.encode(): + output_container.mux(pkt) + + def __extract_images_button(self, batch_mode: bool): + self.__run_in_thread(self.__extract_images_multi, batch_mode) + + def __extract_images_multi(self, batch_mode : bool): + if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + # validate numeric inputs + try: + capture_rate = float(self.capture_rate_entry.get()) + blur_threshold = float(self.blur_threshold_entry.get()) + crop_variation = float(self.image_crop_entry.get()) + input_single_entry = self.image_single_entry.get() + input_multiple_entry = self.image_list_entry.get() + output_entry = self.image_output_entry.get() + except ValueError: + self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") + return + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + return + if not (0.0 <= blur_threshold < 1.0): + self.__update_status("Blur Removal must be between 0.0 and 1.0.") + return + if not (0.0 <= crop_variation < 1.0): + self.__update_status("Crop Variation must be between 0.0 and 1.0.") + return + + input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + if not input_videos: + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for video_path in input_videos: + output_directory = self.__get_output_dir( + self.output_subdir_img_entry.get(), batch_mode, + output_entry, video_path, input_multiple_entry) + time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) + time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + executor.submit(self.__save_frames, + str(video_path), time_start, time_end, capture_rate, + blur_threshold, self.image_bordercrop.get(), + crop_variation, output_directory) + if batch_mode: + self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + else: + self.__update_status(f'Image extraction from "{input_single_entry}" complete') + + def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, + blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): + video = cv2.VideoCapture(video_path) + vid_fps = self.__get_safe_fps(video, video_path) + if capture_rate <= 0: + self.__update_status("Images/sec must be > 0.") + video.release() + return + image_rate = max(int(vid_fps / capture_rate), 1) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 + timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) + timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) + frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) + frame_list = [] + + for n in frame_range: + #pick frame from random triangular distribution around center of each "chunk" of the video + frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) + frame = max(0, min(frame, max(total_frames - 1, 0))) + frame_list.append(frame) + + self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') + + output_list = [] + for f in frame_list: + video.set(cv2.CAP_PROP_POS_FRAMES, f) + success, frame = video.read() + if success and frame is not None: + frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() + output_list.append((f, frame_sharpness)) + + if not output_list: + self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') + video.release() + return + + output_list_sorted = sorted(output_list, key=lambda x: x[1]) + cutoff = int(blur_threshold * len(output_list_sorted)) + output_list_cut = output_list_sorted[cutoff:] + self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') + + basename, ext = os.path.splitext(os.path.basename(video_path)) + os.makedirs(output_dir, exist_ok=True) + + for f in output_list_cut: + filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' + video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) + success, frame = video.read() + + #crop out borders of frame + if remove_borders and success and frame is not None: + x1, y1, w1, h1 = self.find_main_contour(frame) + frame_cropped = frame[y1:y1+h1, x1:x1+w1] + else: + frame_cropped = frame if success and frame is not None else None + if frame_cropped is not None: + x1 = 0 + y1 = 0 + h1, w1, _ = frame_cropped.shape + + y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) + + if success and frame is not None and frame_cropped is not None: + print(end='\x1b[2K') #clear terminal so next line can overwrite it + print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") + try: + preview = Image.fromarray( + cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) + preview.thumbnail((150, 150)) + filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext + self.preview_image.configure(light_image=preview, size=preview.size) + self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + except Exception: + pass # preview update is non-critical + + cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) + video.release() + + def __download_button(self, batch_mode: bool): + self.__run_in_thread(self.__download_multi, batch_mode) + + def __update_status(self, status_text: str): + print(status_text) + self.status_label.configure(state="normal") + self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.configure(state="disabled") + + def __download_multi(self, batch_mode: bool): + if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + self.__update_status("Invalid output directory!") + return + + if not batch_mode: + ydl_urls = [self.download_link_entry.get()] + elif batch_mode: + ydl_path = pathlib.Path(self.download_list_entry.get()) + if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": + with open(ydl_path) as file: + ydl_urls = file.readlines() + else: + self.__update_status("Invalid link list!") + return + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + for url in ydl_urls: + executor.submit(self.__download_video, + url.strip(), self.download_output_entry.get(), + self.download_args_entry.get("0.0", ctk.END)) + + self.__update_status(f'Completed {len(ydl_urls)} downloads.') + + def __download_video(self, url: str, output_dir: str, output_args: str): + url = (url or "").strip() + if not url: + self.__update_status("Empty URL, skipping download.") + return + + #Respect quotes and split into list to run as yt-dlp command + additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] + cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] + + self.__update_status(f'Downloading {url}') + subprocess.run(cmd) + self.__update_status(f'Download {url} done!') diff --git a/modules/util/ui/ctk_validation.py b/modules/util/ui/ctk_validation.py new file mode 100644 index 000000000..9f44de8eb --- /dev/null +++ b/modules/util/ui/ctk_validation.py @@ -0,0 +1,501 @@ +from __future__ import annotations + +import contextlib +import os +import re +import sys +import tkinter as tk +from collections import deque +from collections.abc import Callable +from pathlib import PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse + +from modules.util.enum.ModelFormat import ModelFormat +from modules.util.enum.PathIOType import PathIOType + +if TYPE_CHECKING: + from modules.util.ui.UIState import UIState + + import customtkinter as ctk + + +DEBOUNCE_TYPING_MS = 250 +UNDO_DEBOUNCE_MS = 500 +ERROR_BORDER_COLOR = "#dc3545" + +_active_validators: set[FieldValidator] = set() + +TRAILING_SLASH_RE = re.compile(r"[\\/]$") +ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") +HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") + +_INVALID_CHARS = {chr(c) for c in range(32)} +_IS_WINDOWS = sys.platform == "win32" +if _IS_WINDOWS: + _INVALID_CHARS |= set('<>"|?*') + + +def _is_huggingface_repo_or_file(value: str) -> bool: + trimmed = value.strip() + + if trimmed.startswith("https://"): + parsed = urlparse(trimmed) + if parsed.netloc not in {"huggingface.co", "huggingface.com"}: + return False + parts = parsed.path.strip("/").split("/") + if len(parts) >= 5 and parts[2] in {"resolve", "blob"}: + return bool(ENDS_WITH_EXT.search(parts[-1])) + return False + + if len(trimmed) > 96: + return False + if " " in trimmed or "\t" in trimmed: + return False + if "—" in trimmed or ".." in trimmed: + return False + if trimmed.startswith(("\\\\", "//", "/")): + return False + if len(trimmed) >= 2 and trimmed[1] == ":" and trimmed[0].isalpha(): + return False + if trimmed.count("/") != 1: + return False + + return bool(HUGGINGFACE_REPO_RE.match(trimmed)) + + +def _has_invalid_chars(value: str) -> bool: + return bool(_INVALID_CHARS.intersection(value)) + + +def _check_overwrite(path: str, *, is_dir: bool, prevent: bool) -> str | None: + if not prevent: + return None + abs_path = os.path.abspath(path) + if is_dir and os.path.isdir(abs_path): + return "Output folder already exists (overwrite prevented)" + if not is_dir and os.path.isfile(abs_path): + return "Output file already exists (overwrite prevented)" + return None + + +def validate_path( + value: str, + io_type: PathIOType = PathIOType.INPUT, + *, + prevent_overwrites: bool = False, + output_format: str | None = None, +) -> str | None: + """Return an error string if *value* is an invalid path, else ``None``.""" + trimmed = value.strip() + + if not trimmed: + return "Path is empty" + if TRAILING_SLASH_RE.search(trimmed): + return "Path must not end with a slash" + if _has_invalid_chars(trimmed): + return "Path contains invalid characters" + + if trimmed.startswith("cloud:"): + cloud_path = trimmed[6:] + if not cloud_path: + return "Cloud path is empty" + if cloud_path.startswith(("http://", "https://")): + return "Cloud path cannot be a URL" + if "\\" in cloud_path: + return "Cloud path must use forward slashes (/)" + return None + + if io_type == PathIOType.INPUT and _is_huggingface_repo_or_file(trimmed): + return None + + if io_type == PathIOType.INPUT: + if not os.path.exists(os.path.abspath(trimmed)): + return "Input path does not exist" + + if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): + if not os.path.isdir(os.path.dirname(os.path.abspath(trimmed))): + return "Parent folder does not exist" + + if io_type == PathIOType.MODEL and output_format is not None: + if output_format == "DIFFUSERS": + if ENDS_WITH_EXT.search(trimmed): + return "Diffusers output must be a directory path, not a file" + return _check_overwrite(trimmed, is_dir=True, prevent=prevent_overwrites) + + try: + expected_ext = ModelFormat[output_format].file_extension() + except KeyError: + expected_ext = "" + + if expected_ext: + suffix = (PureWindowsPath(trimmed) if _IS_WINDOWS else PurePosixPath(trimmed)).suffix.lower() + if suffix != expected_ext: + return f"Extension must be '{expected_ext}' for {output_format} format" + return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) + + if io_type == PathIOType.OUTPUT: + return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) + + return None + +DEFAULT_MAX_UNDO = 20 + + +class UndoHistory: + def __init__(self, max_size: int = DEFAULT_MAX_UNDO): + self._stack: deque[str] = deque(maxlen=max_size) + self._redo_stack: list[str] = [] + + def push(self, value: str): + if self._stack and self._stack[-1] == value: + return + self._stack.append(value) + self._redo_stack.clear() + + def undo(self, current: str) -> str | None: + if not self._stack: + return None + top = self._stack[-1] + if top == current and len(self._stack) > 1: + self._redo_stack.append(self._stack.pop()) + return self._stack[-1] + elif top != current: + self._redo_stack.append(current) + return top + return None + + def redo(self) -> str | None: + if not self._redo_stack: + return None + value = self._redo_stack.pop() + self._stack.append(value) + return value + + +class DebounceTimer: + def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): + self.widget = widget + self.delay_ms = delay_ms + self.callback = callback + self._after_id: str | None = None + + def call(self, *args, **kwargs): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + + def fire(): + self._after_id = None + self.callback(*args, **kwargs) + + with contextlib.suppress(tk.TclError): + self._after_id = self.widget.after(self.delay_ms, fire) + + def cancel(self): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + self._after_id = None + + +class FieldValidator: + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + self.component = component + self.var = var + self.ui_state = ui_state + self.var_name = var_name + self._extra_validate = extra_validate + self._required = required + + try: + self._original_border_color = component.cget("border_color") + except Exception: + self._original_border_color = "gray50" + + self._shadow_var = tk.StringVar(master=component) + self._shadow_trace_name: str | None = None + self._real_var_trace_name: str | None = None + self._syncing = False + self._touched = False + self._bound = False + + self._debounce: DebounceTimer | None = None + self._undo_debounce: DebounceTimer | None = None + self._undo = UndoHistory(max_undo) + + def attach(self) -> None: + self._shadow_var.set(self.var.get()) + self._swap_textvariable(self._shadow_var) + + self._debounce = DebounceTimer( + self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire + ) + self._undo_debounce = DebounceTimer( + self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot + ) + + self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) + self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) + + self.component.bind("", self._on_focus_in) + self.component.bind("", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("", self._on_focus_out) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_enter) + + self._bound = True + _active_validators.add(self) + + def detach(self) -> None: + if not self._bound: + return + self._bound = False + _active_validators.discard(self) + + self._commit() + + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + + if self._shadow_trace_name: + with contextlib.suppress(Exception): + self._shadow_var.trace_remove("write", self._shadow_trace_name) + self._shadow_trace_name = None + + if self._real_var_trace_name: + with contextlib.suppress(Exception): + self.var.trace_remove("write", self._real_var_trace_name) + self._real_var_trace_name = None + + self._swap_textvariable(self.var) + + def _swap_textvariable(self, new_var: tk.Variable) -> None: + comp = self.component + if comp._textvariable_callback_name: + with contextlib.suppress(Exception): + comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] + comp._textvariable_callback_name = "" + + comp.configure(textvariable=new_var) + + if new_var is not None: + comp._textvariable_callback_name = new_var.trace_add( + "write", comp._textvariable_callback + ) + + def _commit(self) -> None: + shadow_val = self._shadow_var.get() + if shadow_val != self.var.get(): + self._syncing = True + self.var.set(shadow_val) + self._syncing = False + + def validate(self, value: str) -> str | None: + """Return an error string if *value* is invalid, else None.""" + meta = self.ui_state.get_field_metadata(self.var_name) + declared_type = meta.type + nullable = meta.nullable + default_val = meta.default + + if value == "": + if self._required: + return "Value required" + if nullable: + return None + if declared_type is str: + if default_val == "": + return None + return "Value required" + return None + + try: + if declared_type is int: + v = int(value) + if v < 0: + return "Value must be non-negative" + elif declared_type is float: + v = float(value) + if v < 0: + return "Value must be non-negative" + elif declared_type is bool: + if value.lower() not in ("true", "false", "0", "1"): + return "Invalid bool" + except ValueError: + return "Invalid value" + + if self._extra_validate is not None: + return self._extra_validate(value) + + return None + + def _apply_error(self) -> None: + self.component.configure(border_color=ERROR_BORDER_COLOR) + + def _clear_error(self) -> None: + self.component.configure(border_color=self._original_border_color) + + def _validate_and_style(self, value: str) -> bool: + error = self.validate(value) + if error is None: + self._clear_error() + return True + else: + self._apply_error() + return False + + def _on_shadow_write(self, *_args) -> None: + if self._syncing: + return + if not self._touched: + # external sync or initial set — commit immediately + self._commit() + if self._debounce: + self._debounce.cancel() + return + if self._debounce: + self._debounce.call() + if self._undo_debounce: + self._undo_debounce.call() + + def _on_real_var_write(self, *_args) -> None: + if self._syncing: + return + # external change (preset load, file dialog, etc) — sync to shadow var + self._syncing = True + self._shadow_var.set(self.var.get()) + self._syncing = False + self._validate_and_style(self._shadow_var.get()) + + def _push_undo_snapshot(self) -> None: + self._undo.push(self._shadow_var.get()) + + def _on_debounce_fire(self) -> None: + val = self._shadow_var.get() + if self._validate_and_style(val): + self._commit() + + def _on_focus_in(self, _e=None) -> None: + self._touched = False + self._undo.push(self._shadow_var.get()) + + def _on_user_input(self, _e=None) -> None: + self._touched = True + + def _on_focus_out(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + self._undo.push(self._shadow_var.get()) + + def _on_enter(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + + def _set_value(self, value: str) -> None: + self._syncing = True + self._shadow_var.set(value) + self._syncing = False + if self._validate_and_style(value): + self._commit() + + def _on_undo(self, _e=None) -> str: + previous = self._undo.undo(self._shadow_var.get()) + if previous is not None: + self._set_value(previous) + return "break" + + def _on_redo(self, _e=None) -> str: + next_val = self._undo.redo() + if next_val is not None: + self._set_value(next_val) + return "break" + + +class PathValidator(FieldValidator): + """FieldValidator with additional path-specific checks.""" + + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + io_type: PathIOType = PathIOType.INPUT, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) + self.io_type = io_type + + def _get_var_safe(self, name: str) -> tk.Variable | None: + try: + return self.ui_state.get_var(name) + except (KeyError, AttributeError): + return None + + def validate(self, value: str) -> str | None: + base_err = super().validate(value) + if base_err is not None: + return base_err + if value == "": + return None + + prevent_var = self._get_var_safe("prevent_overwrites") + format_var = self._get_var_safe("output_model_format") + return validate_path( + value, + io_type=self.io_type, + prevent_overwrites=prevent_var.get() if prevent_var is not None else False, + output_format=format_var.get() if format_var is not None else None, + ) + + def revalidate(self) -> None: + if self.component.winfo_exists(): + self._validate_and_style(self._shadow_var.get()) + + +def flush_and_validate_all() -> list[str]: + invalid: list[str] = [] + + for v in list(_active_validators): + if v._debounce: + v._debounce.cancel() + + value = v._shadow_var.get() + error = v.validate(value) + + if error is not None: + v._apply_error() + invalid.append(f"{v.var_name}: {error}") + else: + v._clear_error() + v._commit() + + return invalid From a88f369790c73bcae356382165a89b90c31f7a63 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:06:56 +0200 Subject: [PATCH 10/67] refactor: abstract CTK views and migrate logic to controllers Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 1 + .../ui/AdditionalEmbeddingsTabController.py | 14 + modules/ui/BaseAdditionalEmbeddingsTabView.py | 127 +- modules/ui/BaseCaptionUIView.py | 570 +------- modules/ui/BaseCloudTabView.py | 233 ++-- modules/ui/BaseConceptTabView.py | 223 +--- modules/ui/BaseConceptWindowView.py | 985 ++++---------- modules/ui/BaseConfigListView.py | 163 +-- modules/ui/BaseConvertModelUIView.py | 137 +- modules/ui/BaseGenerateCaptionsWindowView.py | 134 +- modules/ui/BaseGenerateMasksWindowView.py | 152 +-- modules/ui/BaseLoraTabView.py | 168 +-- modules/ui/BaseModelTabView.py | 476 +++---- modules/ui/BaseMuonAdamWindowView.py | 73 +- modules/ui/BaseOffloadingWindowView.py | 87 +- modules/ui/BaseOptimizerParamsWindowView.py | 171 +-- modules/ui/BaseProfilingWindowView.py | 66 +- modules/ui/BaseSampleFrameView.py | 137 +- modules/ui/BaseSampleParamsWindowView.py | 39 +- modules/ui/BaseSampleWindowView.py | 225 +--- modules/ui/BaseSamplingTabView.py | 121 +- modules/ui/BaseSchedulerParamsWindowView.py | 122 +- .../ui/BaseTimestepDistributionWindowView.py | 176 +-- modules/ui/BaseTopBarView.py | 186 +-- modules/ui/BaseTrainUIView.py | 919 +++---------- modules/ui/BaseTrainingTabView.py | 1148 ++++++++--------- modules/ui/BaseVideoToolUIView.py | 905 ++----------- modules/ui/CaptionUIController.py | 369 ++---- modules/ui/CloudTabController.py | 31 + modules/ui/ConceptTabController.py | 19 + modules/ui/ConceptWindowController.py | 723 +---------- modules/ui/ConvertModelUIController.py | 116 +- modules/ui/CtkAdditionalEmbeddingsTabView.py | 119 +- modules/ui/CtkCaptionUIView.py | 476 +------ modules/ui/CtkCloudTabView.py | 210 +-- modules/ui/CtkConceptTabView.py | 196 +-- modules/ui/CtkConceptWindowView.py | 903 +------------ modules/ui/CtkConfigListView.py | 369 +----- modules/ui/CtkConvertModelUIView.py | 160 +-- modules/ui/CtkGenerateCaptionsWindowView.py | 41 +- modules/ui/CtkGenerateMasksWindowView.py | 44 +- modules/ui/CtkLoraTabView.py | 131 +- modules/ui/CtkModelTabView.py | 674 +--------- modules/ui/CtkMuonAdamWindowView.py | 104 +- modules/ui/CtkOffloadingWindowView.py | 70 +- modules/ui/CtkOptimizerParamsWindowView.py | 259 +--- modules/ui/CtkProfilingWindowView.py | 54 +- modules/ui/CtkSampleFrameView.py | 111 +- modules/ui/CtkSampleParamsWindowView.py | 29 +- modules/ui/CtkSampleWindowView.py | 167 +-- modules/ui/CtkSamplingTabView.py | 116 +- modules/ui/CtkSchedulerParamsWindowView.py | 77 +- .../ui/CtkTimestepDistributionWindowView.py | 146 +-- modules/ui/CtkTopBarView.py | 252 +--- modules/ui/CtkTrainUIView.py | 778 +++-------- modules/ui/CtkTrainingTabView.py | 847 +----------- modules/ui/CtkVideoToolUIView.py | 911 ++----------- .../ui/GenerateCaptionsWindowController.py | 135 +- modules/ui/GenerateMasksWindowController.py | 153 +-- modules/ui/LoraTabController.py | 22 + modules/ui/ModelTabController.py | 13 + modules/ui/MuonAdamWindowController.py | 28 + modules/ui/OffloadingWindowController.py | 6 + modules/ui/OptimizerParamsWindowController.py | 273 +--- modules/ui/ProfilingWindowController.py | 25 + modules/ui/SampleFrameController.py | 17 + modules/ui/SampleParamsWindowController.py | 8 + modules/ui/SampleWindowController.py | 110 +- modules/ui/SamplingTabController.py | 14 + modules/ui/SchedulerParamsWindowController.py | 17 + .../TimestepDistributionWindowController.py | 146 +-- modules/ui/TopBarController.py | 281 +--- modules/ui/TrainUIController.py | 885 ++----------- modules/ui/TrainingTabController.py | 39 + modules/ui/VideoToolUIController.py | 452 ++----- modules/util/path_util.py | 6 + modules/util/ui/CtkUIState.py | 23 + modules/util/ui/UIState.py | 86 +- .../ui/{components.py => ctk_components.py} | 75 +- modules/util/ui/ctk_validation.py | 276 +--- modules/util/ui/validation.py | 300 +---- scripts/train_ui.py | 4 +- 82 files changed, 3652 insertions(+), 16002 deletions(-) create mode 100644 modules/ui/AdditionalEmbeddingsTabController.py create mode 100644 modules/ui/CloudTabController.py create mode 100644 modules/ui/ConceptTabController.py create mode 100644 modules/ui/LoraTabController.py create mode 100644 modules/ui/ModelTabController.py create mode 100644 modules/ui/MuonAdamWindowController.py create mode 100644 modules/ui/OffloadingWindowController.py create mode 100644 modules/ui/ProfilingWindowController.py create mode 100644 modules/ui/SampleFrameController.py create mode 100644 modules/ui/SampleParamsWindowController.py create mode 100644 modules/ui/SamplingTabController.py create mode 100644 modules/ui/SchedulerParamsWindowController.py create mode 100644 modules/ui/TrainingTabController.py create mode 100644 modules/util/ui/CtkUIState.py rename modules/util/ui/{components.py => ctk_components.py} (88%) diff --git a/.gitignore b/.gitignore index dce73072c..da19fd74e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ # development +PLAN.md # claude's planning file, kept local-only .idea *.bak *.pyc diff --git a/modules/ui/AdditionalEmbeddingsTabController.py b/modules/ui/AdditionalEmbeddingsTabController.py new file mode 100644 index 000000000..638df8e68 --- /dev/null +++ b/modules/ui/AdditionalEmbeddingsTabController.py @@ -0,0 +1,14 @@ + +from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig + + +class AdditionalEmbeddingsTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> TrainEmbeddingConfig: + return TrainEmbeddingConfig.default_values() + + def randomize_uuid(self, embedding_config: TrainEmbeddingConfig) -> TrainEmbeddingConfig: + embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid + return embedding_config diff --git a/modules/ui/BaseAdditionalEmbeddingsTabView.py b/modules/ui/BaseAdditionalEmbeddingsTabView.py index 6a5e3fbe7..744f0abdc 100644 --- a/modules/ui/BaseAdditionalEmbeddingsTabView.py +++ b/modules/ui/BaseAdditionalEmbeddingsTabView.py @@ -1,136 +1,75 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util import path_util import customtkinter as ctk -class AdditionalEmbeddingsTab(ConfigList): - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - attr_name="additional_embeddings", - enable_key="train", - from_external_file=False, - add_button_text="add embedding", - is_full_width=True, - show_toggle_button=True - ) +class BaseAdditionalEmbeddingsTabView(BaseConfigListView): def refresh_ui(self): if self.element_list is not None: - self.element_list.destroy() + self._destroy_frame(self.element_list) self.element_list = None self.widgets_initialized = False self._create_element_list() - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return TrainEmbeddingConfig.default_values() - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: pass -class EmbeddingWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) +class BaseEmbeddingWidgetView: + + def __init__(self, components): + self.components = components - self.element = element - self.ui_state = UIState(self, element) + def build_content(self, top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller): + self.ui_state = ui_state self.i = i self.save_command = save_command - self.grid_columnconfigure(0, weight=1) - - top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - top_frame.grid(row=0, column=0, sticky="nsew") - top_frame.grid_columnconfigure(3, weight=1) - top_frame.grid_columnconfigure(5, weight=1) - - bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - bottom_frame.grid(row=1, column=0, sticky="nsew") - bottom_frame.grid_columnconfigure(7, weight=1) - # close button - close_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) + self.components.colored_icon_button(top_frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) # clone button - clone_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_uuid), - ) - clone_button.grid(row=0, column=1, padx=5) + self.components.colored_icon_button(top_frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i, controller.randomize_uuid), padx=5) # embedding model names - components.label(top_frame, 0, 2, "base embedding:", - tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( + self.components.label(top_frame, 0, 2, "base embedding:", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + self.components.path_entry( top_frame, 0, 3, self.ui_state, "model_name", - mode="file", path_modifier=components.json_path_modifier + mode="file", path_modifier=path_util.json_path_modifier ) # placeholder - components.label(top_frame, 0, 4, "placeholder:", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(top_frame, 0, 5, self.ui_state, "placeholder") + self.components.label(top_frame, 0, 4, "placeholder:", + tooltip="The placeholder used when using the embedding in a prompt") + self.components.entry(top_frame, 0, 5, self.ui_state, "placeholder") # token count - components.label(top_frame, 0, 6, "token count:", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") - token_count_entry.configure(width=40) + self.components.label(top_frame, 0, 6, "token count:", + tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") + self.components.entry(top_frame, 0, 7, self.ui_state, "token_count", width=40) # trainable - components.label(bottom_frame, 0, 0, "train:") - trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) - trainable_switch.configure(width=40) + self.components.label(bottom_frame, 0, 0, "train:") + self.components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command, width=40) # output embedding - components.label(bottom_frame, 0, 2, "output embedding:", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") - output_embedding_switch.configure(width=40) + self.components.label(bottom_frame, 0, 2, "output embedding:", + tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") + self.components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding", width=40) # stop training after - components.label(bottom_frame, 0, 4, "stop training after:", - tooltip="When to stop training the embedding") - components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") + self.components.label(bottom_frame, 0, 4, "stop training after:", + tooltip="When to stop training the embedding") + self.components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") # initial embedding text - components.label(bottom_frame, 0, 6, "initial embedding text:", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") - - def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): - embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid - return embedding_config + self.components.label(bottom_frame, 0, 6, "initial embedding text:", + tooltip="The initial embedding text used when creating a new embedding") + self.components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") def configure_element(self): pass - - def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/BaseCaptionUIView.py b/modules/ui/BaseCaptionUIView.py index e6cc0551e..a6561da69 100644 --- a/modules/ui/BaseCaptionUIView.py +++ b/modules/ui/BaseCaptionUIView.py @@ -1,572 +1,32 @@ -import os import platform -import subprocess -import traceback -from tkinter import filedialog -from modules.module.Blip2Model import Blip2Model -from modules.module.BlipModel import BlipModel -from modules.module.ClipSegModel import ClipSegModel -from modules.module.MaskByColor import MaskByColor -from modules.module.RembgHumanModel import RembgHumanModel -from modules.module.RembgModel import RembgModel -from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow -from modules.util import path_util -from modules.util.image_util import load_image -from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState -import torch +class BaseCaptionUIView: + def __init__(self, components): + self.components = components -import customtkinter as ctk -import cv2 -import numpy as np -from customtkinter import ScalingTracker, ThemeManager -from PIL import Image, ImageDraw - - -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - - self.dir = initial_dir - self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) - self.image_size = 850 - self.help_text = """ - Keyboard shortcuts when focusing on the prompt input field: - Up arrow: previous image - Down arrow: next image - Return: save - Ctrl+M: only show the mask - Ctrl+D: draw mask editing mode - Ctrl+F: fill mask editing mode - - When editing masks: - Left click: add mask - Right click: remove mask - Mouse wheel: increase or decrease brush size""" - self.masking_model = None - self.captioning_model = None - self.image_rel_paths = [] - self.current_image_index = -1 - self.file_list = None - self.image_labels = [] - self.pil_image = None - self.image_width = 0 - self.image_height = 0 - self.pil_mask = None - self.mask_draw_x = 0 - self.mask_draw_y = 0 - self.mask_draw_radius = 0.01 - self.display_only_mask = False - self.image = None - self.image_label = None - self.mask_editing_mode = 'draw' - self.enable_mask_editing_var = ctk.BooleanVar() - self.mask_editing_alpha = None - self.prompt_var = None - self.prompt_component = None - - - self.title("OneTrainer") - self.geometry("1280x980") - self.resizable(False, False) - - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) - - - self.top_bar(self) - - self.bottom_frame = ctk.CTkFrame(self) - self.bottom_frame.grid(row=1, column=0, sticky="nsew") - self.bottom_frame.grid_rowconfigure(0, weight=1) - self.bottom_frame.grid_columnconfigure(0, weight=0) - self.bottom_frame.grid_columnconfigure(1, weight=1) - - self.file_list_column(self.bottom_frame) - self.content_column(self.bottom_frame) - self.load_directory() - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, + def build_top_bar(self, frame, controller, ui_state): + self.components.button(frame, 0, 0, "Open", self.open_directory, tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, + self.components.button(frame, 0, 1, "Generate Masks", self.open_mask_window, tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, + self.components.button(frame, 0, 2, "Generate Captions", self.open_caption_window, tooltip="open a dialog to automatically generate captions") if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, + self.components.button(frame, 0, 3, "Open in Explorer", self.open_in_explorer, tooltip="open the current image in Explorer") - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", + self.components.switch(frame, 0, 4, ui_state, "include_subdirectories", text="include subdirectories") - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - - def file_list_column(self, master): - if self.file_list is not None: - self.image_labels = [] - self.file_list.destroy() - - self.file_list = ctk.CTkScrollableFrame(master, width=300) - self.file_list.grid(row=0, column=0, sticky="nsew") - - for i, filename in enumerate(self.image_rel_paths): - def __create_switch_image(index): - def __switch_image(event): - self.switch_image(index) - - return __switch_image + frame.grid_columnconfigure(5, weight=1) - label = ctk.CTkLabel(self.file_list, text=filename) - label.bind("", __create_switch_image(i)) + self.components.button(frame, 0, 6, "Help", controller.print_help, + tooltip=controller.help_text) - self.image_labels.append(label) - label.grid(row=i, column=0, padx=5, sticky="nsw") - - def content_column(self, master): - image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) - - right_frame = ctk.CTkFrame(master, fg_color="transparent") - right_frame.grid(row=0, column=1, sticky="nsew") - - right_frame.grid_columnconfigure(4, weight=1) - right_frame.grid_rowconfigure(1, weight=1) - - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, + def build_mask_buttons(self, right_frame): + self.components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + self.components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, tooltip="draw a mask using a fill tool") - - # checkbox to enable mask editing - self.enable_mask_editing_var = ctk.BooleanVar() - self.enable_mask_editing_var.set(False) - enable_mask_editing_checkbox = ctk.CTkCheckBox( - right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) - enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") - - # mask alpha textbox - self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") - self.mask_editing_alpha.insert(0, "1.0") - self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) - self.bind_key_events(self.mask_editing_alpha) - - mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) - mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") - - # image - self.image = ctk.CTkImage( - light_image=image, - size=(self.image_size, self.image_size) - ) - self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size - ) - self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") - - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) - - # prompt - self.prompt_var = ctk.StringVar() - self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) - self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") - self.bind_key_events(self.prompt_component) - self.prompt_component.focus_set() - - def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) - component.bind("", self.save) - component.bind("", self.toggle_mask) - component.bind("", self.draw_mask_editing_mode) - component.bind("", self.fill_mask_editing_mode) - - def load_directory(self, include_subdirectories: bool = False): - self.scan_directory(include_subdirectories) - self.file_list_column(self.bottom_frame) - - if len(self.image_rel_paths) > 0: - self.switch_image(0) - else: - self.switch_image(-1) - - self.prompt_component.focus_set() - - def scan_directory(self, include_subdirectories: bool = False): - def __is_supported_image_extension(filename): - name, ext = os.path.splitext(filename) - return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") - - self.image_rel_paths = [] - - if not self.dir or not os.path.isdir(self.dir): - return - - if include_subdirectories: - for root, _, files in os.walk(self.dir): - for filename in files: - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(root, filename), self.dir) - ) - else: - for _, filename in enumerate(os.listdir(self.dir)): - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(self.dir, filename), self.dir) - ) - - def load_image(self): - image_name = "resources/icons/icon.png" - - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.join(self.dir, image_name) - - try: - return load_image(image_name, convert_mode="RGB") - except Exception: - print(f'Could not open image {image_name}') - - def load_mask(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - return load_image(mask_name, convert_mode='RGB') - except Exception: - return None - else: - return None - - def load_prompt(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - try: - with open(prompt_name, "r", encoding='utf-8') as f: - return f.readlines()[0].strip() - except Exception: - return "" - else: - return "" - - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( - text_color=ThemeManager.theme["CTkLabel"]["text_color"]) - - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") - - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() - - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) - - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) - - def refresh_image(self): - if self.pil_mask: - resized_pil_mask = self.pil_mask.resize( - (self.pil_image.width, self.pil_image.height), - Image.Resampling.NEAREST - ) - - if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) - else: - np_image = np.array(self.pil_image).astype(np.float32) / 255.0 - np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 - - # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha - norm_min = 0.3 - np_mask_min = np_mask.min() - if np_mask_min == 0: - # optimize for common case - np_mask = np_mask * (1.0 - norm_min) + norm_min - elif np_mask_min < 1: - # note: min of 1 means we get divide by 0 - np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min - - np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) - masked_image = Image.fromarray(np_masked_image, mode='RGB') - - self.image.configure(light_image=masked_image, size=masked_image.size) - else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) - - def draw_mask_radius(self, delta, raw_event): - # Wheel up = Increase radius. Wheel down = Decrease radius. - multiplier = 1.0 + (delta * 0.05) - self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) - - def edit_mask(self, event): - if not self.enable_mask_editing_var.get(): - return - - if event.widget != self.image_label.children["!label"]: - return - - if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): - return - - display_scaling = ScalingTracker.get_window_scaling(self) - - event_x = event.x / display_scaling - event_y = event.y / display_scaling - - start_x = int(event_x / self.pil_image.width * self.image_width) - start_y = int(event_y / self.pil_image.height * self.image_height) - end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) - end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) - - self.mask_draw_x = event_x - self.mask_draw_y = event_y - - is_right = False - is_left = False - if event.state & 0x0100 or event.num == 1: # left mouse button - is_left = True - elif event.state & 0x0400 or event.num == 3: # right mouse button - is_right = True - - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) - if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) - - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) - - draw = ImageDraw.Draw(self.pil_mask) - draw.line((start_x, start_y, end_x, end_y), fill=color, - width=radius + radius + 1) - draw.ellipse((start_x - radius, start_y - radius, - start_x + radius, start_y + radius), fill=color, outline=None) - draw.ellipse((end_x - radius, end_y - radius, end_x + radius, - end_y + radius), fill=color, outline=None) - - self.refresh_image() - - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - np_mask = np.array(self.pil_mask).astype(np.uint8) - cv2.floodFill(np_mask, None, (start_x, start_y), color) - self.pil_mask = Image.fromarray(np_mask, 'RGB') - - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return - - if self.pil_mask: - self.pil_mask.save(mask_name) - - def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' - - if args: - # disable default event - return "break" - return None - - def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' - - def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask - self.refresh_image() - - def open_directory(self): - new_dir = filedialog.askdirectory() - - if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) - - def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() - - def load_masking_model(self, model): - model_type = type(self.masking_model).__name__ if self.masking_model else None - - if model == "ClipSeg" and model_type != "ClipSegModel": - self._release_models() - print("loading ClipSeg model, this may take a while") - self.masking_model = ClipSegModel(default_device, torch.float32) - elif model == "Rembg" and model_type != "RembgModel": - self._release_models() - print("loading Rembg model, this may take a while") - self.masking_model = RembgModel(default_device, torch.float32) - elif model == "Rembg-Human" and model_type != "RembgHumanModel": - self._release_models() - print("loading Rembg-Human model, this may take a while") - self.masking_model = RembgHumanModel(default_device, torch.float32) - elif model == "Hex Color" and model_type != "MaskByColor": - self._release_models() - self.masking_model = MaskByColor(default_device, torch.float32) - - def load_captioning_model(self, model): - model_type = type(self.captioning_model).__name__ if self.captioning_model else None - - if model == "Blip" and model_type != "BlipModel": - self._release_models() - print("loading Blip model, this may take a while") - self.captioning_model = BlipModel(default_device, torch.float16) - elif model == "Blip2" and model_type != "Blip2Model": - self._release_models() - print("loading Blip2 model, this may take a while") - self.captioning_model = Blip2Model(default_device, torch.float16) - elif model == "WD14 VIT v2" and model_type != "WDModel": - self._release_models() - print("loading WD14_VIT_v2 model, this may take a while") - self.captioning_model = WDModel(default_device, torch.float16) - - def print_help(self): - print(self.help_text) - - def _release_models(self): - """Release all models from VRAM""" - freed = False - if self.captioning_model is not None: - self.captioning_model = None - freed = True - if self.masking_model is not None: - self.masking_model = None - freed = True - if freed: - torch_gc() - - def _on_close(self): - self._release_models() - self.destroy() - - def destroy(self): - self._release_models() - super().destroy() diff --git a/modules/ui/BaseCloudTabView.py b/modules/ui/BaseCloudTabView.py index 99057e428..80be17360 100644 --- a/modules/ui/BaseCloudTabView.py +++ b/modules/ui/BaseCloudTabView.py @@ -1,221 +1,184 @@ -import webbrowser +from abc import ABC, abstractmethod -from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.CloudAction import CloudAction from modules.util.enum.CloudFileSync import CloudFileSync from modules.util.enum.CloudType import CloudType -from modules.util.ui import components -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseCloudTabView(ABC): + def __init__(self, components): + self.components = components -class CloudTab: + @property + def reattach(self): + return self.controller.reattach - def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): - super().__init__() + @abstractmethod + def _make_reattach_frame(self, frame): pass - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.parent = parent - self.reattach = False + @abstractmethod + def _make_create_frame(self, frame): pass - self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, weight=0) - self.frame.grid_columnconfigure(3, weight=1) - self.frame.grid_columnconfigure(4, weight=0) - self.frame.grid_columnconfigure(5, weight=1) + @abstractmethod + def _on_set_gpu_types(self): pass - components.label(self.frame, 0, 0, "Enabled", + def build_content(self, frame, controller, ui_state): + self.components.label(frame, 0, 0, "Enabled", tooltip="Enable cloud training") - components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") + self.components.switch(frame, 0, 1, ui_state, "cloud.enabled") - components.label(self.frame, 1, 0, "Type", + self.components.label(frame, 1, 0, "Type", tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") - components.options_kv(self.frame, 1, 1, [ + self.components.options_kv(frame, 1, 1, [ ("RUNPOD", CloudType.RUNPOD), ("LINUX", CloudType.LINUX), - ], self.ui_state, "cloud.type") + ], ui_state, "cloud.type") - components.label(self.frame, 2, 0, "File sync method", + self.components.label(frame, 2, 0, "File sync method", tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") - components.options_kv(self.frame, 2, 1, [ + self.components.options_kv(frame, 2, 1, [ ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), - ], self.ui_state, "cloud.file_sync") + ], ui_state, "cloud.file_sync") - components.label(self.frame, 3, 0, "API key", + self.components.label(frame, 3, 0, "API key", tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") - components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") + self.components.entry(frame, 3, 1, ui_state, "secrets.cloud.api_key") - components.label(self.frame, 4, 0, "Hostname", + self.components.label(frame, 4, 0, "Hostname", tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") + self.components.entry(frame, 4, 1, ui_state, "secrets.cloud.host") - components.label(self.frame, 5, 0, "Port", + self.components.label(frame, 5, 0, "Port", tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") + self.components.entry(frame, 5, 1, ui_state, "secrets.cloud.port") - components.label(self.frame, 6, 0, "User", + self.components.label(frame, 6, 0, "User", tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') - components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") + self.components.entry(frame, 6, 1, ui_state, "secrets.cloud.user") - components.label(self.frame, 7, 0, "SSH keyfile path", + self.components.label(frame, 7, 0, "SSH keyfile path", tooltip="Absolute path to the private key file used for SSH connections. Leave empty to rely on your system SSH configuration.") - components.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="file") + self.components.path_entry(frame, 7, 1, ui_state, "secrets.cloud.key_file", mode="file") - components.label(self.frame, 8, 0, "SSH password", + self.components.label(frame, 8, 0, "SSH password", tooltip="SSH password for password-based authentication. If you try to use native SCP requires sshpass to be installed. Leave empty to use key-based authentication.") - components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") + self.components.entry(frame, 8, 1, ui_state, "secrets.cloud.password") - components.label(self.frame, 9, 0, "Cloud id", + self.components.label(frame, 9, 0, "Cloud id", tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") - components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") + self.components.entry(frame, 9, 1, ui_state, "secrets.cloud.id") - components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", + self.components.label(frame, 10, 0, "Tensorboard TCP tunnel", tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") - components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + self.components.switch(frame, 10, 1, ui_state, "cloud.tensorboard_tunnel") - - - components.label(self.frame, 1, 2, "Remote Directory", + self.components.label(frame, 1, 2, "Remote Directory", tooltip="The directory on the cloud where files will be uploaded and downloaded.") - components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") - components.label(self.frame, 2, 2, "OneTrainer Directory", + self.components.entry(frame, 1, 3, ui_state, "cloud.remote_dir") + self.components.label(frame, 2, 2, "OneTrainer Directory", tooltip="The directory for OneTrainer on the cloud.") - components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") - components.label(self.frame, 3, 2, "Huggingface cache Directory", + self.components.entry(frame, 2, 3, ui_state, "cloud.onetrainer_dir") + self.components.label(frame, 3, 2, "Huggingface cache Directory", tooltip="Huggingface models are downloaded to this remote directory.") - components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") - components.label(self.frame, 4, 2, "Install OneTrainer", + self.components.entry(frame, 3, 3, ui_state, "cloud.huggingface_cache_dir") + self.components.label(frame, 4, 2, "Install OneTrainer", tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") - components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") - components.label(self.frame, 5, 2, "Install command", + self.components.switch(frame, 4, 3, ui_state, "cloud.install_onetrainer") + self.components.label(frame, 5, 2, "Install command", tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") - components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") - components.label(self.frame, 6, 2, "Update OneTrainer", + self.components.entry(frame, 5, 3, ui_state, "cloud.install_cmd") + self.components.label(frame, 6, 2, "Update OneTrainer", tooltip="Update OneTrainer if it already exists on the cloud.") - components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + self.components.switch(frame, 6, 3, ui_state, "cloud.update_onetrainer") - components.label(self.frame, 8, 2, "Detach remote trainer", + self.components.label(frame, 8, 2, "Detach remote trainer", tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") - components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") - components.label(self.frame, 9, 2, "Reattach id", + self.components.switch(frame, 8, 3, ui_state, "cloud.detach_trainer") + self.components.label(frame, 9, 2, "Reattach id", tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") - reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") - reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") - reattach_frame.grid_columnconfigure(0, weight=1) - reattach_frame.grid_columnconfigure(1, weight=1) - components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) - components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) - - components.label(self.frame, 11, 2, "Download samples", + reattach_frame = self._make_reattach_frame(frame) + self.components.entry(reattach_frame, 0, 0, ui_state, "cloud.run_id", width=60) + self.components.button(reattach_frame, 0, 1, "Reattach now", controller.do_reattach) + + self.components.label(frame, 11, 2, "Download samples", tooltip="Download samples from the remote workspace directory to your local machine.") - components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") - components.label(self.frame, 12, 2, "Download output model", + self.components.switch(frame, 11, 3, ui_state, "cloud.download_samples") + self.components.label(frame, 12, 2, "Download output model", tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") - components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") - components.label(self.frame, 13, 2, "Download saved checkpoints", + self.components.switch(frame, 12, 3, ui_state, "cloud.download_output_model") + self.components.label(frame, 13, 2, "Download saved checkpoints", tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") - components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") - components.label(self.frame, 14, 2, "Download backups", + self.components.switch(frame, 13, 3, ui_state, "cloud.download_saves") + self.components.label(frame, 14, 2, "Download backups", tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") - components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") - components.label(self.frame, 15, 2, "Download tensorboard logs", + self.components.switch(frame, 14, 3, ui_state, "cloud.download_backups") + self.components.label(frame, 15, 2, "Download tensorboard logs", tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") - components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") - components.label(self.frame, 16, 2, "Delete remote workspace", + self.components.switch(frame, 15, 3, ui_state, "cloud.download_tensorboard") + self.components.label(frame, 16, 2, "Delete remote workspace", tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") - components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + self.components.switch(frame, 16, 3, ui_state, "cloud.delete_workspace") - components.label(self.frame, 1, 4, "Create cloud via API", + self.components.label(frame, 1, 4, "Create cloud via API", tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") - create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") - create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") - create_frame.grid_columnconfigure(0, weight=0) - create_frame.grid_columnconfigure(1, weight=1) - components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") - components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) - - components.label(self.frame, 2, 4, "Cloud name", + create_frame = self._make_create_frame(frame) + self.components.switch(create_frame, 0, 0, ui_state, "cloud.create") + self.components.button(create_frame, 0, 1, "Create cloud via website", controller.open_create_cloud_url) + + self.components.label(frame, 2, 4, "Cloud name", tooltip="The name of the new cloud instance.") - components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") - components.label(self.frame, 3, 4, "Type", + self.components.entry(frame, 2, 5, ui_state, "cloud.name") + self.components.label(frame, 3, 4, "Type", tooltip="Select the RunPod cloud type. See RunPod's website for details.") - components.options_kv(self.frame, 3, 5, [ + self.components.options_kv(frame, 3, 5, [ ("", ""), ("Community", "COMMUNITY"), ("Secure", "SECURE"), - ], self.ui_state, "cloud.sub_type") - + ], ui_state, "cloud.sub_type") - components.label(self.frame, 4, 4, "GPU", + self.components.label(frame, 4, 4, "GPU", tooltip="Select the GPU type. Enter an API key before pressing the button.") + _, gpu_components = self.components.options_adv(frame, 4, 5, [("")], ui_state, "cloud.gpu_type", adv_command=self._on_set_gpu_types) + self.gpu_types_menu = gpu_components['component'] - _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) - self.gpu_types_menu=gpu_components['component'] - - components.label(self.frame, 5, 4, "Volume size", + self.components.label(frame, 5, 4, "Volume size", tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") - components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") + self.components.entry(frame, 5, 5, ui_state, "cloud.volume_size") - components.label(self.frame, 6, 4, "Min download", + self.components.label(frame, 6, 4, "Min download", tooltip="Set the minimum download speed of the cloud in Mbps.") - components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") + self.components.entry(frame, 6, 5, ui_state, "cloud.min_download") - components.label(self.frame, 8, 4, "Action on finish", + self.components.label(frame, 8, 4, "Action on finish", tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") - components.options_kv(self.frame, 8, 5, [ + self.components.options_kv(frame, 8, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_finish") + ], ui_state, "cloud.on_finish") - components.label(self.frame, 9, 4, "Action on error", + self.components.label(frame, 9, 4, "Action on error", tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") - components.options_kv(self.frame, 9, 5, [ + self.components.options_kv(frame, 9, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_error") + ], ui_state, "cloud.on_error") - components.label(self.frame, 10, 4, "Action on detached finish", + self.components.label(frame, 10, 4, "Action on detached finish", tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 10, 5, [ + self.components.options_kv(frame, 10, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_finish") + ], ui_state, "cloud.on_detached_finish") - components.label(self.frame, 11, 4, "Action on detached error", + self.components.label(frame, 11, 4, "Action on detached error", tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 11, 5, [ + self.components.options_kv(frame, 11, 5, [ ("None", CloudAction.NONE), ("Stop", CloudAction.STOP), ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_error") - - self.frame.pack(fill="both", expand=1) - - def __set_gpu_types(self): - self.gpu_types_menu.configure(values=[]) - if self.train_config.cloud.type == CloudType.RUNPOD: - import runpod - runpod.api_key=self.train_config.secrets.cloud.api_key - gpus=runpod.get_gpus() - self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) - - def __reattach(self): - self.reattach=True - try: - self.parent.start_training() - finally: - self.reattach=False - - def __create_cloud(self): - if self.train_config.cloud.type == CloudType.RUNPOD: - webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) + ], ui_state, "cloud.on_detached_error") diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 0b6505694..077a10c24 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -1,112 +1,24 @@ import os import pathlib -from tkinter import BooleanVar, StringVar -from modules.ui.ConceptWindow import ConceptWindow -from modules.ui.ConfigList import ConfigList +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.ui.ConceptWindowController import ConceptWindowController from modules.util import path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DebounceTimer -import customtkinter as ctk from PIL import Image -class ConceptTab(ConfigList): +class BaseConceptTabView(BaseConfigListView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - self.search_var = StringVar() - self.filter_var = StringVar(value="ALL") - self.show_disabled_var = BooleanVar(value=True) - - super().__init__( - master, - train_config, - ui_state, - from_external_file=True, - attr_name="concept_file_name", - config_dir="training_concepts", - default_config_name="concepts.json", - add_button_text="Add Concept", - add_button_tooltip="Adds a new concept to the current config.", - is_full_width=False, - show_toggle_button=True - ) - self._toolbar = None - self._toolbar_is_wrapped = False - self._add_search_bar() - # wrap toolbar if too narrow - self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) - - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return ConceptConfig.default_values() - - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) - - def _add_search_bar(self): - toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") - toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") - toolbar.grid_columnconfigure(2, weight=1) - self._toolbar = toolbar - - # Search - ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) - self.search_var = StringVar() - self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, - placeholder_text="Filter...", width=200) - self.search_entry.grid(row=0, column=1) - self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) - self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) - - # Spacer - ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) - - # Type filter - ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) - self.filter_var = StringVar(value="ALL") - ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], - variable=self.filter_var, command=lambda x: self._update_filters(), - width=150).grid(row=0, column=4) - - # Show disabled checkbox - self.show_disabled_var = BooleanVar(value=True) - self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, - command=self._update_filters, width=100) - self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) - self._refresh_show_disabled_text() - - # Clear button - ctk.CTkButton(toolbar, text="Clear", width=50, - command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) - - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() - - def _reset_filters(self): - self.search_var.set("") - self.filter_var.set("ALL") - self.show_disabled_var.set(True) - self._update_filters() + _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] def _element_matches_filters(self, element): - # Check enabled status if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: return False - # Search filter search = self.filters.get("search", "").lower() if search: if not hasattr(element, '_search_cache'): @@ -127,7 +39,6 @@ def _element_matches_filters(self, element): if not any(search in text for text in getattr(element, '_search_cache', [])): return False - # Type filter type_filter = self.filters.get("type", "ALL") if type_filter != "ALL": if hasattr(element, 'type') and element.type: @@ -139,101 +50,19 @@ def _element_matches_filters(self, element): return True - def _maybe_reposition_toolbar(self, width): - if not self._toolbar: - return - threshold = 1070 - want_wrapped = width < threshold - if want_wrapped == self._toolbar_is_wrapped: - return - self._toolbar_is_wrapped = want_wrapped - if want_wrapped: - self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) - else: - self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) - - def _refresh_show_disabled_text(self): - try: - disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) - except (AttributeError, TypeError): - disabled_count = 0 - text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" - try: - if getattr(self, 'show_disabled_checkbox', None): - self.show_disabled_checkbox.configure(text=text) - except (AttributeError, RuntimeError): - pass - - -class ConceptWidget(ctk.CTkFrame): - def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, width=150, height=170, corner_radius=10, bg_color="transparent" - ) - - self.concept = concept - self.ui_state = UIState(self, concept) - self.image_ui_state = UIState(self, concept.image) - self.text_ui_state = UIState(self, concept.text) - self.i = i - - self.grid_rowconfigure(1, weight=1) - - # image - self.image = ctk.CTkImage( - light_image=self.__get_preview_image(), - size=(150, 150) - ) - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150) - image_label.grid(row=0, column=0) - - # name - self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) - - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.place(x=0, y=0) - - # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_seed), - ) - clone_button.place(x=25, y=0) + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() - # enabled switch - enabled_switch = ctk.CTkSwitch( - master=self, - width=40, - variable=self.ui_state.get_var("enabled"), - text="", - command=save_command, - ) - enabled_switch.place(x=110, y=0) - image_label.bind( - "", - lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) - ) +class BaseConceptWidgetView: - def __randomize_seed(self, concept: ConceptConfig): - concept.seed = ConceptConfig.default_values().seed - return concept + def __init__(self, components): + self.components = components - def __get_display_name(self): + def _get_display_name(self): if self.concept.name: return self.concept.name elif self.concept.path: @@ -241,20 +70,11 @@ def __get_display_name(self): else: return "" - def configure_element(self): - self.name_label.configure(text=self.__get_display_name()) - self.image.configure(light_image=self.__get_preview_image()) - try: - if hasattr(self.concept, '_search_cache'): - delattr(self.concept, '_search_cache') - except AttributeError: - pass - - def __get_preview_image(self): + def _get_preview_image(self): preview_path = "resources/icons/icon.png" glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" - concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) + concept_path = ConceptWindowController.get_concept_path(getattr(self.concept, 'path', None)) if concept_path: for path in pathlib.Path(concept_path).glob(glob_pattern): if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): @@ -268,7 +88,7 @@ def __get_preview_image(self): break try: image = load_image(preview_path, convert_mode="RGBA") - except (OSError): + except OSError: image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) size = min(image.width, image.height) image = image.crop(( @@ -279,8 +99,9 @@ def __get_preview_image(self): )) return image.resize((150, 150), Image.Resampling.BILINEAR) - def place_in_list(self): - index = getattr(self, 'visible_index', self.i) - x = index % 6 - y = index // 6 - self.grid(row=y, column=x, pady=5, padx=5) + def _clear_search_cache(self): + try: + if hasattr(self.concept, '_search_cache'): + delattr(self.concept, '_search_cache') + except AttributeError: + pass diff --git a/modules/ui/BaseConceptWindowView.py b/modules/ui/BaseConceptWindowView.py index f58879d5f..0f851f5fb 100644 --- a/modules/ui/BaseConceptWindowView.py +++ b/modules/ui/BaseConceptWindowView.py @@ -1,840 +1,405 @@ import fractions import math -import os -import pathlib -import platform -import random -import threading -import time -import traceback - -from modules.util import concept_stats, path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig + +from modules.util import path_util from modules.util.enum.BalancingStrategy import BalancingStrategy from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -from mgds.LoadingPipeline import LoadingPipeline -from mgds.OutputPipelineModule import OutputPipelineModule -from mgds.PipelineModule import PipelineModule -from mgds.pipelineModules.CapitalizeTags import CapitalizeTags -from mgds.pipelineModules.DropTags import DropTags -from mgds.pipelineModules.RandomBrightness import RandomBrightness -from mgds.pipelineModules.RandomCircularMaskShrink import ( - RandomCircularMaskShrink, -) -from mgds.pipelineModules.RandomContrast import RandomContrast -from mgds.pipelineModules.RandomFlip import RandomFlip -from mgds.pipelineModules.RandomHue import RandomHue -from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop -from mgds.pipelineModules.RandomRotate import RandomRotate -from mgds.pipelineModules.RandomSaturation import RandomSaturation -from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) - -import torch -from torchvision.transforms import functional - -import customtkinter as ctk -import huggingface_hub -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -from PIL import Image - - -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - - def get_inputs(self) -> list[str]: - return [] - - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.train_config = train_config - - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - self.text_ui_state = text_ui_state - self.image_preview_file_index = 0 - self.preview_augmentations = ctk.BooleanVar(self, True) - self.bucket_fig = None - - self.title("Concept") - self.geometry("800x700") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) + +class BaseConceptWindowView: + def __init__(self, components): + self.components = components + + def build_general_tab(self, frame, controller, ui_state, text_ui_state): # name - components.label(frame, 0, 0, "Name", + self.components.label(frame, 0, 0, "Name", tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") + self.components.entry(frame, 0, 1, ui_state, "name") # enabled - components.label(frame, 1, 0, "Enabled", + self.components.label(frame, 1, 0, "Enabled", tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") + self.components.switch(frame, 1, 1, ui_state, "enabled") # concept type - components.label(frame, 2, 0, "Concept Type", + self.components.label(frame, 2, 0, "Concept Type", tooltip="STANDARD: Standard finetuning with the sample as training target\n" "VALIDATION: Use concept for validation instead of training\n" "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " "Only implemented for LoRA.", wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") + self.components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], ui_state, "type") # path - components.label(frame, 3, 0, "Path", + self.components.label(frame, 3, 0, "Path", tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, + self.components.path_entry(frame, 3, 1, ui_state, "path", mode="dir") + self.components.button(frame, 3, 2, text="download now", command=controller.download_dataset_threaded, tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.") # prompt source - components.label(frame, 4, 0, "Prompt Source", + self.components.label(frame, 4, 0, "Prompt Source", tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts") - prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") + prompt_path_entry = self.components.path_entry(frame, 4, 2, text_ui_state, "prompt_path", mode="file") def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") + self.components.set_widget_enabled(prompt_path_entry, option == 'concept') - components.options_kv(frame, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("From text file per sample", 'sample'), ("From single text file", 'concept'), ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) + ], text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(controller.concept.text.prompt_source) # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", + self.components.label(frame, 5, 0, "Include Subdirectories", tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") + self.components.switch(frame, 5, 1, ui_state, "include_subdirectories") # image variations - components.label(frame, 6, 0, "Image Variations", + self.components.label(frame, 6, 0, "Image Variations", tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") + self.components.entry(frame, 6, 1, ui_state, "image_variations") # text variations - components.label(frame, 7, 0, "Text Variations", + self.components.label(frame, 7, 0, "Text Variations", tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") + self.components.entry(frame, 7, 1, ui_state, "text_variations") # balancing - components.label(frame, 8, 0, "Balancing", + self.components.label(frame, 8, 0, "Balancing", tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") + self.components.entry(frame, 8, 1, ui_state, "balancing") + self.components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], ui_state, "balancing_strategy") # loss weight - components.label(frame, 9, 0, "Loss Weight", + self.components.label(frame, 9, 0, "Loss Weight", tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 9, 1, ui_state, "loss_weight") + def build_image_augmentation_tab(self, frame, controller, image_ui_state): # header - components.label(frame, 0, 1, "Random", + self.components.label(frame, 0, 1, "Random", tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", + self.components.label(frame, 0, 2, "Fixed", tooltip="Enable this augmentation with fixed values") # crop jitter - components.label(frame, 1, 0, "Crop Jitter", + self.components.label(frame, 1, 0, "Crop Jitter", tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") + self.components.switch(frame, 1, 1, image_ui_state, "enable_crop_jitter") # random flip - components.label(frame, 2, 0, "Random Flip", + self.components.label(frame, 2, 0, "Random Flip", tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") + self.components.switch(frame, 2, 1, image_ui_state, "enable_random_flip") + self.components.switch(frame, 2, 2, image_ui_state, "enable_fixed_flip") # random rotation - components.label(frame, 3, 0, "Random Rotation", + self.components.label(frame, 3, 0, "Random Rotation", tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") + self.components.switch(frame, 3, 1, image_ui_state, "enable_random_rotate") + self.components.switch(frame, 3, 2, image_ui_state, "enable_fixed_rotate") + self.components.entry(frame, 3, 3, image_ui_state, "random_rotate_max_angle") # random brightness - components.label(frame, 4, 0, "Random Brightness", + self.components.label(frame, 4, 0, "Random Brightness", tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") + self.components.switch(frame, 4, 1, image_ui_state, "enable_random_brightness") + self.components.switch(frame, 4, 2, image_ui_state, "enable_fixed_brightness") + self.components.entry(frame, 4, 3, image_ui_state, "random_brightness_max_strength") # random contrast - components.label(frame, 5, 0, "Random Contrast", + self.components.label(frame, 5, 0, "Random Contrast", tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") + self.components.switch(frame, 5, 1, image_ui_state, "enable_random_contrast") + self.components.switch(frame, 5, 2, image_ui_state, "enable_fixed_contrast") + self.components.entry(frame, 5, 3, image_ui_state, "random_contrast_max_strength") # random saturation - components.label(frame, 6, 0, "Random Saturation", + self.components.label(frame, 6, 0, "Random Saturation", tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") + self.components.switch(frame, 6, 1, image_ui_state, "enable_random_saturation") + self.components.switch(frame, 6, 2, image_ui_state, "enable_fixed_saturation") + self.components.entry(frame, 6, 3, image_ui_state, "random_saturation_max_strength") # random hue - components.label(frame, 7, 0, "Random Hue", + self.components.label(frame, 7, 0, "Random Hue", tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") + self.components.switch(frame, 7, 1, image_ui_state, "enable_random_hue") + self.components.switch(frame, 7, 2, image_ui_state, "enable_fixed_hue") + self.components.entry(frame, 7, 3, image_ui_state, "random_hue_max_strength") # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", + self.components.label(frame, 8, 0, "Circular Mask Generation", tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") + self.components.switch(frame, 8, 1, image_ui_state, "enable_random_circular_mask_shrink") # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", + self.components.label(frame, 9, 0, "Random Rotate and Crop", tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") + self.components.switch(frame, 9, 1, image_ui_state, "enable_random_mask_rotate_crop") # circular mask generation - components.label(frame, 10, 0, "Resolution Override", + self.components.label(frame, 10, 0, "Resolution Override", tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") - - # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image = ctk.CTkImage( - light_image=image_preview, - size=image_preview.size, - ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) - image_label.grid(row=0, column=4, rowspan=6) - - # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") - update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") - update_button_frame.grid_columnconfigure(1, weight=1) - - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) - preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) - - prev_preview_button.configure(width=40) - next_preview_button.configure(width=40) - - #caption and filename preview - self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) - self.filename_preview.grid(row=2, column=0, columnspan=3) - self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.switch(frame, 10, 2, image_ui_state, "enable_resolution_override") + self.components.entry(frame, 10, 3, image_ui_state, "resolution_override") + def build_text_augmentation_tab(self, frame, controller, text_ui_state): # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", + self.components.label(frame, 0, 0, "Tag Shuffling", tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") + self.components.switch(frame, 0, 1, text_ui_state, "enable_tag_shuffling") # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", + self.components.label(frame, 1, 0, "Tag Delimiter", tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") + self.components.entry(frame, 1, 1, text_ui_state, "tag_delimiter") # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", + self.components.label(frame, 2, 0, "Keep Tag Count", tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") + self.components.entry(frame, 2, 1, text_ui_state, "keep_tags_count") # tag dropout - components.label(frame, 3, 0, "Tag Dropout", + self.components.label(frame, 3, 0, "Tag Dropout", tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", + self.components.switch(frame, 3, 1, text_ui_state, "tag_dropout_enable") + self.components.label(frame, 4, 0, "Dropout Mode", tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("Full", 'FULL'), ("Random", 'RANDOM'), ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", + ], text_ui_state, "tag_dropout_mode", None) + self.components.label(frame, 4, 2, "Probability", tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") + self.components.entry(frame, 4, 3, text_ui_state, "tag_dropout_probability") - components.label(frame, 5, 0, "Special Dropout Tags", + self.components.label(frame, 5, 0, "Special Dropout Tags", tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ + self.components.options_kv(frame, 5, 1, [ ("None", 'NONE'), ("Blacklist", 'BLACKLIST'), ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", + ], text_ui_state, "tag_dropout_special_tags_mode", None) + self.components.entry(frame, 5, 2, text_ui_state, "tag_dropout_special_tags") + self.components.label(frame, 6, 0, "Special Tags Regex", tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") + self.components.switch(frame, 6, 1, text_ui_state, "tag_dropout_special_tags_regex") #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", + self.components.label(frame, 7, 0, "Randomize Capitalization", tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", + self.components.switch(frame, 7, 1, text_ui_state, "caps_randomize_enable") + self.components.label(frame, 7, 2, "Force Lowercase", tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") + self.components.switch(frame, 7, 3, text_ui_state, "caps_randomize_lowercase") - components.label(frame, 8, 0, "Captialization Mode", + self.components.label(frame, 8, 0, "Captialization Mode", tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", + self.components.entry(frame, 8, 1, text_ui_state, "caps_randomize_mode") + self.components.label(frame, 8, 2, "Probability", tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame + self.components.entry(frame, 8, 3, text_ui_state, "caps_randomize_probability") - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - - self.cancel_scan_flag = threading.Event() + def build_concept_stats_tab(self, frame, controller): + self.concept_stats_tab = frame #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") + self.file_size_label = self.components.label(frame, 1, 0, "Total Size", pad=0, + tooltip="Total size of all image, mask, and caption files in MB", underline=True) + self.file_size_preview = self.components.label(frame, 2, 0, pad=0, text="-") #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") + self.dir_count_label = self.components.label(frame, 1, 1, "Directories", pad=0, + tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory", underline=True) + self.dir_count_preview = self.components.label(frame, 2, 1, pad=0, text="-") #basic img/vid stats - count of each type in the concept #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") + self.image_count_label = self.components.label(frame, 3, 0, "\nTotal Images", pad=0, + tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'", underline=True) + self.image_count_preview = self.components.label(frame, 4, 0, pad=0, text="-") + self.video_count_label = self.components.label(frame, 3, 1, "\nTotal Videos", pad=0, + tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS), underline=True) + self.video_count_preview = self.components.label(frame, 4, 1, pad=0, text="-") + self.mask_count_label = self.components.label(frame, 3, 2, "\nTotal Masks", pad=0, + tooltip="Total number of mask files, any file ending in '-masklabel.png'", underline=True) + self.mask_count_preview = self.components.label(frame, 4, 2, pad=0, text="-") + self.caption_count_label = self.components.label(frame, 3, 3, "\nTotal Captions", pad=0, + tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.", underline=True) + self.caption_count_preview = self.components.label(frame, 4, 3, pad=0, text="-") #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") + self.image_count_mask_label = self.components.label(frame, 5, 0, "\nImages with Masks", pad=0, + tooltip="Total number of image files with an associated mask", underline=True) + self.image_count_mask_preview = self.components.label(frame, 6, 0, pad=0, text="-") + self.mask_count_label_unpaired = self.components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, + tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!", underline=True) + self.mask_count_preview_unpaired = self.components.label(frame, 6, 1, pad=0, text="-") #currently no masks for videos? - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") + self.image_count_caption_label = self.components.label(frame, 7, 0, "\nImages with Captions", pad=0, + tooltip="Total number of image files with an associated caption", underline=True) + self.image_count_caption_preview = self.components.label(frame, 8, 0, pad=0, text="-") + self.video_count_caption_label = self.components.label(frame, 7, 1, "\nVideos with Captions", pad=0, + tooltip="Total number of video files with an associated caption", underline=True) + self.video_count_caption_preview = self.components.label(frame, 8, 1, pad=0, text="-") + self.caption_count_label_unpaired = self.components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, + tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.", underline=True) + self.caption_count_preview_unpaired = self.components.label(frame, 8, 2, pad=0, text="-") #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) + self.pixel_max_label = self.components.label(frame, 9, 0, "\nMax Pixels", pad=0, + tooltip="Largest image in the concept by number of pixels (width * height)", underline=True) + self.pixel_max_preview = self.components.label(frame, 10, 0, pad=0, text="-", wraplength=150) + self.pixel_avg_label = self.components.label(frame, 9, 1, "\nAvg Pixels", pad=0, + tooltip="Average size of images in the concept by number of pixels (width * height)", underline=True) + self.pixel_avg_preview = self.components.label(frame, 10, 1, pad=0, text="-", wraplength=150) + self.pixel_min_label = self.components.label(frame, 9, 2, "\nMin Pixels", pad=0, + tooltip="Smallest image in the concept by number of pixels (width * height)", underline=True) + self.pixel_min_preview = self.components.label(frame, 10, 2, pad=0, text="-", wraplength=150) #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) + self.length_max_label = self.components.label(frame, 11, 0, "\nMax Length", pad=0, + tooltip="Longest video in the concept by number of frames", underline=True) + self.length_max_preview = self.components.label(frame, 12, 0, pad=0, text="-", wraplength=150) + self.length_avg_label = self.components.label(frame, 11, 1, "\nAvg Length", pad=0, + tooltip="Average length of videos in the concept by number of frames", underline=True) + self.length_avg_preview = self.components.label(frame, 12, 1, pad=0, text="-", wraplength=150) + self.length_min_label = self.components.label(frame, 11, 2, "\nMin Length", pad=0, + tooltip="Shortest video in the concept by number of frames", underline=True) + self.length_min_preview = self.components.label(frame, 12, 2, pad=0, text="-", wraplength=150) #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) + self.fps_max_label = self.components.label(frame, 13, 0, "\nMax FPS", pad=0, + tooltip="Video in concept with highest fps", underline=True) + self.fps_max_preview = self.components.label(frame, 14, 0, pad=0, text="-", wraplength=150) + self.fps_avg_label = self.components.label(frame, 13, 1, "\nAvg FPS", pad=0, + tooltip="Average fps of videos in the concept", underline=True) + self.fps_avg_preview = self.components.label(frame, 14, 1, pad=0, text="-", wraplength=150) + self.fps_min_label = self.components.label(frame, 13, 2, "\nMin FPS", pad=0, + tooltip="Video in concept with the lowest fps", underline=True) + self.fps_min_preview = self.components.label(frame, 14, 2, pad=0, text="-", wraplength=150) #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) + self.caption_max_label = self.components.label(frame, 15, 0, "\nMax Caption Length", pad=0, + tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_max_preview = self.components.label(frame, 16, 0, pad=0, text="-", wraplength=150) + self.caption_avg_label = self.components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, + tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_avg_preview = self.components.label(frame, 16, 1, pad=0, text="-", wraplength=150) + self.caption_min_label = self.components.label(frame, 15, 2, "\nMin Caption Length", pad=0, + tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word", underline=True) + self.caption_min_preview = self.components.label(frame, 16, 2, pad=0, text="-", wraplength=150) #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, + self.aspect_bucket_label = self.components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - - #aspect bucketing plot, mostly copied from timestep preview graph - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - plt.set_loglevel('WARNING') #suppress errors about data type in bar chart - - assert self.bucket_fig is None - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) - self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) - - self.bucket_fig.set_facecolor(background_color) - self.bucket_ax.set_facecolor(background_color) - self.bucket_ax.spines['bottom'].set_color(self.text_color) - self.bucket_ax.spines['left'].set_color(self.text_color) - self.bucket_ax.spines['top'].set_visible(False) - self.bucket_ax.spines['right'].set_color(self.text_color) - self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") - self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") - self.bucket_ax.xaxis.label.set_color(self.text_color) - self.bucket_ax.yaxis.label.set_color(self.text_color) + Images which don't match a bucket exactly are cropped to the nearest one.", underline=True) + self.small_bucket_label = self.components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, + tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.", underline=True) + self.small_bucket_preview = self.components.label(frame, 18, 1, pad=0, text="-") #refresh stats - must be after all labels are defined or will give error - self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), + self.refresh_basic_stats_button = self.components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: controller.get_concept_stats_threaded(self, False, 9999), tooltip="Reload basic statistics for the concept directory") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), + self.refresh_advanced_stats_button = self.components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: controller.get_concept_stats_threaded(self, True, 9999), tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), + self.cancel_stats_button = self.components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self._cancel_concept_stats(controller), tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - - frame.pack(fill="both", expand=1) - return frame - - def __prev_image_preview(self): - self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() - - def __next_image_preview(self): - self.image_preview_file_index += 1 - self.__update_image_preview() - - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image.configure(light_image=image_preview, size=image_preview.size) - self.filename_preview.configure(text=filename_preview) - self.caption_preview.configure(state="normal") - self.caption_preview.delete(index1="1.0", index2="end") - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - - @staticmethod - def get_concept_path(path: str) -> str | None: - if os.path.isdir(path): - return path - try: - #don't download, only check if available locally: - return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) - except Exception: - return None - - def __download_dataset(self): - try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) - huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") - except Exception: - traceback.print_exc() - - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) - download_thread.start() - - def _read_text_file_for_preview(self, file_path: str) -> str: - empty_msg = "[Empty prompt]" - try: - with open(file_path, "r") as f: - if self.preview_augmentations.get(): - lines = [line.strip() for line in f if line.strip()] - return random.choice(lines) if lines else empty_msg - content = f.read().strip() - return content if content else empty_msg - except FileNotFoundError: - return "File not found, please check the path" - except IsADirectoryError: - return "[Provided path is a directory, please correct the caption path]" - except PermissionError: - if platform.system() == "Windows": - return "[Permission denied, please check the file permissions or Windows Defender settings]" - else: - return "[Permission denied, please check the file permissions]" - except UnicodeDecodeError: - return "[Invalid file encoding. This should not happen, please report this issue]" - - def __get_preview_image(self): - preview_image_path = "resources/icons/icon.png" - file_index = -1 - glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" - - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if path.is_file() and path_util.is_supported_image_extension(extension) \ - and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): - preview_image_path = path_util.canonical_join(concept_path, path) - file_index += 1 - if file_index == self.image_preview_file_index: - break - - image = load_image(preview_image_path, 'RGB') - image_tensor = functional.to_tensor(image) - - splitext = os.path.splitext(preview_image_path) - preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") - if not os.path.isfile(preview_mask_path): - preview_mask_path = None - - if preview_mask_path: - mask = Image.open(preview_mask_path).convert("L") - mask_tensor = functional.to_tensor(mask) - else: - mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) + self.processing_time = self.components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - source = self.concept.text.prompt_source - preview_p = pathlib.Path(preview_image_path) - if source == "filename": - prompt_output = preview_p.stem or "[Empty prompt]" - else: - file_map = { - "sample": preview_p.with_suffix(".txt"), - "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, - } - file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" - - modules = [] - if self.preview_augmentations.get(): - input_module = InputPipelineModule({ - 'true': True, - 'image': image_tensor, - 'mask': mask_tensor, - 'enable_random_flip': self.concept.image.enable_random_flip, - 'enable_fixed_flip': self.concept.image.enable_fixed_flip, - 'enable_random_rotate': self.concept.image.enable_random_rotate, - 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, - 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, - 'enable_random_brightness': self.concept.image.enable_random_brightness, - 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, - 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, - 'enable_random_contrast': self.concept.image.enable_random_contrast, - 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, - 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, - 'enable_random_saturation': self.concept.image.enable_random_saturation, - 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, - 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, - 'enable_random_hue': self.concept.image.enable_random_hue, - 'enable_fixed_hue': self.concept.image.enable_fixed_hue, - 'random_hue_max_strength': self.concept.image.random_hue_max_strength, - 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, - 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, - - 'prompt' : prompt_output, - 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, - 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, - 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, - 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, - 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, - 'tag_delimiter' : self.concept.text.tag_delimiter, - 'keep_tags_count' : self.concept.text.keep_tags_count, - 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, - 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, - 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, - 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, - 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, - 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, - }) - - circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') - random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') - random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') - random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') - random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') - random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') - random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') - random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') - drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', - special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', - keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') - caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', - capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') - shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') - output_module = OutputPipelineModule(['image', 'mask', 'prompt']) - - modules = [ - input_module, - circular_mask_shrink, - random_mask_rotate_crop, - random_flip, - random_rotate, - random_brightness, - random_contrast, - random_saturation, - random_hue, - drop_tags, - caps_randomize, - shuffle_tags, - output_module, - ] - - pipeline = LoadingPipeline( - device=torch.device('cpu'), - modules=modules, - batch_size=1, - seed=random.randint(0, 2**30), - state=None, - initial_epoch=0, - initial_index=0, - ) - - data = pipeline.__next__() - image_tensor = data['image'] - mask_tensor = data['mask'] - prompt_output = data['prompt'] - - filename_output = os.path.basename(preview_image_path) - - mask_tensor = torch.clamp(mask_tensor, 0.3, 1) - image_tensor = image_tensor * mask_tensor - - image = functional.to_pil_image(image_tensor) - - image.thumbnail((300, 300)) - - return image, filename_output, prompt_output - - def __update_concept_stats(self): + def _update_concept_stats(self, controller): #file size - self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") - self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") + self.components.set_label_text(self.file_size_preview, str(int(controller.concept.concept_stats["file_size"]/1048576)) + " MB") + self.components.set_label_text(self.processing_time, str(round(controller.concept.concept_stats["processing_time"], 2)) + " s") #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) + self.components.set_label_text(self.dir_count_preview, controller.concept.concept_stats["directory_count"]) #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) + self.components.set_label_text(self.image_count_preview, controller.concept.concept_stats["image_count"]) + self.components.set_label_text(self.image_count_mask_preview, controller.concept.concept_stats["image_with_mask_count"]) + self.components.set_label_text(self.image_count_caption_preview, controller.concept.concept_stats["image_with_caption_count"]) #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) + self.components.set_label_text(self.video_count_preview, controller.concept.concept_stats["video_count"]) + #self.components.set_label_text(self.video_count_mask_preview, controller.concept.concept_stats["video_with_mask_count"]) + self.components.set_label_text(self.video_count_caption_preview, controller.concept.concept_stats["video_with_caption_count"]) #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) + self.components.set_label_text(self.mask_count_preview, controller.concept.concept_stats["mask_count"]) + self.components.set_label_text(self.mask_count_preview_unpaired, controller.concept.concept_stats["unpaired_masks"]) #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') + if controller.concept.concept_stats["subcaption_count"] > 0: + self.components.set_label_text(self.caption_count_preview, f'{controller.concept.concept_stats["caption_count"]} ({controller.concept.concept_stats["subcaption_count"]})') else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) + self.components.set_label_text(self.caption_count_preview, controller.concept.concept_stats["caption_count"]) + self.components.set_label_text(self.caption_count_preview_unpaired, controller.concept.concept_stats["unpaired_captions"]) #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") + max_pixels = controller.concept.concept_stats["max_pixels"] + avg_pixels = controller.concept.concept_stats["avg_pixels"] + min_pixels = controller.concept.concept_stats["min_pixels"] + + if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or controller.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.pixel_max_preview, "-") + self.components.set_label_text(self.pixel_avg_preview, "-") + self.components.set_label_text(self.pixel_min_preview, "-") else: #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') + self.components.set_label_text(self.pixel_max_preview, f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') + self.components.set_label_text(self.pixel_avg_preview, f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') + self.components.set_label_text(self.pixel_min_preview, f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") + max_length = controller.concept.concept_stats["max_length"] + avg_length = controller.concept.concept_stats["avg_length"] + min_length = controller.concept.concept_stats["min_length"] + max_fps = controller.concept.concept_stats["max_fps"] + avg_fps = controller.concept.concept_stats["avg_fps"] + min_fps = controller.concept.concept_stats["min_fps"] + + if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or controller.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.length_max_preview, "-") + self.components.set_label_text(self.length_avg_preview, "-") + self.components.set_label_text(self.length_min_preview, "-") + self.components.set_label_text(self.fps_max_preview, "-") + self.components.set_label_text(self.fps_avg_preview, "-") + self.components.set_label_text(self.fps_min_preview, "-") else: #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') + self.components.set_label_text(self.length_max_preview, f'{int(max_length[0])} frames\n{max_length[1]}') + self.components.set_label_text(self.length_avg_preview, f'{int(avg_length)} frames') + self.components.set_label_text(self.length_min_preview, f'{int(min_length[0])} frames\n{min_length[1]}') #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') + self.components.set_label_text(self.fps_max_preview, f'{int(max_fps[0])} fps\n{max_fps[1]}') + self.components.set_label_text(self.fps_avg_preview, f'{int(avg_fps)} fps') + self.components.set_label_text(self.fps_min_preview, f'{int(min_fps[0])} fps\n{min_fps[1]}') #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") + max_caption_length = controller.concept.concept_stats["max_caption_length"] + avg_caption_length = controller.concept.concept_stats["avg_caption_length"] + min_caption_length = controller.concept.concept_stats["min_caption_length"] + + if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or controller.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken + self.components.set_label_text(self.caption_max_preview, "-") + self.components.set_label_text(self.caption_avg_preview, "-") + self.components.set_label_text(self.caption_min_preview, "-") else: #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') + self.components.set_label_text(self.caption_max_preview, f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') + self.components.set_label_text(self.caption_avg_preview, f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') + self.components.set_label_text(self.caption_min_preview, f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') #aspect bucketing - aspect_buckets = self.concept.concept_stats["aspect_buckets"] + aspect_buckets = controller.concept.concept_stats["aspect_buckets"] if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be @@ -846,7 +411,7 @@ def __update_concept_stats(self): for key, val in min_aspect_buckets.items(): min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' min_bucket_str.strip() - self.small_bucket_preview.configure(text=min_bucket_str) + self.components.set_label_text(self.small_bucket_preview, min_bucket_str) self.bucket_ax.cla() aspects = [str(x) for x in list(aspect_buckets.keys())] @@ -866,69 +431,13 @@ def decimal_to_aspect_ratio(self, value : float): aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' return aspect_string - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): - start_time = time.perf_counter() - last_update = time.perf_counter() - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) - concept_path = self.get_concept_path(self.concept.path) - - if not concept_path: - print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - return - subfolders = [concept_path] - - stats_dict = concept_stats.init_concept_stats(advanced_checks) - for path in subfolders: - if self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: - break - stats_dict = concept_stats.folder_scan(path, stats_dict, advanced_checks, self.concept, start_time, wait_time, self.cancel_scan_flag) - if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set(): #add all subfolders of current directory to for loop - subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) - self.concept.concept_stats = stats_dict - #update GUI approx every half second - if time.perf_counter() > (last_update + 0.5): - last_update = time.perf_counter() - self.concept_stats_tab.after(0, self.__update_concept_stats) - - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - self.concept_stats_tab.after(0, self.__update_concept_stats) - - def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): - self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) - self.scan_thread.start() - - def __disable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): - try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError - if self.concept.concept_stats["file_size"] == 0: #force rescan if empty - raise KeyError - except KeyError: - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec - if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s - - def destroy(self): - if self.bucket_fig is not None: - plt.close(self.bucket_fig) - self.bucket_fig = None - - super().destroy() - - def __ok(self): - self.destroy() + def _disable_scan_buttons(self): + self.components.set_widget_enabled(self.refresh_basic_stats_button, False) + self.components.set_widget_enabled(self.refresh_advanced_stats_button, False) + + def _enable_scan_buttons(self): + self.components.set_widget_enabled(self.refresh_basic_stats_button, True) + self.components.set_widget_enabled(self.refresh_advanced_stats_button, True) + + def _cancel_concept_stats(self, controller): + controller.cancel_scan_flag.set() diff --git a/modules/ui/BaseConfigListView.py b/modules/ui/BaseConfigListView.py index 75d69252a..2165e854b 100644 --- a/modules/ui/BaseConfigListView.py +++ b/modules/ui/BaseConfigListView.py @@ -1,27 +1,63 @@ -import contextlib import copy import json import os -import tkinter as tk -from abc import ABCMeta, abstractmethod +from abc import ABC, abstractmethod from modules.util import path_util -from modules.util.config.BaseConfig import BaseConfig -from modules.util.config.TrainConfig import TrainConfig from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState import customtkinter as ctk -class ConfigList(metaclass=ABCMeta): +class BaseConfigListView(ABC): - def __init__( + def __init__(self, components): + self.components = components + + @abstractmethod + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + pass + + @abstractmethod + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + @abstractmethod + def _create_top_frame(self, master): + pass + + @abstractmethod + def _create_element_list_frame(self, master): + pass + + @abstractmethod + def _remove_widget_from_layout(self, widget): + pass + + @abstractmethod + def _destroy_widget(self, widget): + pass + + @abstractmethod + def _destroy_frame(self, frame): + pass + + @abstractmethod + def _wait_for_window(self, window): + pass + + @abstractmethod + def _show_name_dialog(self, callback): + pass + + def _refresh_show_disabled_text(self): + return + + def build( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, from_external_file: bool, attr_name: str = "", enable_key: str = "enabled", @@ -29,11 +65,11 @@ def __init__( default_config_name: str = "", add_button_text: str = "", add_button_tooltip: str = "", - is_full_width: bool = "", + is_full_width: bool = False, show_toggle_button: bool = False, ): self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state self.from_external_file = from_external_file self.attr_name = attr_name @@ -54,13 +90,8 @@ def __init__( self.is_opening_window = False self._is_current_item_enabled = False - self.master.grid_rowconfigure(0, weight=0) - self.master.grid_rowconfigure(1, weight=1) - self.master.grid_columnconfigure(0, weight=1) - if self.from_external_file: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") + self.top_frame = self._create_top_frame(master) self.configs_dropdown = None self.element_list = None @@ -68,59 +99,27 @@ def __init__( self.configs = [] self.__load_available_config_names() - self.current_config = getattr(self.train_config, self.attr_name) + self.current_config = getattr(self.controller.train_config, self.attr_name) self.widgets = [] - self.__load_current_config(getattr(self.train_config, self.attr_name)) + self.__load_current_config(getattr(self.controller.train_config, self.attr_name)) self.__create_configs_dropdown() - components.button(self.top_frame, 0, 1, "Add Config", self.__add_config, tooltip="Adds a new config, which are containers for concepts, which themselves contain your dataset", width=20, padx=5) - components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) + self.components.button(self.top_frame, 0, 1, "Add Config", self.__add_config, tooltip="Adds a new config, which are containers for concepts, which themselves contain your dataset", width=20, padx=5) + self.components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) else: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") - components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, width=20, padx=5) + self.top_frame = self._create_top_frame(master) + self.components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, width=20, padx=5) - self.current_config = getattr(self.train_config, self.attr_name) + self.current_config = getattr(self.controller.train_config, self.attr_name) self.element_list = None self._create_element_list() if show_toggle_button: # tooltips break if you initialize with an empty string, default to a single space - self.toggle_button = components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) + self.toggle_button = self.components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) self._update_toggle_button_text() - - - @abstractmethod - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - pass - - @abstractmethod - def create_new_element(self) -> BaseConfig: - pass - - @abstractmethod - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass - - def _refresh_show_disabled_text(self): - return - - def _reset_filters(self): # pragma: no cover - default noop - search_var = getattr(self, 'search_var', None) - filter_var = getattr(self, 'filter_var', None) - show_disabled_var = getattr(self, 'show_disabled_var', None) - - if search_var: - search_var.set("") - if filter_var: - filter_var.set("ALL") - if show_disabled_var: - show_disabled_var.set(True) - if search_var and hasattr(self, '_update_filters'): - self._update_filters() - def _update_item_enabled_state(self): # Only count items that match current filters self._is_current_item_enabled = any( @@ -154,14 +153,14 @@ def __create_configs_dropdown(self): if self.configs_dropdown is not None: self.configs_dropdown.destroy() - self.configs_dropdown = components.options_kv( + self.configs_dropdown = self.components.options_kv( self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config ) self._update_toggle_button_text() def _create_element_list(self, **filters): if not self.from_external_file: - self.current_config = getattr(self.train_config, self.attr_name) + self.current_config = getattr(self.controller.train_config, self.attr_name) self.filters.update(filters) @@ -175,13 +174,9 @@ def _create_element_list(self, **filters): def _initialize_all_widgets(self): self.widgets = [] if self.element_list is not None: - self.element_list.destroy() + self._destroy_frame(self.element_list) - self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.element_list.grid(row=1, column=0, sticky="nsew") - - if self.is_full_width: - self.element_list.grid_columnconfigure(0, weight=1) + self.element_list = self._create_element_list_frame(self.master) for i, element in enumerate(self.current_config): widget = self.create_widget( @@ -205,7 +200,7 @@ def _update_widget_visibility(self): widget.place_in_list() visible_index += 1 else: - widget.grid_remove() + self._remove_widget_from_layout(widget) def __load_available_config_names(self): if os.path.isdir(self.config_dir): @@ -228,10 +223,10 @@ def __create_config(self, name: str): self.__create_configs_dropdown() def __add_config(self): - dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) + self._show_name_dialog(self.__create_config) def __add_element(self): - new_element = self.create_new_element() + new_element = self.controller.create_new_element() self.current_config.append(new_element) # incremental insertion if widgets already initialized, else fall back to full rebuild if self.widgets_initialized and self.element_list is not None: @@ -276,8 +271,7 @@ def __remove_element(self, remove_i): self.current_config.pop(remove_i) if self.widgets_initialized and 0 <= remove_i < len(self.widgets): removed = self.widgets.pop(remove_i) - with contextlib.suppress(tk.TclError, AttributeError): - removed.destroy() + self._destroy_widget(removed) # Reindex remaining widgets for idx, widget in enumerate(self.widgets): widget.i = idx @@ -294,7 +288,7 @@ def __load_current_config(self, filename): loaded_config_json = json.load(f) for element_json in loaded_config_json: - element = self.create_new_element().from_dict(element_json) + element = self.controller.create_new_element().from_dict(element_json) self.current_config.append(element) except (FileNotFoundError, json.JSONDecodeError) as e: print(f"Failed to load config from {filename}: {e}") @@ -315,7 +309,7 @@ def save_current_config(self): os.makedirs(self.config_dir, exist_ok=True) write_json_atomic( - getattr(self.train_config, self.attr_name), + getattr(self.controller.train_config, self.attr_name), [element.to_dict() for element in self.current_config] ) except (OSError) as e: @@ -324,10 +318,7 @@ def save_current_config(self): self._update_toggle_button_text() if self.widgets_initialized: - try: - self._update_widget_visibility() - except (tk.TclError, AttributeError) as e: - print.debug(f"Widget visibility update failed: {e}") + self._update_widget_visibility() # let subclass refresh any show-disabled UI if hasattr(self, '_refresh_show_disabled_text'): @@ -336,13 +327,27 @@ def save_current_config(self): def _element_matches_filters(self, element): return True # Show all by default + def _reset_filters(self): # pragma: no cover - default noop + search_var = getattr(self, 'search_var', None) + filter_var = getattr(self, 'filter_var', None) + show_disabled_var = getattr(self, 'show_disabled_var', None) + + if search_var: + search_var.set("") + if filter_var: + filter_var.set("ALL") + if show_disabled_var: + show_disabled_var.set(True) + if search_var and hasattr(self, '_update_filters'): + self._update_filters() + def __open_element_window(self, i, ui_state): if self.is_opening_window: return self.is_opening_window = True try: window = self.open_element_window(i, ui_state) - self.master.wait_window(window) + self._wait_for_window(window) try: if self.widgets is not None and 0 <= i < len(self.widgets): self.widgets[i].configure_element() diff --git a/modules/ui/BaseConvertModelUIView.py b/modules/ui/BaseConvertModelUIView.py index 6cb1b507a..7e2883b1e 100644 --- a/modules/ui/BaseConvertModelUIView.py +++ b/modules/ui/BaseConvertModelUIView.py @@ -1,56 +1,20 @@ -import traceback -from uuid import uuid4 - -from modules.util import create -from modules.util.args.ConvertModelArgs import ConvertModelArgs -from modules.util.config.TrainConfig import QuantizationConfig +from modules.util import path_util from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ModelNames import EmbeddingName, ModelNames -from modules.util.torch_util import torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent - self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None - - - self.title("Convert models") - self.geometry("550x350") - self.resizable(True, True) - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.main_frame(self.frame) - self.frame.pack(fill="both", expand=True) +class BaseConvertModelUIView: + def __init__(self, components): + self.components = components - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): + def build_content(self, frame, controller, ui_state): # model type - components.label(master, 0, 0, "Model Type", + self.components.label(frame, 0, 0, "Model Type", tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify + self.components.options_kv(frame, 0, 1, [ #TODO simplify ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), @@ -71,100 +35,49 @@ def main_frame(self, master): ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") + ], ui_state, "model_type") # training method - components.label(master, 1, 0, "Model Type", + self.components.label(frame, 1, 0, "Model Type", tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ + self.components.options_kv(frame, 1, 1, [ ("Base Model", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") + ], ui_state, "training_method") # input name - components.label(master, 2, 0, "Input name", + self.components.label(frame, 2, 0, "Input name", tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, 2, 1, ui_state, "input_name", + mode="file", path_modifier=path_util.json_path_modifier ) # output data type - components.label(master, 3, 0, "Output Data Type", + self.components.label(frame, 3, 0, "Output Data Type", tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ + self.components.options_kv(frame, 3, 1, [ ("float32", DataType.FLOAT_32), ("float16", DataType.FLOAT_16), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") + ], ui_state, "output_dtype") # output format - components.label(master, 4, 0, "Output Format", + self.components.label(frame, 4, 0, "Output Format", tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ + self.components.options_kv(frame, 4, 1, [ ("Safetensors", ModelFormat.SAFETENSORS), ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") + ], ui_state, "output_model_format") # output model destination - components.label(master, 5, 0, "Model Output Destination", + self.components.label(frame, 5, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", + self.components.path_entry( + frame, 5, 1, ui_state, "output_model_destination", mode="file", io_type=PathIOType.MODEL, ) - self.button = components.button(master, 6, 1, "Convert", self.convert_model) - - def convert_model(self): - try: - self.button.configure(state="disabled") - model_loader = create.create_model_loader( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - model_saver = create.create_model_saver( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - - print("Loading model " + self.convert_model_args.input_name) - if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=self.convert_model_args.input_name, - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=None, - lora=self.convert_model_args.input_name, - embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - else: - raise Exception("could not load model: " + self.convert_model_args.input_name) - - print("Saving model " + self.convert_model_args.output_model_destination) - model_saver.save( - model=model, - model_type=self.convert_model_args.model_type, - output_model_format=self.convert_model_args.output_model_format, - output_model_destination=self.convert_model_args.output_model_destination, - dtype=self.convert_model_args.output_dtype.torch_dtype(), - ) - print("Model converted") - except Exception: - traceback.print_exc() - - torch_gc() - self.button.configure(state="normal") + self.button = self.components.button(frame, 6, 1, "Convert", controller.convert_model) diff --git a/modules/ui/BaseGenerateCaptionsWindowView.py b/modules/ui/BaseGenerateCaptionsWindowView.py index 1690879f1..68e24638e 100644 --- a/modules/ui/BaseGenerateCaptionsWindowView.py +++ b/modules/ui/BaseGenerateCaptionsWindowView.py @@ -1,133 +1,7 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - self.parent = parent - - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all captions", "Create if absent", "Add as new line"] - self.model_var = ctk.StringVar(self, "Blip") - self.models = ["Blip", "Blip2", "WD14 VIT v2"] - - self.title("Batch generate captions") - self.geometry("360x360") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) - self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) - self.caption_entry = ctk.CTkEntry(self.frame, width=200) - self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) - self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.prefix_entry = ctk.CTkEntry(self.frame, width=200) - self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - - self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) - self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.postfix_entry = ctk.CTkEntry(self.frame, width=200) - self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) - self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() +class BaseGenerateCaptionsWindowView(ABC): + @abstractmethod def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) - - mode = { - "Replace all captions": "replace", - "Create if absent": "fill", - "Add as new line": "add", - }[self.mode_var.get()] - - self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), - initial_caption=self.caption_entry.get(), - caption_prefix=self.prefix_entry.get(), - caption_postfix=self.postfix_entry.get(), - mode=mode, - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), - ) - self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() + pass diff --git a/modules/ui/BaseGenerateMasksWindowView.py b/modules/ui/BaseGenerateMasksWindowView.py index daff0d3d5..f9e82231e 100644 --- a/modules/ui/BaseGenerateMasksWindowView.py +++ b/modules/ui/BaseGenerateMasksWindowView.py @@ -1,151 +1,7 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating masks for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - - self.parent = parent - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] - self.model_var = ctk.StringVar(self, "ClipSeg") - self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] - - self.title("Batch generate masks") - self.geometry("360x430") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) - self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) - self.prompt_entry = ctk.CTkEntry(self.frame, width=200) - self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) - - self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) - self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") - self.threshold_entry.insert(0, "0.3") - self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) - self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") - self.smooth_entry.insert(0, 5) - self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) - self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") - self.expand_entry.insert(0, 10) - self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) - self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") - self.alpha_entry.insert(0, 1) - self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) - self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() +class BaseGenerateMasksWindowView(ABC): + @abstractmethod def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) - - mode = { - "Replace all masks": "replace", - "Create if absent": "fill", - "Add to existing": "add", - "Subtract from existing": "subtract", - "Blend with existing": "blend", - }[self.mode_var.get()] - - self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], - mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), - ) - self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() + pass diff --git a/modules/ui/BaseLoraTabView.py b/modules/ui/BaseLoraTabView.py index 1c73d90ce..ea6f52d5c 100644 --- a/modules/ui/BaseLoraTabView.py +++ b/modules/ui/BaseLoraTabView.py @@ -1,47 +1,20 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType +from modules.util import path_util from modules.util.enum.ModelType import PeftType -from modules.util.ui import components -from modules.util.ui.UIState import UIState from modules.util.ui.validation_helpers import check_range -import customtkinter as ctk +class BaseLoraTabView: + def __init__(self, components): + self.components = components -class LoraTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() + def build(self, frame, controller, ui_state, setup_lora_callback): + self.components.label(frame, 0, 0, "Type", + tooltip="The type of low-parameter finetuning method.") + self.components.options_kv(frame, 0, 1, controller.get_peft_types(), + ui_state, "peft_type", command=setup_lora_callback) - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - self.scroll_frame = None - self.options_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - - self.scroll_frame.grid_columnconfigure(0, weight=0) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=2) - - components.label(self.scroll_frame, 0, 0, "Type", - tooltip="The type of low-parameter finetuning method.") - # This will instantly call self.setup_lora. - components.options_kv(self.scroll_frame, 0, 1, [ - ("LoRA", PeftType.LORA), - ("LoHa", PeftType.LOHA), - ("OFT v2", PeftType.OFT_2), - ], self.ui_state, "peft_type", command=self.setup_lora) - - def setup_lora(self, peft_type: PeftType): + def build_lora_options(self, master, controller, ui_state, peft_type: PeftType): if peft_type == PeftType.LOHA: name = "LoHa" elif peft_type == PeftType.OFT_2: @@ -49,106 +22,87 @@ def setup_lora(self, peft_type: PeftType): else: name = "LoRA" - if self.options_frame: - self.options_frame.destroy() - self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") - self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") - master = self.options_frame - - master.grid_columnconfigure(0, weight=0, uniform="a") - master.grid_columnconfigure(1, weight=1, uniform="a") - master.grid_columnconfigure(2, minsize=50, uniform="a") - master.grid_columnconfigure(3, weight=0, uniform="a") - master.grid_columnconfigure(4, weight=1, uniform="a") - # lora model name - components.label(master, 0, 0, f"{name} base model", - tooltip=f"The base {name} to train on. Leave empty to create a new {name}") - entry = components.path_entry( - master, 0, 1, self.ui_state, "lora_model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.label(master, 0, 0, f"{name} base model", + tooltip=f"The base {name} to train on. Leave empty to create a new {name}") + self.components.path_entry( + master, 0, 1, ui_state, "lora_model_name", + mode="file", path_modifier=path_util.json_path_modifier, + columnspan=4, ) - entry.grid(row=0, column=1, columnspan=4) - # LoRA decomposition if peft_type == PeftType.LORA: - components.label(master, 1, 3, "Decompose Weights (DoRA)", - tooltip="Decompose LoRA Weights (aka, DoRA).") - components.switch(master, 1, 4, self.ui_state, "lora_decompose") + self.components.label(master, 1, 3, "Decompose Weights (DoRA)", + tooltip="Decompose LoRA Weights (aka, DoRA).") + self.components.switch(master, 1, 4, ui_state, "lora_decompose") - components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", - tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") - components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") - components.label(master, 3, 3, "Apply on output axis (DoRA Only)", - tooltip="Apply the weight decomposition on the output axis instead of the input axis.") - components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") + self.components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", + tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") + self.components.switch(master, 2, 4, ui_state, "lora_decompose_norm_epsilon") + self.components.label(master, 3, 3, "Apply on output axis (DoRA Only)", + tooltip="Apply the weight decomposition on the output axis instead of the input axis.") + self.components.switch(master, 3, 4, ui_state, "lora_decompose_output_axis") # LoRA and LoHA shared settings if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: # rank - components.label(master, 1, 0, f"{name} rank", - tooltip=f"The rank parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) + self.components.label(master, 1, 0, f"{name} rank", + tooltip=f"The rank parameter used when creating a new {name}") + self.components.entry(master, 1, 1, ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) # alpha - components.label(master, 2, 0, f"{name} alpha", - tooltip=f"The alpha parameter used when creating a new {name}") - components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) + self.components.label(master, 2, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + self.components.entry(master, 2, 1, ui_state, "lora_alpha", required=True) # Dropout Percentage - components.label(master, 3, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 3, 1, self.ui_state, "dropout_probability") + self.components.label(master, 3, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + self.components.entry(master, 3, 1, ui_state, "dropout_probability") # weight dtype - components.label(master, 4, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 4, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") + self.components.label(master, 4, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + self.components.options_kv(master, 4, 1, controller.get_lora_weight_dtypes(), ui_state, "lora_weight_dtype") # For use with additional embeddings. - components.label(master, 5, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") + self.components.label(master, 5, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + self.components.switch(master, 5, 1, ui_state, "bundle_additional_embeddings") # OFTv2 elif peft_type == PeftType.OFT_2: # Block Size - components.label(master, 1, 0, f"{name} Block Size", - tooltip=f"The block size parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) + self.components.label(master, 1, 0, f"{name} Block Size", + tooltip=f"The block size parameter used when creating a new {name}") + self.components.entry(master, 1, 1, ui_state, "oft_block_size", required=True) # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") + self.components.label(master, 1, 3, "Constrained OFT (COFT)", + tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") + self.components.switch(master, 1, 4, ui_state, "oft_coft") - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") + self.components.label(master, 2, 3, "COFT Epsilon", + tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") + self.components.entry(master, 2, 4, ui_state, "coft_eps") # Block Share - components.label(master, 3, 3, "Block Share", - tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") + self.components.label(master, 3, 3, "Block Share", + tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") + self.components.switch(master, 3, 4, ui_state, "oft_block_share") # Dropout Percentage - components.label(master, 2, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 2, 1, self.ui_state, "dropout_probability") + self.components.label(master, 2, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") + self.components.entry(master, 2, 1, ui_state, "dropout_probability") # OFT weight dtype - components.label(master, 3, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") + self.components.label(master, 3, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + self.components.options_kv(master, 3, 1, controller.get_lora_weight_dtypes(), ui_state, "lora_weight_dtype") # For use with additional embeddings. - components.label(master, 4, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + self.components.label(master, 4, 0, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + self.components.switch(master, 4, 1, ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index ff17ea3ba..88f223e22 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -1,85 +1,60 @@ -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig +from abc import ABC, abstractmethod + +from modules.util import path_util from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import components -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class ModelTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - - self.scroll_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=1) - - base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) - base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") - - base_frame.grid_columnconfigure(0, weight=0) - base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) - base_frame.grid_columnconfigure(2, minsize=50) - base_frame.grid_columnconfigure(3, weight=0) - base_frame.grid_columnconfigure(4, weight=1) - - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): + + +class BaseModelTabView(ABC): + def __init__(self, components): + self.components = components + + @abstractmethod + def _make_svd_frames(self, parent, row: int): + """Create and place SVDQuant label+entry subframes; return (label_frame, entry_frame).""" + + def build_content(self, frame, controller, ui_state): + if controller.train_config.model_type.is_stable_diffusion(): # TODO simplify + self.__setup_stable_diffusion_ui(frame, controller, ui_state) + if controller.train_config.model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_pixart(): + self.__setup_pixart_alpha_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_flux_1(): + self.__setup_flux_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_flux_2(): + self.__setup_flux_2_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_z_image(): + self.__setup_z_image_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_chroma(): + self.__setup_chroma_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_qwen(): + self.__setup_qwen_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_sana(): + self.__setup_sana_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_hi_dream(): + self.__setup_hi_dream_ui(frame, controller, ui_state) + elif controller.train_config.model_type.is_ernie(): + self.__setup_ernie_ui(frame, controller, ui_state) + + def __setup_stable_diffusion_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_unet=True, has_text_encoder=True, has_vae=True, @@ -87,20 +62,23 @@ def __setup_stable_diffusion_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ + allow_diffusers=controller.train_config.training_method in [ TrainingMethod.FINE_TUNE, TrainingMethod.FINE_TUNE_VAE, ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_stable_diffusion_3_ui(self, frame): + def __setup_stable_diffusion_3_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -110,17 +88,20 @@ def __setup_stable_diffusion_3_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_flux_ui(self, frame): + def __setup_flux_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -130,17 +111,20 @@ def __setup_flux_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_flux_2_ui(self, frame): + def __setup_flux_2_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -149,17 +133,20 @@ def __setup_flux_2_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_z_image_ui(self, frame): + def __setup_z_image_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -168,17 +155,20 @@ def __setup_z_image_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_ernie_ui(self, frame): + def __setup_ernie_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -187,17 +177,20 @@ def __setup_ernie_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_chroma_ui(self, frame): + def __setup_chroma_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -206,17 +199,20 @@ def __setup_chroma_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_qwen_ui(self, frame): + def __setup_qwen_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -225,17 +221,20 @@ def __setup_qwen_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_stable_diffusion_xl_ui(self, frame): + def __setup_stable_diffusion_xl_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_unet=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -244,38 +243,44 @@ def __setup_stable_diffusion_xl_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_wuerstchen_ui(self, frame): + def __setup_wuerstchen_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), + allow_override_prior=controller.train_config.model_type.is_stable_cascade(), has_text_encoder=True, ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) + row = self.__create_effnet_encoder_components(frame, row, ui_state) + row = self.__create_decoder_components(frame, row, ui_state, controller.train_config.model_type.is_wuerstchen_v2()) row = self.__create_output_components( frame, row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ui_state, + allow_safetensors=controller.train_config.training_method != TrainingMethod.FINE_TUNE + or controller.train_config.model_type.is_stable_cascade(), + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_pixart_alpha_ui(self, frame): + def __setup_pixart_alpha_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder=True, has_vae=True, @@ -283,17 +288,20 @@ def __setup_pixart_alpha_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_sana_ui(self, frame): + def __setup_sana_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder=True, has_vae=True, @@ -301,17 +309,20 @@ def __setup_sana_ui(self, frame): row = self.__create_output_components( frame, row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ui_state, + allow_safetensors=controller.train_config.training_method != TrainingMethod.FINE_TUNE, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_hunyuan_video_ui(self, frame): + def __setup_hunyuan_video_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, allow_override_transformer=True, has_text_encoder_1=True, @@ -321,17 +332,20 @@ def __setup_hunyuan_video_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __setup_hi_dream_ui(self, frame): + def __setup_hi_dream_ui(self, frame, controller, ui_state): row = 0 - row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_dtype_components(frame, row, ui_state) row = self.__create_base_components( frame, row, + controller, + ui_state, has_transformer=True, has_text_encoder_1=True, has_text_encoder_2=True, @@ -343,12 +357,13 @@ def __setup_hi_dream_ui(self, frame): row = self.__create_output_components( frame, row, + ui_state, allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, ) - def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: + def __create_dtype_options(self, include_gguf: bool = False, include_a8: bool = False) -> list[tuple[str, DataType]]: options = [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), @@ -373,28 +388,28 @@ def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=Fals return options - def __create_base_dtype_components(self, frame, row: int) -> int: + def __create_base_dtype_components(self, frame, row: int, ui_state) -> int: # huggingface token - components.label(frame, row, 0, "Hugging Face Token", + self.components.label(frame, row, 0, "Hugging Face Token", tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " "Go to https://huggingface.co/settings/tokens to create an access token.", wide_tooltip=True) - components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") + self.components.entry(frame, row, 1, ui_state, "secrets.huggingface_token") row += 1 # base model - components.label(frame, row, 0, "Base Model", + self.components.label(frame, row, 0, "Base Model", tooltip="Filename, directory or Hugging Face repository of the base model") - components.path_entry( - frame, row, 1, self.ui_state, "base_model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "base_model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # compile - components.label(frame, row, 3, "Compile transformer blocks", + self.components.label(frame, row, 3, "Compile transformer blocks", tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") - components.switch(frame, row, 4, self.ui_state, "compile") + self.components.switch(frame, row, 4, ui_state, "compile") row += 1 @@ -404,6 +419,8 @@ def __create_base_components( self, frame, row: int, + controller, + ui_state, has_unet: bool = False, has_prior: bool = False, allow_override_prior: bool = False, @@ -419,54 +436,53 @@ def __create_base_components( ) -> int: if has_unet: # unet weight dtype - components.label(frame, row, 3, "UNet Data Type", + self.components.label(frame, row, 3, "UNet Data Type", tooltip="The unet weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), - self.ui_state, "unet.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), + ui_state, "unet.weight_dtype") row += 1 if has_prior: if allow_override_prior: # prior model - components.label(frame, row, 0, "Prior Model", + self.components.label(frame, row, 0, "Prior Model", tooltip="Filename, directory or Hugging Face repository of the prior model") - components.path_entry( - frame, row, 1, self.ui_state, "prior.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "prior.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # prior weight dtype - components.label(frame, row, 3, "Prior Data Type", + self.components.label(frame, row, 3, "Prior Data Type", tooltip="The prior weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "prior.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "prior.weight_dtype") row += 1 if has_transformer: if allow_override_transformer: # transformer model - components.label(frame, row, 0, "Override Transformer / GGUF", + self.components.label(frame, row, 0, "Override Transformer / GGUF", tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") - components.path_entry( - frame, row, 1, self.ui_state, "transformer.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "transformer.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # transformer weight dtype - components.label(frame, row, 3, "Transformer Data Type", + self.components.label(frame, row, 3, "Transformer Data Type", tooltip="The transformer weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), - self.ui_state, "transformer.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), + ui_state, "transformer.weight_dtype") row += 1 - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} + presets = controller.get_presets() - components.label(frame, row, 0, "Quantization") - components.layer_filter_entry(frame, row, 1, self.ui_state, + self.components.label(frame, row, 0, "Quantization") + self.components.layer_filter_entry(frame, row, 1, ui_state, preset_var_name="quantization.layer_filter_preset", presets=presets, preset_label="Quantization Layer Filter", preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", @@ -478,107 +494,103 @@ def __create_base_components( ) # SVDQuant - create vertical grids to match the size of layer_filter_entry - svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_label_frame.grid(row=row, column=3, sticky="nsew") - svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_entry_frame.grid(row=row, column=4, sticky="nsew") - components.label(svd_label_frame, 0, 0, "SVDQuant", + svd_label_frame, svd_entry_frame = self._make_svd_frames(frame, row) + self.components.label(svd_label_frame, 0, 0, "SVDQuant", tooltip="What datatype to use for SVDQuant weights decomposition.") - components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], - self.ui_state, "quantization.svd_dtype") - components.label(svd_label_frame, 1, 0, "SVDQuant Rank", + self.components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], + ui_state, "quantization.svd_dtype") + self.components.label(svd_label_frame, 1, 0, "SVDQuant Rank", tooltip="Rank for SVDQuant weights decomposition") - components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") + self.components.entry(svd_entry_frame, 1, 0, ui_state, "quantization.svd_rank") row += 1 - if has_text_encoder: # text encoder weight dtype - components.label(frame, row, 3, "Text Encoder Data Type", + self.components.label(frame, row, 3, "Text Encoder Data Type", tooltip="The text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder.weight_dtype") row += 1 if has_text_encoder_1: # text encoder 1 weight dtype - components.label(frame, row, 3, "Text Encoder 1 Data Type", + self.components.label(frame, row, 3, "Text Encoder 1 Data Type", tooltip="The text encoder 1 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder.weight_dtype") row += 1 if has_text_encoder_2: # text encoder 2 weight dtype - components.label(frame, row, 3, "Text Encoder 2 Data Type", + self.components.label(frame, row, 3, "Text Encoder 2 Data Type", tooltip="The text encoder 2 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_2.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_2.weight_dtype") row += 1 if has_text_encoder_3: # text encoder 3 weight dtype - components.label(frame, row, 3, "Text Encoder 3 Data Type", + self.components.label(frame, row, 3, "Text Encoder 3 Data Type", tooltip="The text encoder 3 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_3.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_3.weight_dtype") row += 1 if has_text_encoder_4: if allow_override_text_encoder_4: # text encoder 4 weight dtype - components.label(frame, row, 0, "Text Encoder 4 Override", + self.components.label(frame, row, 0, "Text Encoder 4 Override", tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") - components.path_entry( - frame, row, 1, self.ui_state, "text_encoder_4.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "text_encoder_4.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # text encoder 4 weight dtype - components.label(frame, row, 3, "Text Encoder 4 Data Type", + self.components.label(frame, row, 3, "Text Encoder 4 Data Type", tooltip="The text encoder 4 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_4.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "text_encoder_4.weight_dtype") row += 1 if has_vae: # base model - components.label(frame, row, 0, "VAE Override", + self.components.label(frame, row, 0, "VAE Override", tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. Using a safetensor VAE file will cause an error that the model cannot be loaded.") - components.path_entry( - frame, row, 1, self.ui_state, "vae.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "vae.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # vae weight dtype - components.label(frame, row, 3, "VAE Data Type", + self.components.label(frame, row, 3, "VAE Data Type", tooltip="The vae weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "vae.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "vae.weight_dtype") row += 1 return row - def __create_effnet_encoder_components(self, frame, row: int): + def __create_effnet_encoder_components(self, frame, row: int, ui_state) -> int: # effnet encoder model - components.label(frame, row, 0, "Effnet Encoder Model", + self.components.label(frame, row, 0, "Effnet Encoder Model", tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") - components.path_entry( - frame, row, 1, self.ui_state, "effnet_encoder.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "effnet_encoder.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # effnet encoder weight dtype - components.label(frame, row, 3, "Effnet Encoder Data Type", + self.components.label(frame, row, 3, "Effnet Encoder Data Type", tooltip="The effnet encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "effnet_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "effnet_encoder.weight_dtype") row += 1 @@ -588,38 +600,39 @@ def __create_decoder_components( self, frame, row: int, + ui_state, has_text_encoder: bool, ) -> int: # decoder model - components.label(frame, row, 0, "Decoder Model", + self.components.label(frame, row, 0, "Decoder Model", tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.path_entry( - frame, row, 1, self.ui_state, "decoder.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, row, 1, ui_state, "decoder.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # decoder weight dtype - components.label(frame, row, 3, "Decoder Data Type", + self.components.label(frame, row, 3, "Decoder Data Type", tooltip="The decoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder.weight_dtype") row += 1 if has_text_encoder: # decoder text encoder weight dtype - components.label(frame, row, 3, "Decoder Text Encoder Data Type", + self.components.label(frame, row, 3, "Decoder Text Encoder Data Type", tooltip="The decoder text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_text_encoder.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_text_encoder.weight_dtype") row += 1 # decoder vqgan weight dtype - components.label(frame, row, 3, "Decoder VQGAN Data Type", + self.components.label(frame, row, 3, "Decoder VQGAN Data Type", tooltip="The decoder vqgan weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_vqgan.weight_dtype") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_vqgan.weight_dtype") row += 1 @@ -629,30 +642,31 @@ def __create_output_components( self, frame, row: int, + ui_state, allow_safetensors: bool = False, allow_diffusers: bool = False, allow_legacy_safetensors: bool = False, allow_comfy: bool = False, ) -> int: # output model destination - components.label(frame, row, 0, "Model Output Destination", + self.components.label(frame, row, 0, "Model Output Destination", tooltip="Filename or directory where the output model is saved") - components.path_entry( - frame, row, 1, self.ui_state, "output_model_destination", + self.components.path_entry( + frame, row, 1, ui_state, "output_model_destination", mode="file", io_type=PathIOType.MODEL, ) # output data type - components.label(frame, row, 3, "Output Data Type", + self.components.label(frame, row, 3, "Output Data Type", tooltip="Precision to use when saving the output model") - components.options_kv(frame, row, 4, [ + self.components.options_kv(frame, row, 4, [ ("float16", DataType.FLOAT_16), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), ("float8", DataType.FLOAT_8), ("nfloat4", DataType.NFLOAT_4), - ], self.ui_state, "output_dtype") + ], ui_state, "output_dtype") row += 1 @@ -667,21 +681,21 @@ def __create_output_components( if allow_comfy: formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) - components.label(frame, row, 0, "Output Format", + self.components.label(frame, row, 0, "Output Format", tooltip="Format to use when saving the output model") - components.options_kv(frame, row, 1, formats, self.ui_state, "output_model_format") + self.components.options_kv(frame, row, 1, formats, ui_state, "output_model_format") # include config - components.label(frame, row, 3, "Include Config", + self.components.label(frame, row, 3, "Include Config", tooltip="Include the training configuration in the final model. Only supported for safetensors files. " "None: No config is included. " "Settings: All training settings are included. " "All: All settings, including the samples and concepts are included.") - components.options_kv(frame, row, 4, [ + self.components.options_kv(frame, row, 4, [ ("None", ConfigPart.NONE), ("Settings", ConfigPart.SETTINGS), ("All", ConfigPart.ALL), - ], self.ui_state, "include_train_config") + ], ui_state, "include_train_config") row += 1 diff --git a/modules/ui/BaseMuonAdamWindowView.py b/modules/ui/BaseMuonAdamWindowView.py index 5879ab432..90dc04eee 100644 --- a/modules/ui/BaseMuonAdamWindowView.py +++ b/modules/ui/BaseMuonAdamWindowView.py @@ -1,67 +1,10 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk -MUON_AUX_ADAM_DEFAULTS = { - "beta1": 0.9, - "beta2": 0.999, - "eps": 1e-8, - "weight_decay": 0.0, -} +class BaseMuonAdamWindowView: + def __init__(self, components): + self.components = components -class MuonAdamWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state: UIState, - parent_optimizer_type: Optimizer, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.adam_ui_state = ui_state - self.parent_optimizer_type = parent_optimizer_type - - if self.parent_optimizer_type == Optimizer.MUON: - self.title("Muon's Auxiliary AdamW Settings") - self.adam_params_def = MUON_AUX_ADAM_DEFAULTS - else: - self.title("Muon_adv's Auxiliary AdamW_adv Settings") - self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.destroy) - self.create_adam_params_ui(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def create_adam_params_ui(self, master): + def build_content(self, master, controller, ui_state): # This is a large map, copied from OptimizerParamsWindow for simplicity. # @formatter:off KEY_DETAIL_MAP = { @@ -84,7 +27,7 @@ def create_adam_params_ui(self, master): } # @formatter:on - adam_params = self.adam_params_def + adam_params = controller.get_adam_params_def() for index, key in enumerate(adam_params.keys()): if key not in KEY_DETAIL_MAP: @@ -99,9 +42,9 @@ def create_adam_params_ui(self, master): row = index // 2 col = 3 * (index % 2) - components.label(master, row, col, title, tooltip=tooltip) + self.components.label(master, row, col, title, tooltip=tooltip) if param_type != 'bool': - components.entry(master, row, col + 1, self.adam_ui_state, key) + self.components.entry(master, row, col + 1, ui_state, key) else: - components.switch(master, row, col + 1, self.adam_ui_state, key) + self.components.switch(master, row, col + 1, ui_state, key) diff --git a/modules/ui/BaseOffloadingWindowView.py b/modules/ui/BaseOffloadingWindowView.py index 54035e121..8ec9c0cb1 100644 --- a/modules/ui/BaseOffloadingWindowView.py +++ b/modules/ui/BaseOffloadingWindowView.py @@ -1,75 +1,24 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState +from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod -import customtkinter as ctk +class BaseOffloadingWindowView: + def __init__(self, components): + self.components = components -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) + def build_content(self, frame, controller, ui_state): + self.components.label(frame, 0, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + self.components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, + "gradient_checkpointing") - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None + self.components.label(frame, 1, 0, "Async Offloading", + tooltip="Enables Asynchronous offloading.") + self.components.switch(frame, 1, 1, ui_state, "enable_async_offloading") - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) + self.components.label(frame, 2, 0, "Offload Activations", + tooltip="Enables Activation Offloading") + self.components.switch(frame, 2, 1, ui_state, "enable_activation_offloading") - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() + self.components.label(frame, 3, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + self.components.entry(frame, 3, 1, ui_state, "layer_offload_fraction") diff --git a/modules/ui/BaseOptimizerParamsWindowView.py b/modules/ui/BaseOptimizerParamsWindowView.py index 16063c26c..b5a5c8911 100644 --- a/modules/ui/BaseOptimizerParamsWindowView.py +++ b/modules/ui/BaseOptimizerParamsWindowView.py @@ -1,94 +1,32 @@ -import contextlib -from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow -from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( OPTIMIZER_DEFAULT_PARAMETERS, - change_optimizer, - load_optimizer_defaults, - update_optimizer_config, ) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseOptimizerParamsWindowView: + def __init__(self, components): + self.components = components -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) - self.muon_adam_button = None - - self.title("Optimizer Settings") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): + def build_content(self, frame, controller, ui_state, optimizer_ui_state, + on_optimizer_change_cb, load_defaults_cb): # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") + self.components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) + self.components.options(frame, 0, 1, [str(x) for x in list(Optimizer)], optimizer_ui_state, "optimizer", + command=on_optimizer_change_cb) # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): - with contextlib.suppress(TclError): - for widget in master.winfo_children(): - grid_info = widget.grid_info() - if int(grid_info["row"]) >= 1: - widget.destroy() - - def create_dynamic_ui( - self, - master, - ): + self.components.label(frame, 0, 3, "Optimizer Defaults", + tooltip="Load default settings for the selected optimizer") + self.components.button(frame, 0, 4, "Load Defaults", load_defaults_cb, + tooltip="Load default settings for the selected optimizer") + def build_dynamic_content(self, master, controller, optimizer_ui_state, + update_user_pref_cb, open_muon_adam_cb): # Lookup for the title and tooltip for a key # @formatter:off KEY_DETAIL_MAP = { @@ -197,10 +135,7 @@ def create_dynamic_ui( } # @formatter:on - if not self.winfo_exists(): # check if this window isn't open - return - - selected_optimizer = self.train_config.optimizer.optimizer + selected_optimizer = controller.config.optimizer.optimizer # Extract the keys for the selected optimizer for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): @@ -215,74 +150,20 @@ def create_dynamic_ui( row = (index // 2) + 1 col = 3 * (index % 2) - components.label(master, row, col, title, tooltip=tooltip) + self.components.label(master, row, col, title, tooltip=tooltip) if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) + frame = self.components.inline_frame(master, row, col + 1, columnspan=2) - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) + self.components.switch(frame, 0, 0, optimizer_ui_state, key, command=update_user_pref_cb) - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, + self.muon_adam_button = self.components.button( + frame, 0, 1, "...", open_muon_adam_cb, tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() + width=20, padx=5) elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) + self.components.entry(master, row, col + 1, optimizer_ui_state, key, + command=update_user_pref_cb) else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - - def update_user_pref(self, *args): - update_optimizer_config(self.train_config) - self.toggle_muon_adam_button() - - def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) - - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def on_window_close(self): - self.destroy() - - def toggle_muon_adam_button(self): - if self.muon_adam_button and self.muon_adam_button.winfo_exists(): - muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() - self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") - - def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer - - adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config - - if current_optimizer == Optimizer.MUON: - defaults = MUON_AUX_ADAM_DEFAULTS - else: - defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - if current_state is None: - adam_config.from_dict(defaults) - if current_optimizer != Optimizer.MUON: - adam_config.optimizer = Optimizer.ADAMW_ADV - elif isinstance(current_state, dict): - adam_config.from_dict(current_state) - else: - # Should not happen if TrainConfig defines it as dict, but for safety - adam_config = current_state - - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) - self.wait_window(window) - - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + self.components.switch(master, row, col + 1, optimizer_ui_state, key, + command=update_user_pref_cb) diff --git a/modules/ui/BaseProfilingWindowView.py b/modules/ui/BaseProfilingWindowView.py index 8d298abe3..8e0de3b64 100644 --- a/modules/ui/BaseProfilingWindowView.py +++ b/modules/ui/BaseProfilingWindowView.py @@ -1,57 +1,21 @@ -import faulthandler +from abc import abstractmethod -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -import customtkinter as ctk -from scalene import scalene_profiler +class BaseProfilingWindowView: + def __init__(self, components): + self.components = components - -class ProfilingWindow(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.title("Profiling") - self.geometry("512x512") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=0) - self.grid_rowconfigure(2, weight=1) - self.grid_columnconfigure(0, weight=1) - - components.button(self, 0, 0, "Dump stack", self._dump_stack) - self._profile_button = components.button( - self, 1, 0, "Start Profiling", self._start_profiler, + def build_content(self, frame, bottom_bar, controller): + self.components.button(frame, 0, 0, "Dump stack", controller.dump_stack) + self._profile_button = self.components.button( + frame, 1, 0, "Start Profiling", controller.start_profiler, tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") + self._message_label = self.components.label(bottom_bar, 0, 0, "Inactive") - # Bottom bar - self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) - self._bottom_bar.grid(row=2, column=0, sticky="sew") - self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") - - self.protocol("WM_DELETE_WINDOW", self.withdraw) - self.withdraw() - self.after(200, lambda: set_window_icon(self)) - - def _dump_stack(self): - with open('stacks.txt', 'w') as f: - faulthandler.dump_traceback(f) - self._message_label.configure(text='Stack dumped to stacks.txt') - - def _end_profiler(self): - scalene_profiler.stop() - - self._message_label.configure(text='Inactive') - self._profile_button.configure(text='Start Profiling') - self._profile_button.configure(command=self._start_profiler) - - def _start_profiler(self): - scalene_profiler.start() + @abstractmethod + def set_message(self, text): + pass - self._message_label.configure(text='Profiling active...') - self._profile_button.configure(text='End Profiling') - self._profile_button.configure(command=self._end_profiler) + @abstractmethod + def set_profiling_active(self, active): + pass diff --git a/modules/ui/BaseSampleFrameView.py b/modules/ui/BaseSampleFrameView.py index 297caac29..c3e5860c9 100644 --- a/modules/ui/BaseSampleFrameView.py +++ b/modules/ui/BaseSampleFrameView.py @@ -1,98 +1,59 @@ -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType from modules.util.enum.NoiseScheduler import NoiseScheduler -from modules.util.ui import components -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseSampleFrameView: + def __init__(self, components): + self.components = components -class SampleFrame(ctk.CTkFrame): - def __init__( - self, - parent, - sample: SampleConfig, - ui_state: UIState, - model_type: ModelType, - include_prompt: bool = True, - include_settings: bool = True, - ): - ctk.CTkFrame.__init__(self, parent, fg_color="transparent") - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - is_flow_matching = model_type.is_flow_matching() - is_inpainting_model = model_type.has_conditioning_image_input() - is_video_model = model_type.is_video_model() - - if include_prompt and include_prompt: - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) - - if include_prompt: - top_frame = ctk.CTkFrame(self, fg_color="transparent") - top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - - top_frame.grid_columnconfigure(0, weight=0) - top_frame.grid_columnconfigure(1, weight=1) - - if include_settings: - bottom_frame = ctk.CTkFrame(self, fg_color="transparent") - bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - bottom_frame.grid_columnconfigure(0, weight=0) - bottom_frame.grid_columnconfigure(1, weight=1) - bottom_frame.grid_columnconfigure(2, weight=0) - bottom_frame.grid_columnconfigure(3, weight=1) - + def build_content(self, top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings): + is_flow_matching = controller.is_flow_matching() + is_inpainting_model = controller.is_inpainting_model() + is_video_model = controller.is_video_model() if include_prompt: # prompt - components.label(top_frame, 0, 0, "prompt:") - components.entry(top_frame, 0, 1, self.ui_state, "prompt") + self.components.label(top_frame, 0, 0, "prompt:") + self.components.entry(top_frame, 0, 1, ui_state, "prompt") # negative prompt - components.label(top_frame, 1, 0, "negative prompt:") - components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") + self.components.label(top_frame, 1, 0, "negative prompt:") + self.components.entry(top_frame, 1, 1, ui_state, "negative_prompt") if include_settings: # width - components.label(bottom_frame, 0, 0, "width:") - components.entry(bottom_frame, 0, 1, self.ui_state, "width") + self.components.label(bottom_frame, 0, 0, "width:") + self.components.entry(bottom_frame, 0, 1, ui_state, "width") # height - components.label(bottom_frame, 0, 2, "height:") - components.entry(bottom_frame, 0, 3, self.ui_state, "height") + self.components.label(bottom_frame, 0, 2, "height:") + self.components.entry(bottom_frame, 0, 3, ui_state, "height") if is_video_model: # frames - components.label(bottom_frame, 1, 0, "frames:", - tooltip="Number of frames to generate. Only used when generating videos.") - components.entry(bottom_frame, 1, 1, self.ui_state, "frames") + self.components.label(bottom_frame, 1, 0, "frames:", + tooltip="Number of frames to generate. Only used when generating videos.") + self.components.entry(bottom_frame, 1, 1, ui_state, "frames") # length - components.label(bottom_frame, 1, 2, "length:", - tooltip="Length in seconds of audio output.") - components.entry(bottom_frame, 1, 3, self.ui_state, "length") + self.components.label(bottom_frame, 1, 2, "length:", + tooltip="Length in seconds of audio output.") + self.components.entry(bottom_frame, 1, 3, ui_state, "length") # seed - components.label(bottom_frame, 2, 0, "seed:") - components.entry(bottom_frame, 2, 1, self.ui_state, "seed") + self.components.label(bottom_frame, 2, 0, "seed:") + self.components.entry(bottom_frame, 2, 1, ui_state, "seed") # random seed - components.label(bottom_frame, 2, 2, "random seed:") - components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") + self.components.label(bottom_frame, 2, 2, "random seed:") + self.components.switch(bottom_frame, 2, 3, ui_state, "random_seed") # cfg scale - components.label(bottom_frame, 3, 0, "cfg scale:") - components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") + self.components.label(bottom_frame, 3, 0, "cfg scale:") + self.components.entry(bottom_frame, 3, 1, ui_state, "cfg_scale") # sampler if not is_flow_matching: - components.label(bottom_frame, 4, 2, "sampler:") - components.options_kv(bottom_frame, 4, 3, [ + self.components.label(bottom_frame, 4, 2, "sampler:") + self.components.options_kv(bottom_frame, 4, 3, [ ("DDIM", NoiseScheduler.DDIM), ("Euler", NoiseScheduler.EULER), ("Euler A", NoiseScheduler.EULER_A), @@ -103,32 +64,32 @@ def __init__( ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) - ], self.ui_state, "noise_scheduler") + ], ui_state, "noise_scheduler") # steps - components.label(bottom_frame, 4, 0, "steps:") - components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") + self.components.label(bottom_frame, 4, 0, "steps:") + self.components.entry(bottom_frame, 4, 1, ui_state, "diffusion_steps") # inpainting if is_inpainting_model: - components.label(bottom_frame, 5, 0, "inpainting:", - tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") - components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") + self.components.label(bottom_frame, 5, 0, "inpainting:", + tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") + self.components.switch(bottom_frame, 5, 1, ui_state, "sample_inpainting") # base image path - components.label(bottom_frame, 6, 0, "base image path:", - tooltip="The base image used when inpainting.") - components.file_entry(bottom_frame, 6, 1, self.ui_state, "base_image_path", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.components.label(bottom_frame, 6, 0, "base image path:", + tooltip="The base image used when inpainting.") + self.components.path_entry(bottom_frame, 6, 1, ui_state, "base_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) # mask image path - components.label(bottom_frame, 6, 2, "mask image path:", - tooltip="The mask used when inpainting.") - components.file_entry(bottom_frame, 6, 3, self.ui_state, "mask_image_path", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.components.label(bottom_frame, 6, 2, "mask image path:", + tooltip="The mask used when inpainting.") + self.components.path_entry(bottom_frame, 6, 3, ui_state, "mask_image_path", + mode="file", + allow_model_files=False, + allow_image_files=True, + ) diff --git a/modules/ui/BaseSampleParamsWindowView.py b/modules/ui/BaseSampleParamsWindowView.py index 2b0b3f3f1..bbe410a9a 100644 --- a/modules/ui/BaseSampleParamsWindowView.py +++ b/modules/ui/BaseSampleParamsWindowView.py @@ -1,39 +1,6 @@ -from modules.ui.SampleFrame import SampleFrame -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk -class SampleParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - self.title("Sample") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) - frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __ok(self): - self.destroy() +class BaseSampleParamsWindowView: + def __init__(self, components): + self.components = components diff --git a/modules/ui/BaseSampleWindowView.py b/modules/ui/BaseSampleWindowView.py index 0f91ad2fa..1a9c6c0c6 100644 --- a/modules/ui/BaseSampleWindowView.py +++ b/modules/ui/BaseSampleWindowView.py @@ -1,227 +1,8 @@ -import contextlib -import copy -import os -import tkinter as tk -import traceback -from modules.model.BaseModel import BaseModel -from modules.modelSampler.BaseModelSampler import ( - BaseModelSampler, - ModelSamplerOutput, -) -from modules.ui.SampleFrame import SampleFrame -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.FileType import FileType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.time_util import get_string_timestamp -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import torch -import customtkinter as ctk -from PIL import Image -class SampleWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - use_external_model: bool, - callbacks: TrainCallbacks | None = None, - commands: TrainCommands | None = None, - *args, **kwargs - ): - super().__init__(parent, *args, **kwargs) - - self.title("Sample") - self.geometry("1200x800") - self.resizable(True, True) - - if not use_external_model: - self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) - # remove some settings to speed up model loading for sampling - self.initial_train_config.optimizer.optimizer = None - self.initial_train_config.ema = EMAMode.OFF - else: - self.initial_train_config = None - - #TODO why is there a current_train_config and an initial_train_config? - #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands - - # get model specific defaults - model_type = train_config.model_type - self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: - self.model = None - self.model_sampler = None - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_rowconfigure(3, weight=0) - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1) - - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) - prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") - - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) - settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - # image - self.image = ctk.CTkImage( - light_image=self.__dummy_image(), - size=(512, 512) - ) - - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) - image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def __load_model(self) -> BaseModel: - model_loader = create.create_model_loader( - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - model_setup = create.create_model_setup( - model_type=self.initial_train_config.model_type, - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - training_method=self.initial_train_config.training_method, - ) - - model_names = self.initial_train_config.model_names() - if self.initial_train_config.continue_last_backup: - last_backup_path = self.initial_train_config.get_last_backup_path() - - if last_backup_path: - if self.initial_train_config.training_method == TrainingMethod.LORA: - model_names.lora = last_backup_path - elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: - model_names.embedding.model_name = last_backup_path - else: # fine-tunes - model_names.base_model = last_backup_path - - print(f"Loading from backup '{last_backup_path}'...") - else: - print("No backup found, loading without backup...") - - if self.initial_train_config.quantization.cache_dir is None: - self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" - os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) - - model = model_loader.load( - model_type=self.initial_train_config.model_type, - model_names=model_names, - weight_dtypes=self.initial_train_config.weight_dtypes(), - quantization=self.initial_train_config.quantization, - ) - model.train_config = self.initial_train_config - - model_setup.setup_optimizations(model, self.initial_train_config) - model_setup.setup_train_device(model, self.initial_train_config) - model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) - - return model - - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: - return create.create_model_sampler( - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - model=model, - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - def __update_preview(self, sampler_output: ModelSamplerOutput): - if sampler_output.file_type == FileType.IMAGE: - image = sampler_output.data - self.image.configure( - light_image=image, - size=(image.width, image.height), - ) - - def __update_progress(self, progress: int, max_progress: int): - self.progress.set(progress / max_progress) - self.update() - - def __dummy_image(self) -> Image: - return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - - def __sample(self): - sample = copy.copy(self.sample) - - if self.commands: - self.commands.sample_custom(sample) - else: - if self.model is None: - # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) - - sample.from_train_config(self.current_train_config) - - sample_dir = os.path.join( - self.initial_train_config.workspace_dir, - "samples", - "custom", - ) - - progress = self.model.train_progress - sample_path = os.path.join( - sample_dir, - f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" - ) - - self.model.eval() - - self.model_sampler.sample( - sample_config=sample, - destination=sample_path, - image_format=self.current_train_config.sample_image_format, - video_format=self.current_train_config.sample_video_format, - audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, - ) - - def destroy(self): - try: - if hasattr(self, "_icon_image_ref"): - del self._icon_image_ref - - # Remove any pending after callbacks - for after_id in self.tk.call('after', 'info'): - with contextlib.suppress(tk.TclError, RuntimeError): - self.after_cancel(after_id) - - super().destroy() - except (tk.TclError, RuntimeError) as e: - print(f"Error destroying window: {e}") - except Exception as e: - print(f"Unexpected error destroying window: {e}") - traceback.print_exc() +class BaseSampleWindowView: + def __init__(self, components): + pass diff --git a/modules/ui/BaseSamplingTabView.py b/modules/ui/BaseSamplingTabView.py index 5a3c44f08..69ec082a1 100644 --- a/modules/ui/BaseSamplingTabView.py +++ b/modules/ui/BaseSamplingTabView.py @@ -1,124 +1,67 @@ -from modules.ui.ConfigList import ConfigList -from modules.ui.SampleParamsWindow import SampleParamsWindow -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from abc import ABC, abstractmethod -import customtkinter as ctk +from modules.ui.BaseConfigListView import BaseConfigListView -class SamplingTab(ConfigList): +class BaseSamplingTabView(BaseConfigListView): + pass - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - from_external_file=True, - attr_name="sample_definition_file_name", - config_dir="training_samples", - default_config_name="samples.json", - add_button_text="Add Sample", - add_button_tooltip="Add a new sample configuration.", - is_full_width=True, - show_toggle_button=True - ) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return SampleConfig.default_values(self.train_config.model_type) - - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) - - -class SampleWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) +class BaseSampleWidgetView(ABC): + def __init__(self, components): + self.components = components + def build_content(self, frame, element, ui_state, i, open_command, remove_command, clone_command, save_command): self.element = element - self.ui_state = UIState(self, element) self.i = i self.save_command = save_command - self.grid_columnconfigure(10, weight=1) - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) + self.components.colored_icon_button(frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i), - ) - clone_button.grid(row=0, column=1, padx=5) + self.components.colored_icon_button(frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i), padx=5) # enabled - self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) - self.enabled_switch.configure(width=40) + self.enabled_switch = self.components.switch(frame, 0, 2, ui_state, "enabled", self._switch_enabled, width=40) # width - components.label(self, 0, 3, "width:") - self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") - self.width_entry.bind('', lambda _: save_command()) - self.width_entry.configure(width=50) + self.components.label(frame, 0, 3, "width:") + self.width_entry = self.components.entry(frame, 0, 4, ui_state, "width", width=50) # height - components.label(self, 0, 5, "height:") - self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") - self.height_entry.bind('', lambda _: save_command()) - self.height_entry.configure(width=50) + self.components.label(frame, 0, 5, "height:") + self.height_entry = self.components.entry(frame, 0, 6, ui_state, "height", width=50) # seed - components.label(self, 0, 7, "seed:") - self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") - self.seed_entry.bind('', lambda _: save_command()) - self.seed_entry.configure(width=80) + self.components.label(frame, 0, 7, "seed:") + self.seed_entry = self.components.entry(frame, 0, 8, ui_state, "seed", width=80) # prompt - components.label(self, 0, 9, "prompt:") - self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") - self.prompt_entry.bind('', lambda _: save_command()) + self.components.label(frame, 0, 9, "prompt:") + self.prompt_entry = self.components.entry(frame, 0, 10, ui_state, "prompt") # button - self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) - self.button.configure(width=40) + self.button = self.components.icon_button(frame, 0, 11, "...", lambda: open_command(self.i, ui_state)) - self.__set_enabled() + self._bind_save(save_command) + self._set_enabled() - def __switch_enabled(self): + @abstractmethod + def _bind_save(self, save_command): pass + + # BaseConfigListView calls configure_element() on all widget types generically; + # sampling widgets have no post-window logic, so this is an intentional no-op. + def configure_element(self): pass # noqa: B027 + + def _switch_enabled(self): self.save_command() - self.__set_enabled() + self._set_enabled() - def __set_enabled(self): + def _set_enabled(self): enabled = self.element.enabled self.width_entry.configure(state="normal" if enabled else "disabled") self.height_entry.configure(state="normal" if enabled else "disabled") self.prompt_entry.configure(state="normal" if enabled else "disabled") self.seed_entry.configure(state="normal" if enabled else "disabled") self.button.configure(state="normal" if enabled else "disabled") - - def configure_element(self): - pass - - def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/BaseSchedulerParamsWindowView.py b/modules/ui/BaseSchedulerParamsWindowView.py index f96ed4876..1106d5227 100644 --- a/modules/ui/BaseSchedulerParamsWindowView.py +++ b/modules/ui/BaseSchedulerParamsWindowView.py @@ -1,119 +1,21 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +from modules.ui.BaseConfigListView import BaseConfigListView -class KvParams(ConfigList): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, - attr_name="scheduler_params", - from_external_file=False, - add_button_text="add parameter", - is_full_width=True - ) +class BaseSchedulerParamsWindowView: + def __init__(self, components): + self.components = components - def refresh_ui(self): - self._create_element_list() + def build_content(self, master, controller, ui_state): + if controller.is_custom_scheduler(): + self.components.label(master, 0, 0, "Class Name", + tooltip="Python class module and name for the custom scheduler class, in the form of ..") + self.components.entry(master, 0, 1, ui_state, "custom_learning_rate_scheduler") - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) - def create_new_element(self) -> dict[str, str]: - return {"key": "", "value": ""} +class BaseKvParamsView(BaseConfigListView): + def __init__(self, components): + self.components = components def open_element_window(self, i, ui_state): pass - - -class KvWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__(master=master, bg_color="transparent") - self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command - - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1, uniform=1) - self.grid_columnconfigure(2, weight=1, uniform=1) - - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i)) - close_button.grid(row=0, column=0) - - # Key - tooltip_key = "Key name for an argument in your scheduler" - self.key = components.entry(self, 0, 1, self.ui_state, "key", - tooltip=tooltip_key, wide_tooltip=True) - self.key.bind("", lambda _: save_command()) - self.key.configure(width=50) - - # Value - tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." - self.value = components.entry(self, 0, 2, self.ui_state, "value", - tooltip=tooltip_val, wide_tooltip=True) - self.value.bind("", lambda _: save_command()) - self.value.configure(width=50) - - def place_in_list(self): - self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") - - -class SchedulerParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - - self.title("Learning Rate Scheduler Settings") - self.geometry("800x400") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkFrame(self) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - - self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") - self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: - components.label(master, 0, 0, "Class Name", - tooltip="Python class module and name for the custom scheduler class, in the form of ..") - components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") - - # Any additional parameters, in key-value form. - self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) - - def on_window_close(self): - self.destroy() diff --git a/modules/ui/BaseTimestepDistributionWindowView.py b/modules/ui/BaseTimestepDistributionWindowView.py index 21e41ce3e..19354bf43 100644 --- a/modules/ui/BaseTimestepDistributionWindowView.py +++ b/modules/ui/BaseTimestepDistributionWindowView.py @@ -1,186 +1,46 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import torch -from torch import Tensor -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -class TimestepGenerator(ModelSetupNoiseMixin): - - def __init__( - self, - timestep_distribution: TimestepDistribution, - min_noising_strength: float, - max_noising_strength: float, - noising_weight: float, - noising_bias: float, - timestep_shift: float, - ): - super().__init__() - - self.timestep_distribution = timestep_distribution - self.min_noising_strength = min_noising_strength - self.max_noising_strength = max_noising_strength - self.noising_weight = noising_weight - self.noising_bias = noising_bias - self.timestep_shift = timestep_shift - - def generate(self) -> Tensor: - generator = torch.Generator() - generator.seed() - - config = TrainConfig.default_values() - config.timestep_distribution = self.timestep_distribution - config.min_noising_strength = self.min_noising_strength - config.max_noising_strength = self.max_noising_strength - config.noising_weight = self.noising_weight - config.noising_bias = self.noising_bias - config.timestep_shift = self.timestep_shift - - - return self._get_timestep_discrete( - num_train_timesteps=1000, - deterministic=False, - generator=generator, - batch_size=1000000, - config=config, - ) - - -class TimestepDistributionWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.title("Timestep Distribution") - self.geometry("900x600") - self.resizable(True, True) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - frame.grid_rowconfigure(7, weight=1) +class BaseTimestepDistributionWindowView: + def __init__(self, components): + self.components = components + def build_content(self, frame, controller, ui_state): # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", + self.components.label(frame, 0, 0, "Timestep Distribution", tooltip="Selects the function to sample timesteps during training", wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, + self.components.options(frame, 0, 1, controller.get_distribution_options(), ui_state, "timestep_distribution") # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", + self.components.label(frame, 1, 0, "Min Noising Strength", tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") + self.components.entry(frame, 1, 1, ui_state, "min_noising_strength") # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", + self.components.label(frame, 2, 0, "Max Noising Strength", tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") + self.components.entry(frame, 2, 1, ui_state, "max_noising_strength") # noising weight - components.label(frame, 3, 0, "Noising Weight", + self.components.label(frame, 3, 0, "Noising Weight", tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") + self.components.entry(frame, 3, 1, ui_state, "noising_weight") # noising bias - components.label(frame, 4, 0, "Noising Bias", + self.components.label(frame, 4, 0, "Noising Bias", tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") + self.components.entry(frame, 4, 1, ui_state, "noising_bias") # timestep shift - components.label(frame, 5, 0, "Timestep Shift", + self.components.label(frame, 5, 0, "Timestep Shift", tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") + self.components.entry(frame, 5, 1, ui_state, "timestep_shift") # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", + self.components.label(frame, 6, 0, "Dynamic Timestep Shifting", tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - fig, ax = plt.subplots() - self.ax = ax - self.canvas = FigureCanvasTkAgg(fig, master=frame) - self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) - - fig.set_facecolor(background_color) - ax.set_facecolor(background_color) - ax.spines['bottom'].set_color(text_color) - ax.spines['left'].set_color(text_color) - ax.spines['top'].set_color(text_color) - ax.spines['right'].set_color(text_color) - ax.tick_params(axis='x', colors=text_color, which="both") - ax.tick_params(axis='y', colors=text_color, which="both") - ax.xaxis.label.set_color(text_color) - ax.yaxis.label.set_color(text_color) - - self.__update_preview() - - # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) - - frame.pack(fill="both", expand=1) - return frame - - def __update_preview(self): - generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, - ) - - self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) - self.canvas.draw() - - def __ok(self): - self.destroy() + self.components.switch(frame, 6, 1, ui_state, "dynamic_timestep_shifting") diff --git a/modules/ui/BaseTopBarView.py b/modules/ui/BaseTopBarView.py index 820fdb71a..f50c1ce14 100644 --- a/modules/ui/BaseTopBarView.py +++ b/modules/ui/BaseTopBarView.py @@ -1,7 +1,7 @@ import json import os import traceback -import webbrowser +from abc import abstractmethod from collections.abc import Callable from contextlib import suppress @@ -11,25 +11,45 @@ from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.optimizer_util import change_optimizer -from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class BaseTopBarView: + def __init__(self, components): + self.components = components -class TopBar: - def __init__( + @abstractmethod + def _make_config_ui_state(self, master, data): + pass + + @abstractmethod + def _get_dropdown_text(self, widget) -> str: + pass + + @abstractmethod + def _setup_frame_column_weight(self): + pass + + @abstractmethod + def _forget_dropdown(self): + pass + + @abstractmethod + def _show_save_dialog(self, default_value: str, callback): + pass + + def build( self, + frame, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, change_model_type_callback: Callable[[ModelType], None], change_training_method_callback: Callable[[TrainingMethod], None], load_preset_callback: Callable[[], None], ): + self.controller = controller + self.frame = frame self.master = master - self.train_config = train_config self.ui_state = ui_state self.change_model_type_callback = change_model_type_callback self.change_training_method_callback = change_training_method_callback @@ -40,20 +60,16 @@ def __init__( self.config_ui_data = { "config_name": path_util.canonical_join(self.dir, "#.json") } - self.config_ui_state = UIState(master, self.config_ui_data) + self.config_ui_state = self._make_config_ui_state(master, self.config_ui_data) - self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() + self.configs = controller.load_available_config_names(self.dir) self.current_config = [] - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - self.training_method = None # title - components.app_title(self.frame, 0, 0) + self.components.app_title(self.frame, 0, 0) # dropdown self.configs_dropdown = None @@ -61,49 +77,25 @@ def __init__( # remove button # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + # self.components.icon_button(self.frame, 0, 2, "-", self.__remove_config) # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + self.components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) + self.components.button(self.frame, 0, 3, "Save config", self.__save_config, + tooltip="Save the current configuration in a custom preset", width=90) # padding - self.frame.grid_columnconfigure(5, weight=1) + self._setup_frame_column_weight() # model type - components.options_kv( + self.components.options_kv( master=self.frame, row=0, column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, + values=controller.get_model_types(), + ui_state=ui_state, var_name="model_type", command=self.__change_model_type, ) @@ -112,40 +104,9 @@ def __create_training_method(self): if self.training_method: self.training_method.destroy() - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] - - # training method - self.training_method = components.options_kv( + values = self.controller.get_training_methods(self.controller.train_config.model_type) + + self.training_method = self.components.options_kv( master=self.frame, row=0, column=7, @@ -161,40 +122,21 @@ def __change_model_type(self, model_type: ModelType): def __create_configs_dropdown(self): if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() + self._forget_dropdown() - self.configs_dropdown = components.options_kv( + self.configs_dropdown = self.components.options_kv( self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config ) - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): - if path != "#.json": - path = path_util.canonical_join(self.dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() - - def __save_to_file(self, name) -> str: - name = path_util.safe_filename(name) - path = path_util.canonical_join("training_presets", f"{name}.json") - - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - - return path - - def __save_secrets(self, path) -> str: - write_json_atomic(path, self.train_config.secrets.to_dict()) - return path + def __save_config(self): + default_value = self._get_dropdown_text(self.configs_dropdown) + while default_value.startswith('#'): + default_value = default_value[1:] - def open_wiki(self): - webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) + self._show_save_dialog(default_value, self.__save_new_config) def __save_new_config(self, name): - path = self.__save_to_file(name) + path = self.controller.save_to_file(name) is_new_config = name not in [x[0] for x in self.configs] @@ -208,20 +150,6 @@ def __save_new_config(self, name): if is_new_config: self.__create_configs_dropdown() - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] - - dialogs.StringInputDialog( - parent=self.master, - title="name", - question="Config Name", - callback=self.__save_new_config, - default_value=default_value, - validate_callback=lambda x: not x.startswith("#") - ) - def __load_current_config(self, filename): try: basename = os.path.basename(filename) @@ -239,10 +167,10 @@ def __load_current_config(self, filename): secrets_dict=json.load(f) loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - self.train_config.from_dict(loaded_config.to_dict()) + self.controller.train_config.from_dict(loaded_config.to_dict()) self.ui_state.update(loaded_config) - optimizer_config = change_optimizer(self.train_config) + optimizer_config = change_optimizer(self.controller.train_config) self.ui_state.get_var("optimizer").update(optimizer_config) self.load_preset_callback() @@ -255,6 +183,8 @@ def __remove_config(self): # TODO pass + def open_wiki(self): + self.controller.open_wiki() + def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") + self.controller.save_default() diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index ba90d2e64..7acbf0005 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -1,889 +1,356 @@ -import ctypes -import datetime -import json -import os -import platform -import subprocess -import sys -import threading -import time -import traceback -import webbrowser +from abc import ABC, abstractmethod from collections.abc import Callable -from contextlib import suppress -from pathlib import Path -from tkinter import filedialog, messagebox - -import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.TrainConfig import TrainConfig + +from modules.util import path_util from modules.util.enum.DataType import DataType from modules.util.enum.GradientReducePrecision import GradientReducePrecision from modules.util.enum.ImageFormat import ImageFormat -from modules.util.enum.ModelType import ModelType from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.torch_util import torch_gc -from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import flush_and_validate_all - -import torch - -import customtkinter as ctk -from customtkinter import AppearanceModeTracker - -# chunk for forcing Windows to ignore DPI scaling when moving between monitors -# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 -if platform.system() == "Windows": - with suppress(Exception): - # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically - ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE - -class TrainUI(ctk.CTk): - set_step_progress: Callable[[int, int], None] - set_epoch_progress: Callable[[int, int], None] - - status_label: ctk.CTkLabel | None - training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None - - _TRAIN_BUTTON_STYLES = { - "idle": { - "text": "Start Training", - "state": "normal", - "fg_color": "#198754", - "hover_color": "#146c43", - "text_color": "white", - "text_color_disabled": "white", - }, - "running": { - "text": "Stop Training", - "state": "normal", - "fg_color": "#dc3545", - "hover_color": "#bb2d3b", - "text_color": "white", - }, - "stopping": { - "text": "Stopping...", - "state": "disabled", - "fg_color": "#dc3545", - "hover_color": "#dc3545", - "text_color": "white", - "text_color_disabled": "white", - }, - } - - def __init__(self): - super().__init__() - - self.title("OneTrainer") - self.geometry("1100x740") - - self.after(100, lambda: self._set_icon()) - - # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration - ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") - ctk.set_default_color_theme("blue") - - self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.status_label = None - self.eta_label = None - self.training_button = None - self.export_button = None - self.tabview = None - - self.model_tab = None - self.training_tab = None - self.lora_tab = None - self.cloud_tab = None - self.additional_embeddings_tab = None - - self.top_bar_component = self.top_bar(self) - self.content_frame(self) - self.bottom_bar(self) - - self.training_thread = None - self.training_callbacks = None - self.training_commands = None - - self.always_on_tensorboard_subprocess = None - self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() - - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) - - # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) - - self.protocol("WM_DELETE_WINDOW", self.__close) - - def __close(self): - self.top_bar_component.save_default() - self._stop_always_on_tensorboard() - if hasattr(self, 'workspace_dir_trace_id'): - self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) - self.quit() - - def top_bar(self, master): - return TopBar( - master, - self.train_config, - self.ui_state, - self.change_model_type, - self.change_training_method, - self.load_preset, - ) - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - def bottom_bar(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=2, column=0, sticky="nsew") +class BaseTrainUIView(ABC): + def __init__(self, components): + self.components = components - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") + # --- Abstract callbacks (controller calls into view) --- - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) + @abstractmethod + def on_update_status(self, status: str): pass - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + @abstractmethod + def on_training_started(self): pass - # padding - frame.grid_columnconfigure(2, weight=1) + @abstractmethod + def on_training_stopped(self, error_caught: bool): pass + @abstractmethod + def on_training_stopping(self): pass - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") + @abstractmethod + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): pass - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + @abstractmethod + def schedule_on_main_thread(self, fn: Callable): pass - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) + @abstractmethod + def get_cloud_reattach(self) -> bool: pass - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) - self._set_training_button_style("idle") # centralized styling + @abstractmethod + def save_default(self): pass + + @abstractmethod + def show_validation_errors(self, errors: list[str]): pass + + @abstractmethod + def wait_window(self, window): pass + + @abstractmethod + def show_window(self, window): pass + + @abstractmethod + def connect_window_closed(self, window, callback): pass + + def sync_cloud_secrets(self): + self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - return frame + def start_training(self): + self.controller.start_training() + + def open_tensorboard(self): + self.controller.open_tensorboard() - def content_frame(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") + def sample_now(self): + self.controller.sample_now() - frame.grid_rowconfigure(0, weight=1) - frame.grid_columnconfigure(0, weight=1) + def backup_now(self): + self.controller.backup_now() - self.tabview = ctk.CTkTabview(frame) - self.tabview.grid(row=0, column=0, sticky="nsew") + def save_now(self): + self.controller.save_now() - self.general_tab = self.create_general_tab(self.tabview.add("general")) - self.model_tab = self.create_model_tab(self.tabview.add("model")) - self.data_tab = self.create_data_tab(self.tabview.add("data")) - self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) - self.training_tab = self.create_training_tab(self.tabview.add("training")) - self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) - self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) - self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) - self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) - self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) + @abstractmethod + def open_dataset_tool(self): pass - self.change_training_method(self.train_config.training_method) + @abstractmethod + def open_video_tool(self): pass - return frame + @abstractmethod + def open_convert_model_tool(self): pass - def create_general_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + @abstractmethod + def open_sampling_tool(self): pass + + @abstractmethod + def open_manual_sample_window(self): pass + + # --- Content builders (components calls; called by CTK view after frame creation) --- + + def build_bottom_bar_content(self, frame, status_frame, controller, ui_state): + self.set_step_progress, self.set_epoch_progress = self.components.double_progress(frame, 0, 0, "step", "epoch") + + self.status_label = self.components.label(status_frame, 0, 0, "", pad=0, + tooltip="Current status of the training run") + self.eta_label = self.components.label(status_frame, 1, 0, "", pad=0) + self.export_button = self.components.button(frame, 0, 3, "Export", self.export_training, + width=60, padx=5, pady=(15, 0), + tooltip="Export the current configuration as a script to run without a UI") + + self.components.button(frame, 0, 4, "Debug", self.generate_debug_package, + width=60, padx=(5, 25), pady=(15, 0), + tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") + + self.components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, + width=100, padx=(0, 5), pady=(15, 0)) + + self.training_button = self.components.button(frame, 0, 6, "Start Training", self.start_training, + padx=(5, 20), pady=(15, 0)) + + def build_general_tab_content(self, frame, controller, ui_state): # workspace dir - components.label(frame, 0, 0, "Workspace Directory", + self.components.label(frame, 0, 0, "Workspace Directory", tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) + self.components.path_entry(frame, 0, 1, ui_state, "workspace_dir", mode="dir", command=controller._on_workspace_dir_change) # cache dir - components.label(frame, 0, 2, "Cache Directory", + self.components.label(frame, 0, 2, "Cache Directory", tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") + self.components.path_entry(frame, 0, 3, ui_state, "cache_dir", mode="dir") # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", + self.components.label(frame, 2, 0, "Continue from last backup", tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") + self.components.switch(frame, 2, 1, ui_state, "continue_last_backup") # only cache - components.label(frame, 2, 2, "Only Cache", + self.components.label(frame, 2, 2, "Only Cache", tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") + self.components.switch(frame, 2, 3, ui_state, "only_cache") # TODO: In Phase 4 rework the general tab. # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", + self.components.label(frame, 3, 0, "Prevent Overwrites", tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") + self.components.switch(frame, 3, 1, ui_state, "prevent_overwrites") # debug - components.label(frame, 4, 0, "Debug mode", + self.components.label(frame, 4, 0, "Debug mode", tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") + self.components.switch(frame, 4, 1, ui_state, "debug_mode") - components.label(frame, 4, 2, "Debug Directory", + self.components.label(frame, 4, 2, "Debug Directory", tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + self.components.path_entry(frame, 4, 3, ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) # tensorboard - components.label(frame, 6, 0, "Tensorboard", + self.components.label(frame, 6, 0, "Tensorboard", tooltip="Starts the Tensorboard Web UI during training") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") + self.components.switch(frame, 6, 1, ui_state, "tensorboard") - components.label(frame, 6, 2, "Always-On Tensorboard", + self.components.label(frame, 6, 2, "Always-On Tensorboard", tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) + self.components.switch(frame, 6, 3, ui_state, "tensorboard_always_on", command=controller._on_always_on_tensorboard_toggle) - components.label(frame, 7, 0, "Expose Tensorboard", + self.components.label(frame, 7, 0, "Expose Tensorboard", tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") - components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") - components.label(frame, 7, 2, "Tensorboard Port", + self.components.switch(frame, 7, 1, ui_state, "tensorboard_expose") + self.components.label(frame, 7, 2, "Tensorboard Port", tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - + self.components.entry(frame, 7, 3, ui_state, "tensorboard_port") # validation - components.label(frame, 8, 0, "Validation", + self.components.label(frame, 8, 0, "Validation", tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") + self.components.switch(frame, 8, 1, ui_state, "validation") - components.label(frame, 8, 2, "Validate after", + self.components.label(frame, 8, 2, "Validate after", tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") + self.components.time_entry(frame, 8, 3, ui_state, "validate_after", "validate_after_unit") # device - components.label(frame, 10, 0, "Dataloader Threads", + self.components.label(frame, 10, 0, "Dataloader Threads", tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) + self.components.entry(frame, 10, 1, ui_state, "dataloader_threads", required=True) - components.label(frame, 11, 0, "Train Device", + self.components.label(frame, 11, 0, "Train Device", tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") - components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + self.components.entry(frame, 11, 1, ui_state, "train_device", required=True) - components.label(frame, 12, 0, "Multi-GPU", + self.components.label(frame, 12, 0, "Multi-GPU", tooltip="Enable multi-GPU training") - components.switch(frame, 12, 1, self.ui_state, "multi_gpu") - components.label(frame, 12, 2, "Device Indexes", + self.components.switch(frame, 12, 1, ui_state, "multi_gpu") + self.components.label(frame, 12, 2, "Device Indexes", tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") - components.entry(frame, 12, 3, self.ui_state, "device_indexes") + self.components.entry(frame, 12, 3, ui_state, "device_indexes") - components.label(frame, 13, 0, "Gradient Reduce Precision", + self.components.label(frame, 13, 0, "Gradient Reduce Precision", tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" "FLOAT_32: Reduce gradients in float32\n" "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", wide_tooltip=True) - components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + self.components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], ui_state, "gradient_reduce_precision") - components.label(frame, 13, 2, "Fused Gradient Reduce", + self.components.label(frame, 13, 2, "Fused Gradient Reduce", tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") + self.components.switch(frame, 13, 3, ui_state, "fused_gradient_reduce") - components.label(frame, 14, 0, "Async Gradient Reduce", + self.components.label(frame, 14, 0, "Async Gradient Reduce", tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 14, 2, "Buffer size (MB)", + self.components.switch(frame, 14, 1, ui_state, "async_gradient_reduce") + self.components.label(frame, 14, 2, "Buffer size (MB)", tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") + self.components.entry(frame, 14, 3, ui_state, "async_gradient_reduce_buffer") - components.label(frame, 15, 0, "Temp Device", + self.components.label(frame, 15, 0, "Temp Device", tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 15, 1, self.ui_state, "temp_device") - - frame.pack(fill="both", expand=1) - return frame - - def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) - - def create_data_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.entry(frame, 15, 1, ui_state, "temp_device") + def build_data_tab_content(self, frame, controller, ui_state): # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", + self.components.label(frame, 0, 0, "Aspect Ratio Bucketing", tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") + self.components.switch(frame, 0, 1, ui_state, "aspect_ratio_bucketing") # latent caching - components.label(frame, 1, 0, "Latent Caching", + self.components.label(frame, 1, 0, "Latent Caching", tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") + self.components.switch(frame, 1, 1, ui_state, "latent_caching") # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", + self.components.label(frame, 2, 0, "Clear cache before training", tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + self.components.switch(frame, 2, 1, ui_state, "clear_cache_before_training") - frame.pack(fill="both", expand=1) - return frame - - def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) - - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) - - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) - - def create_sampling_tab(self, master): - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) - - # sample after - top_frame = ctk.CTkFrame(master=master, corner_radius=0) - top_frame.grid(row=0, column=0, sticky="nsew") - sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") - sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - - components.label(top_frame, 0, 0, "Sample After", + def build_sampling_tab_header(self, top_frame, sub_frame, controller, ui_state): + self.components.label(top_frame, 0, 0, "Sample After", tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") + self.components.time_entry(top_frame, 0, 1, ui_state, "sample_after", "sample_after_unit") - components.label(top_frame, 0, 2, "Skip First", + self.components.label(top_frame, 0, 2, "Skip First", tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") + self.components.entry(top_frame, 0, 3, ui_state, "sample_skip_first", width=50, sticky="nw") - components.label(top_frame, 0, 4, "Format", + self.components.label(top_frame, 0, 4, "Format", tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ + self.components.options_kv(top_frame, 0, 5, [ ("PNG", ImageFormat.PNG), ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") + ], ui_state, "sample_image_format") - components.button(top_frame, 0, 6, "sample now", self.sample_now) + self.components.button(top_frame, 0, 6, "sample now", self.sample_now) - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) + self.components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window) - components.label(sub_frame, 0, 0, "Non-EMA Sampling", + self.components.label(sub_frame, 0, 0, "Non-EMA Sampling", tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") + self.components.switch(sub_frame, 0, 1, ui_state, "non_ema_sampling") - components.label(sub_frame, 0, 2, "Samples to Tensorboard", + self.components.label(sub_frame, 0, 2, "Samples to Tensorboard", tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") - - # table - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - return SamplingTab(frame, self.train_config, self.ui_state) - - def create_backup_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.switch(sub_frame, 0, 3, ui_state, "samples_to_tensorboard") + def build_backup_tab_content(self, frame, controller, ui_state): # backup after - components.label(frame, 0, 0, "Backup After", + self.components.label(frame, 0, 0, "Backup After", tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") + self.components.time_entry(frame, 0, 1, ui_state, "backup_after", "backup_after_unit") # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) + self.components.button(frame, 0, 3, "backup now", self.backup_now) # rolling backup - components.label(frame, 1, 0, "Rolling Backup", + self.components.label(frame, 1, 0, "Rolling Backup", tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") + self.components.switch(frame, 1, 1, ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + self.components.label(frame, 1, 3, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + self.components.entry(frame, 1, 4, ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + self.components.label(frame, 2, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + self.components.switch(frame, 2, 1, ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + self.components.label(frame, 3, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + self.components.time_entry(frame, 3, 1, ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + self.components.button(frame, 3, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + self.components.label(frame, 4, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + self.components.entry(frame, 4, 1, ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + self.components.label(frame, 5, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - - frame.pack(fill="both", expand=1) - return frame - - def embedding_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.entry(frame, 5, 1, ui_state, "save_filename_prefix") + def build_embedding_tab_content(self, frame, controller, ui_state): # embedding model name - components.label(frame, 0, 0, "Base embedding", + self.components.label(frame, 0, 0, "Base embedding", tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier + self.components.path_entry( + frame, 0, 1, ui_state, "embedding.model_name", + mode="file", path_modifier=path_util.json_path_modifier ) # token count - components.label(frame, 1, 0, "Token count", + self.components.label(frame, 1, 0, "Token count", tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") + self.components.entry(frame, 1, 1, ui_state, "embedding.token_count") # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", + self.components.label(frame, 2, 0, "Initial embedding text", tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") + self.components.entry(frame, 2, 1, ui_state, "embedding.initial_embedding_text") # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", + self.components.label(frame, 3, 0, "Embedding Weight Data Type", tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ + self.components.options_kv(frame, 3, 1, [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") + ], ui_state, "embedding_weight_dtype") # placeholder - components.label(frame, 4, 0, "Placeholder", + self.components.label(frame, 4, 0, "Placeholder", tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") + self.components.entry(frame, 4, 1, ui_state, "embedding.placeholder") # output embedding - components.label(frame, 5, 0, "Output embedding", + self.components.label(frame, 5, 0, "Output embedding", tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - - frame.pack(fill="both", expand=1) - return frame - - def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) - - def create_tools_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) + self.components.switch(frame, 5, 1, ui_state, "embedding.is_output_embedding") + def build_tools_tab_content(self, frame, controller, ui_state): # dataset - components.label(frame, 0, 0, "Dataset Tools", + self.components.label(frame, 0, 0, "Dataset Tools", tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) + self.components.button(frame, 0, 1, "Open", self.open_dataset_tool) # video tools - components.label(frame, 1, 0, "Video Tools", + self.components.label(frame, 1, 0, "Video Tools", tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) + self.components.button(frame, 1, 1, "Open", self.open_video_tool) # convert model - components.label(frame, 2, 0, "Convert Model Tools", + self.components.label(frame, 2, 0, "Convert Model Tools", tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + self.components.button(frame, 2, 1, "Open", self.open_convert_model_tool) # sample - components.label(frame, 3, 0, "Sampling Tool", + self.components.label(frame, 3, 0, "Sampling Tool", tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) + self.components.button(frame, 3, 1, "Open", self.open_sampling_tool) - components.label(frame, 4, 0, "Profiling Tool", + self.components.label(frame, 4, 0, "Profiling Tool", tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - - frame.pack(fill="both", expand=1) - return frame - - def change_model_type(self, model_type: ModelType): - if self.model_tab: - self.model_tab.refresh_ui() - - if self.training_tab: - self.training_tab.refresh_ui() - - if self.lora_tab: - self.lora_tab.refresh_ui() - - def change_training_method(self, training_method: TrainingMethod): - if not self.tabview: - return - - if self.model_tab: - self.model_tab.refresh_ui() - - if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: - self.tabview.delete("LoRA") - self.lora_tab = None - if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: - self.tabview.delete("embedding") - - if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) - if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: - self.embedding_tab(self.tabview.add("embedding")) - - def load_preset(self): - if not self.tabview: - return - - if self.additional_embeddings_tab: - self.additional_embeddings_tab.refresh_ui() - - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) - - def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: - spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step - remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - - if train_progress.global_step <= 30: - return "Estimating ..." - - td = datetime.timedelta(seconds=total_eta) - days = td.days - hours, remainder = divmod(td.seconds, 3600) - minutes, seconds = divmod(remainder, 60) - if days > 0: - return f"{days}d {hours}h" - elif hours > 0: - return f"{hours}h {minutes}m" - elif minutes > 0: - return f"{minutes}m {seconds}s" - else: - return f"{seconds}s" - - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") - - def delete_eta_label(self): - self.eta_label.configure(text="") - - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) - - def on_update_status(self, status: str): - self.status_label.configure(text=status) - - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) - - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) - - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) - - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() - - def open_profiling_tool(self): - self.profiling_window.deiconify() - - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) - - if not zip_path: - return - - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" - - self.on_update_status("Generating debug package...") - - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") - - - def open_manual_sample_window (self): - training_callbacks = self.training_callbacks - training_commands = self.training_commands - - if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, - use_external_model=True, - callbacks=training_callbacks, - commands=training_commands, - ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() - - def __training_thread_function(self): - error_caught = False - - self.training_callbacks = TrainCallbacks( - on_update_train_progress=self.on_update_train_progress, - on_update_status=self.on_update_status, - ) - - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) - try: - trainer.start() - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - - self.start_time = time.monotonic() - trainer.train() - except Exception: - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - error_caught = True - traceback.print_exc() - - trainer.end() - - # clear gpu memory - del trainer - - self.training_thread = None - self.training_commands = None - torch.clear_autocast_cache() - torch_gc() - - if error_caught: - self.on_update_status("Error: check the console for details") - else: - self.on_update_status("Stopped") - self.delete_eta_label() - - # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash - self.after(0, self._set_training_button_idle) - - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self.after(0, self._start_always_on_tensorboard) - - def start_training(self): - if self.training_thread is None: - self.save_default() - - # --- pre-training validation gate --- - errors = flush_and_validate_all() - - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) - return - - self._set_training_button_running() - - if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - self.training_commands = TrainCommands() - torch_gc() - - self.training_thread = threading.Thread(target=self.__training_thread_function) - self.training_thread.start() - else: - self._set_training_button_stopping() - self.on_update_status("Stopping ...") - self.training_commands.stop() - - def save_default(self): - self.top_bar_component.save_default() - self.concepts_tab.save_current_config() - self.sampling_tab.save_current_config() - self.additional_embeddings_tab.save_current_config() - - def export_training(self): - file_path = filedialog.asksaveasfilename(filetypes=[ - ("All Files", "*.*"), - ("json", "*.json"), - ], initialdir=".", initialfile="config.json") - - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.save() - - def _check_start_always_on_tensorboard(self): - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _start_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") - - os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - - tensorboard_args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.train_config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - try: - self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) - except Exception: - self.always_on_tensorboard_subprocess = None - - def _stop_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - try: - self.always_on_tensorboard_subprocess.terminate() - self.always_on_tensorboard_subprocess.wait(timeout=5) - except subprocess.TimeoutExpired: - self.always_on_tensorboard_subprocess.kill() - except Exception: - pass - finally: - self.always_on_tensorboard_subprocess = None - - def _on_workspace_dir_change(self, new_workspace_dir: str): - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_workspace_dir_change_trace(self, *args): - new_workspace_dir = self.train_config.workspace_dir - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_always_on_tensorboard_toggle(self): - if self.train_config.tensorboard_always_on: - if not (self.training_thread and self.train_config.tensorboard): - self._start_always_on_tensorboard() - else: - if not (self.training_thread and self.train_config.tensorboard): - self._stop_always_on_tensorboard() - - def _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: - return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") + self.components.button(frame, 4, 1, "Open", self.open_profiling_tool) diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index bcca11ae9..22a6af34e 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -1,9 +1,5 @@ -from modules.ui.OffloadingWindow import OffloadingWindow -from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow -from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow -from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig +from abc import ABC + from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod @@ -13,844 +9,750 @@ from modules.util.enum.LossWeight import LossWeight from modules.util.enum.Optimizer import Optimizer from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.optimizer_util import change_optimizer -from modules.util.ui import components -from modules.util.ui.UIState import UIState from modules.util.ui.validation_helpers import check_range, validate_resolution -import customtkinter as ctk - - -class TrainingTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - - self.master = master - self.train_config = train_config - self.ui_state = ui_state - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - - self.scroll_frame = None - - self.refresh_ui() - - def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - - self.scroll_frame.grid_columnconfigure(0, weight=1) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=1) - - column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_0.grid(row=0, column=0, sticky="nsew") - column_0.grid_columnconfigure(0, weight=1) - - column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_1.grid(row=0, column=1, sticky="nsew") - column_1.grid_columnconfigure(0, weight=1) - - column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_2.grid(row=0, column=2, sticky="nsew") - column_2.grid_columnconfigure(0, weight=1) - - if self.train_config.model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2) - - - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1) - self.__create_text_encoder_n_frame(column_0, 2, i=2) - self.__create_embedding_frame(column_0, 3) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_wuerstchen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_prior_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 0) - self.__create_loss_frame(column_2, 1) - self.__create_layer_frame(column_2, 2) - - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2, supports_vb_loss=True) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_2_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_chroma_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_qwen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_z_image_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_ernie_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_sana_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hi_dream_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) - self.__create_embedding_frame(column_0, 5) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __create_base_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + +class BaseTrainingTabView(ABC): + def __init__(self, components): + self.components = components + + def build(self, column_0, column_1, column_2, controller, ui_state, callbacks: dict): + model_type = controller.config.model_type + if model_type.is_stable_diffusion(): + self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + if model_type.is_stable_diffusion_3(): + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_stable_diffusion_xl(): + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_wuerstchen(): + self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_pixart(): + self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_flux_1(): + self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_flux_2(): + self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_chroma(): + self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_qwen(): + self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_sana(): + self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_hunyuan_video(): + self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_hi_dream(): + self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_z_image(): + self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + elif model_type.is_ernie(): + self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2) + self.__create_embedding_frame(column_0, 3, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_prior_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 0, ui_state) + self.__create_loss_frame(column_2, 1, controller, ui_state) + self.__create_layer_frame(column_2, 2, controller, ui_state) + + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state, supports_vb_loss=True) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True, supports_sequence_length=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_frame(column_0, 1, ui_state) + self.__create_embedding_frame(column_0, 2, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_embedding_frame(column_0, 4, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): + self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) + self.__create_text_encoder_n_frame(column_0, 4, ui_state, i=4, supports_include=True, supports_layer_skip=False) + self.__create_embedding_frame(column_0, 5, ui_state) + + self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, callbacks) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + + def __create_base_frame(self, master, row, controller, ui_state, callbacks): + frame = self.components.section_frame(master, row) # optimizer - components.label(frame, 0, 0, "Optimizer", - tooltip="The type of optimizer") - components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", - command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) + self.components.label(frame, 0, 0, "Optimizer", + tooltip="The type of optimizer") + self.components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], ui_state, "optimizer.optimizer", + command=callbacks.get('restore_optimizer'), + adv_command=callbacks.get('open_optimizer_params')) # learning rate scheduler # Wackiness will ensue when reloading configs if we don't check and clear this first. if hasattr(self, "lr_scheduler_comp"): delattr(self, "lr_scheduler_comp") delattr(self, "lr_scheduler_adv_comp") - components.label(frame, 1, 0, "Learning Rate Scheduler", - tooltip="Learning rate scheduler that automatically changes the learning rate during training") - _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, - "learning_rate_scheduler", command=self.__restore_scheduler_config, - adv_command=self.__open_scheduler_params_window) + self.components.label(frame, 1, 0, "Learning Rate Scheduler", + tooltip="Learning rate scheduler that automatically changes the learning rate during training") + _, d = self.components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], ui_state, + "learning_rate_scheduler", + command=callbacks.get('restore_scheduler'), + adv_command=callbacks.get('open_scheduler_params')) self.lr_scheduler_comp = d['component'] self.lr_scheduler_adv_comp = d['button_component'] # Initial call requires the presence of self.lr_scheduler_adv_comp. - self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) + restore_scheduler = callbacks.get('restore_scheduler') + if restore_scheduler: + restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) # learning rate - components.label(frame, 2, 0, "Learning Rate", - tooltip="The base learning rate") - components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) + self.components.label(frame, 2, 0, "Learning Rate", + tooltip="The base learning rate") + self.components.entry(frame, 2, 1, ui_state, "learning_rate", required=True) # learning rate warmup steps - components.label(frame, 3, 0, "Learning Rate Warmup Steps", - tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") - components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") + self.components.label(frame, 3, 0, "Learning Rate Warmup Steps", + tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") + self.components.entry(frame, 3, 1, ui_state, "learning_rate_warmup_steps") # learning rate min factor - components.label(frame, 4, 0, "Learning Rate Min Factor", - tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") - components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", - extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) + self.components.label(frame, 4, 0, "Learning Rate Min Factor", + tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") + self.components.entry(frame, 4, 1, ui_state, "learning_rate_min_factor", + extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) # learning rate cycles - components.label(frame, 5, 0, "Learning Rate Cycles", - tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") - components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") + self.components.label(frame, 5, 0, "Learning Rate Cycles", + tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") + self.components.entry(frame, 5, 1, ui_state, "learning_rate_cycles") # epochs - components.label(frame, 6, 0, "Epochs", - tooltip="The number of epochs for a full training run") - components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) + self.components.label(frame, 6, 0, "Epochs", + tooltip="The number of epochs for a full training run") + self.components.entry(frame, 6, 1, ui_state, "epochs", required=True) # batch size - components.label(frame, 7, 0, "Local Batch Size", - tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") - components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) + self.components.label(frame, 7, 0, "Local Batch Size", + tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") + self.components.entry(frame, 7, 1, ui_state, "batch_size", required=True) # accumulation steps - components.label(frame, 8, 0, "Accumulation Steps", - tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") - components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) + self.components.label(frame, 8, 0, "Accumulation Steps", + tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") + self.components.entry(frame, 8, 1, ui_state, "gradient_accumulation_steps", required=True) # Learning Rate Scaler - components.label(frame, 9, 0, "Learning Rate Scaler", - tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") - components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, - "learning_rate_scaler") + self.components.label(frame, 9, 0, "Learning Rate Scaler", + tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") + self.components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], ui_state, + "learning_rate_scaler") # clip grad norm - components.label(frame, 10, 0, "Clip Grad Norm", - tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") - components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") - - def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + self.components.label(frame, 10, 0, "Clip Grad Norm", + tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") + self.components.entry(frame, 10, 1, ui_state, "clip_grad_norm") + + def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_enabled: bool = False, + supports_circular_padding: bool = False): + frame = self.components.section_frame(master, row) row = 0 # ema - components.label(frame, row, 0, "EMA", - tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") - components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") + self.components.label(frame, row, 0, "EMA", + tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") + self.components.options(frame, row, 1, [str(x) for x in list(EMAMode)], ui_state, "ema") row += 1 # ema decay - components.label(frame, row, 0, "EMA Decay", - tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") - components.entry(frame, row, 1, self.ui_state, "ema_decay", - extra_validate=check_range(lower=0.5, upper=1, - message="EMA decay must be between 0.5 and 1")) + self.components.label(frame, row, 0, "EMA Decay", + tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") + self.components.entry(frame, row, 1, ui_state, "ema_decay", + extra_validate=check_range(lower=0.5, upper=1, + message="EMA decay must be between 0.5 and 1")) row += 1 # ema update step interval - components.label(frame, row, 0, "EMA Update Step Interval", - tooltip="Number of steps between EMA update steps") - components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") + self.components.label(frame, row, 0, "EMA Update Step Interval", + tooltip="Number of steps between EMA update steps") + self.components.entry(frame, row, 1, ui_state, "ema_update_step_interval") row += 1 # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) + self.components.label(frame, row, 0, "Gradient checkpointing", + tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") + self.components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, + "gradient_checkpointing", + adv_command=callbacks.get('open_offloading')) row += 1 # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") + self.components.label(frame, row, 0, "Layer offload fraction", + tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") + self.components.entry(frame, row, 1, ui_state, "layer_offload_fraction") row += 1 # train dtype - components.label(frame, row, 0, "Train Data Type", - tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ + self.components.label(frame, row, 0, "Train Data Type", + tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") + self.components.options_kv(frame, row, 1, [ ("float32", DataType.FLOAT_32), ("float16", DataType.FLOAT_16), ("bfloat16", DataType.BFLOAT_16), ("tfloat32", DataType.TFLOAT_32), - ], self.ui_state, "train_dtype") + ], ui_state, "train_dtype") row += 1 # fallback train dtype - components.label(frame, row, 0, "Fallback Train Data Type", - tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ + self.components.label(frame, row, 0, "Fallback Train Data Type", + tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") + self.components.options_kv(frame, row, 1, [ ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "fallback_train_dtype") + ], ui_state, "fallback_train_dtype") row += 1 # autocast cache - components.label(frame, row, 0, "Autocast Cache", - tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") - components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") + self.components.label(frame, row, 0, "Autocast Cache", + tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") + self.components.switch(frame, row, 1, ui_state, "enable_autocast_cache") row += 1 # resolution - components.label(frame, row, 0, "Resolution", - tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.entry(frame, row, 1, self.ui_state, "resolution", required=True, - extra_validate=validate_resolution()) + self.components.label(frame, row, 0, "Resolution", + tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") + self.components.entry(frame, row, 1, ui_state, "resolution", required=True, + extra_validate=validate_resolution()) row += 1 # frames if video_training_enabled: - components.label(frame, row, 0, "Frames", - tooltip="The number of frames used for training.") - components.entry(frame, row, 1, self.ui_state, "frames", required=True) + self.components.label(frame, row, 0, "Frames", + tooltip="The number of frames used for training.") + self.components.entry(frame, row, 1, ui_state, "frames", required=True) row += 1 # force circular padding if supports_circular_padding: - components.label(frame, row, 0, "Force Circular Padding", - tooltip="Enables circular padding for all conv layers to better train seamless images") - components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + self.components.label(frame, row, 0, "Force Circular Padding", + tooltip="Enables circular padding for all conv layers to better train seamless images") + self.components.switch(frame, row, 1, ui_state, "force_circular_padding") - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_text_encoder_frame(self, master, row, ui_state, supports_clip_skip=True, supports_training=True, + supports_sequence_length=False): + frame = self.components.section_frame(master, row) if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", - tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + self.components.label(frame, 0, 0, "Train Text Encoder", + tooltip="Enables training the text encoder model") + self.components.switch(frame, 0, 1, ui_state, "text_encoder.train") # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + self.components.label(frame, 1, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + self.components.entry(frame, 1, 1, ui_state, "text_encoder.dropout_probability") if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", - tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", - "text_encoder.stop_training_after_unit", supports_time_units=False) + self.components.label(frame, 2, 0, "Stop Training After", + tooltip="When to stop training the text encoder") + self.components.time_entry(frame, 2, 1, ui_state, "text_encoder.stop_training_after", + "text_encoder.stop_training_after_unit", supports_time_units=False) # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", - tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + self.components.label(frame, 3, 0, "Text Encoder Learning Rate", + tooltip="The learning rate of the text encoder. Overrides the base learning rate") + self.components.entry(frame, 3, 1, ui_state, "text_encoder.learning_rate") if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + self.components.label(frame, 4, 0, "Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + self.components.entry(frame, 4, 1, ui_state, "text_encoder_layer_skip") if supports_sequence_length: # text encoder sequence length - components.label(frame, row, 0, "Text Encoder Sequence Length", - tooltip="Number of tokens for captions") - components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") + self.components.label(frame, row, 0, "Text Encoder Sequence Length", + tooltip="Number of tokens for captions") + self.components.entry(frame, row, 1, ui_state, "text_encoder_sequence_length") row += 1 def __create_text_encoder_n_frame( self, master, row: int, + ui_state, i: int, supports_include: bool = False, supports_layer_skip: bool = True, supports_sequence_length: bool = False, ): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + frame = self.components.section_frame(master, row) row = 0 suffix = f"_{i}" if i > 1 else "" if supports_include: # include text encoder - components.label(frame, row, 0, f"Include Text Encoder {i}", - tooltip=f"Includes text encoder {i} in the training run") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") + self.components.label(frame, row, 0, f"Include Text Encoder {i}", + tooltip=f"Includes text encoder {i} in the training run") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.include") row += 1 # train text encoder - components.label(frame, row, 0, f"Train Text Encoder {i}", - tooltip=f"Enables training the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") + self.components.label(frame, row, 0, f"Train Text Encoder {i}", + tooltip=f"Enables training the text encoder {i} model") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.train") row += 1 # train text encoder embedding - components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", - tooltip=f"Enables training embeddings for the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") + self.components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", + tooltip=f"Enables training embeddings for the text encoder {i} model") + self.components.switch(frame, row, 1, ui_state, f"text_encoder{suffix}.train_embedding") row += 1 # dropout - components.label(frame, row, 0, "Dropout Probability", - tooltip=f"The Probability for dropping the text encoder {i} conditioning") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") + self.components.label(frame, row, 0, "Dropout Probability", + tooltip=f"The Probability for dropping the text encoder {i} conditioning") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}.dropout_probability") row += 1 # train text encoder epochs - components.label(frame, row, 0, "Stop Training After", - tooltip=f"When to stop training the text encoder {i}") - components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", - f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) + self.components.label(frame, row, 0, "Stop Training After", + tooltip=f"When to stop training the text encoder {i}") + self.components.time_entry(frame, row, 1, ui_state, f"text_encoder{suffix}.stop_training_after", + f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) row += 1 # text encoder learning rate - components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", - tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") + self.components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", + tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}.learning_rate") row += 1 if supports_layer_skip: # text encoder layer skip (clip skip) - components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") + self.components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}_layer_skip") row += 1 if supports_sequence_length: # text encoder sequence length - components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", - tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") + self.components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", + tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") + self.components.entry(frame, row, 1, ui_state, f"text_encoder{suffix}_sequence_length") row += 1 - def __create_embedding_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + def __create_embedding_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # embedding learning rate - components.label(frame, 0, 0, "Embeddings Learning Rate", - tooltip="The learning rate of embeddings. Overrides the base learning rate") - components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") + self.components.label(frame, 0, 0, "Embeddings Learning Rate", + tooltip="The learning rate of embeddings. Overrides the base learning rate") + self.components.entry(frame, 0, 1, ui_state, "embedding_learning_rate") # preserve embedding norm - components.label(frame, 1, 0, "Preserve Embedding Norm", - tooltip="Rescales each trained embedding to the median embedding norm") - components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") + self.components.label(frame, 1, 0, "Preserve Embedding Norm", + tooltip="Rescales each trained embedding to the median embedding norm") + self.components.switch(frame, 1, 1, ui_state, "preserve_embedding_norm") - def __create_unet_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_unet_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # train unet - components.label(frame, 0, 0, "Train UNet", - tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") + self.components.label(frame, 0, 0, "Train UNet", + tooltip="Enables training the UNet model") + self.components.switch(frame, 0, 1, ui_state, "unet.train") # train unet epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the UNet") + self.components.time_entry(frame, 1, 1, ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + supports_time_units=False) # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", - tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + self.components.label(frame, 2, 0, "UNet Learning Rate", + tooltip="The learning rate of the UNet. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "unet.learning_rate") # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", - tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") - rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + self.components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target", + wraplength=130) + self.components.switch(frame, 3, 1, ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") - def __create_prior_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_prior_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # train prior - components.label(frame, 0, 0, "Train Prior", - tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") + self.components.label(frame, 0, 0, "Train Prior", + tooltip="Enables training the Prior model") + self.components.switch(frame, 0, 1, ui_state, "prior.train") # train prior epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Prior") + self.components.time_entry(frame, 1, 1, ui_state, "prior.stop_training_after", + "prior.stop_training_after_unit", supports_time_units=False) # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", - tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + self.components.label(frame, 2, 0, "Prior Learning Rate", + tooltip="The learning rate of the Prior. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "prior.learning_rate") - def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_transformer_frame(self, master, row, ui_state, supports_guidance_scale: bool = False, + supports_force_attention_mask: bool = True): + frame = self.components.section_frame(master, row) # train transformer - components.label(frame, 0, 0, "Train Transformer", - tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") + self.components.label(frame, 0, 0, "Train Transformer", + tooltip="Enables training the Transformer model") + self.components.switch(frame, 0, 1, ui_state, "transformer.train") # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", - supports_time_units=False) + self.components.label(frame, 1, 0, "Stop Training After", + tooltip="When to stop training the Transformer") + self.components.time_entry(frame, 1, 1, ui_state, "transformer.stop_training_after", + "transformer.stop_training_after_unit", supports_time_units=False) # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", - tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + self.components.label(frame, 2, 0, "Transformer Learning Rate", + tooltip="The learning rate of the Transformer. Overrides the base learning rate") + self.components.entry(frame, 2, 1, ui_state, "transformer.learning_rate") if supports_force_attention_mask: # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", - tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + self.components.label(frame, 3, 0, "Force Attention Mask", + tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") + self.components.switch(frame, 3, 1, ui_state, "transformer.attention_mask") if supports_guidance_scale: # guidance scale - components.label(frame, 4, 0, "Guidance Scale", - tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + self.components.label(frame, 4, 0, "Guidance Scale", + tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") + self.components.entry(frame, 4, 1, ui_state, "transformer.guidance_scale") - def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_noise_frame(self, master, row, ui_state, callbacks, + supports_generalized_offset_noise: bool = False, + supports_dynamic_timestep_shifting: bool = False): + frame = self.components.section_frame(master, row) # offset noise weight - components.label(frame, 0, 0, "Offset Noise Weight", - tooltip="The weight of offset noise added to each training step") - components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") + self.components.label(frame, 0, 0, "Offset Noise Weight", + tooltip="The weight of offset noise added to each training step") + self.components.entry(frame, 0, 1, ui_state, "offset_noise_weight") if supports_generalized_offset_noise: # generalized offset noise weight - generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", - tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") - generalised_offset_label.configure(wraplength=130, justify="left") - components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") + self.components.label(frame, 1, 0, "Generalized Offset Noise", + tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.", + wraplength=130) + self.components.switch(frame, 1, 1, ui_state, "generalized_offset_noise") # perturbation noise weight - components.label(frame, 2, 0, "Perturbation Noise Weight", - tooltip="The weight of perturbation noise added to each training step") - components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") + self.components.label(frame, 2, 0, "Perturbation Noise Weight", + tooltip="The weight of perturbation noise added to each training step") + self.components.entry(frame, 2, 1, ui_state, "perturbation_noise_weight") # timestep distribution - components.label(frame, 3, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", - adv_command=self.__open_timestep_distribution_window) + self.components.label(frame, 3, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + self.components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], ui_state, + "timestep_distribution", + adv_command=callbacks.get('open_timestep_distribution')) # min noising strength - components.label(frame, 4, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) + self.components.label(frame, 4, 0, "Min Noising Strength", + tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") + self.components.entry(frame, 4, 1, ui_state, "min_noising_strength", required=True) # max noising strength - components.label(frame, 5, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) + self.components.label(frame, 5, 0, "Max Noising Strength", + tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") + self.components.entry(frame, 5, 1, ui_state, "max_noising_strength", required=True) # noising weight - components.label(frame, 6, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) + self.components.label(frame, 6, 0, "Noising Weight", + tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") + self.components.entry(frame, 6, 1, ui_state, "noising_weight", required=True) # noising bias - components.label(frame, 7, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) + self.components.label(frame, 7, 0, "Noising Bias", + tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") + self.components.entry(frame, 7, 1, ui_state, "noising_bias", required=True) # timestep shift - components.label(frame, 8, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) + self.components.label(frame, 8, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + self.components.entry(frame, 8, 1, ui_state, "timestep_shift", required=True) if supports_dynamic_timestep_shifting: # dynamic timestep shifting - components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - - + self.components.label(frame, 9, 0, "Dynamic Timestep Shifting", + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + self.components.switch(frame, 9, 1, ui_state, "dynamic_timestep_shifting") - def __create_masked_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_masked_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) # Masked Training - components.label(frame, 0, 0, "Masked Training", - tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") - components.switch(frame, 0, 1, self.ui_state, "masked_training") + self.components.label(frame, 0, 0, "Masked Training", + tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") + self.components.switch(frame, 0, 1, ui_state, "masked_training") # unmasked probability - components.label(frame, 1, 0, "Unmasked Probability", - tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") - components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", - extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) + self.components.label(frame, 1, 0, "Unmasked Probability", + tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") + self.components.entry(frame, 1, 1, ui_state, "unmasked_probability", + extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) # unmasked weight - components.label(frame, 2, 0, "Unmasked Weight", - tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") - components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", - extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) + self.components.label(frame, 2, 0, "Unmasked Weight", + tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") + self.components.entry(frame, 2, 1, ui_state, "unmasked_weight", + extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) # normalize masked area loss - components.label(frame, 3, 0, "Normalize Masked Area Loss", - tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") - components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") + self.components.label(frame, 3, 0, "Normalize Masked Area Loss", + tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") + self.components.switch(frame, 3, 1, ui_state, "normalize_masked_area_loss") # masked prior preservation - components.label(frame, 4, 0, "Masked Prior Preservation Weight", - tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") - components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", - extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) + self.components.label(frame, 4, 0, "Masked Prior Preservation Weight", + tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") + self.components.entry(frame, 4, 1, ui_state, "masked_prior_preservation_weight", + extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) # use custom conditioning image - components.label(frame, 5, 0, "Custom Conditioning Image", - tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") - components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") + self.components.label(frame, 5, 0, "Custom Conditioning Image", + tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") + self.components.switch(frame, 5, 1, ui_state, "custom_conditioning_image") - def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) + def __create_loss_frame(self, master, row, controller, ui_state, + supports_vb_loss: bool = False): + frame = self.components.section_frame(master, row) # MSE Strength - components.label(frame, 0, 0, "MSE Strength", - tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) + self.components.label(frame, 0, 0, "MSE Strength", + tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 0, 1, ui_state, "mse_strength", required=True) # MAE Strength - components.label(frame, 1, 0, "MAE Strength", - tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) + self.components.label(frame, 1, 0, "MAE Strength", + tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 1, 1, ui_state, "mae_strength", required=True) # log-cosh Strength - components.label(frame, 2, 0, "log-cosh Strength", - tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) + self.components.label(frame, 2, 0, "log-cosh Strength", + tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") + self.components.entry(frame, 2, 1, ui_state, "log_cosh_strength", required=True) # Huber Strength - components.label(frame, 3, 0, "Huber Strength", - tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") - components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) + self.components.label(frame, 3, 0, "Huber Strength", + tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") + self.components.entry(frame, 3, 1, ui_state, "huber_strength", required=True) # Huber Delta - components.label(frame, 4, 0, "Huber Delta", - tooltip="Delta parameter for huber loss") - components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) + self.components.label(frame, 4, 0, "Huber Delta", + tooltip="Delta parameter for huber loss") + self.components.entry(frame, 4, 1, ui_state, "huber_delta", required=True) if supports_vb_loss: # VB Strength - components.label(frame, 5, 0, "VB Strength", - tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") - components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) + self.components.label(frame, 5, 0, "VB Strength", + tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") + self.components.entry(frame, 5, 1, ui_state, "vb_loss_strength", required=True) # Loss Weight function - components.label(frame, 6, 0, "Loss Weight Function", - tooltip="Choice of loss weight function. Can help the model learn details more accurately.") - components.options(frame, 6, 1, [str(x) for x in list(LossWeight) - if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() - or x == LossWeight.CONSTANT - ], - self.ui_state, "loss_weight_fn") + self.components.label(frame, 6, 0, "Loss Weight Function", + tooltip="Choice of loss weight function. Can help the model learn details more accurately.") + self.components.options(frame, 6, 1, [str(x) for x in list(LossWeight) + if x.supports_flow_matching() == controller.is_flow_matching() + or x == LossWeight.CONSTANT + ], + ui_state, "loss_weight_fn") row = 7 # Loss weight strength - if not self.train_config.model_type.is_flow_matching(): - components.label(frame, row, 0, "Gamma", - tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") - components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", - extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) + if not controller.is_flow_matching(): + self.components.label(frame, row, 0, "Gamma", + tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") + self.components.entry(frame, row, 1, ui_state, "loss_weight_strength", + extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) row += 1 # Loss Scaler - components.label(frame, row, 0, "Loss Scaler", - tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") - components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") + self.components.label(frame, row, 0, "Loss Scaler", + tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") + self.components.options(frame, row, 1, [str(x) for x in list(LossScaler)], ui_state, "loss_scaler") row += 1 - def __create_layer_frame(self, master, row): - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - components.layer_filter_entry(master, row, 0, self.ui_state, - preset_var_name="layer_filter_preset", presets=presets, - preset_label="Layer Filter", - preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", - entry_var_name="layer_filter", - entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", - regex_var_name="layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - ) - - - def __on_layer_filter_preset_change(self): - if not self.layer_selector: - return - selected = self.ui_state.get_var("layer_filter_preset").get() - self.__preset_set_layer_choice(selected) - - def __hide_layer_entry(self): - if self.layer_entry and self.layer_entry.winfo_manager(): - self.layer_entry.grid_remove() - - def __show_layer_entry(self): - if self.layer_entry and not self.layer_entry.winfo_manager(): - self.layer_entry.grid() - - def __open_optimizer_params_window(self): - window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_scheduler_params_window(self): - window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_timestep_distribution_window(self): - window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __restore_optimizer_config(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def __restore_scheduler_config(self, variable): - if not hasattr(self, 'lr_scheduler_adv_comp'): - return - - if variable == "CUSTOM": - self.lr_scheduler_adv_comp.configure(state="normal") - else: - self.lr_scheduler_adv_comp.configure(state="disabled") + def __create_layer_frame(self, master, row, controller, ui_state): + presets = controller.get_layer_presets() + self.components.layer_filter_entry(master, row, 0, ui_state, + preset_var_name="layer_filter_preset", presets=presets, + preset_label="Layer Filter", + preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", + entry_var_name="layer_filter", + entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", + regex_var_name="layer_filter_regex", + regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", + ) diff --git a/modules/ui/BaseVideoToolUIView.py b/modules/ui/BaseVideoToolUIView.py index c3291e6ea..b82ba5701 100644 --- a/modules/ui/BaseVideoToolUIView.py +++ b/modules/ui/BaseVideoToolUIView.py @@ -1,877 +1,196 @@ -import concurrent.futures -import math -import os -import pathlib -import random -import shlex -import subprocess -import threading import webbrowser -from fractions import Fraction -from tkinter import filedialog +from abc import ABC, abstractmethod -from modules.util.image_util import load_image from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components - -import av -import customtkinter as ctk -import cv2 -import scenedetect -from PIL import Image - - -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - - self.title("Video Tools") - self.geometry("600x720") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=1, column=0) - frame.grid_columnconfigure(0, weight=0, minsize=160) - frame.grid_columnconfigure(1, weight=0, minsize=300) - frame.grid_columnconfigure(2, weight=1) - - #create preview image - preview_path = "resources/icons/icon.png" - preview = load_image(preview_path, 'RGB') - preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) - self.preview_image_label = ctk.CTkLabel( - master=frame, text="Preview image", image=self.preview_image, height=150, width=150, - compound="top") - self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - - #displays progress and messages that also go to terminal - self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) - self.status_label.insert(index="1.0", text="Current status") - self.status_label.configure(state="disabled") - self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + +class BaseVideoToolUIView(ABC): + def __init__(self, components): + self.components = components + + def build_clip_extract_tab(self, frame, controller, ui_state): # single video - components.label(frame, 0, 0, "Single Video", + self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) + self.components.entry(frame, 0, 1, ui_state, "clip_single", width=190) + self._create_browse_file_button(frame, 0, ui_state, "clip_single", + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.button(frame, 0, 2, "Extract Single", + command=lambda: self._extract_clips(False, controller)) # time range - components.label(frame, 1, 0, " Time Range", + self.components.label(frame, 1, 0, " Time Range", tooltip="Time range to limit selection for single video, \ format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") + self.components.entry(frame, 1, 1, ui_state, "clip_time_start", width=100, sticky="w") + self.components.entry(frame, 1, 1, ui_state, "clip_time_end", width=100, sticky="e") # directory of videos - components.label(frame, 2, 0, "Directory", + self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) + self.components.entry(frame, 2, 1, ui_state, "clip_list", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "clip_list") + self.components.button(frame, 2, 2, "Extract Directory", + command=lambda: self._extract_clips(True, controller)) # output directory - components.label(frame, 3, 0, "Output", + self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 3, 1, ui_state, "clip_output", width=190) + self._create_browse_dir_button(frame, 3, ui_state, "clip_output") # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", + self.components.label(frame, 4, 0, "Output to\nSubdirectories", tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 4, 1, ui_state, "output_subdir_clip") # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", + self.components.label(frame, 5, 0, "Split at Cuts", tooltip="If enabled, detect cuts in the input video and split at those points. \ Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 5, 1, ui_state, "split_cuts") # maximum length - components.label(frame, 6, 0, "Max Length (s)", + self.components.label(frame, 6, 0, "Max Length (s)", tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") + self.components.entry(frame, 6, 1, ui_state, "clip_length", width=220) # Set FPS - components.label(frame, 7, 0, "Set FPS", + self.components.label(frame, 7, 0, "Set FPS", tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") + self.components.entry(frame, 7, 1, ui_state, "clip_fps", width=220) # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", + self.components.label(frame, 8, 0, "Remove Borders", tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 8, 1, ui_state, "clip_bordercrop") # Crop Variation - components.label(frame, 9, 0, "Crop Variation", + self.components.label(frame, 9, 0, "Crop Variation", tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 9, 1, ui_state, "clip_crop", width=220) + def build_image_extract_tab(self, frame, controller, ui_state): # single video - components.label(frame, 0, 0, "Single Video", + self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) + self.components.entry(frame, 0, 1, ui_state, "image_single", width=190) + self._create_browse_file_button(frame, 0, ui_state, "image_single", + [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.button(frame, 0, 2, "Extract Single", + command=lambda: self._extract_images(False, controller)) # time range - components.label(frame, 1, 0, " Time Range", + self.components.label(frame, 1, 0, " Time Range", tooltip="Time range to limit selection for single video, \ format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") + self.components.entry(frame, 1, 1, ui_state, "image_time_start", width=100, sticky="w") + self.components.entry(frame, 1, 1, ui_state, "image_time_end", width=100, sticky="e") # directory of videos - components.label(frame, 2, 0, "Directory", + self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) + self.components.entry(frame, 2, 1, ui_state, "image_list", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "image_list") + self.components.button(frame, 2, 2, "Extract Directory", + command=lambda: self._extract_images(True, controller)) # output directory - components.label(frame, 3, 0, "Output", + self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 3, 1, ui_state, "image_output", width=190) + self._create_browse_dir_button(frame, 3, ui_state, "image_output") # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", + self.components.label(frame, 4, 0, "Output to\nSubdirectories", tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 4, 1, ui_state, "output_subdir_img") # image capture rate - components.label(frame, 5, 0, "Images/sec", + self.components.label(frame, 5, 0, "Images/sec", tooltip="Number of images to capture per second of video. \ Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") + self.components.entry(frame, 5, 1, ui_state, "capture_rate", width=220) # blur removal - components.label(frame, 6, 0, "Blur Removal", + self.components.label(frame, 6, 0, "Blur Removal", tooltip="Threshold for removal of blurry images, relative to all others. \ For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") + self.components.entry(frame, 6, 1, ui_state, "blur_threshold", width=220) # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", + self.components.label(frame, 7, 0, "Remove Borders", tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + self.components.switch(frame, 7, 1, ui_state, "image_bordercrop") # Crop Variation - components.label(frame, 8, 0, "Crop Variation", + self.components.label(frame, 8, 0, "Crop Variation", tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) + self.components.entry(frame, 8, 1, ui_state, "image_crop", width=220) + def build_video_download_tab(self, frame, controller, ui_state): # link - components.label(frame, 0, 0, "Single Link", + self.components.label(frame, 0, 0, "Single Link", tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) + self.components.entry(frame, 0, 1, ui_state, "download_link", width=220) + self.components.button(frame, 0, 2, "Download Link", + command=lambda: self._download(False, controller)) # link list - components.label(frame, 1, 0, "Link List", + self.components.label(frame, 1, 0, "Link List", tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) + self.components.entry(frame, 1, 1, ui_state, "download_list", width=190) + self._create_browse_file_button(frame, 1, ui_state, "download_list", [("Text file", ".txt")]) + self.components.button(frame, 1, 2, "Download List", + command=lambda: self._download(True, controller)) # output directory - components.label(frame, 2, 0, "Output", + self.components.label(frame, 2, 0, "Output", tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) + self.components.entry(frame, 2, 1, ui_state, "download_output", width=190) + self._create_browse_dir_button(frame, 2, ui_state, "download_output") # additional args - components.label(frame, 3, 0, "Additional Args", + self.components.label(frame, 3, 0, "Additional Args", tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ Default args will hide most terminal outputs.") - self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) - self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) - self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress --format mp4") - components.button(frame, 3, 2, "yt-dlp info", + self._create_textbox(frame, 3, 1, 220, 90, ui_state, "download_args") + self.components.button(frame, 3, 2, "yt-dlp info", command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): - input_videos = [] - if not batch_mode: - path = pathlib.Path(input_path_single) - if path.is_file(): - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - return [path] - else: - self.__update_status("Invalid video file!") - return [] - else: - self.__update_status("No file specified, or invalid file path!") - return [] - else: - input_videos = [] - if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": - self.__update_status("Invalid input directory!") - return [] - # Only traverse supported extensions to avoid opening every file. - lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} - for path in pathlib.Path(input_path_dir).rglob("*"): - if path.is_file() and path.suffix.lower() in lower_exts: - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - input_videos.append(path) - self.__update_status(f'Found {len(input_videos)} videos to process') - return input_videos - - def __run_in_thread(self, target, *args): - """Clear status box and run target function in a daemon thread.""" - self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") - self.status_label.configure(state="disabled") - t = threading.Thread(target=target, args=args) - t.daemon = True - t.start() - - @staticmethod - def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: - return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) - - def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: - fps = video.get(cv2.CAP_PROP_FPS) or 0.0 - if fps <= 0: - self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') - return 30.0 - return fps - - @staticmethod - def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, - video_path, input_dir: str) -> str: - if use_subdir and batch_mode: - return os.path.join(output_entry, - os.path.splitext(os.path.relpath(video_path, input_dir))[0]) - elif use_subdir: - return os.path.join(output_entry, - os.path.splitext(os.path.basename(video_path))[0]) - return output_entry - - def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: - # Return original dimensions and no offset if variation is zero - if variation == 0: - return 0, height, 0, width - - old_aspect = height/width - variation_scaled = old_aspect*variation - if old_aspect > 1.2: #tall image - new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) - elif old_aspect < 0.85: #wide image - new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) - else: #square image - new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) - - new_aspect = round(new_aspect, 2) - #keep the height the same if reducing width, and vice versa - if new_aspect > old_aspect: - new_height = int(height) - new_width = int(width*(old_aspect/new_aspect)) - elif new_aspect < old_aspect: - new_height = int(height*(new_aspect/old_aspect)) - new_width = int(width) - else: - new_height = int(height) - new_width = int(width) - - #random offset in dimension that was cropped - position_x = random.randint(0, width-new_width) - position_y = random.randint(0, height-new_height) - return position_y, new_height, position_x, new_width - - def find_main_contour(self, frame): - #outline image to find main content and exclude black bars often present on letterboxed videos - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) - frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - if frame_contours: - #select largest contour by area - frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) - x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) - else: #fallback if no contours detected - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - #if bounding box did not detect the correct area, likely due to all-black frame - if not frame_contours or h1 < 10 or w1 < 10: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - return x1, y1, w1, h1 - - def __extract_clips_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_clips_multi, batch_mode) - - def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") - return - if max_length <= 0.25: - self.__update_status("Max Length of clips must be > 0.25 seconds.") - return - if target_fps < 0: - self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if len(input_videos) == 0: # exit if no paths found - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) - executor.submit(self.__extract_clips, - str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), - crop_variation, target_fps, output_directory) - - if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') - else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') - - def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, - split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - max_length_frames = int(max_length * vid_fps) - min_length_frames = max(int(0.25 * vid_fps), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - - if split_at_cuts: - #use scenedetect to find cuts, based on start/end frame number - self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') - timecode_list = scenedetect.detect( - video_path=str(video_path), - detector=scenedetect.AdaptiveDetector(), - start_time=int(timestamp_min_frame), - end_time=int(timestamp_max_frame)) - scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] - if not scene_list: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - else: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - - scene_list_split = [] - for scene in scene_list: - length = scene[1]-scene[0] - if length > max_length_frames: #check for any scenes longer than max length - n = math.ceil(length/max_length_frames) #divide into n new scenes - new_length = int(length/n) - new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length - for i, _n in enumerate(new_splits[:-1]): - if new_splits[i + 1] - new_splits[i] > min_length_frames: - scene_list_split.append((new_splits[i], new_splits[i + 1])) - elif length > (min_length_frames + 2): - # Trim first/last frame to avoid transition artifacts - scene_list_split.append((scene[0] + 1, scene[1] - 1)) - - self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - futures = [ - executor.submit(self.__save_clip, scene, video_path, target_fps, - remove_borders, crop_variation, output_dir) - for scene in scene_list_split - ] - for future in concurrent.futures.as_completed(futures): - exc = future.exception() - if exc is not None: - self.__update_status(f'Error saving clip: {exc}') - - video.release() - - def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, - remove_borders: bool, crop_variation: float, output_dir: str): - basename, ext = os.path.splitext(os.path.basename(video_path)) - video = cv2.VideoCapture(str(video_path)) - fps = self.__get_safe_fps(video, video_path) - os.makedirs(output_dir, exist_ok=True) - output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' - output_ext = ".mp4" - - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene - frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) - success, frame = video.read() - if not success or frame is None: - self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') - video.release() - return - - # Blend random frames to detect borders, avoiding incorrect crop from black frames - if remove_borders: - frame_blend = frame - for i in range(5): - random_frame = random.randint(scene[0], scene[1]) - video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) - success, frame = video.read() - if not success or frame is None: - continue - a = 1/(i+1) - b = 1-a - frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) - x1, y1, w1, h1 = self.find_main_contour(frame_blend) - else: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - # Ensure dimensions are even, required - h2 -= h2 % 2 - w2 -= w2 % 2 - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) - success, frame = video.read() - if success: - try: - preview = Image.fromarray( - cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) - #truncate filename of long files so UI doesn't shift around - filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') - except Exception: - pass - video.release() - - if target_fps <= 0: - target_fps = fps - - output_path = f'{output_name}{output_ext}' - self.__write_clip_av(video_path, output_path, scene, fps, target_fps, - x1 + x2, y1 + y2, w2, h2) - - @staticmethod - def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], - src_fps: float, target_fps: float, - crop_x: int, crop_y: int, crop_w: int, crop_h: int): - start_sec = scene[0] / src_fps - end_sec = scene[1] / src_fps - rate_frac = Fraction(target_fps).limit_denominator(10000) - stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) - - with av.open(video_path) as input_container: - in_video = input_container.streams.video[0] - in_video.thread_type = 'AUTO' - in_audio = input_container.streams.audio[0] if input_container.streams.audio else None - - with av.open(output_path, mode='w') as output_container: - out_video = output_container.add_stream('libx264', rate=rate_frac) - out_video.width = crop_w - out_video.height = crop_h - out_video.pix_fmt = 'yuv420p' - out_video.time_base = stream_time_base - - out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None - - input_container.seek(int(start_sec * 1_000_000)) - - out_frame_idx = 0 - out_time_step = 1.0 / target_fps - video_done = False - decode_streams = [s for s in (in_video, in_audio) if s is not None] - - for packet in input_container.demux(decode_streams): - if packet.stream == in_video: - if video_done: - continue - for frame in packet.decode(): - if frame.time is None or frame.time < start_sec: - continue - if frame.time >= end_sec: - video_done = True - break - - # FPS conversion: skip frames when source fps > target fps - if frame.time < start_sec + out_frame_idx * out_time_step: - continue - - img = frame.to_ndarray(format='bgr24') - cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] - out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') - out_frame.pts = out_frame_idx - out_frame.time_base = stream_time_base - - for out_pkt in out_video.encode(out_frame): - output_container.mux(out_pkt) - out_frame_idx += 1 - - elif packet.stream == in_audio and out_audio is not None: - if packet.dts is None: - continue - pkt_time = float(packet.pts * packet.time_base) - if pkt_time < start_sec or pkt_time >= end_sec: - continue - # Re-timestamp audio relative to clip start - packet.pts = int((pkt_time - start_sec) / packet.time_base) - packet.dts = packet.pts - packet.stream = out_audio - output_container.mux(packet) - - # Flush video encoder - for pkt in out_video.encode(): - output_container.mux(pkt) - - def __extract_images_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_images_multi, batch_mode) - - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") - return - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - return - if not (0.0 <= blur_threshold < 1.0): - self.__update_status("Blur Removal must be between 0.0 and 1.0.") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if not input_videos: - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) - executor.submit(self.__save_frames, - str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), - crop_variation, output_directory) - if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') - else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') - - def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, - blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - video.release() - return - image_rate = max(int(vid_fps / capture_rate), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) - frame_list = [] - - for n in frame_range: - #pick frame from random triangular distribution around center of each "chunk" of the video - frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) - frame = max(0, min(frame, max(total_frames - 1, 0))) - frame_list.append(frame) - - self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') - - output_list = [] - for f in frame_list: - video.set(cv2.CAP_PROP_POS_FRAMES, f) - success, frame = video.read() - if success and frame is not None: - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() - output_list.append((f, frame_sharpness)) - - if not output_list: - self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') - video.release() - return - - output_list_sorted = sorted(output_list, key=lambda x: x[1]) - cutoff = int(blur_threshold * len(output_list_sorted)) - output_list_cut = output_list_sorted[cutoff:] - self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') - - basename, ext = os.path.splitext(os.path.basename(video_path)) - os.makedirs(output_dir, exist_ok=True) - - for f in output_list_cut: - filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' - video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) - success, frame = video.read() - - #crop out borders of frame - if remove_borders and success and frame is not None: - x1, y1, w1, h1 = self.find_main_contour(frame) - frame_cropped = frame[y1:y1+h1, x1:x1+w1] - else: - frame_cropped = frame if success and frame is not None else None - if frame_cropped is not None: - x1 = 0 - y1 = 0 - h1, w1, _ = frame_cropped.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - - if success and frame is not None and frame_cropped is not None: - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") - try: - preview = Image.fromarray( - cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') - except Exception: - pass # preview update is non-critical - - cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) - video.release() - - def __download_button(self, batch_mode: bool): - self.__run_in_thread(self.__download_multi, batch_mode) - - def __update_status(self, status_text: str): - print(status_text) - self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") - self.status_label.configure(state="disabled") - - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - if not batch_mode: - ydl_urls = [self.download_link_entry.get()] - elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) - if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": - with open(ydl_path) as file: - ydl_urls = file.readlines() - else: - self.__update_status("Invalid link list!") - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for url in ydl_urls: - executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) - - self.__update_status(f'Completed {len(ydl_urls)} downloads.') - - def __download_video(self, url: str, output_dir: str, output_args: str): - url = (url or "").strip() - if not url: - self.__update_status("Empty URL, skipping download.") - return - - #Respect quotes and split into list to run as yt-dlp command - additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] - cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] - - self.__update_status(f'Downloading {url}') - subprocess.run(cmd) - self.__update_status(f'Download {url} done!') + @abstractmethod + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + pass + + @abstractmethod + def _create_browse_dir_button(self, master, row, ui_state, var_name): + pass + + @abstractmethod + def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): + pass + + @abstractmethod + def update_status(self, status_text: str): + pass + + @abstractmethod + def clear_status(self): + pass + + @abstractmethod + def update_preview(self, preview_image, label_text: str): + pass + + def _extract_clips(self, batch_mode: bool, controller): + controller.extract_clips_button(batch_mode) + + def _extract_images(self, batch_mode: bool, controller): + controller.extract_images_button(batch_mode) + + def _download(self, batch_mode: bool, controller): + controller.download_button(batch_mode) diff --git a/modules/ui/CaptionUIController.py b/modules/ui/CaptionUIController.py index e6cc0551e..0495322a3 100644 --- a/modules/ui/CaptionUIController.py +++ b/modules/ui/CaptionUIController.py @@ -1,8 +1,6 @@ import os -import platform import subprocess import traceback -from tkinter import filedialog from modules.module.Blip2Model import Blip2Model from modules.module.BlipModel import BlipModel @@ -11,39 +9,22 @@ from modules.module.RembgHumanModel import RembgHumanModel from modules.module.RembgModel import RembgModel from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow +from modules.ui.GenerateCaptionsWindowController import GenerateCaptionsWindowController +from modules.ui.GenerateMasksWindowController import GenerateMasksWindowController from modules.util import path_util from modules.util.image_util import load_image from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState import torch -import customtkinter as ctk -import cv2 import numpy as np -from customtkinter import ScalingTracker, ThemeManager from PIL import Image, ImageDraw -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - +class CaptionUIController: + def __init__(self, initial_dir: str | None, initial_include_subdirectories: bool): self.dir = initial_dir self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) self.image_size = 850 self.help_text = """ Keyboard shortcuts when focusing on the prompt input field: @@ -62,8 +43,6 @@ def __init__( self.captioning_model = None self.image_rel_paths = [] self.current_image_index = -1 - self.file_list = None - self.image_labels = [] self.pil_image = None self.image_width = 0 self.image_height = 0 @@ -72,155 +51,67 @@ def __init__( self.mask_draw_y = 0 self.mask_draw_radius = 0.01 self.display_only_mask = False - self.image = None - self.image_label = None self.mask_editing_mode = 'draw' - self.enable_mask_editing_var = ctk.BooleanVar() - self.mask_editing_alpha = None - self.prompt_var = None - self.prompt_component = None - + self.view = None - self.title("OneTrainer") - self.geometry("1280x980") - self.resizable(False, False) + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + def open_mask_window(self, parent_window, view_cls): + controller = GenerateMasksWindowController(self) + return controller.create_window(parent_window, self.dir, self.config_ui_data["include_subdirectories"], view_cls) - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) + def open_caption_window(self, parent_window, view_cls): + controller = GenerateCaptionsWindowController(self) + return controller.create_window(parent_window, self.dir, self.config_ui_data["include_subdirectories"], view_cls) + def open_in_explorer(self): + try: + image_name = self.image_rel_paths[self.current_image_index] + image_name = os.path.realpath(os.path.join(self.dir, image_name)) + subprocess.Popen(f"explorer /select,{image_name}") + except Exception: + traceback.print_exc() - self.top_bar(self) + def switch_image(self, index): + old_index = self.current_image_index + self.current_image_index = index + if index >= 0: + self.pil_image = self.load_image() + self.pil_mask = self.load_mask() + prompt = self.load_prompt() - self.bottom_frame = ctk.CTkFrame(self) - self.bottom_frame.grid(row=1, column=0, sticky="nsew") - self.bottom_frame.grid_rowconfigure(0, weight=1) - self.bottom_frame.grid_columnconfigure(0, weight=0) - self.bottom_frame.grid_columnconfigure(1, weight=1) + self.image_width = self.pil_image.width + self.image_height = self.pil_image.height + scale = self.image_size / max(self.pil_image.height, self.pil_image.width) + height = int(self.pil_image.height * scale) + width = int(self.pil_image.width * scale) - self.file_list_column(self.bottom_frame) - self.content_column(self.bottom_frame) - self.load_directory() + self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, - tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, - tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, - tooltip="open a dialog to automatically generate captions") - - if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, - tooltip="open the current image in Explorer") - - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", - text="include subdirectories") - - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - - def file_list_column(self, master): - if self.file_list is not None: - self.image_labels = [] - self.file_list.destroy() - - self.file_list = ctk.CTkScrollableFrame(master, width=300) - self.file_list.grid(row=0, column=0, sticky="nsew") - - for i, filename in enumerate(self.image_rel_paths): - def __create_switch_image(index): - def __switch_image(event): - self.switch_image(index) - - return __switch_image - - label = ctk.CTkLabel(self.file_list, text=filename) - label.bind("", __create_switch_image(i)) - - self.image_labels.append(label) - label.grid(row=i, column=0, padx=5, sticky="nsw") - - def content_column(self, master): - image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) - - right_frame = ctk.CTkFrame(master, fg_color="transparent") - right_frame.grid(row=0, column=1, sticky="nsew") - - right_frame.grid_columnconfigure(4, weight=1) - right_frame.grid_rowconfigure(1, weight=1) - - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, - tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, - tooltip="draw a mask using a fill tool") - - # checkbox to enable mask editing - self.enable_mask_editing_var = ctk.BooleanVar() - self.enable_mask_editing_var.set(False) - enable_mask_editing_checkbox = ctk.CTkCheckBox( - right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) - enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") - - # mask alpha textbox - self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") - self.mask_editing_alpha.insert(0, "1.0") - self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) - self.bind_key_events(self.mask_editing_alpha) + self.view.on_image_switched(old_index, index, prompt) + else: + self.view.on_image_cleared() - mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) - mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + def previous_image(self): + if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: + self.view.switch_image(self.current_image_index - 1) - # image - self.image = ctk.CTkImage( - light_image=image, - size=(self.image_size, self.image_size) - ) - self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size - ) - self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") - - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) - - # prompt - self.prompt_var = ctk.StringVar() - self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) - self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") - self.bind_key_events(self.prompt_component) - self.prompt_component.focus_set() - - def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) - component.bind("", self.save) - component.bind("", self.toggle_mask) - component.bind("", self.draw_mask_editing_mode) - component.bind("", self.fill_mask_editing_mode) + def next_image(self): + if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): + self.view.switch_image(self.current_image_index + 1) def load_directory(self, include_subdirectories: bool = False): self.scan_directory(include_subdirectories) - self.file_list_column(self.bottom_frame) + self.view.refresh_file_list() if len(self.image_rel_paths) > 0: self.switch_image(0) else: self.switch_image(-1) - self.prompt_component.focus_set() + self.view.focus_prompt() def scan_directory(self, include_subdirectories: bool = False): def __is_supported_image_extension(filename): @@ -285,42 +176,26 @@ def load_prompt(self): else: return "" - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( - text_color=ThemeManager.theme["CTkLabel"]["text_color"]) - - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") + def save(self, prompt_text): + if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): + image_name = self.image_rel_paths[self.current_image_index] - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() + prompt_name = os.path.splitext(image_name)[0] + ".txt" + prompt_name = os.path.join(self.dir, prompt_name) - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) + mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" + mask_name = os.path.join(self.dir, mask_name) - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) + try: + with open(prompt_name, "w", encoding='utf-8') as f: + f.write(prompt_text) + except Exception: + return - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) + if self.pil_mask: + self.pil_mask.save(mask_name) - def refresh_image(self): + def get_display_image(self): if self.pil_mask: resized_pil_mask = self.pil_mask.resize( (self.pil_image.width, self.pil_image.height), @@ -328,7 +203,7 @@ def refresh_image(self): ) if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) + return resized_pil_mask, resized_pil_mask.size else: np_image = np.array(self.pil_image).astype(np.float32) / 255.0 np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 @@ -346,29 +221,26 @@ def refresh_image(self): np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) masked_image = Image.fromarray(np_masked_image, mode='RGB') - self.image.configure(light_image=masked_image, size=masked_image.size) + return masked_image, masked_image.size else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + return self.pil_image, self.pil_image.size + + def toggle_mask(self): + self.display_only_mask = not self.display_only_mask + + def set_mask_editing_mode(self, mode): + self.mask_editing_mode = mode - def draw_mask_radius(self, delta, raw_event): + def update_mask_draw_radius(self, delta): # Wheel up = Increase radius. Wheel down = Decrease radius. multiplier = 1.0 + (delta * 0.05) self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) - def edit_mask(self, event): - if not self.enable_mask_editing_var.get(): - return - - if event.widget != self.image_label.children["!label"]: - return - + def handle_edit_mask(self, event_x, event_y, is_left, is_right, alpha): if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): return - - display_scaling = ScalingTracker.get_window_scaling(self) - - event_x = event.x / display_scaling - event_y = event.y / display_scaling + if self.pil_image is None: + return start_x = int(event_x / self.pil_image.width * self.image_width) start_y = int(event_y / self.pil_image.height * self.image_height) @@ -378,27 +250,16 @@ def edit_mask(self, event): self.mask_draw_x = event_x self.mask_draw_y = event_y - is_right = False - is_left = False - if event.state & 0x0100 or event.num == 1: # left mouse button - is_left = True - elif event.state & 0x0400 or event.num == 3: # right mouse button - is_right = True - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) + self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right, alpha) if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) + self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right, alpha) - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right, alpha): color = None adding_to_mask = True if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range color = (rgb_value, rgb_value, rgb_value) @@ -423,17 +284,13 @@ def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): draw.ellipse((end_x - radius, end_y - radius, end_x + radius, end_y + radius), fill=color, outline=None) - self.refresh_image() + self.view.refresh_image() - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): + def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right, alpha): color = None adding_to_mask = True if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range color = (rgb_value, rgb_value, rgb_value) @@ -449,69 +306,11 @@ def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) np_mask = np.array(self.pil_mask).astype(np.uint8) + import cv2 cv2.floodFill(np_mask, None, (start_x, start_y), color) self.pil_mask = Image.fromarray(np_mask, 'RGB') - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return - - if self.pil_mask: - self.pil_mask.save(mask_name) - - def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' - - if args: - # disable default event - return "break" - return None - - def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' - - def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask - self.refresh_image() - - def open_directory(self): - new_dir = filedialog.askdirectory() - - if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) - - def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) - - def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() + self.view.refresh_image() def load_masking_model(self, model): model_type = type(self.masking_model).__name__ if self.masking_model else None @@ -563,10 +362,6 @@ def _release_models(self): if freed: torch_gc() - def _on_close(self): - self._release_models() - self.destroy() - - def destroy(self): + def on_close(self): self._release_models() - super().destroy() + self.view.destroy() diff --git a/modules/ui/CloudTabController.py b/modules/ui/CloudTabController.py new file mode 100644 index 000000000..d21dda6ce --- /dev/null +++ b/modules/ui/CloudTabController.py @@ -0,0 +1,31 @@ + +import webbrowser + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.CloudType import CloudType + + +class CloudTabController: + def __init__(self, config: TrainConfig, parent): + self.config = config + self.parent = parent + self.reattach = False + + def do_reattach(self): + self.reattach = True + try: + self.parent.start_training() + finally: + self.reattach = False + + def get_gpu_types(self) -> list[str]: + if self.config.cloud.type == CloudType.RUNPOD: + import runpod + runpod.api_key = self.config.secrets.cloud.api_key + gpus = runpod.get_gpus() + return [gpu['id'] for gpu in gpus] + return [] + + def open_create_cloud_url(self): + if self.config.cloud.type == CloudType.RUNPOD: + webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) diff --git a/modules/ui/ConceptTabController.py b/modules/ui/ConceptTabController.py new file mode 100644 index 000000000..b90b7a4f6 --- /dev/null +++ b/modules/ui/ConceptTabController.py @@ -0,0 +1,19 @@ + +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.config.ConceptConfig import ConceptConfig +from modules.util.config.TrainConfig import TrainConfig + + +class ConceptTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> ConceptConfig: + return ConceptConfig.default_values() + + def randomize_seed(self, concept: ConceptConfig) -> ConceptConfig: + concept.seed = ConceptConfig.default_values().seed + return concept + + def open_element_window(self, parent, concept_config, ui_state, image_ui_state, text_ui_state, view_cls): + return view_cls(parent, ConceptWindowController(self.train_config, concept_config), ui_state, image_ui_state, text_ui_state) diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py index f58879d5f..c2c486961 100644 --- a/modules/ui/ConceptWindowController.py +++ b/modules/ui/ConceptWindowController.py @@ -1,5 +1,3 @@ -import fractions -import math import os import pathlib import platform @@ -11,12 +9,7 @@ from modules.util import concept_stats, path_util from modules.util.config.ConceptConfig import ConceptConfig from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.BalancingStrategy import BalancingStrategy -from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState from mgds.LoadingPipeline import LoadingPipeline from mgds.OutputPipelineModule import OutputPipelineModule @@ -34,534 +27,21 @@ from mgds.pipelineModules.RandomRotate import RandomRotate from mgds.pipelineModules.RandomSaturation import RandomSaturation from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) +from mgds.pipelineModuleTypes.RandomAccessPipelineModule import RandomAccessPipelineModule import torch from torchvision.transforms import functional -import customtkinter as ctk import huggingface_hub -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg from PIL import Image -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - - def get_inputs(self) -> list[str]: - return [] - - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - +class ConceptWindowController: + def __init__(self, train_config: TrainConfig, concept: ConceptConfig): self.train_config = train_config - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - self.text_ui_state = text_ui_state - self.image_preview_file_index = 0 - self.preview_augmentations = ctk.BooleanVar(self, True) - self.bucket_fig = None - - self.title("Concept") - self.geometry("800x700") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) - - # name - components.label(frame, 0, 0, "Name", - tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") - - # enabled - components.label(frame, 1, 0, "Enabled", - tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") - - # concept type - components.label(frame, 2, 0, "Concept Type", - tooltip="STANDARD: Standard finetuning with the sample as training target\n" - "VALIDATION: Use concept for validation instead of training\n" - "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " - "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " - "Only implemented for LoRA.", - wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") - - # path - components.label(frame, 3, 0, "Path", - tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, - tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.") - - # prompt source - components.label(frame, 4, 0, "Prompt Source", - tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts") - prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") - - def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") - - components.options_kv(frame, 4, 1, [ - ("From text file per sample", 'sample'), - ("From single text file", 'concept'), - ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) - - # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", - tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") - - # image variations - components.label(frame, 6, 0, "Image Variations", - tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") - - # text variations - components.label(frame, 7, 0, "Text Variations", - tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") - - # balancing - components.label(frame, 8, 0, "Balancing", - tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") - - # loss weight - components.label(frame, 9, 0, "Loss Weight", - tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # header - components.label(frame, 0, 1, "Random", - tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", - tooltip="Enable this augmentation with fixed values") - - # crop jitter - components.label(frame, 1, 0, "Crop Jitter", - tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") - - # random flip - components.label(frame, 2, 0, "Random Flip", - tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") - - # random rotation - components.label(frame, 3, 0, "Random Rotation", - tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") - - # random brightness - components.label(frame, 4, 0, "Random Brightness", - tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") - - # random contrast - components.label(frame, 5, 0, "Random Contrast", - tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") - - # random saturation - components.label(frame, 6, 0, "Random Saturation", - tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") - - # random hue - components.label(frame, 7, 0, "Random Hue", - tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") - - # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", - tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") - - # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", - tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") - - # circular mask generation - components.label(frame, 10, 0, "Resolution Override", - tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") - - # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image = ctk.CTkImage( - light_image=image_preview, - size=image_preview.size, - ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) - image_label.grid(row=0, column=4, rowspan=6) - - # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") - update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") - update_button_frame.grid_columnconfigure(1, weight=1) - - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) - preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) - - prev_preview_button.configure(width=40) - next_preview_button.configure(width=40) - - #caption and filename preview - self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) - self.filename_preview.grid(row=2, column=0, columnspan=3) - self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", - tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") - - # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", - tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") - - # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", - tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") - - # tag dropout - components.label(frame, 3, 0, "Tag Dropout", - tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", - tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ - ("Full", 'FULL'), - ("Random", 'RANDOM'), - ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", - tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") - - components.label(frame, 5, 0, "Special Dropout Tags", - tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ - ("None", 'NONE'), - ("Blacklist", 'BLACKLIST'), - ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", - tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") - - #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", - tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", - tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") - - components.label(frame, 8, 0, "Captialization Mode", - tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", - tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame - - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - self.cancel_scan_flag = threading.Event() - - #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") - - #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") - - #basic img/vid stats - count of each type in the concept - #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") - - #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") - #currently no masks for videos? - - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") - - #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) - - #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) - - #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) - - #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) - - #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, - tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") - - #aspect bucketing plot, mostly copied from timestep preview graph - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - plt.set_loglevel('WARNING') #suppress errors about data type in bar chart - - assert self.bucket_fig is None - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) - self.bucket_fig.tight_layout() - self.bucket_fig.subplots_adjust(bottom=0.15) - - self.bucket_fig.set_facecolor(background_color) - self.bucket_ax.set_facecolor(background_color) - self.bucket_ax.spines['bottom'].set_color(self.text_color) - self.bucket_ax.spines['left'].set_color(self.text_color) - self.bucket_ax.spines['top'].set_visible(False) - self.bucket_ax.spines['right'].set_color(self.text_color) - self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") - self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") - self.bucket_ax.xaxis.label.set_color(self.text_color) - self.bucket_ax.yaxis.label.set_color(self.text_color) - - #refresh stats - must be after all labels are defined or will give error - self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), - tooltip="Reload basic statistics for the concept directory") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), - tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), - tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") - - frame.pack(fill="both", expand=1) - return frame - - def __prev_image_preview(self): - self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() - - def __next_image_preview(self): - self.image_preview_file_index += 1 - self.__update_image_preview() - - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() - self.image.configure(light_image=image_preview, size=image_preview.size) - self.filename_preview.configure(text=filename_preview) - self.caption_preview.configure(state="normal") - self.caption_preview.delete(index1="1.0", index2="end") - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") + self.scan_thread = None @staticmethod def get_concept_path(path: str) -> str | None: @@ -573,22 +53,22 @@ def get_concept_path(path: str) -> str | None: except Exception: return None - def __download_dataset(self): + def download_dataset(self): try: huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) + def download_dataset_threaded(self): + download_thread = threading.Thread(target=self.download_dataset, daemon=True) download_thread.start() - def _read_text_file_for_preview(self, file_path: str) -> str: + def _read_text_file_for_preview(self, file_path: str, preview_augmentations: bool) -> str: empty_msg = "[Empty prompt]" try: with open(file_path, "r") as f: - if self.preview_augmentations.get(): + if preview_augmentations: lines = [line.strip() for line in f if line.strip()] return random.choice(lines) if lines else empty_msg content = f.read().strip() @@ -605,7 +85,7 @@ def _read_text_file_for_preview(self, file_path: str) -> str: except UnicodeDecodeError: return "[Invalid file encoding. This should not happen, please report this issue]" - def __get_preview_image(self): + def get_preview_image(self, image_preview_file_index: int, preview_augmentations: bool): preview_image_path = "resources/icons/icon.png" file_index = -1 glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" @@ -620,7 +100,7 @@ def __get_preview_image(self): and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): preview_image_path = path_util.canonical_join(concept_path, path) file_index += 1 - if file_index == self.image_preview_file_index: + if file_index == image_preview_file_index: break image = load_image(preview_image_path, 'RGB') @@ -647,10 +127,10 @@ def __get_preview_image(self): "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, } file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" + prompt_output = self._read_text_file_for_preview(str(file_path), preview_augmentations) if file_path else "[Empty prompt]" modules = [] - if self.preview_augmentations.get(): + if preview_augmentations: input_module = InputPipelineModule({ 'true': True, 'image': image_tensor, @@ -749,133 +229,16 @@ def __get_preview_image(self): return image, filename_output, prompt_output - def __update_concept_stats(self): - #file size - self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") - self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") - - #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) - - #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) - - #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) - - #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) - - #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') - else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) - - #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") - else: - #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') - - #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") - else: - #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') - #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') - - #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") - else: - #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') - - #aspect bucketing - aspect_buckets = self.concept.concept_stats["aspect_buckets"] - if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero - min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values - if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be - min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket - else: - min_val2 = min_val #if no second smallest bucket exists set to min_val - min_aspect_buckets = {key: val for key,val in aspect_buckets.items() if val in (min_val, min_val2)} - min_bucket_str = "" - for key, val in min_aspect_buckets.items(): - min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' - min_bucket_str.strip() - self.small_bucket_preview.configure(text=min_bucket_str) - - self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = self.bucket_ax.bar(aspect_ratios, counts) - self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() - - def decimal_to_aspect_ratio(self, value : float): - #find closest fraction to decimal aspect value and convert to a:b format - aspect_fraction = fractions.Fraction(value).limit_denominator(16) - aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' - return aspect_string - - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): + def get_concept_stats(self, view, advanced_checks: bool, wait_time: float): start_time = time.perf_counter() last_update = time.perf_counter() self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, view._disable_scan_buttons) concept_path = self.get_concept_path(self.concept.path) if not concept_path: print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, view._enable_scan_buttons) return subfolders = [concept_path] @@ -890,45 +253,45 @@ def __get_concept_stats(self, advanced_checks: bool, wait_time: float): #update GUI approx every half second if time.perf_counter() > (last_update + 0.5): last_update = time.perf_counter() - self.concept_stats_tab.after(0, self.__update_concept_stats) + view.components.call_after(view.concept_stats_tab, 0, lambda: view._update_concept_stats(self)) self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - self.concept_stats_tab.after(0, self.__update_concept_stats) + view.components.call_after(view.concept_stats_tab, 0, view._enable_scan_buttons) + view.components.call_after(view.concept_stats_tab, 0, lambda: view._update_concept_stats(self)) - def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): - self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) + def get_concept_stats_threaded(self, view, advanced_checks: bool, waittime: float): + self.scan_thread = threading.Thread(target=self.get_concept_stats, args=[view, advanced_checks, waittime], daemon=True) self.scan_thread.start() - def __disable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): + def auto_update_concept_stats(self, view): try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError + view._update_concept_stats(self) #load stats from config if available, else raises KeyError if self.concept.concept_stats["file_size"] == 0: #force rescan if empty raise KeyError except KeyError: concept_path = self.get_concept_path(self.concept.path) if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec + self.get_concept_stats(view, False, 2) #force rescan if config is empty, timeout of 2 sec if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s + self.get_concept_stats(view, True, 2) #do advanced scan automatically if basic took <0.1s + + +class InputPipelineModule( + PipelineModule, + RandomAccessPipelineModule, +): + def __init__(self, data: dict): + super().__init__() + self.data = data - def destroy(self): - if self.bucket_fig is not None: - plt.close(self.bucket_fig) - self.bucket_fig = None + def length(self) -> int: + return 1 - super().destroy() + def get_inputs(self) -> list[str]: + return [] - def __ok(self): - self.destroy() + def get_outputs(self) -> list[str]: + return list(self.data.keys()) + + def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: + return self.data diff --git a/modules/ui/ConvertModelUIController.py b/modules/ui/ConvertModelUIController.py index 6cb1b507a..c3cdebd31 100644 --- a/modules/ui/ConvertModelUIController.py +++ b/modules/ui/ConvertModelUIController.py @@ -4,123 +4,23 @@ from modules.util import create from modules.util.args.ConvertModelArgs import ConvertModelArgs from modules.util.config.TrainConfig import QuantizationConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.ModelNames import EmbeddingName, ModelNames from modules.util.torch_util import torch_gc -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk - -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent +class ConvertModelUIController: + def __init__(self): self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None - - - self.title("Convert models") - self.geometry("550x350") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - - self.main_frame(self.frame) - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - # model type - components.label(master, 0, 0, "Model Type", - tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify - ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), - ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), - ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), - ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), - ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), - ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2", ModelType.FLUX_2), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here - ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here - ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") - - # training method - components.label(master, 1, 0, "Model Type", - tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ - ("Base Model", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") - - # input name - components.label(master, 2, 0, "Input name", - tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # output data type - components.label(master, 3, 0, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") - - # output format - components.label(master, 4, 0, "Output Format", - tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ - ("Safetensors", ModelFormat.SAFETENSORS), - ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") - - # output model destination - components.label(master, 5, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) + self.view = None - self.button = components.button(master, 6, 1, "Convert", self.convert_model) + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view def convert_model(self): try: - self.button.configure(state="disabled") + self.view.set_converting(True) model_loader = create.create_model_loader( model_type=self.convert_model_args.model_type, training_method=self.convert_model_args.training_method @@ -167,4 +67,4 @@ def convert_model(self): traceback.print_exc() torch_gc() - self.button.configure(state="normal") + self.view.set_converting(False) diff --git a/modules/ui/CtkAdditionalEmbeddingsTabView.py b/modules/ui/CtkAdditionalEmbeddingsTabView.py index 6a5e3fbe7..fc24c61d1 100644 --- a/modules/ui/CtkAdditionalEmbeddingsTabView.py +++ b/modules/ui/CtkAdditionalEmbeddingsTabView.py @@ -1,54 +1,38 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseAdditionalEmbeddingsTabView import BaseAdditionalEmbeddingsTabView, BaseEmbeddingWidgetView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class AdditionalEmbeddingsTab(ConfigList): +class CtkAdditionalEmbeddingsTabView(CtkConfigListView, BaseAdditionalEmbeddingsTabView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, + def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, attr_name="additional_embeddings", enable_key="train", from_external_file=False, add_button_text="add embedding", is_full_width=True, - show_toggle_button=True + show_toggle_button=True, ) - def refresh_ui(self): - if self.element_list is not None: - self.element_list.destroy() - self.element_list = None - self.widgets_initialized = False - self._create_element_list() - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return EmbeddingWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return TrainEmbeddingConfig.default_values() + return CtkEmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass +class CtkEmbeddingWidgetView(BaseEmbeddingWidgetView, ctk.CTkFrame): -class EmbeddingWidget(ctk.CTkFrame): - def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseEmbeddingWidgetView.__init__(self, ctk_components) self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command + ui_state = CtkUIState(self, element) self.grid_columnconfigure(0, weight=1) @@ -61,76 +45,7 @@ def __init__(self, master, element, i, open_command, remove_command, clone_comma bottom_frame.grid(row=1, column=0, sticky="nsew") bottom_frame.grid_columnconfigure(7, weight=1) - # close button - close_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) - - # clone button - clone_button = ctk.CTkButton( - master=top_frame, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_uuid), - ) - clone_button.grid(row=0, column=1, padx=5) - - # embedding model names - components.label(top_frame, 0, 2, "base embedding:", - tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( - top_frame, 0, 3, self.ui_state, "model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # placeholder - components.label(top_frame, 0, 4, "placeholder:", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(top_frame, 0, 5, self.ui_state, "placeholder") - - # token count - components.label(top_frame, 0, 6, "token count:", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - token_count_entry = components.entry(top_frame, 0, 7, self.ui_state, "token_count") - token_count_entry.configure(width=40) - - # trainable - components.label(bottom_frame, 0, 0, "train:") - trainable_switch = components.switch(bottom_frame, 0, 1, self.ui_state, "train", command=save_command) - trainable_switch.configure(width=40) - - # output embedding - components.label(bottom_frame, 0, 2, "output embedding:", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - output_embedding_switch = components.switch(bottom_frame, 0, 3, self.ui_state, "is_output_embedding") - output_embedding_switch.configure(width=40) - - # stop training after - components.label(bottom_frame, 0, 4, "stop training after:", - tooltip="When to stop training the embedding") - components.time_entry(bottom_frame, 0, 5, self.ui_state, "stop_training_after", "stop_training_after_unit") - - # initial embedding text - components.label(bottom_frame, 0, 6, "initial embedding text:", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(bottom_frame, 0, 7, self.ui_state, "initial_embedding_text") - - def __randomize_uuid(self, embedding_config: TrainEmbeddingConfig): - embedding_config.uuid = TrainEmbeddingConfig.default_values().uuid - return embedding_config - - def configure_element(self): - pass + self.build_content(top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller) def place_in_list(self): self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkCaptionUIView.py b/modules/ui/CtkCaptionUIView.py index e6cc0551e..281036912 100644 --- a/modules/ui/CtkCaptionUIView.py +++ b/modules/ui/CtkCaptionUIView.py @@ -1,97 +1,47 @@ -import os -import platform -import subprocess -import traceback from tkinter import filedialog -from modules.module.Blip2Model import Blip2Model -from modules.module.BlipModel import BlipModel -from modules.module.ClipSegModel import ClipSegModel -from modules.module.MaskByColor import MaskByColor -from modules.module.RembgHumanModel import RembgHumanModel -from modules.module.RembgModel import RembgModel -from modules.module.WDModel import WDModel -from modules.ui.GenerateCaptionsWindow import GenerateCaptionsWindow -from modules.ui.GenerateMasksWindow import GenerateMasksWindow -from modules.util import path_util -from modules.util.image_util import load_image -from modules.util.torch_util import default_device, torch_gc -from modules.util.ui import components +from modules.ui.BaseCaptionUIView import BaseCaptionUIView +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.CtkGenerateCaptionsWindowView import CtkGenerateCaptionsWindowView +from modules.ui.CtkGenerateMasksWindowView import CtkGenerateMasksWindowView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon -from modules.util.ui.UIState import UIState - -import torch import customtkinter as ctk -import cv2 -import numpy as np from customtkinter import ScalingTracker, ThemeManager -from PIL import Image, ImageDraw - - -class CaptionUI(ctk.CTkToplevel): - def __init__( - self, - parent, - initial_dir: str | None, - initial_include_subdirectories: bool, - *args, - **kwargs, - ) -> None: - super().__init__(parent, *args, **kwargs) - self.protocol("WM_DELETE_WINDOW", self._on_close) - - self.dir = initial_dir - self.config_ui_data = {"include_subdirectories": initial_include_subdirectories} - self.config_ui_state = UIState(self, self.config_ui_data) - self.image_size = 850 - self.help_text = """ - Keyboard shortcuts when focusing on the prompt input field: - Up arrow: previous image - Down arrow: next image - Return: save - Ctrl+M: only show the mask - Ctrl+D: draw mask editing mode - Ctrl+F: fill mask editing mode - - When editing masks: - Left click: add mask - Right click: remove mask - Mouse wheel: increase or decrease brush size""" - self.masking_model = None - self.captioning_model = None - self.image_rel_paths = [] - self.current_image_index = -1 - self.file_list = None - self.image_labels = [] - self.pil_image = None - self.image_width = 0 - self.image_height = 0 - self.pil_mask = None - self.mask_draw_x = 0 - self.mask_draw_y = 0 - self.mask_draw_radius = 0.01 - self.display_only_mask = False - self.image = None - self.image_label = None - self.mask_editing_mode = 'draw' +from PIL import Image + + +class CtkCaptionUIView(BaseCaptionUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: CaptionUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseCaptionUIView.__init__(self, ctk_components) + self.protocol("WM_DELETE_WINDOW", controller.on_close) + + self.controller = controller + controller.view = self + self.config_ui_state = CtkUIState(self, controller.config_ui_data) self.enable_mask_editing_var = ctk.BooleanVar() self.mask_editing_alpha = None self.prompt_var = None self.prompt_component = None - + self.image = None + self.image_label = None + self.file_list = None + self.image_labels = [] self.title("OneTrainer") self.geometry("1280x980") self.resizable(False, False) - self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_columnconfigure(0, weight=1) - - self.top_bar(self) + top_frame = ctk.CTkFrame(self) + top_frame.grid(row=0, column=0, sticky="nsew") + self.build_top_bar(top_frame, controller, self.config_ui_state) self.bottom_frame = ctk.CTkFrame(self) self.bottom_frame.grid(row=1, column=0, sticky="nsew") @@ -101,35 +51,12 @@ def __init__( self.file_list_column(self.bottom_frame) self.content_column(self.bottom_frame) - self.load_directory() + self.controller.load_directory() self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - def top_bar(self, master): - top_frame = ctk.CTkFrame(master) - top_frame.grid(row=0, column=0, sticky="nsew") - - components.button(top_frame, 0, 0, "Open", self.open_directory, - tooltip="open a new directory") - components.button(top_frame, 0, 1, "Generate Masks", self.open_mask_window, - tooltip="open a dialog to automatically generate masks") - components.button(top_frame, 0, 2, "Generate Captions", self.open_caption_window, - tooltip="open a dialog to automatically generate captions") - - if platform.system() == "Windows": - components.button(top_frame, 0, 3, "Open in Explorer", self.open_in_explorer, - tooltip="open the current image in Explorer") - - components.switch(top_frame, 0, 4, self.config_ui_state, "include_subdirectories", - text="include subdirectories") - - top_frame.grid_columnconfigure(5, weight=1) - - components.button(top_frame, 0, 6, "Help", self.print_help, - tooltip=self.help_text) - def file_list_column(self, master): if self.file_list is not None: self.image_labels = [] @@ -138,10 +65,10 @@ def file_list_column(self, master): self.file_list = ctk.CTkScrollableFrame(master, width=300) self.file_list.grid(row=0, column=0, sticky="nsew") - for i, filename in enumerate(self.image_rel_paths): + for i, filename in enumerate(self.controller.image_rel_paths): def __create_switch_image(index): def __switch_image(event): - self.switch_image(index) + self.controller.switch_image(index) return __switch_image @@ -160,10 +87,7 @@ def content_column(self, master): right_frame.grid_columnconfigure(4, weight=1) right_frame.grid_rowconfigure(1, weight=1) - components.button(right_frame, 0, 0, "Draw", self.draw_mask_editing_mode, - tooltip="draw a mask using a brush") - components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, - tooltip="draw a mask using a fill tool") + self.build_mask_buttons(right_frame) # checkbox to enable mask editing self.enable_mask_editing_var = ctk.BooleanVar() @@ -184,10 +108,11 @@ def content_column(self, master): # image self.image = ctk.CTkImage( light_image=image, - size=(self.image_size, self.image_size) + size=(self.controller.image_size, self.controller.image_size) ) self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, height=self.image_size, width=self.image_size + master=right_frame, text="", image=self.image, + height=self.controller.image_size, width=self.controller.image_size ) self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") @@ -204,156 +129,37 @@ def content_column(self, master): self.prompt_component.focus_set() def bind_key_events(self, component): - component.bind("", self.next_image) - component.bind("", self.previous_image) + component.bind("", lambda e: self.controller.next_image()) + component.bind("", lambda e: self.controller.previous_image()) component.bind("", self.save) component.bind("", self.toggle_mask) component.bind("", self.draw_mask_editing_mode) component.bind("", self.fill_mask_editing_mode) - def load_directory(self, include_subdirectories: bool = False): - self.scan_directory(include_subdirectories) + def refresh_file_list(self): self.file_list_column(self.bottom_frame) - if len(self.image_rel_paths) > 0: - self.switch_image(0) - else: - self.switch_image(-1) - + def focus_prompt(self): self.prompt_component.focus_set() - def scan_directory(self, include_subdirectories: bool = False): - def __is_supported_image_extension(filename): - name, ext = os.path.splitext(filename) - return path_util.is_supported_image_extension(ext) and not name.endswith("-masklabel") and not name.endswith("-condlabel") - - self.image_rel_paths = [] - - if not self.dir or not os.path.isdir(self.dir): - return - - if include_subdirectories: - for root, _, files in os.walk(self.dir): - for filename in files: - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(root, filename), self.dir) - ) - else: - for _, filename in enumerate(os.listdir(self.dir)): - if __is_supported_image_extension(filename): - self.image_rel_paths.append( - os.path.relpath(os.path.join(self.dir, filename), self.dir) - ) - - def load_image(self): - image_name = "resources/icons/icon.png" - - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.join(self.dir, image_name) - - try: - return load_image(image_name, convert_mode="RGB") - except Exception: - print(f'Could not open image {image_name}') - - def load_mask(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) - - try: - return load_image(mask_name, convert_mode='RGB') - except Exception: - return None - else: - return None - - def load_prompt(self): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - try: - with open(prompt_name, "r", encoding='utf-8') as f: - return f.readlines()[0].strip() - except Exception: - return "" - else: - return "" - - def previous_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index - 1) >= 0: - self.switch_image(self.current_image_index - 1) - - def next_image(self, event): - if len(self.image_rel_paths) > 0 and (self.current_image_index + 1) < len(self.image_rel_paths): - self.switch_image(self.current_image_index + 1) - - def switch_image(self, index): - if len(self.image_labels) > 0 and self.current_image_index < len(self.image_labels): - self.image_labels[self.current_image_index].configure( + def on_image_switched(self, old_index, new_index, prompt): + if len(self.image_labels) > 0 and old_index < len(self.image_labels): + self.image_labels[old_index].configure( text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + self.image_labels[new_index].configure(text_color="#FF0000") + self.refresh_image() + self.prompt_var.set(prompt) - self.current_image_index = index - if index >= 0: - self.image_labels[index].configure(text_color="#FF0000") - - self.pil_image = self.load_image() - self.pil_mask = self.load_mask() - prompt = self.load_prompt() - - self.image_width = self.pil_image.width - self.image_height = self.pil_image.height - scale = self.image_size / max(self.pil_image.height, self.pil_image.width) - height = int(self.pil_image.height * scale) - width = int(self.pil_image.width * scale) - - self.pil_image = self.pil_image.resize((width, height), Image.Resampling.LANCZOS) - - self.refresh_image() - self.prompt_var.set(prompt) - else: - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) + def on_image_cleared(self): + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) def refresh_image(self): - if self.pil_mask: - resized_pil_mask = self.pil_mask.resize( - (self.pil_image.width, self.pil_image.height), - Image.Resampling.NEAREST - ) - - if self.display_only_mask: - self.image.configure(light_image=resized_pil_mask, size=resized_pil_mask.size) - else: - np_image = np.array(self.pil_image).astype(np.float32) / 255.0 - np_mask = np.array(resized_pil_mask).astype(np.float32) / 255.0 - - # normalize mask between 0.3 - 1.0 so we can see image underneath and gauge strength of the alpha - norm_min = 0.3 - np_mask_min = np_mask.min() - if np_mask_min == 0: - # optimize for common case - np_mask = np_mask * (1.0 - norm_min) + norm_min - elif np_mask_min < 1: - # note: min of 1 means we get divide by 0 - np_mask = (np_mask - np_mask_min) / (1.0 - np_mask_min) * (1.0 - norm_min) + norm_min - - np_masked_image = (np_image * np_mask * 255.0).astype(np.uint8) - masked_image = Image.fromarray(np_masked_image, mode='RGB') - - self.image.configure(light_image=masked_image, size=masked_image.size) - else: - self.image.configure(light_image=self.pil_image, size=self.pil_image.size) + pil_image, size = self.controller.get_display_image() + self.image.configure(light_image=pil_image, size=size) def draw_mask_radius(self, delta, raw_event): - # Wheel up = Increase radius. Wheel down = Decrease radius. - multiplier = 1.0 + (delta * 0.05) - self.mask_draw_radius = max(0.0025, self.mask_draw_radius * multiplier) + self.controller.update_mask_draw_radius(delta) def edit_mask(self, event): if not self.enable_mask_editing_var.get(): @@ -362,22 +168,11 @@ def edit_mask(self, event): if event.widget != self.image_label.children["!label"]: return - if len(self.image_rel_paths) == 0 or self.current_image_index >= len(self.image_rel_paths): - return - display_scaling = ScalingTracker.get_window_scaling(self) event_x = event.x / display_scaling event_y = event.y / display_scaling - start_x = int(event_x / self.pil_image.width * self.image_width) - start_y = int(event_y / self.pil_image.height * self.image_height) - end_x = int(self.mask_draw_x / self.pil_image.width * self.image_width) - end_y = int(self.mask_draw_y / self.pil_image.height * self.image_height) - - self.mask_draw_x = event_x - self.mask_draw_y = event_y - is_right = False is_left = False if event.state & 0x0100 or event.num == 1: # left mouse button @@ -385,96 +180,18 @@ def edit_mask(self, event): elif event.state & 0x0400 or event.num == 3: # right mouse button is_right = True - if self.mask_editing_mode == 'draw': - self.draw_mask(start_x, start_y, end_x, end_y, is_left, is_right) - if self.mask_editing_mode == 'fill': - self.fill_mask(start_x, start_y, end_x, end_y, is_left, is_right) - - def draw_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - radius = int(self.mask_draw_radius * max(self.pil_mask.width, self.pil_mask.height)) - - draw = ImageDraw.Draw(self.pil_mask) - draw.line((start_x, start_y, end_x, end_y), fill=color, - width=radius + radius + 1) - draw.ellipse((start_x - radius, start_y - radius, - start_x + radius, start_y + radius), fill=color, outline=None) - draw.ellipse((end_x - radius, end_y - radius, end_x + radius, - end_y + radius), fill=color, outline=None) - - self.refresh_image() - - def fill_mask(self, start_x, start_y, end_x, end_y, is_left, is_right): - color = None - - adding_to_mask = True - if is_left: - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - rgb_value = int(max(0.0, min(alpha, 1.0)) * 255) # max/min stuff to clamp to 0 - 255 range - color = (rgb_value, rgb_value, rgb_value) - - elif is_right: - color = (0, 0, 0) - adding_to_mask = False - - if color is not None: - if self.pil_mask is None: - if adding_to_mask: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(0, 0, 0)) - else: - self.pil_mask = Image.new('RGB', size=(self.image_width, self.image_height), color=(255, 255, 255)) - - np_mask = np.array(self.pil_mask).astype(np.uint8) - cv2.floodFill(np_mask, None, (start_x, start_y), color) - self.pil_mask = Image.fromarray(np_mask, 'RGB') - - self.refresh_image() - - def save(self, event): - if len(self.image_rel_paths) > 0 and self.current_image_index < len(self.image_rel_paths): - image_name = self.image_rel_paths[self.current_image_index] - - prompt_name = os.path.splitext(image_name)[0] + ".txt" - prompt_name = os.path.join(self.dir, prompt_name) - - mask_name = os.path.splitext(image_name)[0] + "-masklabel.png" - mask_name = os.path.join(self.dir, mask_name) + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 - try: - with open(prompt_name, "w", encoding='utf-8') as f: - f.write(self.prompt_var.get()) - except Exception: - return + self.controller.handle_edit_mask(event_x, event_y, is_left, is_right, alpha) - if self.pil_mask: - self.pil_mask.save(mask_name) + def save(self, event): + self.controller.save(self.prompt_var.get()) def draw_mask_editing_mode(self, *args): - self.mask_editing_mode = 'draw' + self.controller.set_mask_editing_mode('draw') if args: # disable default event @@ -482,91 +199,30 @@ def draw_mask_editing_mode(self, *args): return None def fill_mask_editing_mode(self, *args): - self.mask_editing_mode = 'fill' + self.controller.set_mask_editing_mode('fill') def toggle_mask(self, *args): - self.display_only_mask = not self.display_only_mask + self.controller.toggle_mask() self.refresh_image() def open_directory(self): new_dir = filedialog.askdirectory() if new_dir: - self.dir = new_dir - self.load_directory(include_subdirectories=self.config_ui_data["include_subdirectories"]) + self.controller.dir = new_dir + self.controller.load_directory(include_subdirectories=self.controller.config_ui_data["include_subdirectories"]) def open_mask_window(self): - dialog = GenerateMasksWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) + self.wait_window(self.controller.open_mask_window(self, CtkGenerateMasksWindowView)) + self.controller.switch_image(self.controller.current_image_index) def open_caption_window(self): - dialog = GenerateCaptionsWindow(self, self.dir, self.config_ui_data["include_subdirectories"]) - self.wait_window(dialog) - self.switch_image(self.current_image_index) + self.wait_window(self.controller.open_caption_window(self, CtkGenerateCaptionsWindowView)) + self.controller.switch_image(self.controller.current_image_index) def open_in_explorer(self): - try: - image_name = self.image_rel_paths[self.current_image_index] - image_name = os.path.realpath(os.path.join(self.dir, image_name)) - subprocess.Popen(f"explorer /select,{image_name}") - except Exception: - traceback.print_exc() - - def load_masking_model(self, model): - model_type = type(self.masking_model).__name__ if self.masking_model else None - - if model == "ClipSeg" and model_type != "ClipSegModel": - self._release_models() - print("loading ClipSeg model, this may take a while") - self.masking_model = ClipSegModel(default_device, torch.float32) - elif model == "Rembg" and model_type != "RembgModel": - self._release_models() - print("loading Rembg model, this may take a while") - self.masking_model = RembgModel(default_device, torch.float32) - elif model == "Rembg-Human" and model_type != "RembgHumanModel": - self._release_models() - print("loading Rembg-Human model, this may take a while") - self.masking_model = RembgHumanModel(default_device, torch.float32) - elif model == "Hex Color" and model_type != "MaskByColor": - self._release_models() - self.masking_model = MaskByColor(default_device, torch.float32) - - def load_captioning_model(self, model): - model_type = type(self.captioning_model).__name__ if self.captioning_model else None - - if model == "Blip" and model_type != "BlipModel": - self._release_models() - print("loading Blip model, this may take a while") - self.captioning_model = BlipModel(default_device, torch.float16) - elif model == "Blip2" and model_type != "Blip2Model": - self._release_models() - print("loading Blip2 model, this may take a while") - self.captioning_model = Blip2Model(default_device, torch.float16) - elif model == "WD14 VIT v2" and model_type != "WDModel": - self._release_models() - print("loading WD14_VIT_v2 model, this may take a while") - self.captioning_model = WDModel(default_device, torch.float16) - - def print_help(self): - print(self.help_text) - - def _release_models(self): - """Release all models from VRAM""" - freed = False - if self.captioning_model is not None: - self.captioning_model = None - freed = True - if self.masking_model is not None: - self.masking_model = None - freed = True - if freed: - torch_gc() - - def _on_close(self): - self._release_models() - self.destroy() + self.controller.open_in_explorer() def destroy(self): - self._release_models() + self.controller._release_models() super().destroy() diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py index 99057e428..0a5249069 100644 --- a/modules/ui/CtkCloudTabView.py +++ b/modules/ui/CtkCloudTabView.py @@ -1,26 +1,18 @@ -import webbrowser -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.CloudAction import CloudAction -from modules.util.enum.CloudFileSync import CloudFileSync -from modules.util.enum.CloudType import CloudType -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseCloudTabView import BaseCloudTabView +from modules.ui.CloudTabController import CloudTabController +from modules.util.ui import ctk_components import customtkinter as ctk -class CloudTab: - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent): - super().__init__() - +class CtkCloudTabView(BaseCloudTabView): + def __init__(self, master, controller: CloudTabController, ui_state): + BaseCloudTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state - self.parent = parent - self.reattach = False self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") self.frame.grid_columnconfigure(0, weight=0) @@ -30,192 +22,24 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState, parent) self.frame.grid_columnconfigure(4, weight=0) self.frame.grid_columnconfigure(5, weight=1) - components.label(self.frame, 0, 0, "Enabled", - tooltip="Enable cloud training") - components.switch(self.frame, 0, 1, self.ui_state, "cloud.enabled") - - components.label(self.frame, 1, 0, "Type", - tooltip="Choose LINUX to connect to a linux machine via SSH. Choose RUNPOD for additional functionality such as automatically creating and deleting pods.") - components.options_kv(self.frame, 1, 1, [ - ("RUNPOD", CloudType.RUNPOD), - ("LINUX", CloudType.LINUX), - ], self.ui_state, "cloud.type") - - components.label(self.frame, 2, 0, "File sync method", - tooltip="Choose NATIVE_SCP to use scp.exe to transfer files. FABRIC_SFTP uses the Paramiko/Fabric SFTP implementation for file transfers instead.") - components.options_kv(self.frame, 2, 1, [ - ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), - ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), - ], self.ui_state, "cloud.file_sync") - - components.label(self.frame, 3, 0, "API key", - tooltip="Cloud service API key for RUNPOD. Leave empty for LINUX. This value is stored separately, not saved to your configuration file. ") - components.entry(self.frame, 3, 1, self.ui_state, "secrets.cloud.api_key") - - components.label(self.frame, 4, 0, "Hostname", - tooltip="SSH server hostname or IP. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 4, 1, self.ui_state, "secrets.cloud.host") - - components.label(self.frame, 5, 0, "Port", - tooltip="SSH server port. Leave empty if you have a Cloud ID or want to automatically create a new cloud.") - components.entry(self.frame, 5, 1, self.ui_state, "secrets.cloud.port") - - components.label(self.frame, 6, 0, "User", - tooltip='SSH username. Use "root" for RUNPOD. Your SSH client must be set up to connect to the cloud using a public key, without a password. For RUNPOD, create an ed25519 key locally, and copy the contents of the public keyfile to your "SSH Public Keys" on the RunPod website.') - components.entry(self.frame, 6, 1, self.ui_state, "secrets.cloud.user") - - components.label(self.frame, 7, 0, "SSH keyfile path", - tooltip="Absolute path to the private key file used for SSH connections. Leave empty to rely on your system SSH configuration.") - components.path_entry(self.frame, 7, 1, self.ui_state, "secrets.cloud.key_file", mode="file") - - components.label(self.frame, 8, 0, "SSH password", - tooltip="SSH password for password-based authentication. If you try to use native SCP requires sshpass to be installed. Leave empty to use key-based authentication.") - components.entry(self.frame, 8, 1, self.ui_state, "secrets.cloud.password") - - components.label(self.frame, 9, 0, "Cloud id", - tooltip="RUNPOD Cloud ID. The cloud service must have a public IP and SSH service. Leave empty if you want to automatically create a new RUNPOD cloud, or if you're connecting to another cloud provider via SSH Hostname and Port.") - components.entry(self.frame, 9, 1, self.ui_state, "secrets.cloud.id") - - components.label(self.frame, 10, 0, "Tensorboard TCP tunnel", - tooltip="Instead of starting tensorboard locally, make a TCP tunnel to a tensorboard on the cloud") - components.switch(self.frame, 10, 1, self.ui_state, "cloud.tensorboard_tunnel") + self.build_content(self.frame, controller, ui_state) + self.frame.pack(fill="both", expand=1) - components.label(self.frame, 1, 2, "Remote Directory", - tooltip="The directory on the cloud where files will be uploaded and downloaded.") - components.entry(self.frame, 1, 3, self.ui_state, "cloud.remote_dir") - components.label(self.frame, 2, 2, "OneTrainer Directory", - tooltip="The directory for OneTrainer on the cloud.") - components.entry(self.frame, 2, 3, self.ui_state, "cloud.onetrainer_dir") - components.label(self.frame, 3, 2, "Huggingface cache Directory", - tooltip="Huggingface models are downloaded to this remote directory.") - components.entry(self.frame, 3, 3, self.ui_state, "cloud.huggingface_cache_dir") - components.label(self.frame, 4, 2, "Install OneTrainer", - tooltip="Automatically install OneTrainer from GitHub if the directory doesn't already exist.") - components.switch(self.frame, 4, 3, self.ui_state, "cloud.install_onetrainer") - components.label(self.frame, 5, 2, "Install command", - tooltip="The command for installing OneTrainer. Leave the default, unless you want to use a development branch of OneTrainer.") - components.entry(self.frame, 5, 3, self.ui_state, "cloud.install_cmd") - components.label(self.frame, 6, 2, "Update OneTrainer", - tooltip="Update OneTrainer if it already exists on the cloud.") - components.switch(self.frame, 6, 3, self.ui_state, "cloud.update_onetrainer") + def _on_set_gpu_types(self): + self.gpu_types_menu.configure(values=self.controller.get_gpu_types()) - components.label(self.frame, 8, 2, "Detach remote trainer", - tooltip="Allows the trainer to keep running even if your connection to the cloud is lost.") - components.switch(self.frame, 8, 3, self.ui_state, "cloud.detach_trainer") - components.label(self.frame, 9, 2, "Reattach id", - tooltip="An id identifying the remotely running trainer. In case you have lost connection or closed OneTrainer, it will try to reattach to this id instead of starting a new remote trainer.") - reattach_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + def _make_reattach_frame(self, frame): + reattach_frame = ctk.CTkFrame(frame, fg_color="transparent") reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") reattach_frame.grid_columnconfigure(0, weight=1) reattach_frame.grid_columnconfigure(1, weight=1) - components.entry(reattach_frame, 0, 0, self.ui_state, "cloud.run_id", width=60) - components.button(reattach_frame, 0, 1, "Reattach now", self.__reattach) - - components.label(self.frame, 11, 2, "Download samples", - tooltip="Download samples from the remote workspace directory to your local machine.") - components.switch(self.frame, 11, 3, self.ui_state, "cloud.download_samples") - components.label(self.frame, 12, 2, "Download output model", - tooltip="Download the final model after training. You can disable this if you plan to use an automatically saved checkpoint instead.") - components.switch(self.frame, 12, 3, self.ui_state, "cloud.download_output_model") - components.label(self.frame, 13, 2, "Download saved checkpoints", - tooltip="Download the automatically saved training checkpoints from the remote workspace directory to your local machine.") - components.switch(self.frame, 13, 3, self.ui_state, "cloud.download_saves") - components.label(self.frame, 14, 2, "Download backups", - tooltip="Download backups from the remote workspace directory to your local machine. It's usually not necessary to download them, because as long as the backups are still available on the cloud, the training can be restarted using one of the cloud's backups.") - components.switch(self.frame, 14, 3, self.ui_state, "cloud.download_backups") - components.label(self.frame, 15, 2, "Download tensorboard logs", - tooltip="Download TensorBoard event logs from the remote workspace directory to your local machine. They can then be viewed locally in TensorBoard. It is recommended to disable \"Sample to TensorBoard\" to reduce the event log size.") - components.switch(self.frame, 15, 3, self.ui_state, "cloud.download_tensorboard") - components.label(self.frame, 16, 2, "Delete remote workspace", - tooltip="Delete the workspace directory on the cloud after training has finished successfully and data has been downloaded.") - components.switch(self.frame, 16, 3, self.ui_state, "cloud.delete_workspace") + return reattach_frame - components.label(self.frame, 1, 4, "Create cloud via API", - tooltip="Automatically creates a new cloud instance if both Host:Port and Cloud ID are empty. Currently supported for RUNPOD.") - create_frame = ctk.CTkFrame(self.frame, fg_color="transparent") + def _make_create_frame(self, frame): + create_frame = ctk.CTkFrame(frame, fg_color="transparent") create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") create_frame.grid_columnconfigure(0, weight=0) create_frame.grid_columnconfigure(1, weight=1) - components.switch(create_frame, 0, 0, self.ui_state, "cloud.create") - components.button(create_frame, 0, 1, "Create cloud via website", self.__create_cloud) - - components.label(self.frame, 2, 4, "Cloud name", - tooltip="The name of the new cloud instance.") - components.entry(self.frame, 2, 5, self.ui_state, "cloud.name") - components.label(self.frame, 3, 4, "Type", - tooltip="Select the RunPod cloud type. See RunPod's website for details.") - components.options_kv(self.frame, 3, 5, [ - ("", ""), - ("Community", "COMMUNITY"), - ("Secure", "SECURE"), - ], self.ui_state, "cloud.sub_type") - - - components.label(self.frame, 4, 4, "GPU", - tooltip="Select the GPU type. Enter an API key before pressing the button.") - - _,gpu_components=components.options_adv(self.frame, 4, 5, [("")], self.ui_state, "cloud.gpu_type",adv_command=self.__set_gpu_types) - self.gpu_types_menu=gpu_components['component'] - - components.label(self.frame, 5, 4, "Volume size", - tooltip="Set the storage volume size in GB. This volume persists only until the cloud is deleted - not a RunPod network volume") - components.entry(self.frame, 5, 5, self.ui_state, "cloud.volume_size") - - components.label(self.frame, 6, 4, "Min download", - tooltip="Set the minimum download speed of the cloud in Mbps.") - components.entry(self.frame, 6, 5, self.ui_state, "cloud.min_download") - - components.label(self.frame, 8, 4, "Action on finish", - tooltip="What to do when training finishes and the data has been fully downloaded: Stop or delete the cloud, or do nothing.") - components.options_kv(self.frame, 8, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_finish") - - components.label(self.frame, 9, 4, "Action on error", - tooltip="What to do if training stops due to an error: Stop or delete the cloud, or do nothing. Data may be lost.") - components.options_kv(self.frame, 9, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_error") - - components.label(self.frame, 10, 4, "Action on detached finish", - tooltip="What to do when training finishes, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 10, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_finish") - - components.label(self.frame, 11, 4, "Action on detached error", - tooltip="What to if training stops due to an error, but the client has been detached and cannot download data. Data may be lost.") - components.options_kv(self.frame, 11, 5, [ - ("None", CloudAction.NONE), - ("Stop", CloudAction.STOP), - ("Delete", CloudAction.DELETE), - ], self.ui_state, "cloud.on_detached_error") - - self.frame.pack(fill="both", expand=1) - - def __set_gpu_types(self): - self.gpu_types_menu.configure(values=[]) - if self.train_config.cloud.type == CloudType.RUNPOD: - import runpod - runpod.api_key=self.train_config.secrets.cloud.api_key - gpus=runpod.get_gpus() - self.gpu_types_menu.configure(values=[gpu['id'] for gpu in gpus]) - - def __reattach(self): - self.reattach=True - try: - self.parent.start_training() - finally: - self.reattach=False - - def __create_cloud(self): - if self.train_config.cloud.type == CloudType.RUNPOD: - webbrowser.open("https://www.runpod.io/console/deploy?template=1a33vbssq9&type=gpu", new=0, autoraise=False) + return create_frame diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py index 0b6505694..5b3e86ac9 100644 --- a/modules/ui/CtkConceptTabView.py +++ b/modules/ui/CtkConceptTabView.py @@ -1,33 +1,27 @@ -import os -import pathlib from tkinter import BooleanVar, StringVar -from modules.ui.ConceptWindow import ConceptWindow -from modules.ui.ConfigList import ConfigList -from modules.util import path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DebounceTimer +from modules.ui.BaseConceptTabView import BaseConceptTabView, BaseConceptWidgetView +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkConceptWindowView import CtkConceptWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.ctk_validation import DebounceTimer +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -from PIL import Image -class ConceptTab(ConfigList): +class CtkConceptTabView(CtkConfigListView, BaseConceptTabView): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): + def __init__(self, master, controller: ConceptTabController, ui_state): + # Pre-initialize before CtkConfigListView.__init__ because _reset_filters is + # called during build() via options_kv's immediate update_var() call. self.search_var = StringVar() self.filter_var = StringVar(value="ALL") self.show_disabled_var = BooleanVar(value=True) - super().__init__( - master, - train_config, - ui_state, + CtkConfigListView.__init__( + self, master, controller, ui_state, from_external_file=True, attr_name="concept_file_name", config_dir="training_concepts", @@ -35,22 +29,18 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): add_button_text="Add Concept", add_button_tooltip="Adds a new concept to the current config.", is_full_width=False, - show_toggle_button=True + show_toggle_button=True, ) self._toolbar = None self._toolbar_is_wrapped = False self._add_search_bar() - # wrap toolbar if too narrow self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return ConceptWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return ConceptConfig.default_values() - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return ConceptWindow(self.master, self.train_config, self.current_config[i], ui_state[0], ui_state[1], ui_state[2]) + return self.controller.open_element_window(self.master, self.current_config[i], ui_state[0], ui_state[1], ui_state[2], CtkConceptWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) def _add_search_bar(self): toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") @@ -58,8 +48,7 @@ def _add_search_bar(self): toolbar.grid_columnconfigure(2, weight=1) self._toolbar = toolbar - # Search - ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0,5)) + ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0, 5)) self.search_var = StringVar() self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, placeholder_text="Filter...", width=200) @@ -67,77 +56,22 @@ def _add_search_bar(self): self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) - # Spacer ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) - # Type filter - ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0,5)) + ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0, 5)) self.filter_var = StringVar(value="ALL") - ctk.CTkOptionMenu(toolbar, values=["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"], + ctk.CTkOptionMenu(toolbar, values=self._FILTER_TYPES, variable=self.filter_var, command=lambda x: self._update_filters(), width=150).grid(row=0, column=4) - # Show disabled checkbox self.show_disabled_var = BooleanVar(value=True) self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, command=self._update_filters, width=100) - self.show_disabled_checkbox.grid(row=0, column=5, padx=(10,0)) + self.show_disabled_checkbox.grid(row=0, column=5, padx=(10, 0)) self._refresh_show_disabled_text() - # Clear button ctk.CTkButton(toolbar, text="Clear", width=50, - command=self._reset_filters).grid(row=0, column=6, padx=(10,0)) - - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() - - def _reset_filters(self): - self.search_var.set("") - self.filter_var.set("ALL") - self.show_disabled_var.set(True) - self._update_filters() - - def _element_matches_filters(self, element): - # Check enabled status - if not self.filters.get("show_disabled", True): - if hasattr(element, 'enabled') and not element.enabled: - return False - - # Search filter - search = self.filters.get("search", "").lower() - if search: - if not hasattr(element, '_search_cache'): - cache = [] - try: - if getattr(element, 'name', None): - cache.append(element.name.lower()) - p = getattr(element, 'path', None) - if p: - try: - cache.append(os.path.basename(p).lower()) - cache.append(p.lower()) - except (TypeError, AttributeError): - pass - except (AttributeError, TypeError): - pass - element._search_cache = cache - if not any(search in text for text in getattr(element, '_search_cache', [])): - return False - - # Type filter - type_filter = self.filters.get("type", "ALL") - if type_filter != "ALL": - if hasattr(element, 'type') and element.type: - try: - return ConceptType(element.type).value == type_filter - except (ValueError, AttributeError): - return False - return False - - return True + command=self._reset_filters).grid(row=0, column=6, padx=(10, 0)) def _maybe_reposition_toolbar(self, width): if not self._toolbar: @@ -152,6 +86,12 @@ def _maybe_reposition_toolbar(self, width): else: self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + def _reset_filters(self): + self.search_var.set("") + self.filter_var.set("ALL") + self.show_disabled_var.set(True) + self._update_filters() + def _refresh_show_disabled_text(self): try: disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) @@ -165,32 +105,29 @@ def _refresh_show_disabled_text(self): pass -class ConceptWidget(ctk.CTkFrame): - def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, width=150, height=170, corner_radius=10, bg_color="transparent" - ) +class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): + + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") + BaseConceptWidgetView.__init__(self, ctk_components) self.concept = concept - self.ui_state = UIState(self, concept) - self.image_ui_state = UIState(self, concept.image) - self.text_ui_state = UIState(self, concept.text) + self.ui_state = CtkUIState(self, concept) + self.image_ui_state = CtkUIState(self, concept.image) + self.text_ui_state = CtkUIState(self, concept.text) self.i = i self.grid_rowconfigure(1, weight=1) - # image self.image = ctk.CTkImage( - light_image=self.__get_preview_image(), + light_image=self._get_preview_image(), size=(150, 150) ) image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150) image_label.grid(row=0, column=0) - # name - self.name_label = components.label(self, 1, 0, self.__get_display_name(), pad=5, wraplength=140) + self.name_label = self.components.label(self, 1, 0, self._get_display_name(), pad=5, wraplength=140) - # close button close_button = ctk.CTkButton( master=self, width=20, @@ -202,7 +139,6 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma ) close_button.place(x=0, y=0) - # clone button clone_button = ctk.CTkButton( master=self, width=20, @@ -210,11 +146,10 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma text="+", corner_radius=2, fg_color="#00C000", - command=lambda: clone_command(self.i, self.__randomize_seed), + command=lambda: clone_command(self.i, controller.randomize_seed), ) clone_button.place(x=25, y=0) - # enabled switch enabled_switch = ctk.CTkSwitch( master=self, width=40, @@ -229,55 +164,10 @@ def __init__(self, master, concept, i, open_command, remove_command, clone_comma lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) ) - def __randomize_seed(self, concept: ConceptConfig): - concept.seed = ConceptConfig.default_values().seed - return concept - - def __get_display_name(self): - if self.concept.name: - return self.concept.name - elif self.concept.path: - return os.path.basename(self.concept.path) - else: - return "" - def configure_element(self): - self.name_label.configure(text=self.__get_display_name()) - self.image.configure(light_image=self.__get_preview_image()) - try: - if hasattr(self.concept, '_search_cache'): - delattr(self.concept, '_search_cache') - except AttributeError: - pass - - def __get_preview_image(self): - preview_path = "resources/icons/icon.png" - glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" - - concept_path = ConceptWindow.get_concept_path(getattr(self.concept, 'path', None)) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if (path.is_file() - and path_util.is_supported_image_extension(extension) - and not path.name.endswith("-masklabel.png") - and not path.name.endswith("-condlabel.png")): - preview_path = path_util.canonical_join(concept_path, path) - break - try: - image = load_image(preview_path, convert_mode="RGBA") - except (OSError): - image = Image.new("RGBA", (150, 150), (200, 200, 200, 255)) - size = min(image.width, image.height) - image = image.crop(( - (image.width - size) // 2, - (image.height - size) // 2, - (image.width - size) // 2 + size, - (image.height - size) // 2 + size, - )) - return image.resize((150, 150), Image.Resampling.BILINEAR) + self.name_label.configure(text=self._get_display_name()) + self.image.configure(light_image=self._get_preview_image()) + self._clear_search_cache() def place_in_list(self): index = getattr(self, 'visible_index', self.i) diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index f58879d5f..60c0f57fe 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -1,94 +1,30 @@ -import fractions -import math -import os -import pathlib -import platform -import random import threading -import time -import traceback - -from modules.util import concept_stats, path_util -from modules.util.config.ConceptConfig import ConceptConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.BalancingStrategy import BalancingStrategy -from modules.util.enum.ConceptType import ConceptType -from modules.util.image_util import load_image -from modules.util.ui import components + +from modules.ui.BaseConceptWindowView import BaseConceptWindowView +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -from mgds.LoadingPipeline import LoadingPipeline -from mgds.OutputPipelineModule import OutputPipelineModule -from mgds.PipelineModule import PipelineModule -from mgds.pipelineModules.CapitalizeTags import CapitalizeTags -from mgds.pipelineModules.DropTags import DropTags -from mgds.pipelineModules.RandomBrightness import RandomBrightness -from mgds.pipelineModules.RandomCircularMaskShrink import ( - RandomCircularMaskShrink, -) -from mgds.pipelineModules.RandomContrast import RandomContrast -from mgds.pipelineModules.RandomFlip import RandomFlip -from mgds.pipelineModules.RandomHue import RandomHue -from mgds.pipelineModules.RandomMaskRotateCrop import RandomMaskRotateCrop -from mgds.pipelineModules.RandomRotate import RandomRotate -from mgds.pipelineModules.RandomSaturation import RandomSaturation -from mgds.pipelineModules.ShuffleTags import ShuffleTags -from mgds.pipelineModuleTypes.RandomAccessPipelineModule import ( - RandomAccessPipelineModule, -) - -import torch -from torchvision.transforms import functional import customtkinter as ctk -import huggingface_hub from customtkinter import AppearanceModeTracker, ThemeManager from matplotlib import pyplot as plt from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -from PIL import Image - - -class InputPipelineModule( - PipelineModule, - RandomAccessPipelineModule, -): - def __init__(self, data: dict): - super().__init__() - self.data = data - - def length(self) -> int: - return 1 - def get_inputs(self) -> list[str]: - return [] - def get_outputs(self) -> list[str]: - return list(self.data.keys()) - - def get_item(self, variation: int, index: int, requested_name: str = None) -> dict: - return self.data - - -class ConceptWindow(ctk.CTkToplevel): +class CtkConceptWindowView(BaseConceptWindowView, ctk.CTkToplevel): def __init__( self, parent, - train_config: TrainConfig, - concept: ConceptConfig, - ui_state: UIState, - image_ui_state: UIState, - text_ui_state: UIState, + controller: ConceptWindowController, + ui_state, + image_ui_state, + text_ui_state, *args, **kwargs, ): - super().__init__(parent, *args, **kwargs) - - self.train_config = train_config + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConceptWindowView.__init__(self, ctk_components) - self.concept = concept - self.ui_state = ui_state - self.image_ui_state = image_ui_state - self.text_ui_state = text_ui_state + self.controller = controller self.image_preview_file_index = 0 self.preview_augmentations = ctk.BooleanVar(self, True) self.bucket_fig = None @@ -103,197 +39,40 @@ def __init__( tabview = ctk.CTkTabview(self) tabview.grid(row=0, column=0, sticky="nsew") - self.general_tab = self.__general_tab(tabview.add("general"), concept) - self.image_augmentation_tab = self.__image_augmentation_tab(tabview.add("image augmentation")) - self.text_augmentation_tab = self.__text_augmentation_tab(tabview.add("text augmentation")) - self.concept_stats_tab = self.__concept_stats_tab(tabview.add("statistics")) - - #automatic concept scan - self.scan_thread = threading.Thread(target=self.__auto_update_concept_stats, daemon=True) - self.scan_thread.start() - - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __general_tab(self, master, concept: ConceptConfig): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=1) - - # name - components.label(frame, 0, 0, "Name", - tooltip="Name of the concept") - components.entry(frame, 0, 1, self.ui_state, "name") - - # enabled - components.label(frame, 1, 0, "Enabled", - tooltip="Enable or disable this concept") - components.switch(frame, 1, 1, self.ui_state, "enabled") - - # concept type - components.label(frame, 2, 0, "Concept Type", - tooltip="STANDARD: Standard finetuning with the sample as training target\n" - "VALIDATION: Use concept for validation instead of training\n" - "PRIOR_PREDICTION: Use the sample to make a prediction using the model as it was before training. This prediction is then used as the training target " - "for the model in training. This can be used as regularisation and to preserve prior model knowledge while finetuning the model on other concepts. " - "Only implemented for LoRA.", - wide_tooltip=True) - components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], self.ui_state, "type") - - # path - components.label(frame, 3, 0, "Path", - tooltip="Path where the training data is located") - components.path_entry(frame, 3, 1, self.ui_state, "path", mode="dir") - components.button(frame, 3, 2, text="download now", command=self.__download_dataset_threaded, - tooltip="Download dataset from Huggingface now, for the purpose of previewing and statistics. Otherwise, it will be downloaded when you start training. Path must be a Huggingface repository.") - - # prompt source - components.label(frame, 4, 0, "Prompt Source", - tooltip="The source for prompts used during training. When selecting \"From single text file\", select a text file that contains a list of prompts") - prompt_path_entry = components.path_entry(frame, 4, 2, self.text_ui_state, "prompt_path", mode="file") - - def set_prompt_path_entry_enabled(option: str): - if option == 'concept': - for child in prompt_path_entry.children.values(): - child.configure(state="normal") - else: - for child in prompt_path_entry.children.values(): - child.configure(state="disabled") - - components.options_kv(frame, 4, 1, [ - ("From text file per sample", 'sample'), - ("From single text file", 'concept'), - ("From image file name", 'filename'), - ], self.text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) - set_prompt_path_entry_enabled(concept.text.prompt_source) - - # include subdirectories - components.label(frame, 5, 0, "Include Subdirectories", - tooltip="Includes images from subdirectories into the dataset") - components.switch(frame, 5, 1, self.ui_state, "include_subdirectories") - - # image variations - components.label(frame, 6, 0, "Image Variations", - tooltip="The number of different image versions to cache if latent caching is enabled.") - components.entry(frame, 6, 1, self.ui_state, "image_variations") - - # text variations - components.label(frame, 7, 0, "Text Variations", - tooltip="The number of different text versions to cache if latent caching is enabled.") - components.entry(frame, 7, 1, self.ui_state, "text_variations") - - # balancing - components.label(frame, 8, 0, "Balancing", - tooltip="The number of samples used during training. Use repeats to multiply the concept, or samples to specify an exact number of samples used in each epoch.") - components.entry(frame, 8, 1, self.ui_state, "balancing") - components.options(frame, 8, 2, [str(x) for x in list(BalancingStrategy)], self.ui_state, "balancing_strategy") - - # loss weight - components.label(frame, 9, 0, "Loss Weight", - tooltip="The loss multiplyer for this concept.") - components.entry(frame, 9, 1, self.ui_state, "loss_weight") - - frame.pack(fill="both", expand=1) - return frame - - def __image_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # header - components.label(frame, 0, 1, "Random", - tooltip="Enable this augmentation with random values") - components.label(frame, 0, 2, "Fixed", - tooltip="Enable this augmentation with fixed values") - - # crop jitter - components.label(frame, 1, 0, "Crop Jitter", - tooltip="Enables random cropping of samples") - components.switch(frame, 1, 1, self.image_ui_state, "enable_crop_jitter") - - # random flip - components.label(frame, 2, 0, "Random Flip", - tooltip="Randomly flip the sample during training") - components.switch(frame, 2, 1, self.image_ui_state, "enable_random_flip") - components.switch(frame, 2, 2, self.image_ui_state, "enable_fixed_flip") - - # random rotation - components.label(frame, 3, 0, "Random Rotation", - tooltip="Randomly rotates the sample during training") - components.switch(frame, 3, 1, self.image_ui_state, "enable_random_rotate") - components.switch(frame, 3, 2, self.image_ui_state, "enable_fixed_rotate") - components.entry(frame, 3, 3, self.image_ui_state, "random_rotate_max_angle") - - # random brightness - components.label(frame, 4, 0, "Random Brightness", - tooltip="Randomly adjusts the brightness of the sample during training") - components.switch(frame, 4, 1, self.image_ui_state, "enable_random_brightness") - components.switch(frame, 4, 2, self.image_ui_state, "enable_fixed_brightness") - components.entry(frame, 4, 3, self.image_ui_state, "random_brightness_max_strength") - - # random contrast - components.label(frame, 5, 0, "Random Contrast", - tooltip="Randomly adjusts the contrast of the sample during training") - components.switch(frame, 5, 1, self.image_ui_state, "enable_random_contrast") - components.switch(frame, 5, 2, self.image_ui_state, "enable_fixed_contrast") - components.entry(frame, 5, 3, self.image_ui_state, "random_contrast_max_strength") - - # random saturation - components.label(frame, 6, 0, "Random Saturation", - tooltip="Randomly adjusts the saturation of the sample during training") - components.switch(frame, 6, 1, self.image_ui_state, "enable_random_saturation") - components.switch(frame, 6, 2, self.image_ui_state, "enable_fixed_saturation") - components.entry(frame, 6, 3, self.image_ui_state, "random_saturation_max_strength") - - # random hue - components.label(frame, 7, 0, "Random Hue", - tooltip="Randomly adjusts the hue of the sample during training") - components.switch(frame, 7, 1, self.image_ui_state, "enable_random_hue") - components.switch(frame, 7, 2, self.image_ui_state, "enable_fixed_hue") - components.entry(frame, 7, 3, self.image_ui_state, "random_hue_max_strength") - - # random circular mask shrink - components.label(frame, 8, 0, "Circular Mask Generation", - tooltip="Automatically create circular masks for masked training") - components.switch(frame, 8, 1, self.image_ui_state, "enable_random_circular_mask_shrink") - - # random rotate and crop - components.label(frame, 9, 0, "Random Rotate and Crop", - tooltip="Randomly rotate the training samples and crop to the masked region") - components.switch(frame, 9, 1, self.image_ui_state, "enable_random_mask_rotate_crop") - - # circular mask generation - components.label(frame, 10, 0, "Resolution Override", - tooltip="Override the resolution for this concept. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.switch(frame, 10, 2, self.image_ui_state, "enable_resolution_override") - components.entry(frame, 10, 3, self.image_ui_state, "resolution_override") + # general tab + general_frame = ctk.CTkScrollableFrame(tabview.add("general"), fg_color="transparent") + general_frame.grid_columnconfigure(1, weight=1) + general_frame.grid_columnconfigure(2, weight=1) + self.build_general_tab(general_frame, controller, ui_state, text_ui_state) + general_frame.pack(fill="both", expand=1) + + # image augmentation tab + image_aug_master = tabview.add("image augmentation") + image_aug_frame = ctk.CTkScrollableFrame(image_aug_master, fg_color="transparent") + image_aug_frame.grid_columnconfigure(0, weight=0) + image_aug_frame.grid_columnconfigure(1, weight=0) + image_aug_frame.grid_columnconfigure(2, weight=0) + image_aug_frame.grid_columnconfigure(3, weight=1) + self.build_image_augmentation_tab(image_aug_frame, controller, image_ui_state) # image - image_preview, filename_preview, caption_preview = self.__get_preview_image() + image_preview, filename_preview, caption_preview = controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) self.image = ctk.CTkImage( light_image=image_preview, size=image_preview.size, ) - image_label = ctk.CTkLabel(master=frame, text="", image=self.image, height=300, width=300) + image_label = ctk.CTkLabel(master=image_aug_frame, text="", image=self.image, height=300, width=300) image_label.grid(row=0, column=4, rowspan=6) # refresh preview - update_button_frame = ctk.CTkFrame(master=frame, corner_radius=0, fg_color="transparent") + update_button_frame = ctk.CTkFrame(master=image_aug_frame, corner_radius=0, fg_color="transparent") update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") update_button_frame.grid_columnconfigure(1, weight=1) - prev_preview_button = components.button(update_button_frame, 0, 0, "<", command=self.__prev_image_preview) - components.button(update_button_frame, 0, 1, "Update Preview", command=self.__update_image_preview) - next_preview_button = components.button(update_button_frame, 0, 2, ">", command=self.__next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self.__update_image_preview) + prev_preview_button = self.components.button(update_button_frame, 0, 0, "<", command=self._prev_image_preview) + self.components.button(update_button_frame, 0, 1, "Update Preview", command=self._update_image_preview) + next_preview_button = self.components.button(update_button_frame, 0, 2, ">", command=self._next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self._update_image_preview) preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) prev_preview_button.configure(width=40) @@ -307,205 +86,24 @@ def __image_augmentation_tab(self, master): self.caption_preview.configure(state="disabled") self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - frame.pack(fill="both", expand=1) - return frame - - def __text_augmentation_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # tag shuffling - components.label(frame, 0, 0, "Tag Shuffling", - tooltip="Enables tag shuffling") - components.switch(frame, 0, 1, self.text_ui_state, "enable_tag_shuffling") - - # keep tag count - components.label(frame, 1, 0, "Tag Delimiter", - tooltip="The delimiter between tags") - components.entry(frame, 1, 1, self.text_ui_state, "tag_delimiter") - - # keep tag count - components.label(frame, 2, 0, "Keep Tag Count", - tooltip="The number of tags at the start of the caption that are not shuffled or dropped") - components.entry(frame, 2, 1, self.text_ui_state, "keep_tags_count") - - # tag dropout - components.label(frame, 3, 0, "Tag Dropout", - tooltip="Enables random dropout for tags in the captions.") - components.switch(frame, 3, 1, self.text_ui_state, "tag_dropout_enable") - components.label(frame, 4, 0, "Dropout Mode", - tooltip="Method used to drop captions. 'Full' will drop the entire caption past the 'kept' tags with a certain probability, 'Random' will drop individual tags with the set probability, and 'Random Weighted' will linearly increase the probability of dropping tags, more likely to preseve tags near the front with full probability to drop at the end.") - components.options_kv(frame, 4, 1, [ - ("Full", 'FULL'), - ("Random", 'RANDOM'), - ("Random Weighted", 'RANDOM WEIGHTED'), - ], self.text_ui_state, "tag_dropout_mode", None) - components.label(frame, 4, 2, "Probability", - tooltip="Probability to drop tags, from 0 to 1.") - components.entry(frame, 4, 3, self.text_ui_state, "tag_dropout_probability") - - components.label(frame, 5, 0, "Special Dropout Tags", - tooltip="List of tags which will be whitelisted/blacklisted by dropout. 'Whitelist' tags will never be dropped but all others may be, 'Blacklist' tags may be dropped but all others will never be, 'None' may drop any tags. Can specify either a delimiter-separated list in the field, or a file path to a .txt or .csv file with entries separated by newlines.") - components.options_kv(frame, 5, 1, [ - ("None", 'NONE'), - ("Blacklist", 'BLACKLIST'), - ("Whitelist", 'WHITELIST'), - ], self.text_ui_state, "tag_dropout_special_tags_mode", None) - components.entry(frame, 5, 2, self.text_ui_state, "tag_dropout_special_tags") - components.label(frame, 6, 0, "Special Tags Regex", - tooltip="Interpret special tags with regex, such as 'photo.*' to match 'photo, photograph, photon' but not 'telephoto'. Includes exception for '/(' and '/)' syntax found in many booru/e6 tags.") - components.switch(frame, 6, 1, self.text_ui_state, "tag_dropout_special_tags_regex") - - #capitalization randomization - components.label(frame, 7, 0, "Randomize Capitalization", - tooltip="Enables randomization of capitalization for tags in the caption.") - components.switch(frame, 7, 1, self.text_ui_state, "caps_randomize_enable") - components.label(frame, 7, 2, "Force Lowercase", - tooltip="If enabled, converts the caption to lowercase before any further processing.") - components.switch(frame, 7, 3, self.text_ui_state, "caps_randomize_lowercase") - - components.label(frame, 8, 0, "Captialization Mode", - tooltip="Comma-separated list of types of capitalization randomization to perform. 'capslock' for ALL CAPS, 'title' for First Letter Of Every Word, 'first' for First word only, 'random' for rAndOMiZeD lEtTERs.") - components.entry(frame, 8, 1, self.text_ui_state, "caps_randomize_mode") - components.label(frame, 8, 2, "Probability", - tooltip="Probability to randomize capitialization of each tag, from 0 to 1.") - components.entry(frame, 8, 3, self.text_ui_state, "caps_randomize_probability") - - frame.pack(fill="both", expand=1) - return frame - - def __concept_stats_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=150) - frame.grid_columnconfigure(1, weight=0, minsize=150) - frame.grid_columnconfigure(2, weight=0, minsize=150) - frame.grid_columnconfigure(3, weight=0, minsize=150) - - self.cancel_scan_flag = threading.Event() - - #file size - self.file_size_label = components.label(frame, 1, 0, "Total Size", pad=0, - tooltip="Total size of all image, mask, and caption files in MB") - self.file_size_label.configure(font=ctk.CTkFont(underline=True)) - self.file_size_preview = components.label(frame, 2, 0, pad=0, text="-") - - #subdirectory count - self.dir_count_label = components.label(frame, 1, 1, "Directories", pad=0, - tooltip="Total number of directories including and under (if 'include subdirectories' is enabled) the main concept directory") - self.dir_count_label.configure(font=ctk.CTkFont(underline=True)) - self.dir_count_preview = components.label(frame, 2, 1, pad=0, text="-") - - #basic img/vid stats - count of each type in the concept - #the \n at the start of the label gives it better vertical spacing with other rows - self.image_count_label = components.label(frame, 3, 0, "\nTotal Images", pad=0, - tooltip="Total number of image files, any of the extensions " + str(path_util.SUPPORTED_IMAGE_EXTENSIONS) + ", excluding '-masklabel.png and -condlabel.png'") - self.image_count_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_preview = components.label(frame, 4, 0, pad=0, text="-") - self.video_count_label = components.label(frame, 3, 1, "\nTotal Videos", pad=0, - tooltip="Total number of video files, any of the extensions " + str(path_util.SUPPORTED_VIDEO_EXTENSIONS)) - self.video_count_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_preview = components.label(frame, 4, 1, pad=0, text="-") - self.mask_count_label = components.label(frame, 3, 2, "\nTotal Masks", pad=0, - tooltip="Total number of mask files, any file ending in '-masklabel.png'") - self.mask_count_label.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview = components.label(frame, 4, 2, pad=0, text="-") - self.caption_count_label = components.label(frame, 3, 3, "\nTotal Captions", pad=0, - tooltip="Total number of caption files, any .txt file. With advanced scan, includes the total number of captions on separate lines across all files in parentheses.") - self.caption_count_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview = components.label(frame, 4, 3, pad=0, text="-") - - #advanced img/vid stats - how many img/vid files have a mask or caption of the same name - self.image_count_mask_label = components.label(frame, 5, 0, "\nImages with Masks", pad=0, - tooltip="Total number of image files with an associated mask") - self.image_count_mask_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_mask_preview = components.label(frame, 6, 0, pad=0, text="-") - self.mask_count_label_unpaired = components.label(frame, 5, 1, "\nUnpaired Masks", pad=0, - tooltip="Total number of mask files which lack a corresponding image file - if >0, check your data set!") - self.mask_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.mask_count_preview_unpaired = components.label(frame, 6, 1, pad=0, text="-") - #currently no masks for videos? - - self.image_count_caption_label = components.label(frame, 7, 0, "\nImages with Captions", pad=0, - tooltip="Total number of image files with an associated caption") - self.image_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.image_count_caption_preview = components.label(frame, 8, 0, pad=0, text="-") - self.video_count_caption_label = components.label(frame, 7, 1, "\nVideos with Captions", pad=0, - tooltip="Total number of video files with an associated caption") - self.video_count_caption_label.configure(font=ctk.CTkFont(underline=True)) - self.video_count_caption_preview = components.label(frame, 8, 1, pad=0, text="-") - self.caption_count_label_unpaired = components.label(frame, 7, 2, "\nUnpaired Captions", pad=0, - tooltip="Total number of caption files which lack a corresponding image file - if >0, check your data set! If using 'from file name' or 'from single text file' then this can be ignored.") - self.caption_count_label_unpaired.configure(font=ctk.CTkFont(underline=True)) - self.caption_count_preview_unpaired = components.label(frame, 8, 2, pad=0, text="-") - - #resolution info - self.pixel_max_label = components.label(frame, 9, 0, "\nMax Pixels", pad=0, - tooltip="Largest image in the concept by number of pixels (width * height)") - self.pixel_max_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_max_preview = components.label(frame, 10, 0, pad=0, text="-", wraplength=150) - self.pixel_avg_label = components.label(frame, 9, 1, "\nAvg Pixels", pad=0, - tooltip="Average size of images in the concept by number of pixels (width * height)") - self.pixel_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_avg_preview = components.label(frame, 10, 1, pad=0, text="-", wraplength=150) - self.pixel_min_label = components.label(frame, 9, 2, "\nMin Pixels", pad=0, - tooltip="Smallest image in the concept by number of pixels (width * height)") - self.pixel_min_label.configure(font=ctk.CTkFont(underline=True)) - self.pixel_min_preview = components.label(frame, 10, 2, pad=0, text="-", wraplength=150) - - #video length info - self.length_max_label = components.label(frame, 11, 0, "\nMax Length", pad=0, - tooltip="Longest video in the concept by number of frames") - self.length_max_label.configure(font=ctk.CTkFont(underline=True)) - self.length_max_preview = components.label(frame, 12, 0, pad=0, text="-", wraplength=150) - self.length_avg_label = components.label(frame, 11, 1, "\nAvg Length", pad=0, - tooltip="Average length of videos in the concept by number of frames") - self.length_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.length_avg_preview = components.label(frame, 12, 1, pad=0, text="-", wraplength=150) - self.length_min_label = components.label(frame, 11, 2, "\nMin Length", pad=0, - tooltip="Shortest video in the concept by number of frames") - self.length_min_label.configure(font=ctk.CTkFont(underline=True)) - self.length_min_preview = components.label(frame, 12, 2, pad=0, text="-", wraplength=150) - - #video fps info - self.fps_max_label = components.label(frame, 13, 0, "\nMax FPS", pad=0, - tooltip="Video in concept with highest fps") - self.fps_max_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_max_preview = components.label(frame, 14, 0, pad=0, text="-", wraplength=150) - self.fps_avg_label = components.label(frame, 13, 1, "\nAvg FPS", pad=0, - tooltip="Average fps of videos in the concept") - self.fps_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_avg_preview = components.label(frame, 14, 1, pad=0, text="-", wraplength=150) - self.fps_min_label = components.label(frame, 13, 2, "\nMin FPS", pad=0, - tooltip="Video in concept with the lowest fps") - self.fps_min_label.configure(font=ctk.CTkFont(underline=True)) - self.fps_min_preview = components.label(frame, 14, 2, pad=0, text="-", wraplength=150) - - #caption info - self.caption_max_label = components.label(frame, 15, 0, "\nMax Caption Length", pad=0, - tooltip="Largest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_max_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_max_preview = components.label(frame, 16, 0, pad=0, text="-", wraplength=150) - self.caption_avg_label = components.label(frame, 15, 1, "\nAvg Caption Length", pad=0, - tooltip="Average length of caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_avg_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_avg_preview = components.label(frame, 16, 1, pad=0, text="-", wraplength=150) - self.caption_min_label = components.label(frame, 15, 2, "\nMin Caption Length", pad=0, - tooltip="Smallest caption in concept by character count. For token count, assume ~2 tokens/word") - self.caption_min_label.configure(font=ctk.CTkFont(underline=True)) - self.caption_min_preview = components.label(frame, 16, 2, pad=0, text="-", wraplength=150) - - #aspect bucket info - self.aspect_bucket_label = components.label(frame, 17, 0, "\nAspect Bucketing", pad=0, - tooltip="Graph of all possible buckets and the number of images in each one, defined as height/width. Buckets range from 0.25 (4:1 extremely wide) to 4 (1:4 extremely tall). \ - Images which don't match a bucket exactly are cropped to the nearest one.") - self.aspect_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_label = components.label(frame, 17, 1, "\nSmallest Buckets", pad=0, - tooltip="Image buckets with the least nonzero total images - if 'batch size' is larger than this, these images will be ignored during training! See the wiki for more details.") - self.small_bucket_label.configure(font=ctk.CTkFont(underline=True)) - self.small_bucket_preview = components.label(frame, 18, 1, pad=0, text="-") + image_aug_frame.pack(fill="both", expand=1) + + # text augmentation tab + text_aug_frame = ctk.CTkScrollableFrame(tabview.add("text augmentation"), fg_color="transparent") + text_aug_frame.grid_columnconfigure(0, weight=0) + text_aug_frame.grid_columnconfigure(1, weight=0) + text_aug_frame.grid_columnconfigure(2, weight=0) + text_aug_frame.grid_columnconfigure(3, weight=1) + self.build_text_augmentation_tab(text_aug_frame, controller, text_ui_state) + text_aug_frame.pack(fill="both", expand=1) + + # statistics tab + stats_frame = ctk.CTkScrollableFrame(tabview.add("statistics"), fg_color="transparent") + stats_frame.grid_columnconfigure(0, weight=0, minsize=150) + stats_frame.grid_columnconfigure(1, weight=0, minsize=150) + stats_frame.grid_columnconfigure(2, weight=0, minsize=150) + stats_frame.grid_columnconfigure(3, weight=0, minsize=150) + self.build_concept_stats_tab(stats_frame, controller) #aspect bucketing plot, mostly copied from timestep preview graph appearance_mode = AppearanceModeTracker.get_mode() @@ -518,7 +116,7 @@ def __concept_stats_tab(self, master): assert self.bucket_fig is None self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=frame) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=stats_frame) self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) self.bucket_fig.tight_layout() self.bucket_fig.subplots_adjust(bottom=0.15) @@ -534,28 +132,29 @@ def __concept_stats_tab(self, master): self.bucket_ax.xaxis.label.set_color(self.text_color) self.bucket_ax.yaxis.label.set_color(self.text_color) - #refresh stats - must be after all labels are defined or will give error - self.refresh_basic_stats_button = components.button(master=frame, row=0, column=0, text="Refresh Basic", command=lambda: self.__get_concept_stats_threaded(False, 9999), - tooltip="Reload basic statistics for the concept directory") - self.refresh_advanced_stats_button = components.button(master=frame, row=0, column=1, text="Refresh Advanced", command=lambda: self.__get_concept_stats_threaded(True, 9999), - tooltip="Reload advanced statistics for the concept directory") #run "basic" scan first before "advanced", seems to help the system cache the directories and run faster - self.cancel_stats_button = components.button(master=frame, row=0, column=2, text="Abort Scan", command=lambda: self.__cancel_concept_stats(), - tooltip="Stop the currently running scan if it's taking a long time - advanced scan will be slow on large folders and on HDDs") - self.processing_time = components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + stats_frame.pack(fill="both", expand=1) - frame.pack(fill="both", expand=1) - return frame + #automatic concept scan + self.scan_thread = threading.Thread(target=controller.auto_update_concept_stats, args=[self], daemon=True) + self.scan_thread.start() - def __prev_image_preview(self): + self.components.button(self, 1, 0, "ok", self._ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def _prev_image_preview(self): self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) - self.__update_image_preview() + self._update_image_preview() - def __next_image_preview(self): + def _next_image_preview(self): self.image_preview_file_index += 1 - self.__update_image_preview() + self._update_image_preview() - def __update_image_preview(self): - image_preview, filename_preview, caption_preview = self.__get_preview_image() + def _update_image_preview(self): + image_preview, filename_preview, caption_preview = self.controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) self.image.configure(light_image=image_preview, size=image_preview.size) self.filename_preview.configure(text=filename_preview) self.caption_preview.configure(state="normal") @@ -563,366 +162,6 @@ def __update_image_preview(self): self.caption_preview.insert(index="1.0", text=caption_preview) self.caption_preview.configure(state="disabled") - @staticmethod - def get_concept_path(path: str) -> str | None: - if os.path.isdir(path): - return path - try: - #don't download, only check if available locally: - return huggingface_hub.snapshot_download(repo_id=path, repo_type="dataset", local_files_only=True) - except Exception: - return None - - def __download_dataset(self): - try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) - huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") - except Exception: - traceback.print_exc() - - def __download_dataset_threaded(self): - download_thread = threading.Thread(target=self.__download_dataset, daemon=True) - download_thread.start() - - def _read_text_file_for_preview(self, file_path: str) -> str: - empty_msg = "[Empty prompt]" - try: - with open(file_path, "r") as f: - if self.preview_augmentations.get(): - lines = [line.strip() for line in f if line.strip()] - return random.choice(lines) if lines else empty_msg - content = f.read().strip() - return content if content else empty_msg - except FileNotFoundError: - return "File not found, please check the path" - except IsADirectoryError: - return "[Provided path is a directory, please correct the caption path]" - except PermissionError: - if platform.system() == "Windows": - return "[Permission denied, please check the file permissions or Windows Defender settings]" - else: - return "[Permission denied, please check the file permissions]" - except UnicodeDecodeError: - return "[Invalid file encoding. This should not happen, please report this issue]" - - def __get_preview_image(self): - preview_image_path = "resources/icons/icon.png" - file_index = -1 - glob_pattern = "**/*.*" if self.concept.include_subdirectories else "*.*" - - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - for path in pathlib.Path(concept_path).glob(glob_pattern): - if any(part.startswith('.') for part in path.relative_to(concept_path).parent.parts): - continue - extension = os.path.splitext(path)[1] - if path.is_file() and path_util.is_supported_image_extension(extension) \ - and not path.name.endswith("-masklabel.png") and not path.name.endswith("-condlabel.png"): - preview_image_path = path_util.canonical_join(concept_path, path) - file_index += 1 - if file_index == self.image_preview_file_index: - break - - image = load_image(preview_image_path, 'RGB') - image_tensor = functional.to_tensor(image) - - splitext = os.path.splitext(preview_image_path) - preview_mask_path = path_util.canonical_join(splitext[0] + "-masklabel.png") - if not os.path.isfile(preview_mask_path): - preview_mask_path = None - - if preview_mask_path: - mask = Image.open(preview_mask_path).convert("L") - mask_tensor = functional.to_tensor(mask) - else: - mask_tensor = torch.ones((1, image_tensor.shape[1], image_tensor.shape[2])) - - source = self.concept.text.prompt_source - preview_p = pathlib.Path(preview_image_path) - if source == "filename": - prompt_output = preview_p.stem or "[Empty prompt]" - else: - file_map = { - "sample": preview_p.with_suffix(".txt"), - "concept": pathlib.Path(self.concept.text.prompt_path) if self.concept.text.prompt_path else None, - } - file_path = file_map.get(source) - prompt_output = self._read_text_file_for_preview(str(file_path)) if file_path else "[Empty prompt]" - - modules = [] - if self.preview_augmentations.get(): - input_module = InputPipelineModule({ - 'true': True, - 'image': image_tensor, - 'mask': mask_tensor, - 'enable_random_flip': self.concept.image.enable_random_flip, - 'enable_fixed_flip': self.concept.image.enable_fixed_flip, - 'enable_random_rotate': self.concept.image.enable_random_rotate, - 'enable_fixed_rotate': self.concept.image.enable_fixed_rotate, - 'random_rotate_max_angle': self.concept.image.random_rotate_max_angle, - 'enable_random_brightness': self.concept.image.enable_random_brightness, - 'enable_fixed_brightness': self.concept.image.enable_fixed_brightness, - 'random_brightness_max_strength': self.concept.image.random_brightness_max_strength, - 'enable_random_contrast': self.concept.image.enable_random_contrast, - 'enable_fixed_contrast': self.concept.image.enable_fixed_contrast, - 'random_contrast_max_strength': self.concept.image.random_contrast_max_strength, - 'enable_random_saturation': self.concept.image.enable_random_saturation, - 'enable_fixed_saturation': self.concept.image.enable_fixed_saturation, - 'random_saturation_max_strength': self.concept.image.random_saturation_max_strength, - 'enable_random_hue': self.concept.image.enable_random_hue, - 'enable_fixed_hue': self.concept.image.enable_fixed_hue, - 'random_hue_max_strength': self.concept.image.random_hue_max_strength, - 'enable_random_circular_mask_shrink': self.concept.image.enable_random_circular_mask_shrink, - 'enable_random_mask_rotate_crop': self.concept.image.enable_random_mask_rotate_crop, - - 'prompt' : prompt_output, - 'tag_dropout_enable' : self.concept.text.tag_dropout_enable, - 'tag_dropout_probability' : self.concept.text.tag_dropout_probability, - 'tag_dropout_mode' : self.concept.text.tag_dropout_mode, - 'tag_dropout_special_tags' : self.concept.text.tag_dropout_special_tags, - 'tag_dropout_special_tags_mode' : self.concept.text.tag_dropout_special_tags_mode, - 'tag_delimiter' : self.concept.text.tag_delimiter, - 'keep_tags_count' : self.concept.text.keep_tags_count, - 'tag_dropout_special_tags_regex' : self.concept.text.tag_dropout_special_tags_regex, - 'caps_randomize_enable' : self.concept.text.caps_randomize_enable, - 'caps_randomize_probability' : self.concept.text.caps_randomize_probability, - 'caps_randomize_mode' : self.concept.text.caps_randomize_mode, - 'caps_randomize_lowercase' : self.concept.text.caps_randomize_lowercase, - 'enable_tag_shuffling' : self.concept.text.enable_tag_shuffling, - }) - - circular_mask_shrink = RandomCircularMaskShrink(mask_name='mask', shrink_probability=1.0, shrink_factor_min=0.2, shrink_factor_max=1.0, enabled_in_name='enable_random_circular_mask_shrink') - random_mask_rotate_crop = RandomMaskRotateCrop(mask_name='mask', additional_names=['image'], min_size=512, min_padding_percent=10, max_padding_percent=30, max_rotate_angle=20, enabled_in_name='enable_random_mask_rotate_crop') - random_flip = RandomFlip(names=['image', 'mask'], enabled_in_name='enable_random_flip', fixed_enabled_in_name='enable_fixed_flip') - random_rotate = RandomRotate(names=['image', 'mask'], enabled_in_name='enable_random_rotate', fixed_enabled_in_name='enable_fixed_rotate', max_angle_in_name='random_rotate_max_angle') - random_brightness = RandomBrightness(names=['image'], enabled_in_name='enable_random_brightness', fixed_enabled_in_name='enable_fixed_brightness', max_strength_in_name='random_brightness_max_strength') - random_contrast = RandomContrast(names=['image'], enabled_in_name='enable_random_contrast', fixed_enabled_in_name='enable_fixed_contrast', max_strength_in_name='random_contrast_max_strength') - random_saturation = RandomSaturation(names=['image'], enabled_in_name='enable_random_saturation', fixed_enabled_in_name='enable_fixed_saturation', max_strength_in_name='random_saturation_max_strength') - random_hue = RandomHue(names=['image'], enabled_in_name='enable_random_hue', fixed_enabled_in_name='enable_fixed_hue', max_strength_in_name='random_hue_max_strength') - drop_tags = DropTags(text_in_name='prompt', enabled_in_name='tag_dropout_enable', probability_in_name='tag_dropout_probability', dropout_mode_in_name='tag_dropout_mode', - special_tags_in_name='tag_dropout_special_tags', special_tag_mode_in_name='tag_dropout_special_tags_mode', delimiter_in_name='tag_delimiter', - keep_tags_count_in_name='keep_tags_count', text_out_name='prompt', regex_enabled_in_name='tag_dropout_special_tags_regex') - caps_randomize = CapitalizeTags(text_in_name='prompt', enabled_in_name='caps_randomize_enable', probability_in_name='caps_randomize_probability', - capitalize_mode_in_name='caps_randomize_mode', delimiter_in_name='tag_delimiter', convert_lowercase_in_name='caps_randomize_lowercase', text_out_name='prompt') - shuffle_tags = ShuffleTags(text_in_name='prompt', enabled_in_name='enable_tag_shuffling', delimiter_in_name='tag_delimiter', keep_tags_count_in_name='keep_tags_count', text_out_name='prompt') - output_module = OutputPipelineModule(['image', 'mask', 'prompt']) - - modules = [ - input_module, - circular_mask_shrink, - random_mask_rotate_crop, - random_flip, - random_rotate, - random_brightness, - random_contrast, - random_saturation, - random_hue, - drop_tags, - caps_randomize, - shuffle_tags, - output_module, - ] - - pipeline = LoadingPipeline( - device=torch.device('cpu'), - modules=modules, - batch_size=1, - seed=random.randint(0, 2**30), - state=None, - initial_epoch=0, - initial_index=0, - ) - - data = pipeline.__next__() - image_tensor = data['image'] - mask_tensor = data['mask'] - prompt_output = data['prompt'] - - filename_output = os.path.basename(preview_image_path) - - mask_tensor = torch.clamp(mask_tensor, 0.3, 1) - image_tensor = image_tensor * mask_tensor - - image = functional.to_pil_image(image_tensor) - - image.thumbnail((300, 300)) - - return image, filename_output, prompt_output - - def __update_concept_stats(self): - #file size - self.file_size_preview.configure(text=str(int(self.concept.concept_stats["file_size"]/1048576)) + " MB") - self.processing_time.configure(text=str(round(self.concept.concept_stats["processing_time"], 2)) + " s") - - #directory count - self.dir_count_preview.configure(text=self.concept.concept_stats["directory_count"]) - - #image count - self.image_count_preview.configure(text=self.concept.concept_stats["image_count"]) - self.image_count_mask_preview.configure(text=self.concept.concept_stats["image_with_mask_count"]) - self.image_count_caption_preview.configure(text=self.concept.concept_stats["image_with_caption_count"]) - - #video count - self.video_count_preview.configure(text=self.concept.concept_stats["video_count"]) - #self.video_count_mask_preview.configure(text=self.concept.concept_stats["video_with_mask_count"]) - self.video_count_caption_preview.configure(text=self.concept.concept_stats["video_with_caption_count"]) - - #mask count - self.mask_count_preview.configure(text=self.concept.concept_stats["mask_count"]) - self.mask_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_masks"]) - - #caption count - if self.concept.concept_stats["subcaption_count"] > 0: - self.caption_count_preview.configure(text=f'{self.concept.concept_stats["caption_count"]} ({self.concept.concept_stats["subcaption_count"]})') - else: - self.caption_count_preview.configure(text=self.concept.concept_stats["caption_count"]) - self.caption_count_preview_unpaired.configure(text=self.concept.concept_stats["unpaired_captions"]) - - #resolution info - max_pixels = self.concept.concept_stats["max_pixels"] - avg_pixels = self.concept.concept_stats["avg_pixels"] - min_pixels = self.concept.concept_stats["min_pixels"] - - if any(isinstance(x, str) for x in [max_pixels, avg_pixels, min_pixels]) or self.concept.concept_stats["image_count"] == 0: #will be str if adv stats were not taken - self.pixel_max_preview.configure(text="-") - self.pixel_avg_preview.configure(text="-") - self.pixel_min_preview.configure(text="-") - else: - #formatted as (#pixels/1000000) MP, width x height, \n filename - self.pixel_max_preview.configure(text=f'{str(round(max_pixels[0]/1000000, 2))} MP, {max_pixels[2]}\n{max_pixels[1]}') - self.pixel_avg_preview.configure(text=f'{str(round(avg_pixels/1000000, 2))} MP, ~{int(math.sqrt(avg_pixels))}w x {int(math.sqrt(avg_pixels))}h') - self.pixel_min_preview.configure(text=f'{str(round(min_pixels[0]/1000000, 2))} MP, {min_pixels[2]}\n{min_pixels[1]}') - - #video length and fps info - max_length = self.concept.concept_stats["max_length"] - avg_length = self.concept.concept_stats["avg_length"] - min_length = self.concept.concept_stats["min_length"] - max_fps = self.concept.concept_stats["max_fps"] - avg_fps = self.concept.concept_stats["avg_fps"] - min_fps = self.concept.concept_stats["min_fps"] - - if any(isinstance(x, str) for x in [max_length, avg_length, min_length]) or self.concept.concept_stats["video_count"] == 0: #will be str if adv stats were not taken - self.length_max_preview.configure(text="-") - self.length_avg_preview.configure(text="-") - self.length_min_preview.configure(text="-") - self.fps_max_preview.configure(text="-") - self.fps_avg_preview.configure(text="-") - self.fps_min_preview.configure(text="-") - else: - #formatted as (#frames) frames \n filename - self.length_max_preview.configure(text=f'{int(max_length[0])} frames\n{max_length[1]}') - self.length_avg_preview.configure(text=f'{int(avg_length)} frames') - self.length_min_preview.configure(text=f'{int(min_length[0])} frames\n{min_length[1]}') - #formatted as (#fps) fps \n filename - self.fps_max_preview.configure(text=f'{int(max_fps[0])} fps\n{max_fps[1]}') - self.fps_avg_preview.configure(text=f'{int(avg_fps)} fps') - self.fps_min_preview.configure(text=f'{int(min_fps[0])} fps\n{min_fps[1]}') - - #caption info - max_caption_length = self.concept.concept_stats["max_caption_length"] - avg_caption_length = self.concept.concept_stats["avg_caption_length"] - min_caption_length = self.concept.concept_stats["min_caption_length"] - - if any(isinstance(x, str) for x in [max_caption_length, avg_caption_length, min_caption_length]) or self.concept.concept_stats["caption_count"] == 0: #will be str if adv stats were not taken - self.caption_max_preview.configure(text="-") - self.caption_avg_preview.configure(text="-") - self.caption_min_preview.configure(text="-") - else: - #formatted as (#chars) chars, (#words) words, \n filename - self.caption_max_preview.configure(text=f'{max_caption_length[0]} chars, {max_caption_length[2]} words\n{max_caption_length[1]}') - self.caption_avg_preview.configure(text=f'{int(avg_caption_length[0])} chars, {int(avg_caption_length[1])} words') - self.caption_min_preview.configure(text=f'{min_caption_length[0]} chars, {min_caption_length[2]} words\n{min_caption_length[1]}') - - #aspect bucketing - aspect_buckets = self.concept.concept_stats["aspect_buckets"] - if len(aspect_buckets) != 0 and max(val for val in aspect_buckets.values()) > 0: #check aspect_bucket data exists and is not all zero - min_val = min(val for val in aspect_buckets.values() if val > 0) #smallest nonzero values - if max(val for val in aspect_buckets.values()) > min_val: #check if any buckets larger than min_val exist - if all images are same aspect then there won't be - min_val2 = min(val for val in aspect_buckets.values() if (val > 0 and val != min_val)) #second smallest bucket - else: - min_val2 = min_val #if no second smallest bucket exists set to min_val - min_aspect_buckets = {key: val for key,val in aspect_buckets.items() if val in (min_val, min_val2)} - min_bucket_str = "" - for key, val in min_aspect_buckets.items(): - min_bucket_str += f'aspect {self.decimal_to_aspect_ratio(key)} : {val} img\n' - min_bucket_str.strip() - self.small_bucket_preview.configure(text=min_bucket_str) - - self.bucket_ax.cla() - aspects = [str(x) for x in list(aspect_buckets.keys())] - aspect_ratios = [self.decimal_to_aspect_ratio(x) for x in list(aspect_buckets.keys())] - counts = list(aspect_buckets.values()) - b = self.bucket_ax.bar(aspect_ratios, counts) - self.bucket_ax.bar_label(b, color=self.text_color) - sec = self.bucket_ax.secondary_xaxis(location=-0.1) - sec.spines["bottom"].set_linewidth(0) - sec.set_xticks([0, (len(aspects)-1)/2, len(aspects)-1], labels=["Wide", "Square", "Tall"]) - sec.tick_params('x', length=0) - self.canvas.draw() - - def decimal_to_aspect_ratio(self, value : float): - #find closest fraction to decimal aspect value and convert to a:b format - aspect_fraction = fractions.Fraction(value).limit_denominator(16) - aspect_string = f'{aspect_fraction.denominator}:{aspect_fraction.numerator}' - return aspect_string - - def __get_concept_stats(self, advanced_checks: bool, wait_time: float): - start_time = time.perf_counter() - last_update = time.perf_counter() - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__disable_scan_buttons) - concept_path = self.get_concept_path(self.concept.path) - - if not concept_path: - print(f"Unable to get statistics for concept path: {self.concept.path}") - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - return - subfolders = [concept_path] - - stats_dict = concept_stats.init_concept_stats(advanced_checks) - for path in subfolders: - if self.cancel_scan_flag.is_set() or time.perf_counter() - start_time > wait_time: - break - stats_dict = concept_stats.folder_scan(path, stats_dict, advanced_checks, self.concept, start_time, wait_time, self.cancel_scan_flag) - if self.concept.include_subdirectories and not self.cancel_scan_flag.is_set(): #add all subfolders of current directory to for loop - subfolders.extend([f for f in os.scandir(path) if f.is_dir() and not f.name.startswith('.')]) - self.concept.concept_stats = stats_dict - #update GUI approx every half second - if time.perf_counter() > (last_update + 0.5): - last_update = time.perf_counter() - self.concept_stats_tab.after(0, self.__update_concept_stats) - - self.cancel_scan_flag.clear() - self.concept_stats_tab.after(0, self.__enable_scan_buttons) - self.concept_stats_tab.after(0, self.__update_concept_stats) - - def __get_concept_stats_threaded(self, advanced_checks : bool, waittime : float): - self.scan_thread = threading.Thread(target=self.__get_concept_stats, args=[advanced_checks, waittime], daemon=True) - self.scan_thread.start() - - def __disable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="disabled") - self.refresh_advanced_stats_button.configure(state="disabled") - - def __enable_scan_buttons(self): - self.refresh_basic_stats_button.configure(state="normal") - self.refresh_advanced_stats_button.configure(state="normal") - - def __cancel_concept_stats(self): - self.cancel_scan_flag.set() - - def __auto_update_concept_stats(self): - try: - self.__update_concept_stats() #load stats from config if available, else raises KeyError - if self.concept.concept_stats["file_size"] == 0: #force rescan if empty - raise KeyError - except KeyError: - concept_path = self.get_concept_path(self.concept.path) - if concept_path: - self.__get_concept_stats(False, 2) #force rescan if config is empty, timeout of 2 sec - if self.concept.concept_stats["processing_time"] < 0.1: - self.__get_concept_stats(True, 2) #do advanced scan automatically if basic took <0.1s - def destroy(self): if self.bucket_fig is not None: plt.close(self.bucket_fig) @@ -930,5 +169,5 @@ def destroy(self): super().destroy() - def __ok(self): + def _ok(self): self.destroy() diff --git a/modules/ui/CtkConfigListView.py b/modules/ui/CtkConfigListView.py index 75d69252a..72995bfcc 100644 --- a/modules/ui/CtkConfigListView.py +++ b/modules/ui/CtkConfigListView.py @@ -1,27 +1,19 @@ import contextlib -import copy -import json -import os -import tkinter as tk -from abc import ABCMeta, abstractmethod +from abc import ABC -from modules.util import path_util -from modules.util.config.BaseConfig import BaseConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util.ui import ctk_components, dialogs import customtkinter as ctk -class ConfigList(metaclass=ABCMeta): +class CtkConfigListView(BaseConfigListView, ABC): def __init__( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller, + ui_state, from_external_file: bool, attr_name: str = "", enable_key: str = "enabled", @@ -29,326 +21,51 @@ def __init__( default_config_name: str = "", add_button_text: str = "", add_button_tooltip: str = "", - is_full_width: bool = "", + is_full_width: bool = False, show_toggle_button: bool = False, ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.from_external_file = from_external_file - self.attr_name = attr_name - self.enable_key = enable_key - - self.config_dir = config_dir - self.default_config_name = default_config_name - - self.is_full_width = is_full_width - - # From search-concepts - self.filters = {"search": "", "type": "ALL", "show_disabled": True} - self.widgets_initialized = False - - # From master - self.toggle_button = None - self.show_toggle_button = show_toggle_button - self.is_opening_window = False - self._is_current_item_enabled = False - - self.master.grid_rowconfigure(0, weight=0) - self.master.grid_rowconfigure(1, weight=1) - self.master.grid_columnconfigure(0, weight=1) - - if self.from_external_file: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") - - self.configs_dropdown = None - self.element_list = None - - self.configs = [] - self.__load_available_config_names() - - self.current_config = getattr(self.train_config, self.attr_name) - self.widgets = [] - self.__load_current_config(getattr(self.train_config, self.attr_name)) - - self.__create_configs_dropdown() - components.button(self.top_frame, 0, 1, "Add Config", self.__add_config, tooltip="Adds a new config, which are containers for concepts, which themselves contain your dataset", width=20, padx=5) - components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, tooltip=add_button_tooltip, width=30, padx=5) - else: - self.top_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.top_frame.grid(row=0, column=0, sticky="nsew") - components.button(self.top_frame, 0, 2, add_button_text, self.__add_element, width=20, padx=5) - - self.current_config = getattr(self.train_config, self.attr_name) - - self.element_list = None - self._create_element_list() - - if show_toggle_button: - # tooltips break if you initialize with an empty string, default to a single space - self.toggle_button = components.button(self.top_frame, 0, 3, " ", self._toggle, tooltip="Disables/Enables all visible items in the current view", width=30, padx=5) - self._update_toggle_button_text() - - - - @abstractmethod - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - pass - - @abstractmethod - def create_new_element(self) -> BaseConfig: - pass - - @abstractmethod - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - pass - - def _refresh_show_disabled_text(self): - return - - def _reset_filters(self): # pragma: no cover - default noop - search_var = getattr(self, 'search_var', None) - filter_var = getattr(self, 'filter_var', None) - show_disabled_var = getattr(self, 'show_disabled_var', None) - - if search_var: - search_var.set("") - if filter_var: - filter_var.set("ALL") - if show_disabled_var: - show_disabled_var.set(True) - if search_var and hasattr(self, '_update_filters'): - self._update_filters() - - def _update_item_enabled_state(self): - # Only count items that match current filters - self._is_current_item_enabled = any( - item.ui_state.get_var(self.enable_key).get() - for i, item in enumerate(self.widgets) - if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]) + BaseConfigListView.__init__(self, ctk_components) + + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.build( + master, controller, ui_state, from_external_file, + attr_name=attr_name, + enable_key=enable_key, + config_dir=config_dir, + default_config_name=default_config_name, + add_button_text=add_button_text, + add_button_tooltip=add_button_tooltip, + is_full_width=is_full_width, + show_toggle_button=show_toggle_button, ) - def _update_toggle_button_text(self): - if not self.show_toggle_button: - return - self._update_item_enabled_state() - if self.toggle_button is not None: - self.toggle_button.configure(text="Disable" if self._is_current_item_enabled else "Enable") - - def _toggle(self): - self._toggle_items() - - def _toggle_items(self): - enable_state = not self._is_current_item_enabled - - # Only toggle items that match current filters - for i, widget in enumerate(self.widgets): - if i < len(self.current_config) and self._element_matches_filters(self.current_config[i]): - widget.ui_state.get_var(self.enable_key).set(enable_state) - self.save_current_config() - - self._update_widget_visibility() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.destroy() - - self.configs_dropdown = components.options_kv( - self.top_frame, 0, 0, self.configs, self.ui_state, self.attr_name, self.__load_current_config - ) - self._update_toggle_button_text() - - def _create_element_list(self, **filters): - if not self.from_external_file: - self.current_config = getattr(self.train_config, self.attr_name) - - self.filters.update(filters) - - if not self.widgets_initialized: - self._initialize_all_widgets() - self.widgets_initialized = True - - self._update_widget_visibility() - self._update_toggle_button_text() - - def _initialize_all_widgets(self): - self.widgets = [] - if self.element_list is not None: - self.element_list.destroy() - - self.element_list = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.element_list.grid(row=1, column=0, sticky="nsew") + def _create_top_frame(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew") + return frame + def _create_element_list_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid(row=1, column=0, sticky="nsew") if self.is_full_width: - self.element_list.grid_columnconfigure(0, weight=1) - - for i, element in enumerate(self.current_config): - widget = self.create_widget( - self.element_list, element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - - def _update_widget_visibility(self): - visible_index = 0 - - for i, widget in enumerate(self.widgets): - if i < len(self.current_config): - element = self.current_config[i] - - if self._element_matches_filters(element): - widget.visible_index = visible_index - widget.place_in_list() - visible_index += 1 - else: - widget.grid_remove() - - def __load_available_config_names(self): - if os.path.isdir(self.config_dir): - for path in os.listdir(self.config_dir): - path = path_util.canonical_join(self.config_dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - - if len(self.configs) == 0: - name = self.default_config_name.removesuffix(".json") - self.__create_config(name) - self.save_current_config() - - def __create_config(self, name: str): - name = path_util.safe_filename(name) - path = path_util.canonical_join(self.config_dir, f"{name}.json") - self.configs.append((name, path)) - self.__create_configs_dropdown() - - def __add_config(self): - dialogs.StringInputDialog(self.master, "name", "Name", self.__create_config) - - def __add_element(self): - new_element = self.create_new_element() - self.current_config.append(new_element) - # incremental insertion if widgets already initialized, else fall back to full rebuild - if self.widgets_initialized and self.element_list is not None: - i = len(self.current_config) - 1 - widget = self.create_widget( - self.element_list, new_element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - self._update_widget_visibility() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - def __clone_element(self, clone_i, modify_element_fun=None): - new_element = copy.deepcopy(self.current_config[clone_i]) - - if modify_element_fun is not None: - new_element = modify_element_fun(new_element) - self.current_config.append(new_element) - if self.widgets_initialized and self.element_list is not None: - i = len(self.current_config) - 1 - widget = self.create_widget( - self.element_list, new_element, i, - self.__open_element_window, - self.__remove_element, - self.__clone_element, - self.save_current_config - ) - self.widgets.append(widget) - self._update_widget_visibility() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - def __remove_element(self, remove_i): - self.current_config.pop(remove_i) - if self.widgets_initialized and 0 <= remove_i < len(self.widgets): - removed = self.widgets.pop(remove_i) - with contextlib.suppress(tk.TclError, AttributeError): - removed.destroy() - # Reindex remaining widgets - for idx, widget in enumerate(self.widgets): - widget.i = idx - self._update_widget_visibility() - else: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - - def __load_current_config(self, filename): - try: - with open(filename, "r") as f: - self.current_config = [] - - loaded_config_json = json.load(f) - for element_json in loaded_config_json: - element = self.create_new_element().from_dict(element_json) - self.current_config.append(element) - except (FileNotFoundError, json.JSONDecodeError) as e: - print(f"Failed to load config from {filename}: {e}") - self.current_config = [] - - # reset filters when switching configs - if hasattr(self, '_reset_filters') and self.widgets_initialized: - self._reset_filters() - - self.widgets_initialized = False - self._create_element_list() - self._update_toggle_button_text() - - def save_current_config(self): - if self.from_external_file: - try: - if not os.path.exists(self.config_dir): - os.makedirs(self.config_dir, exist_ok=True) - - write_json_atomic( - getattr(self.train_config, self.attr_name), - [element.to_dict() for element in self.current_config] - ) - except (OSError) as e: - print(f"Failed to save config: {e}") + frame.grid_columnconfigure(0, weight=1) + return frame - self._update_toggle_button_text() + def _wait_for_window(self, window): + self.master.wait_window(window) - if self.widgets_initialized: - try: - self._update_widget_visibility() - except (tk.TclError, AttributeError) as e: - print.debug(f"Widget visibility update failed: {e}") + def _remove_widget_from_layout(self, widget): + widget.grid_remove() - # let subclass refresh any show-disabled UI - if hasattr(self, '_refresh_show_disabled_text'): - self._refresh_show_disabled_text() + def _destroy_widget(self, widget): + with contextlib.suppress(AttributeError): + widget.destroy() - def _element_matches_filters(self, element): - return True # Show all by default + def _destroy_frame(self, frame): + frame.destroy() - def __open_element_window(self, i, ui_state): - if self.is_opening_window: - return - self.is_opening_window = True - try: - window = self.open_element_window(i, ui_state) - self.master.wait_window(window) - try: - if self.widgets is not None and 0 <= i < len(self.widgets): - self.widgets[i].configure_element() - except Exception: - self.widgets_initialized = False - self._create_element_list() - self.save_current_config() - finally: - self.is_opening_window = False + def _show_name_dialog(self, callback): + dialogs.StringInputDialog(self.master, "name", "Name", callback) diff --git a/modules/ui/CtkConvertModelUIView.py b/modules/ui/CtkConvertModelUIView.py index 6cb1b507a..782637348 100644 --- a/modules/ui/CtkConvertModelUIView.py +++ b/modules/ui/CtkConvertModelUIView.py @@ -1,33 +1,18 @@ -import traceback -from uuid import uuid4 - -from modules.util import create -from modules.util.args.ConvertModelArgs import ConvertModelArgs -from modules.util.config.TrainConfig import QuantizationConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ModelNames import EmbeddingName, ModelNames -from modules.util.torch_util import torch_gc -from modules.util.ui import components +from modules.ui.BaseConvertModelUIView import BaseConvertModelUIView +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class ConvertModelUI(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent - - self.parent = parent - self.convert_model_args = ConvertModelArgs.default_values() - self.ui_state = UIState(self, self.convert_model_args) - self.button = None +class CtkConvertModelUIView(BaseConvertModelUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: ConvertModelUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConvertModelUIView.__init__(self, ctk_components) + ui_state = CtkUIState(self, controller.convert_model_args) self.title("Convert models") self.geometry("550x350") @@ -38,133 +23,12 @@ def __init__(self, parent, *args, **kwargs): self.frame.grid_columnconfigure(0, weight=0) self.frame.grid_columnconfigure(1, weight=1) - self.main_frame(self.frame) + self.build_content(self.frame, controller, ui_state) self.frame.pack(fill="both", expand=True) self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def main_frame(self, master): - # model type - components.label(master, 0, 0, "Model Type", - tooltip="Type of the model") - components.options_kv(master, 0, 1, [ #TODO simplify - ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15), - ("Stable Diffusion 1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("Stable Diffusion 2.0", ModelType.STABLE_DIFFUSION_20), - ("Stable Diffusion 2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("Stable Diffusion 2.1", ModelType.STABLE_DIFFUSION_21), - ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3), - ("Stable Diffusion 3.5", ModelType.STABLE_DIFFUSION_35), - ("Stable Diffusion XL 1.0 Base", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("Stable Diffusion XL 1.0 Base Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2", ModelType.FLUX_2), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here - ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here - ("ZImage", ModelType.Z_IMAGE), - ], self.ui_state, "model_type") - - # training method - components.label(master, 1, 0, "Model Type", - tooltip="The type of model to convert") - components.options_kv(master, 1, 1, [ - ("Base Model", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ], self.ui_state, "training_method") - - # input name - components.label(master, 2, 0, "Input name", - tooltip="Filename, directory or hugging face repository of the base model") - components.path_entry( - master, 2, 1, self.ui_state, "input_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # output data type - components.label(master, 3, 0, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "output_dtype") - - # output format - components.label(master, 4, 0, "Output Format", - tooltip="Format to use when saving the output model") - components.options_kv(master, 4, 1, [ - ("Safetensors", ModelFormat.SAFETENSORS), - ("Diffusers", ModelFormat.DIFFUSERS), - ], self.ui_state, "output_model_format") - - # output model destination - components.label(master, 5, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - master, 5, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) - - self.button = components.button(master, 6, 1, "Convert", self.convert_model) - - def convert_model(self): - try: - self.button.configure(state="disabled") - model_loader = create.create_model_loader( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - model_saver = create.create_model_saver( - model_type=self.convert_model_args.model_type, - training_method=self.convert_model_args.training_method - ) - - print("Loading model " + self.convert_model_args.input_name) - if self.convert_model_args.training_method in [TrainingMethod.FINE_TUNE]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=self.convert_model_args.input_name, - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - elif self.convert_model_args.training_method in [TrainingMethod.LORA, TrainingMethod.EMBEDDING]: - model = model_loader.load( - model_type=self.convert_model_args.model_type, - model_names=ModelNames( - base_model=None, - lora=self.convert_model_args.input_name, - embedding=EmbeddingName(str(uuid4()), self.convert_model_args.input_name), - ), - weight_dtypes=self.convert_model_args.weight_dtypes(), - quantization=QuantizationConfig.default_values(), - ) - else: - raise Exception("could not load model: " + self.convert_model_args.input_name) - - print("Saving model " + self.convert_model_args.output_model_destination) - model_saver.save( - model=model, - model_type=self.convert_model_args.model_type, - output_model_format=self.convert_model_args.output_model_format, - output_model_destination=self.convert_model_args.output_model_destination, - dtype=self.convert_model_args.output_dtype.torch_dtype(), - ) - print("Model converted") - except Exception: - traceback.print_exc() - - torch_gc() - self.button.configure(state="normal") + def set_converting(self, active): + self.button.configure(state="disabled" if active else "normal") diff --git a/modules/ui/CtkGenerateCaptionsWindowView.py b/modules/ui/CtkGenerateCaptionsWindowView.py index 1690879f1..09d82f74b 100644 --- a/modules/ui/CtkGenerateCaptionsWindowView.py +++ b/modules/ui/CtkGenerateCaptionsWindowView.py @@ -2,27 +2,22 @@ import tkinter as tk from tkinter import filedialog +from modules.ui.BaseGenerateCaptionsWindowView import BaseGenerateCaptionsWindowView +from modules.ui.GenerateCaptionsWindowController import GenerateCaptionsWindowController from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - self.parent = parent +class CtkGenerateCaptionsWindowView(BaseGenerateCaptionsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateCaptionsWindowController, path, parent_include_subdirectories, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) if path is None: path = "" + self.controller = controller + self.mode_var = ctk.StringVar(self, "Create if absent") self.modes = ["Replace all captions", "Create if absent", "Add as new line"] self.model_var = ctk.StringVar(self, "Blip") @@ -79,7 +74,7 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self._on_create_captions) self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) self.frame.pack(fill="both", expand=True) @@ -89,7 +84,6 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.focus_set() self.after(200, lambda: set_window_icon(self)) - def browse_for_path(self, entry_box): # get the path from the user path = filedialog.askdirectory() @@ -106,25 +100,16 @@ def set_progress(self, value, max_value): self.progress_label.configure(text=f"{value}/{max_value}") self.progress.update() - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) - - mode = { - "Replace all captions": "replace", - "Create if absent": "fill", - "Add as new line": "add", - }[self.mode_var.get()] - - self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), + def _on_create_captions(self): + self.controller.create_captions( + model_name=self.model_var.get(), + path=self.path_entry.get(), initial_caption=self.caption_entry.get(), caption_prefix=self.prefix_entry.get(), caption_postfix=self.postfix_entry.get(), - mode=mode, - progress_callback=self.set_progress, + mode_str=self.mode_var.get(), include_subdirectories=self.include_subdirectories_var.get(), ) - self.parent.load_image() def destroy(self): with contextlib.suppress(tk.TclError): diff --git a/modules/ui/CtkGenerateMasksWindowView.py b/modules/ui/CtkGenerateMasksWindowView.py index daff0d3d5..631179fac 100644 --- a/modules/ui/CtkGenerateMasksWindowView.py +++ b/modules/ui/CtkGenerateMasksWindowView.py @@ -2,13 +2,15 @@ import tkinter as tk from tkinter import filedialog +from modules.ui.BaseGenerateMasksWindowView import BaseGenerateMasksWindowView +from modules.ui.GenerateMasksWindowController import GenerateMasksWindowController from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): +class CtkGenerateMasksWindowView(BaseGenerateMasksWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateMasksWindowController, path, parent_include_subdirectories, *args, **kwargs): """ Window for generating masks for a folder of images @@ -17,9 +19,9 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) path (`str`): the path to the folder parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox """ - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - self.parent = parent + self.controller = controller if path is None: path = "" @@ -93,7 +95,7 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self._on_create_masks) self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) self.frame.pack(fill="both", expand=True) @@ -103,7 +105,6 @@ def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs) self.focus_set() self.after(200, lambda: set_window_icon(self)) - def browse_for_path(self, entry_box): # get the path from the user path = filedialog.askdirectory() @@ -120,29 +121,18 @@ def set_progress(self, value, max_value): self.progress_label.configure(text=f"{value}/{max_value}") self.progress.update() - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) - - mode = { - "Replace all masks": "replace", - "Create if absent": "fill", - "Add to existing": "add", - "Subtract from existing": "subtract", - "Blend with existing": "blend", - }[self.mode_var.get()] - - self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], - mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, + def _on_create_masks(self): + self.controller.create_masks( + model_name=self.model_var.get(), + path=self.path_entry.get(), + prompt=self.prompt_entry.get(), + mode_str=self.mode_var.get(), + alpha_str=self.alpha_entry.get(), + threshold_str=self.threshold_entry.get(), + smooth_str=self.smooth_entry.get(), + expand_str=self.expand_entry.get(), include_subdirectories=self.include_subdirectories_var.get(), ) - self.parent.load_image() def destroy(self): with contextlib.suppress(tk.TclError): diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index 1c73d90ce..8caa1f171 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -1,25 +1,20 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType +from modules.ui.BaseLoraTabView import BaseLoraTabView +from modules.ui.LoraTabController import LoraTabController from modules.util.enum.ModelType import PeftType -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation_helpers import check_range +from modules.util.ui import ctk_components import customtkinter as ctk -class LoraTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() - +class CtkLoraTabView(BaseLoraTabView): + def __init__(self, master, controller: LoraTabController, ui_state): + BaseLoraTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state - self.scroll_frame = None self.options_frame = None - self.refresh_ui() def refresh_ui(self): @@ -27,128 +22,20 @@ def refresh_ui(self): self.scroll_frame.destroy() self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=0) self.scroll_frame.grid_columnconfigure(1, weight=1) self.scroll_frame.grid_columnconfigure(2, weight=2) - - components.label(self.scroll_frame, 0, 0, "Type", - tooltip="The type of low-parameter finetuning method.") - # This will instantly call self.setup_lora. - components.options_kv(self.scroll_frame, 0, 1, [ - ("LoRA", PeftType.LORA), - ("LoHa", PeftType.LOHA), - ("OFT v2", PeftType.OFT_2), - ], self.ui_state, "peft_type", command=self.setup_lora) + self.build(self.scroll_frame, self.controller, self.ui_state, self.setup_lora) def setup_lora(self, peft_type: PeftType): - if peft_type == PeftType.LOHA: - name = "LoHa" - elif peft_type == PeftType.OFT_2: - name = "OFT v2" - else: - name = "LoRA" - if self.options_frame: self.options_frame.destroy() self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") master = self.options_frame - master.grid_columnconfigure(0, weight=0, uniform="a") master.grid_columnconfigure(1, weight=1, uniform="a") master.grid_columnconfigure(2, minsize=50, uniform="a") master.grid_columnconfigure(3, weight=0, uniform="a") master.grid_columnconfigure(4, weight=1, uniform="a") - - # lora model name - components.label(master, 0, 0, f"{name} base model", - tooltip=f"The base {name} to train on. Leave empty to create a new {name}") - entry = components.path_entry( - master, 0, 1, self.ui_state, "lora_model_name", - mode="file", path_modifier=components.json_path_modifier - ) - entry.grid(row=0, column=1, columnspan=4) - - - # LoRA decomposition - if peft_type == PeftType.LORA: - components.label(master, 1, 3, "Decompose Weights (DoRA)", - tooltip="Decompose LoRA Weights (aka, DoRA).") - components.switch(master, 1, 4, self.ui_state, "lora_decompose") - - components.label(master, 2, 3, "Use Norm Epsilon (DoRA Only)", - tooltip="Add an epsilon to the norm divison calculation in DoRA. Can aid in training stability, and also acts as regularization.") - components.switch(master, 2, 4, self.ui_state, "lora_decompose_norm_epsilon") - components.label(master, 3, 3, "Apply on output axis (DoRA Only)", - tooltip="Apply the weight decomposition on the output axis instead of the input axis.") - components.switch(master, 3, 4, self.ui_state, "lora_decompose_output_axis") - - # LoRA and LoHA shared settings - if peft_type == PeftType.LORA or peft_type == PeftType.LOHA: - # rank - components.label(master, 1, 0, f"{name} rank", - tooltip=f"The rank parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "lora_rank", required=True, extra_validate=check_range(lower=1, message="Rank must be at least 1")) - - # alpha - components.label(master, 2, 0, f"{name} alpha", - tooltip=f"The alpha parameter used when creating a new {name}") - components.entry(master, 2, 1, self.ui_state, "lora_alpha", required=True) - - # Dropout Percentage - components.label(master, 3, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 3, 1, self.ui_state, "dropout_probability") - - # weight dtype - components.label(master, 4, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 4, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") - - # For use with additional embeddings. - components.label(master, 5, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 5, 1, self.ui_state, "bundle_additional_embeddings") - - # OFTv2 - elif peft_type == PeftType.OFT_2: - # Block Size - components.label(master, 1, 0, f"{name} Block Size", - tooltip=f"The block size parameter used when creating a new {name}") - components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) - - # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") - - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") - - # Block Share - components.label(master, 3, 3, "Block Share", - tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") - - # Dropout Percentage - components.label(master, 2, 0, "Dropout Probability", - tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") - components.entry(master, 2, 1, self.ui_state, "dropout_probability") - - # OFT weight dtype - components.label(master, 3, 0, f"{name} Weight Data Type", - tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(master, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "lora_weight_dtype") - - # For use with additional embeddings. - components.label(master, 4, 0, "Bundle Embeddings", - tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") - components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + self.build_lora_options(master, self.controller, self.ui_state, peft_type) diff --git a/modules/ui/CtkModelTabView.py b/modules/ui/CtkModelTabView.py index ff17ea3ba..5b43b7dca 100644 --- a/modules/ui/CtkModelTabView.py +++ b/modules/ui/CtkModelTabView.py @@ -1,24 +1,17 @@ -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.ConfigPart import ConfigPart -from modules.util.enum.DataType import DataType -from modules.util.enum.ModelFormat import ModelFormat -from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import components -from modules.util.ui.UIState import UIState - -import customtkinter as ctk +from modules.ui.BaseModelTabView import BaseModelTabView +from modules.ui.ModelTabController import ModelTabController +from modules.util.ui import ctk_components -class ModelTab: +import customtkinter as ctk - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() +class CtkModelTabView(BaseModelTabView): + def __init__(self, master, controller: ModelTabController, ui_state): + BaseModelTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state master.grid_rowconfigure(0, weight=1) @@ -28,6 +21,13 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): self.refresh_ui() + def _make_svd_frames(self, parent, row: int): + svd_label_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_label_frame.grid(row=row, column=3, sticky="nsew") + svd_entry_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_entry_frame.grid(row=row, column=4, sticky="nsew") + return svd_label_frame, svd_entry_frame + def refresh_ui(self): if self.scroll_frame: self.scroll_frame.destroy() @@ -40,649 +40,9 @@ def refresh_ui(self): base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") base_frame.grid_columnconfigure(0, weight=0) - base_frame.grid_columnconfigure(1, weight=10)#, minsize=500) + base_frame.grid_columnconfigure(1, weight=10) # , minsize=500) base_frame.grid_columnconfigure(2, minsize=50) base_frame.grid_columnconfigure(3, weight=0) base_frame.grid_columnconfigure(4, weight=1) - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ - TrainingMethod.FINE_TUNE, - TrainingMethod.FINE_TUNE_VAE, - ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_3_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_2_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_z_image_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_ernie_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_chroma_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_qwen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_xl_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_wuerstchen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), - has_text_encoder=True, - ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_pixart_alpha_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_sana_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hunyuan_video_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hi_dream_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_text_encoder_4=True, - allow_override_text_encoder_4=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: - options = [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ("float16", DataType.FLOAT_16), - ("float8 (W8)", DataType.FLOAT_8), - # ("int8", DataType.INT_8), # TODO: reactivate when the int8 implementation is fixed in bitsandbytes: https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1332 - ("nfloat4", DataType.NFLOAT_4), - ] - if include_a8: - options += [ - ("float W8A8", DataType.FLOAT_W8A8), - ("int W8A8", DataType.INT_W8A8), - ] - - if include_gguf: - options.append(("GGUF", DataType.GGUF)) - if include_a8: - options += [ - ("GGUF A8 float", DataType.GGUF_A8_FLOAT), - ("GGUF A8 int", DataType.GGUF_A8_INT), - ] - - return options - - def __create_base_dtype_components(self, frame, row: int) -> int: - # huggingface token - components.label(frame, row, 0, "Hugging Face Token", - tooltip="Enter your Hugging Face access token if you have used a protected Hugging Face repository below.\nThis value is stored separately, not saved to your configuration file. " - "Go to https://huggingface.co/settings/tokens to create an access token.", - wide_tooltip=True) - components.entry(frame, row, 1, self.ui_state, "secrets.huggingface_token") - - row += 1 - - # base model - components.label(frame, row, 0, "Base Model", - tooltip="Filename, directory or Hugging Face repository of the base model") - components.path_entry( - frame, row, 1, self.ui_state, "base_model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # compile - components.label(frame, row, 3, "Compile transformer blocks", - tooltip="Uses torch.compile and Triton to significantly speed up training. Only applies to transformer/unet. Disable in case of compatibility issues.") - components.switch(frame, row, 4, self.ui_state, "compile") - - row += 1 - - return row - - def __create_base_components( - self, - frame, - row: int, - has_unet: bool = False, - has_prior: bool = False, - allow_override_prior: bool = False, - has_transformer: bool = False, - allow_override_transformer: bool = False, - allow_override_text_encoder_4: bool = False, - has_text_encoder: bool = False, - has_text_encoder_1: bool = False, - has_text_encoder_2: bool = False, - has_text_encoder_3: bool = False, - has_text_encoder_4: bool = False, - has_vae: bool = False, - ) -> int: - if has_unet: - # unet weight dtype - components.label(frame, row, 3, "UNet Data Type", - tooltip="The unet weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), - self.ui_state, "unet.weight_dtype") - - row += 1 - - if has_prior: - if allow_override_prior: - # prior model - components.label(frame, row, 0, "Prior Model", - tooltip="Filename, directory or Hugging Face repository of the prior model") - components.path_entry( - frame, row, 1, self.ui_state, "prior.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # prior weight dtype - components.label(frame, row, 3, "Prior Data Type", - tooltip="The prior weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "prior.weight_dtype") - - row += 1 - - if has_transformer: - if allow_override_transformer: - # transformer model - components.label(frame, row, 0, "Override Transformer / GGUF", - tooltip="Can be used to override the transformer in the base model. Safetensors and GGUF files are supported, local and on Huggingface. If a GGUF file is used, the DataType must also be set to GGUF") - components.path_entry( - frame, row, 1, self.ui_state, "transformer.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # transformer weight dtype - components.label(frame, row, 3, "Transformer Data Type", - tooltip="The transformer weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), - self.ui_state, "transformer.weight_dtype") - - row += 1 - - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - - components.label(frame, row, 0, "Quantization") - components.layer_filter_entry(frame, row, 1, self.ui_state, - preset_var_name="quantization.layer_filter_preset", presets=presets, - preset_label="Quantization Layer Filter", - preset_tooltip="Select a preset defining which layers to quantize. Quantization of certain layers can decrease model quality. Only applies to the transformer/unet", - entry_var_name="quantization.layer_filter", - entry_tooltip="Comma-separated list of layers to quantize. Regular expressions (if toggled) are supported. Any model layer with a matching name will be quantized", - regex_var_name="quantization.layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - frame_color="transparent", - ) - - # SVDQuant - create vertical grids to match the size of layer_filter_entry - svd_label_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_label_frame.grid(row=row, column=3, sticky="nsew") - svd_entry_frame = ctk.CTkFrame(frame, fg_color="transparent") - svd_entry_frame.grid(row=row, column=4, sticky="nsew") - components.label(svd_label_frame, 0, 0, "SVDQuant", - tooltip="What datatype to use for SVDQuant weights decomposition.") - components.options_kv(svd_entry_frame, 0, 0, [("disabled", DataType.NONE), ("float32", DataType.FLOAT_32), ("bfloat16", DataType.BFLOAT_16)], - self.ui_state, "quantization.svd_dtype") - components.label(svd_label_frame, 1, 0, "SVDQuant Rank", - tooltip="Rank for SVDQuant weights decomposition") - components.entry(svd_entry_frame, 1, 0, self.ui_state, "quantization.svd_rank") - row += 1 - - - if has_text_encoder: - # text encoder weight dtype - components.label(frame, row, 3, "Text Encoder Data Type", - tooltip="The text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") - - row += 1 - - if has_text_encoder_1: - # text encoder 1 weight dtype - components.label(frame, row, 3, "Text Encoder 1 Data Type", - tooltip="The text encoder 1 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder.weight_dtype") - - row += 1 - - if has_text_encoder_2: - # text encoder 2 weight dtype - components.label(frame, row, 3, "Text Encoder 2 Data Type", - tooltip="The text encoder 2 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_2.weight_dtype") - - row += 1 - - if has_text_encoder_3: - # text encoder 3 weight dtype - components.label(frame, row, 3, "Text Encoder 3 Data Type", - tooltip="The text encoder 3 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_3.weight_dtype") - - row += 1 - - if has_text_encoder_4: - if allow_override_text_encoder_4: - # text encoder 4 weight dtype - components.label(frame, row, 0, "Text Encoder 4 Override", - tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") - components.path_entry( - frame, row, 1, self.ui_state, "text_encoder_4.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # text encoder 4 weight dtype - components.label(frame, row, 3, "Text Encoder 4 Data Type", - tooltip="The text encoder 4 weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "text_encoder_4.weight_dtype") - - row += 1 - - if has_vae: - # base model - components.label(frame, row, 0, "VAE Override", - tooltip="Directory or Hugging Face repository of a VAE model in diffusers format. Can be used to override the VAE included in the base model. Using a safetensor VAE file will cause an error that the model cannot be loaded.") - components.path_entry( - frame, row, 1, self.ui_state, "vae.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # vae weight dtype - components.label(frame, row, 3, "VAE Data Type", - tooltip="The vae weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "vae.weight_dtype") - - row += 1 - - return row - - def __create_effnet_encoder_components(self, frame, row: int): - # effnet encoder model - components.label(frame, row, 0, "Effnet Encoder Model", - tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") - components.path_entry( - frame, row, 1, self.ui_state, "effnet_encoder.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # effnet encoder weight dtype - components.label(frame, row, 3, "Effnet Encoder Data Type", - tooltip="The effnet encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "effnet_encoder.weight_dtype") - - row += 1 - - return row - - def __create_decoder_components( - self, - frame, - row: int, - has_text_encoder: bool, - ) -> int: - # decoder model - components.label(frame, row, 0, "Decoder Model", - tooltip="Filename, directory or Hugging Face repository of the decoder model") - components.path_entry( - frame, row, 1, self.ui_state, "decoder.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # decoder weight dtype - components.label(frame, row, 3, "Decoder Data Type", - tooltip="The decoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder.weight_dtype") - - row += 1 - - if has_text_encoder: - # decoder text encoder weight dtype - components.label(frame, row, 3, "Decoder Text Encoder Data Type", - tooltip="The decoder text encoder weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_text_encoder.weight_dtype") - - row += 1 - - # decoder vqgan weight dtype - components.label(frame, row, 3, "Decoder VQGAN Data Type", - tooltip="The decoder vqgan weight data type") - components.options_kv(frame, row, 4, self.__create_dtype_options(), - self.ui_state, "decoder_vqgan.weight_dtype") - - row += 1 - - return row - - def __create_output_components( - self, - frame, - row: int, - allow_safetensors: bool = False, - allow_diffusers: bool = False, - allow_legacy_safetensors: bool = False, - allow_comfy: bool = False, - ) -> int: - # output model destination - components.label(frame, row, 0, "Model Output Destination", - tooltip="Filename or directory where the output model is saved") - components.path_entry( - frame, row, 1, self.ui_state, "output_model_destination", - mode="file", - io_type=PathIOType.MODEL, - ) - - # output data type - components.label(frame, row, 3, "Output Data Type", - tooltip="Precision to use when saving the output model") - components.options_kv(frame, row, 4, [ - ("float16", DataType.FLOAT_16), - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ("float8", DataType.FLOAT_8), - ("nfloat4", DataType.NFLOAT_4), - ], self.ui_state, "output_dtype") - - row += 1 - - # output format - formats = [] - if allow_safetensors: - formats.append(("Safetensors", ModelFormat.SAFETENSORS)) - if allow_diffusers: - formats.append(("Diffusers", ModelFormat.DIFFUSERS)) - # if allow_legacy_safetensors: - # formats.append(("Legacy Safetensors", ModelFormat.LEGACY_SAFETENSORS)) - if allow_comfy: - formats.append(("Comfy LoRA", ModelFormat.COMFY_LORA)) - - components.label(frame, row, 0, "Output Format", - tooltip="Format to use when saving the output model") - components.options_kv(frame, row, 1, formats, self.ui_state, "output_model_format") - - # include config - components.label(frame, row, 3, "Include Config", - tooltip="Include the training configuration in the final model. Only supported for safetensors files. " - "None: No config is included. " - "Settings: All training settings are included. " - "All: All settings, including the samples and concepts are included.") - components.options_kv(frame, row, 4, [ - ("None", ConfigPart.NONE), - ("Settings", ConfigPart.SETTINGS), - ("All", ConfigPart.ALL), - ], self.ui_state, "include_train_config") - - row += 1 - - return row + self.build_content(base_frame, self.controller, self.ui_state) diff --git a/modules/ui/CtkMuonAdamWindowView.py b/modules/ui/CtkMuonAdamWindowView.py index 5879ab432..3dc48ab74 100644 --- a/modules/ui/CtkMuonAdamWindowView.py +++ b/modules/ui/CtkMuonAdamWindowView.py @@ -1,42 +1,17 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS -from modules.util.ui import components +from modules.ui.BaseMuonAdamWindowView import BaseMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -MUON_AUX_ADAM_DEFAULTS = { - "beta1": 0.9, - "beta2": 0.999, - "eps": 1e-8, - "weight_decay": 0.0, -} -class MuonAdamWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state: UIState, - parent_optimizer_type: Optimizer, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.adam_ui_state = ui_state - self.parent_optimizer_type = parent_optimizer_type - - if self.parent_optimizer_type == Optimizer.MUON: - self.title("Muon's Auxiliary AdamW Settings") - self.adam_params_def = MUON_AUX_ADAM_DEFAULTS - else: - self.title("Muon_adv's Auxiliary AdamW_adv Settings") - self.adam_params_def = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] +class CtkMuonAdamWindowView(BaseMuonAdamWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: MuonAdamWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseMuonAdamWindowView.__init__(self, ctk_components) + self.title(controller.get_title()) self.geometry("800x500") self.resizable(True, True) @@ -44,64 +19,19 @@ def __init__( self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) - components.button(self, 1, 0, "ok", command=self.destroy) - self.create_adam_params_ui(self.frame) + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def create_adam_params_ui(self, master): - # This is a large map, copied from OptimizerParamsWindow for simplicity. - # @formatter:off - KEY_DETAIL_MAP = { - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - } - # @formatter:on - - adam_params = self.adam_params_def - - for index, key in enumerate(adam_params.keys()): - if key not in KEY_DETAIL_MAP: - continue - - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - param_type = arg_info['type'] - - row = index // 2 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if param_type != 'bool': - components.entry(master, row, col + 1, self.adam_ui_state, key) - else: - components.switch(master, row, col + 1, self.adam_ui_state, key) diff --git a/modules/ui/CtkOffloadingWindowView.py b/modules/ui/CtkOffloadingWindowView.py index 54035e121..b752f1602 100644 --- a/modules/ui/CtkOffloadingWindowView.py +++ b/modules/ui/CtkOffloadingWindowView.py @@ -1,75 +1,31 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components +from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None +class CtkOffloadingWindowView(BaseOffloadingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OffloadingWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOffloadingWindowView.__init__(self, ctk_components) self.title("Offloading") self.geometry("800x400") self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) self.grid_columnconfigure(0, weight=1) - frame = self.__content_frame(self) + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid_columnconfigure(0, weight=1) + frame.grid_columnconfigure(1, weight=1) + self.build_content(frame, controller, ui_state) + frame.pack(fill="both", expand=1) frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) + self.components.button(self, 1, 0, "ok", self.destroy) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkOptimizerParamsWindowView.py b/modules/ui/CtkOptimizerParamsWindowView.py index 16063c26c..ecc1f6a38 100644 --- a/modules/ui/CtkOptimizerParamsWindowView.py +++ b/modules/ui/CtkOptimizerParamsWindowView.py @@ -1,38 +1,27 @@ import contextlib from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow -from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig -from modules.util.enum.Optimizer import Optimizer -from modules.util.optimizer_util import ( - OPTIMIZER_DEFAULT_PARAMETERS, - change_optimizer, - load_optimizer_defaults, - update_optimizer_config, -) -from modules.util.ui import components +from modules.ui.BaseOptimizerParamsWindowView import BaseOptimizerParamsWindowView +from modules.ui.CtkMuonAdamWindowView import CtkMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) +class CtkOptimizerParamsWindowView(BaseOptimizerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OptimizerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOptimizerParamsWindowView.__init__(self, ctk_components) - self.parent = parent - self.train_config = train_config + self.controller = controller self.ui_state = ui_state self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) self.muon_adam_button = None + self.protocol("WM_DELETE_WINDOW", self.on_window_close) self.title("Optimizer Settings") self.geometry("800x500") @@ -51,206 +40,40 @@ def __init__( self.frame.grid_columnconfigure(3, weight=0) self.frame.grid_columnconfigure(4, weight=1) - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) + self.components.button(self, 1, 0, "ok", command=self.on_window_close) + self.build_content(self.frame, controller, ui_state, self.optimizer_ui_state, + self.on_optimizer_change, self._load_defaults) + self._rebuild_dynamic_ui() self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - def main_frame(self, master): - # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") - - # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) - - # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): + def _rebuild_dynamic_ui(self): with contextlib.suppress(TclError): - for widget in master.winfo_children(): + for widget in self.frame.winfo_children(): grid_info = widget.grid_info() if int(grid_info["row"]) >= 1: widget.destroy() - def create_dynamic_ui( - self, - master, - ): - - # Lookup for the title and tooltip for a key - # @formatter:off - KEY_DETAIL_MAP = { - 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, - 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, - 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, - 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, - 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, - 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, - 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, - 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, - 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, - 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, - 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, - 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, - 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, - 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, - 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, - 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, - 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, - 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, - 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, - 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, - 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, - 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, - 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, - 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, - 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, - 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, - 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, - 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, - 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, - 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, - 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, - 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, - 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, - 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, - 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, - 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, - 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, - 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, - 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, - 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, - 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, - 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, - 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, - 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, - 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, - 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, - 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, - 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, - 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, - 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, - 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, - 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, - 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, - 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, - 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, - 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, - 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, - 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, - 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, - 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, - 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, - 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, - 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, - 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, - 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, - 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, - 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, - 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, - 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, - 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, - 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, - 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, - 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, - 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, - 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, - 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, - } - # @formatter:on - - if not self.winfo_exists(): # check if this window isn't open + if not self.winfo_exists(): return - selected_optimizer = self.train_config.optimizer.optimizer - - # Extract the keys for the selected optimizer - for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): - if key not in KEY_DETAIL_MAP: - continue - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - type = arg_info['type'] - - row = (index // 2) + 1 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) - - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, - tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() - elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) + self.build_dynamic_content(self.frame, self.controller, self.optimizer_ui_state, + self.update_user_pref, self.open_muon_adam_window) + self.toggle_muon_adam_button() def update_user_pref(self, *args): - update_optimizer_config(self.train_config) + self.controller.on_close() self.toggle_muon_adam_button() def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) + self.controller.restore_optimizer_config(self.ui_state) + self._rebuild_dynamic_ui() - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) + def _load_defaults(self, *args): + self.controller.load_defaults(self.ui_state) def on_window_close(self): self.destroy() @@ -261,28 +84,8 @@ def toggle_muon_adam_button(self): self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer - - adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config - - if current_optimizer == Optimizer.MUON: - defaults = MUON_AUX_ADAM_DEFAULTS - else: - defaults = OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] - - if current_state is None: - adam_config.from_dict(defaults) - if current_optimizer != Optimizer.MUON: - adam_config.optimizer = Optimizer.ADAMW_ADV - elif isinstance(current_state, dict): - adam_config.from_dict(current_state) - else: - # Should not happen if TrainConfig defines it as dict, but for safety - adam_config = current_state - - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) + adam_config, current_optimizer = self.controller.prepare_muon_adam_config() + temp_adam_ui_state = CtkUIState(self, adam_config) + window = CtkMuonAdamWindowView(self, MuonAdamWindowController(self.controller.config, current_optimizer), temp_adam_ui_state) self.wait_window(window) - - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + self.controller.save_muon_adam_config(adam_config) diff --git a/modules/ui/CtkProfilingWindowView.py b/modules/ui/CtkProfilingWindowView.py index 8d298abe3..15d5055a0 100644 --- a/modules/ui/CtkProfilingWindowView.py +++ b/modules/ui/CtkProfilingWindowView.py @@ -1,16 +1,18 @@ -import faulthandler -from modules.util.ui import components +from modules.ui.BaseProfilingWindowView import BaseProfilingWindowView +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon import customtkinter as ctk -from scalene import scalene_profiler -class ProfilingWindow(ctk.CTkToplevel): - def __init__(self, parent, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - self.parent = parent +class CtkProfilingWindowView(BaseProfilingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: ProfilingWindowController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseProfilingWindowView.__init__(self, ctk_components) + + self._controller = controller self.title("Profiling") self.geometry("512x512") @@ -23,35 +25,25 @@ def __init__(self, parent, *args, **kwargs): self.grid_rowconfigure(2, weight=1) self.grid_columnconfigure(0, weight=1) - components.button(self, 0, 0, "Dump stack", self._dump_stack) - self._profile_button = components.button( - self, 1, 0, "Start Profiling", self._start_profiler, - tooltip="Turns on/off Scalene profiling. Only works when OneTrainer is launched with Scalene!") - # Bottom bar self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) self._bottom_bar.grid(row=2, column=0, sticky="sew") - self._message_label = components.label(self._bottom_bar, 0, 0, "Inactive") + + self.build_content(self, self._bottom_bar, controller) self.protocol("WM_DELETE_WINDOW", self.withdraw) self.withdraw() self.after(200, lambda: set_window_icon(self)) - def _dump_stack(self): - with open('stacks.txt', 'w') as f: - faulthandler.dump_traceback(f) - self._message_label.configure(text='Stack dumped to stacks.txt') - - def _end_profiler(self): - scalene_profiler.stop() - - self._message_label.configure(text='Inactive') - self._profile_button.configure(text='Start Profiling') - self._profile_button.configure(command=self._start_profiler) - - def _start_profiler(self): - scalene_profiler.start() - - self._message_label.configure(text='Profiling active...') - self._profile_button.configure(text='End Profiling') - self._profile_button.configure(command=self._end_profiler) + def set_message(self, text): + self._message_label.configure(text=text) + + def set_profiling_active(self, active): + if active: + self._message_label.configure(text='Profiling active...') + self._profile_button.configure(text='End Profiling') + self._profile_button.configure(command=self._controller.end_profiler) + else: + self._message_label.configure(text='Inactive') + self._profile_button.configure(text='Start Profiling') + self._profile_button.configure(command=self._controller.start_profiler) diff --git a/modules/ui/CtkSampleFrameView.py b/modules/ui/CtkSampleFrameView.py index 297caac29..167b25692 100644 --- a/modules/ui/CtkSampleFrameView.py +++ b/modules/ui/CtkSampleFrameView.py @@ -1,37 +1,28 @@ -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.enum.NoiseScheduler import NoiseScheduler -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseSampleFrameView import BaseSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.util.ui import ctk_components import customtkinter as ctk -class SampleFrame(ctk.CTkFrame): +class CtkSampleFrameView(BaseSampleFrameView, ctk.CTkFrame): def __init__( self, parent, - sample: SampleConfig, - ui_state: UIState, - model_type: ModelType, + controller: SampleFrameController, + ui_state, include_prompt: bool = True, include_settings: bool = True, ): ctk.CTkFrame.__init__(self, parent, fg_color="transparent") - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type - - is_flow_matching = model_type.is_flow_matching() - is_inpainting_model = model_type.has_conditioning_image_input() - is_video_model = model_type.is_video_model() + BaseSampleFrameView.__init__(self, ctk_components) if include_prompt and include_prompt: self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_columnconfigure(0, weight=1) + top_frame = None if include_prompt: top_frame = ctk.CTkFrame(self, fg_color="transparent") top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") @@ -39,6 +30,7 @@ def __init__( top_frame.grid_columnconfigure(0, weight=0) top_frame.grid_columnconfigure(1, weight=1) + bottom_frame = None if include_settings: bottom_frame = ctk.CTkFrame(self, fg_color="transparent") bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") @@ -48,87 +40,4 @@ def __init__( bottom_frame.grid_columnconfigure(2, weight=0) bottom_frame.grid_columnconfigure(3, weight=1) - if include_prompt: - # prompt - components.label(top_frame, 0, 0, "prompt:") - components.entry(top_frame, 0, 1, self.ui_state, "prompt") - - # negative prompt - components.label(top_frame, 1, 0, "negative prompt:") - components.entry(top_frame, 1, 1, self.ui_state, "negative_prompt") - - if include_settings: - # width - components.label(bottom_frame, 0, 0, "width:") - components.entry(bottom_frame, 0, 1, self.ui_state, "width") - - # height - components.label(bottom_frame, 0, 2, "height:") - components.entry(bottom_frame, 0, 3, self.ui_state, "height") - - if is_video_model: - # frames - components.label(bottom_frame, 1, 0, "frames:", - tooltip="Number of frames to generate. Only used when generating videos.") - components.entry(bottom_frame, 1, 1, self.ui_state, "frames") - - # length - components.label(bottom_frame, 1, 2, "length:", - tooltip="Length in seconds of audio output.") - components.entry(bottom_frame, 1, 3, self.ui_state, "length") - - # seed - components.label(bottom_frame, 2, 0, "seed:") - components.entry(bottom_frame, 2, 1, self.ui_state, "seed") - - # random seed - components.label(bottom_frame, 2, 2, "random seed:") - components.switch(bottom_frame, 2, 3, self.ui_state, "random_seed") - - # cfg scale - components.label(bottom_frame, 3, 0, "cfg scale:") - components.entry(bottom_frame, 3, 1, self.ui_state, "cfg_scale") - - # sampler - if not is_flow_matching: - components.label(bottom_frame, 4, 2, "sampler:") - components.options_kv(bottom_frame, 4, 3, [ - ("DDIM", NoiseScheduler.DDIM), - ("Euler", NoiseScheduler.EULER), - ("Euler A", NoiseScheduler.EULER_A), - # ("DPM++", NoiseScheduler.DPMPP), # TODO: produces noisy samples - # ("DPM++ SDE", NoiseScheduler.DPMPP_SDE), # TODO: produces noisy samples - ("UniPC", NoiseScheduler.UNIPC), - ("Euler Karras", NoiseScheduler.EULER_KARRAS), - ("DPM++ Karras", NoiseScheduler.DPMPP_KARRAS), - ("DPM++ SDE Karras", NoiseScheduler.DPMPP_SDE_KARRAS), - ("UniPC Karras", NoiseScheduler.UNIPC_KARRAS) - ], self.ui_state, "noise_scheduler") - - # steps - components.label(bottom_frame, 4, 0, "steps:") - components.entry(bottom_frame, 4, 1, self.ui_state, "diffusion_steps") - - # inpainting - if is_inpainting_model: - components.label(bottom_frame, 5, 0, "inpainting:", - tooltip="Enables inpainting sampling. Only available when sampling from an inpainting model.") - components.switch(bottom_frame, 5, 1, self.ui_state, "sample_inpainting") - - # base image path - components.label(bottom_frame, 6, 0, "base image path:", - tooltip="The base image used when inpainting.") - components.file_entry(bottom_frame, 6, 1, self.ui_state, "base_image_path", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) - - # mask image path - components.label(bottom_frame, 6, 2, "mask image path:", - tooltip="The mask used when inpainting.") - components.file_entry(bottom_frame, 6, 3, self.ui_state, "mask_image_path", - mode="file", - allow_model_files=False, - allow_image_files=True, - ) + self.build_content(top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings) diff --git a/modules/ui/CtkSampleParamsWindowView.py b/modules/ui/CtkSampleParamsWindowView.py index 2b0b3f3f1..8229a19f6 100644 --- a/modules/ui/CtkSampleParamsWindowView.py +++ b/modules/ui/CtkSampleParamsWindowView.py @@ -1,20 +1,17 @@ -from modules.ui.SampleFrame import SampleFrame -from modules.util.config.SampleConfig import SampleConfig -from modules.util.enum.ModelType import ModelType -from modules.util.ui import components +from modules.ui.BaseSampleParamsWindowView import BaseSampleParamsWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class SampleParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: ModelType | None = None, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.sample = sample - self.ui_state = ui_state - self.model_type = model_type +class CtkSampleParamsWindowView(BaseSampleParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SampleParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleParamsWindowView.__init__(self, ctk_components) self.title("Sample") self.geometry("800x500") @@ -24,16 +21,12 @@ def __init__(self, parent, sample: SampleConfig, ui_state: UIState, model_type: self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - frame = SampleFrame(self, self.sample, self.ui_state, model_type=model_type) + frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - components.button(self, 1, 0, "ok", self.__ok) + self.components.button(self, 1, 0, "ok", self.destroy) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkSampleWindowView.py b/modules/ui/CtkSampleWindowView.py index 0f91ad2fa..3b67f03d5 100644 --- a/modules/ui/CtkSampleWindowView.py +++ b/modules/ui/CtkSampleWindowView.py @@ -1,75 +1,43 @@ import contextlib -import copy -import os import tkinter as tk import traceback -from modules.model.BaseModel import BaseModel from modules.modelSampler.BaseModelSampler import ( - BaseModelSampler, ModelSamplerOutput, ) -from modules.ui.SampleFrame import SampleFrame -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.EMAMode import EMAMode +from modules.ui.BaseSampleWindowView import BaseSampleWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleWindowController import SampleWindowController from modules.util.enum.FileType import FileType -from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.time_util import get_string_timestamp -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import torch import customtkinter as ctk from PIL import Image -class SampleWindow(ctk.CTkToplevel): +class CtkSampleWindowView(BaseSampleWindowView, ctk.CTkToplevel): def __init__( self, parent, - train_config: TrainConfig, - use_external_model: bool, - callbacks: TrainCallbacks | None = None, - commands: TrainCommands | None = None, + controller: SampleWindowController, *args, **kwargs ): - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleWindowView.__init__(self, ctk_components) self.title("Sample") self.geometry("1200x800") self.resizable(True, True) - if not use_external_model: - self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) - # remove some settings to speed up model loading for sampling - self.initial_train_config.optimizer.optimizer = None - self.initial_train_config.ema = EMAMode.OFF - else: - self.initial_train_config = None - - #TODO why is there a current_train_config and an initial_train_config? - #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands - - # get model specific defaults - model_type = train_config.model_type - self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: - self.model = None - self.model_sampler = None + model_type = controller.get_model_type() + self.ui_state = CtkUIState(self, controller.sample) + + if controller.use_external_model: + controller.callbacks.set_on_sample_custom(self.__update_preview) + controller.callbacks.set_on_update_sample_custom_progress(self.__update_progress) self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) @@ -78,10 +46,10 @@ def __init__( self.grid_columnconfigure(0, weight=0) self.grid_columnconfigure(1, weight=1) - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) + prompt_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_settings=False) prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) + settings_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_prompt=False) settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") # image @@ -93,70 +61,14 @@ def __init__( image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) + self.progress = self.components.progress(self, 2, 0) + self.components.button(self, 3, 0, "sample", + lambda: controller.do_sample(self.__update_preview, self.__update_progress)) self.wait_visibility() self.focus_set() self.after(200, lambda: set_window_icon(self)) - def __load_model(self) -> BaseModel: - model_loader = create.create_model_loader( - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - - model_setup = create.create_model_setup( - model_type=self.initial_train_config.model_type, - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - training_method=self.initial_train_config.training_method, - ) - - model_names = self.initial_train_config.model_names() - if self.initial_train_config.continue_last_backup: - last_backup_path = self.initial_train_config.get_last_backup_path() - - if last_backup_path: - if self.initial_train_config.training_method == TrainingMethod.LORA: - model_names.lora = last_backup_path - elif self.initial_train_config.training_method == TrainingMethod.EMBEDDING: - model_names.embedding.model_name = last_backup_path - else: # fine-tunes - model_names.base_model = last_backup_path - - print(f"Loading from backup '{last_backup_path}'...") - else: - print("No backup found, loading without backup...") - - if self.initial_train_config.quantization.cache_dir is None: - self.initial_train_config.quantization.cache_dir = self.initial_train_config.cache_dir + "/quantization" - os.makedirs(self.initial_train_config.quantization.cache_dir, exist_ok=True) - - model = model_loader.load( - model_type=self.initial_train_config.model_type, - model_names=model_names, - weight_dtypes=self.initial_train_config.weight_dtypes(), - quantization=self.initial_train_config.quantization, - ) - model.train_config = self.initial_train_config - - model_setup.setup_optimizations(model, self.initial_train_config) - model_setup.setup_train_device(model, self.initial_train_config) - model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) - - return model - - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: - return create.create_model_sampler( - train_device=torch.device(self.initial_train_config.train_device), - temp_device=torch.device(self.initial_train_config.temp_device), - model=model, - model_type=self.initial_train_config.model_type, - training_method=self.initial_train_config.training_method, - ) - def __update_preview(self, sampler_output: ModelSamplerOutput): if sampler_output.file_type == FileType.IMAGE: image = sampler_output.data @@ -172,43 +84,6 @@ def __update_progress(self, progress: int, max_progress: int): def __dummy_image(self) -> Image: return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - def __sample(self): - sample = copy.copy(self.sample) - - if self.commands: - self.commands.sample_custom(sample) - else: - if self.model is None: - # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) - - sample.from_train_config(self.current_train_config) - - sample_dir = os.path.join( - self.initial_train_config.workspace_dir, - "samples", - "custom", - ) - - progress = self.model.train_progress - sample_path = os.path.join( - sample_dir, - f"{get_string_timestamp()}-training-sample-{progress.filename_string()}" - ) - - self.model.eval() - - self.model_sampler.sample( - sample_config=sample, - destination=sample_path, - image_format=self.current_train_config.sample_image_format, - video_format=self.current_train_config.sample_video_format, - audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, - ) - def destroy(self): try: if hasattr(self, "_icon_image_ref"): diff --git a/modules/ui/CtkSamplingTabView.py b/modules/ui/CtkSamplingTabView.py index 5a3c44f08..dfe1a704a 100644 --- a/modules/ui/CtkSamplingTabView.py +++ b/modules/ui/CtkSamplingTabView.py @@ -1,20 +1,17 @@ -from modules.ui.ConfigList import ConfigList -from modules.ui.SampleParamsWindow import SampleParamsWindow -from modules.util.config.SampleConfig import SampleConfig -from modules.util.config.TrainConfig import TrainConfig -from modules.util.ui import components -from modules.util.ui.UIState import UIState +from modules.ui.BaseSamplingTabView import BaseSampleWidgetView, BaseSamplingTabView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.CtkSampleParamsWindowView import CtkSampleParamsWindowView +from modules.ui.SamplingTabController import SamplingTabController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class SamplingTab(ConfigList): - - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, +class CtkSamplingTabView(CtkConfigListView, BaseSamplingTabView): + def __init__(self, master, controller: SamplingTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, from_external_file=True, attr_name="sample_definition_file_name", config_dir="training_samples", @@ -22,103 +19,32 @@ def __init__(self, master, train_config: TrainConfig, ui_state: UIState): add_button_text="Add Sample", add_button_tooltip="Add a new sample configuration.", is_full_width=True, - show_toggle_button=True + show_toggle_button=True, ) - def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return SampleWidget(master, element, i, open_command, remove_command, clone_command, save_command) - - def create_new_element(self) -> dict: - return SampleConfig.default_values(self.train_config.model_type) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return SampleParamsWindow(self.master, self.current_config[i], ui_state, model_type=self.train_config.model_type) + return self.controller.open_element_window(self.master, self.current_config[i], ui_state, CtkSampleParamsWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkSampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) -class SampleWidget(ctk.CTkFrame): +class CtkSampleWidgetView(BaseSampleWidgetView, ctk.CTkFrame): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__( - master=master, corner_radius=10, bg_color="transparent" - ) + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseSampleWidgetView.__init__(self, ctk_components) - self.element = element - self.ui_state = UIState(self, element) - self.i = i - self.save_command = save_command + self.ui_state = CtkUIState(self, element) self.grid_columnconfigure(10, weight=1) - # close button - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), - ) - close_button.grid(row=0, column=0) - - # clone button - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i), - ) - clone_button.grid(row=0, column=1, padx=5) + self.build_content(self, element, self.ui_state, i, open_command, remove_command, clone_command, save_command) - # enabled - self.enabled_switch = components.switch(self, 0, 2, self.ui_state, "enabled", self.__switch_enabled) - self.enabled_switch.configure(width=40) - - # width - components.label(self, 0, 3, "width:") - self.width_entry = components.entry(self, 0, 4, self.ui_state, "width") + def _bind_save(self, save_command): self.width_entry.bind('', lambda _: save_command()) - self.width_entry.configure(width=50) - - # height - components.label(self, 0, 5, "height:") - self.height_entry = components.entry(self, 0, 6, self.ui_state, "height") self.height_entry.bind('', lambda _: save_command()) - self.height_entry.configure(width=50) - - # seed - components.label(self, 0, 7, "seed:") - self.seed_entry = components.entry(self, 0, 8, self.ui_state, "seed") self.seed_entry.bind('', lambda _: save_command()) - self.seed_entry.configure(width=80) - - # prompt - components.label(self, 0, 9, "prompt:") - self.prompt_entry = components.entry(self, 0, 10, self.ui_state, "prompt") self.prompt_entry.bind('', lambda _: save_command()) - # button - self.button = components.icon_button(self, 0, 11, "...", lambda: open_command(self.i, self.ui_state)) - self.button.configure(width=40) - - self.__set_enabled() - - def __switch_enabled(self): - self.save_command() - self.__set_enabled() - - def __set_enabled(self): - enabled = self.element.enabled - self.width_entry.configure(state="normal" if enabled else "disabled") - self.height_entry.configure(state="normal" if enabled else "disabled") - self.prompt_entry.configure(state="normal" if enabled else "disabled") - self.seed_entry.configure(state="normal" if enabled else "disabled") - self.button.configure(state="normal" if enabled else "disabled") - - def configure_element(self): - pass - def place_in_list(self): self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/CtkSchedulerParamsWindowView.py b/modules/ui/CtkSchedulerParamsWindowView.py index f96ed4876..9a7c41b96 100644 --- a/modules/ui/CtkSchedulerParamsWindowView.py +++ b/modules/ui/CtkSchedulerParamsWindowView.py @@ -1,24 +1,23 @@ -from modules.ui.ConfigList import ConfigList -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.ui import components +from modules.ui.BaseSchedulerParamsWindowView import BaseKvParamsView, BaseSchedulerParamsWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.SchedulerParamsWindowController import KvParamsController, SchedulerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import customtkinter as ctk -class KvParams(ConfigList): - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__( - master, - train_config, - ui_state, +class CtkKvParamsView(CtkConfigListView, BaseKvParamsView): + def __init__(self, master, controller: KvParamsController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, attr_name="scheduler_params", from_external_file=False, add_button_text="add parameter", - is_full_width=True + is_full_width=True, ) + BaseKvParamsView.__init__(self, ctk_components) def refresh_ui(self): self._create_element_list() @@ -26,18 +25,12 @@ def refresh_ui(self): def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) - def create_new_element(self) -> dict[str, str]: - return {"key": "", "value": ""} - - def open_element_window(self, i, ui_state): - pass - class KvWidget(ctk.CTkFrame): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): super().__init__(master=master, bg_color="transparent") self.element = element - self.ui_state = UIState(self, element) + self.ui_state = CtkUIState(self, element) self.i = i self.save_command = save_command @@ -57,14 +50,14 @@ def __init__(self, master, element, i, open_command, remove_command, clone_comma # Key tooltip_key = "Key name for an argument in your scheduler" - self.key = components.entry(self, 0, 1, self.ui_state, "key", + self.key = ctk_components.entry(self, 0, 1, self.ui_state, "key", tooltip=tooltip_key, wide_tooltip=True) self.key.bind("", lambda _: save_command()) self.key.configure(width=50) # Value tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." - self.value = components.entry(self, 0, 2, self.ui_state, "value", + self.value = ctk_components.entry(self, 0, 2, self.ui_state, "value", tooltip=tooltip_val, wide_tooltip=True) self.value.bind("", lambda _: save_command()) self.value.configure(width=50) @@ -73,47 +66,31 @@ def place_in_list(self): self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") -class SchedulerParamsWindow(ctk.CTkToplevel): - def __init__(self, parent, train_config: TrainConfig, ui_state, *args, **kwargs): - super().__init__(parent, *args, **kwargs) - - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state +class CtkSchedulerParamsWindowView(BaseSchedulerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSchedulerParamsWindowView.__init__(self, ctk_components) self.title("Learning Rate Scheduler Settings") self.geometry("800x400") self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) self.grid_rowconfigure(1, weight=0) self.grid_columnconfigure(0, weight=1) - self.frame = ctk.CTkFrame(self) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) + frame = ctk.CTkFrame(self) + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) - self.expand_frame = ctk.CTkFrame(self.frame, bg_color="transparent") - self.expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + expand_frame = ctk.CTkFrame(frame, bg_color="transparent") + expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) + CtkKvParamsView(expand_frame, KvParamsController(controller.config), ui_state) self.wait_visibility() self.grab_set() self.focus_set() self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - if self.train_config.learning_rate_scheduler is LearningRateScheduler.CUSTOM: - components.label(master, 0, 0, "Class Name", - tooltip="Python class module and name for the custom scheduler class, in the form of ..") - components.entry(master, 0, 1, self.ui_state, "custom_learning_rate_scheduler") - - # Any additional parameters, in key-value form. - self.params = KvParams(self.expand_frame, self.train_config, self.ui_state) - - def on_window_close(self): - self.destroy() diff --git a/modules/ui/CtkTimestepDistributionWindowView.py b/modules/ui/CtkTimestepDistributionWindowView.py index 21e41ce3e..69f2ae7d3 100644 --- a/modules/ui/CtkTimestepDistributionWindowView.py +++ b/modules/ui/CtkTimestepDistributionWindowView.py @@ -1,15 +1,8 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components +from modules.ui.BaseTimestepDistributionWindowView import BaseTimestepDistributionWindowView +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util.ui import ctk_components from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import torch -from torch import Tensor import customtkinter as ctk from customtkinter import AppearanceModeTracker, ThemeManager @@ -17,127 +10,38 @@ from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg -class TimestepGenerator(ModelSetupNoiseMixin): - - def __init__( - self, - timestep_distribution: TimestepDistribution, - min_noising_strength: float, - max_noising_strength: float, - noising_weight: float, - noising_bias: float, - timestep_shift: float, - ): - super().__init__() - - self.timestep_distribution = timestep_distribution - self.min_noising_strength = min_noising_strength - self.max_noising_strength = max_noising_strength - self.noising_weight = noising_weight - self.noising_bias = noising_bias - self.timestep_shift = timestep_shift - - def generate(self) -> Tensor: - generator = torch.Generator() - generator.seed() - - config = TrainConfig.default_values() - config.timestep_distribution = self.timestep_distribution - config.min_noising_strength = self.min_noising_strength - config.max_noising_strength = self.max_noising_strength - config.noising_weight = self.noising_weight - config.noising_bias = self.noising_bias - config.timestep_shift = self.timestep_shift - - - return self._get_timestep_discrete( - num_train_timesteps=1000, - deterministic=False, - generator=generator, - batch_size=1000000, - config=config, - ) - - -class TimestepDistributionWindow(ctk.CTkToplevel): +class CtkTimestepDistributionWindowView(BaseTimestepDistributionWindowView, ctk.CTkToplevel): def __init__( self, parent, - config: TrainConfig, - ui_state: UIState, + controller: TimestepDistributionWindowController, + ui_state, *args, **kwargs, ): - super().__init__(parent, *args, **kwargs) + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseTimestepDistributionWindowView.__init__(self, ctk_components) self.title("Timestep Distribution") self.geometry("900x600") self.resizable(True, True) - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 + self.controller = controller self.ax = None self.canvas = None self.grid_rowconfigure(0, weight=1) self.grid_columnconfigure(0, weight=1) - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") frame.grid_columnconfigure(0, weight=0) frame.grid_columnconfigure(1, weight=0) frame.grid_columnconfigure(2, weight=0) frame.grid_columnconfigure(3, weight=1) frame.grid_rowconfigure(7, weight=1) - # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, - "timestep_distribution") - - # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") - - # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") - - # noising weight - components.label(frame, 3, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") - - # noising bias - components.label(frame, 4, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") - - # timestep shift - components.label(frame, 5, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") - - # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot + self.build_content(frame, controller, ui_state) + + # matplotlib chart (CTK-only: needs winfo_rgb from the toplevel) appearance_mode = AppearanceModeTracker.get_mode() background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) @@ -163,24 +67,18 @@ def __content_frame(self, master): self.__update_preview() # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + ctk_components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) frame.pack(fill="both", expand=1) - return frame + frame.grid(row=0, column=0, sticky='nsew') + ctk_components.button(self, 1, 0, "ok", self.destroy) - def __update_preview(self): - generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, - ) + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + def __update_preview(self): self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) + self.ax.hist(self.controller.generate_preview_data(), bins=1000, range=(0, 999)) self.canvas.draw() - - def __ok(self): - self.destroy() diff --git a/modules/ui/CtkTopBarView.py b/modules/ui/CtkTopBarView.py index 820fdb71a..ee1bcfa75 100644 --- a/modules/ui/CtkTopBarView.py +++ b/modules/ui/CtkTopBarView.py @@ -1,260 +1,50 @@ -import json -import os -import traceback -import webbrowser from collections.abc import Callable -from contextlib import suppress -from modules.util import path_util -from modules.util.config.SecretsConfig import SecretsConfig -from modules.util.config.TrainConfig import TrainConfig +from modules.ui.BaseTopBarView import BaseTopBarView +from modules.ui.TopBarController import TopBarController from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.optimizer_util import change_optimizer -from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState +from modules.util.ui import ctk_components, dialogs +from modules.util.ui.CtkUIState import CtkUIState import customtkinter as ctk -class TopBar: +class CtkTopBarView(BaseTopBarView): def __init__( self, master, - train_config: TrainConfig, - ui_state: UIState, + controller: TopBarController, + ui_state, change_model_type_callback: Callable[[ModelType], None], change_training_method_callback: Callable[[TrainingMethod], None], load_preset_callback: Callable[[], None], ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.change_model_type_callback = change_model_type_callback - self.change_training_method_callback = change_training_method_callback - self.load_preset_callback = load_preset_callback + BaseTopBarView.__init__(self, ctk_components) - self.dir = "training_presets" + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=0, column=0, sticky="nsew") - self.config_ui_data = { - "config_name": path_util.canonical_join(self.dir, "#.json") - } - self.config_ui_state = UIState(master, self.config_ui_data) + self.build(frame, master, controller, ui_state, change_model_type_callback, change_training_method_callback, load_preset_callback) - self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() + def _make_config_ui_state(self, master, data): + return CtkUIState(master, data) - self.current_config = [] + def _get_dropdown_text(self, widget) -> str: + return widget.get() - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - - self.training_method = None - - # title - components.app_title(self.frame, 0, 0) - - # dropdown - self.configs_dropdown = None - self.__create_configs_dropdown() - - # remove button - # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) - - # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) - - # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) - - # padding + def _setup_frame_column_weight(self): self.frame.grid_columnconfigure(5, weight=1) - # model type - components.options_kv( - master=self.frame, - row=0, - column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, - var_name="model_type", - command=self.__change_model_type, - ) - - def __create_training_method(self): - if self.training_method: - self.training_method.destroy() - - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] - - # training method - self.training_method = components.options_kv( - master=self.frame, - row=0, - column=7, - values=values, - ui_state=self.ui_state, - var_name="training_method", - command=self.change_training_method_callback, - ) - - def __change_model_type(self, model_type: ModelType): - self.change_model_type_callback(model_type) - self.__create_training_method() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() - - self.configs_dropdown = components.options_kv( - self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config - ) - - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): - if path != "#.json": - path = path_util.canonical_join(self.dir, path) - if path.endswith(".json") and os.path.isfile(path): - name = os.path.basename(path) - name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() - - def __save_to_file(self, name) -> str: - name = path_util.safe_filename(name) - path = path_util.canonical_join("training_presets", f"{name}.json") - - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - - return path - - def __save_secrets(self, path) -> str: - write_json_atomic(path, self.train_config.secrets.to_dict()) - return path - - def open_wiki(self): - webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) - - def __save_new_config(self, name): - path = self.__save_to_file(name) - - is_new_config = name not in [x[0] for x in self.configs] - - if is_new_config: - self.configs.append((name, path)) - self.configs.sort() - - if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): - self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) - - if is_new_config: - self.__create_configs_dropdown() - - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] + def _forget_dropdown(self): + self.configs_dropdown.grid_forget() + def _show_save_dialog(self, default_value: str, callback): dialogs.StringInputDialog( parent=self.master, title="name", question="Config Name", - callback=self.__save_new_config, + callback=callback, default_value=default_value, - validate_callback=lambda x: not x.startswith("#") + validate_callback=lambda x: not x.startswith("#"), ) - - def __load_current_config(self, filename): - try: - basename = os.path.basename(filename) - is_built_in_preset = basename.startswith("#") and basename != "#.json" - - with open(filename, "r") as f: - loaded_dict = json.load(f) - default_config = TrainConfig.default_values() - if is_built_in_preset: - # always assume built-in configs are saved in the most recent version - loaded_dict["__version"] = default_config.config_version - loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() - - with suppress(FileNotFoundError), open("secrets.json", "r") as f: - secrets_dict=json.load(f) - loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - - self.train_config.from_dict(loaded_config.to_dict()) - self.ui_state.update(loaded_config) - - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.load_preset_callback() - except FileNotFoundError: - pass - except Exception: - print(traceback.format_exc()) - - def __remove_config(self): - # TODO - pass - - def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index ba90d2e64..d3c1a70b5 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -1,51 +1,40 @@ import ctypes -import datetime -import json -import os import platform -import subprocess -import sys -import threading -import time -import traceback -import webbrowser from collections.abc import Callable from contextlib import suppress from pathlib import Path from tkinter import filedialog, messagebox -import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI -from modules.util import create -from modules.util.callbacks.TrainCallbacks import TrainCallbacks -from modules.util.commands.TrainCommands import TrainCommands +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseTrainUIView import BaseTrainUIView +from modules.ui.CloudTabController import CloudTabController +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkAdditionalEmbeddingsTabView import CtkAdditionalEmbeddingsTabView +from modules.ui.CtkCaptionUIView import CtkCaptionUIView +from modules.ui.CtkCloudTabView import CtkCloudTabView +from modules.ui.CtkConceptTabView import CtkConceptTabView +from modules.ui.CtkConvertModelUIView import CtkConvertModelUIView +from modules.ui.CtkLoraTabView import CtkLoraTabView +from modules.ui.CtkModelTabView import CtkModelTabView +from modules.ui.CtkProfilingWindowView import CtkProfilingWindowView +from modules.ui.CtkSampleWindowView import CtkSampleWindowView +from modules.ui.CtkSamplingTabView import CtkSamplingTabView +from modules.ui.CtkTopBarView import CtkTopBarView +from modules.ui.CtkTrainingTabView import CtkTrainingTabView +from modules.ui.CtkVideoToolUIView import CtkVideoToolUIView +from modules.ui.LoraTabController import LoraTabController +from modules.ui.ModelTabController import ModelTabController +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.ui.SamplingTabController import SamplingTabController +from modules.ui.TopBarController import TopBarController +from modules.ui.TrainingTabController import TrainingTabController +from modules.ui.TrainUIController import TrainUIController from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.GradientReducePrecision import GradientReducePrecision -from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.torch_util import torch_gc -from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import flush_and_validate_all - -import torch import customtkinter as ctk from customtkinter import AppearanceModeTracker @@ -57,14 +46,13 @@ # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE -class TrainUI(ctk.CTk): + +class CtkTrainUIView(BaseTrainUIView, ctk.CTk): set_step_progress: Callable[[int, int], None] set_epoch_progress: Callable[[int, int], None] status_label: ctk.CTkLabel | None training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None _TRAIN_BUTTON_STYLES = { "idle": { @@ -93,7 +81,8 @@ class TrainUI(ctk.CTk): } def __init__(self): - super().__init__() + ctk.CTk.__init__(self) + BaseTrainUIView.__init__(self, ctk_components) self.title("OneTrainer") self.geometry("1100x740") @@ -105,7 +94,10 @@ def __init__(self): ctk.set_default_color_theme("blue") self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) + self.ui_state = CtkUIState(self, self.train_config) + + self.controller = TrainUIController(self.train_config) + self.controller.view = self self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) @@ -128,80 +120,120 @@ def __init__(self): self.content_frame(self) self.bottom_bar(self) - self.training_thread = None - self.training_callbacks = None - self.training_commands = None + self.controller._check_start_always_on_tensorboard() - self.always_on_tensorboard_subprocess = None - self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() - - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) + self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self.controller._on_workspace_dir_change_trace) # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) + self._profiling_controller = ProfilingWindowController() + self.profiling_window = self._profiling_controller.create_window(self, CtkProfilingWindowView) self.protocol("WM_DELETE_WINDOW", self.__close) def __close(self): self.top_bar_component.save_default() - self._stop_always_on_tensorboard() + self.controller._stop_always_on_tensorboard() if hasattr(self, 'workspace_dir_trace_id'): self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) self.quit() + # --- BaseTrainUIView abstract method implementations --- + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def on_training_started(self): + self._set_training_button_style("running") + + def on_training_stopped(self, error_caught: bool): + self.eta_label.configure(text="") + self._set_training_button_style("idle") + + def on_training_stopping(self): + self._set_training_button_style("stopping") + + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): + self.set_step_progress(epoch_step, max_step) + self.set_epoch_progress(epoch, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def schedule_on_main_thread(self, fn: Callable): + self.after(0, fn) + + def get_cloud_reattach(self) -> bool: + return self.cloud_tab.reattach + + def save_default(self): + self.top_bar_component.save_default() + self.concepts_tab.save_current_config() + self.sampling_tab.save_current_config() + self.additional_embeddings_tab.save_current_config() + + def show_validation_errors(self, errors: list[str]): + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + + def open_dataset_tool(self): + self.wait_window(self.controller.open_dataset_tool(self, CtkCaptionUIView)) + + def open_video_tool(self): + self.wait_window(self.controller.open_video_tool(self, CtkVideoToolUIView)) + + def open_convert_model_tool(self): + self.wait_window(self.controller.open_convert_model_tool(self, CtkConvertModelUIView)) + + def open_sampling_tool(self): + self.controller.open_sampling_tool(self, CtkSampleWindowView) + + def open_manual_sample_window(self): + self.controller.open_manual_sample_window(self, CtkSampleWindowView) + + def wait_window(self, window): + ctk.CTk.wait_window(self, window) + + def show_window(self, window): + window.focus_set() + + def connect_window_closed(self, window, callback): + window.bind("", lambda _: callback()) + + # --- CTK layout and frame builders --- + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + def top_bar(self, master): - return TopBar( + return CtkTopBarView( master, - self.train_config, + TopBarController(self.train_config), self.ui_state, self.change_model_type, self.change_training_method, self.load_preset, ) - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - def bottom_bar(self, master): frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=2, column=0, sticky="nsew") - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) - - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) + status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + status_frame.grid(row=0, column=1, sticky="w") + status_frame.grid_rowconfigure(0, weight=0) + status_frame.grid_rowconfigure(1, weight=0) + status_frame.grid_columnconfigure(0, weight=1) # padding frame.grid_columnconfigure(2, weight=1) - - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") - - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") - - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) - - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) + self.build_bottom_bar_content(frame, status_frame, self.controller, self.ui_state) self._set_training_button_style("idle") # centralized styling return frame @@ -237,113 +269,12 @@ def create_general_tab(self, master): frame.grid_columnconfigure(1, weight=1) frame.grid_columnconfigure(2, weight=0) frame.grid_columnconfigure(3, weight=1) - - # workspace dir - components.label(frame, 0, 0, "Workspace Directory", - tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) - - # cache dir - components.label(frame, 0, 2, "Cache Directory", - tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") - - # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", - tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") - - # only cache - components.label(frame, 2, 2, "Only Cache", - tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") - - # TODO: In Phase 4 rework the general tab. - # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", - tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") - - # debug - components.label(frame, 4, 0, "Debug mode", - tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") - - components.label(frame, 4, 2, "Debug Directory", - tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) - - # tensorboard - components.label(frame, 6, 0, "Tensorboard", - tooltip="Starts the Tensorboard Web UI during training") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") - - components.label(frame, 6, 2, "Always-On Tensorboard", - tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) - - components.label(frame, 7, 0, "Expose Tensorboard", - tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") - components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") - components.label(frame, 7, 2, "Tensorboard Port", - tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - - - # validation - components.label(frame, 8, 0, "Validation", - tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") - - components.label(frame, 8, 2, "Validate after", - tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") - - # device - components.label(frame, 10, 0, "Dataloader Threads", - tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) - - components.label(frame, 11, 0, "Train Device", - tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") - components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) - - components.label(frame, 12, 0, "Multi-GPU", - tooltip="Enable multi-GPU training") - components.switch(frame, 12, 1, self.ui_state, "multi_gpu") - components.label(frame, 12, 2, "Device Indexes", - tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") - components.entry(frame, 12, 3, self.ui_state, "device_indexes") - - components.label(frame, 13, 0, "Gradient Reduce Precision", - tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" - "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" - "FLOAT_32: Reduce gradients in float32\n" - "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", - wide_tooltip=True) - components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, - "gradient_reduce_precision") - - components.label(frame, 13, 2, "Fused Gradient Reduce", - tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") - - components.label(frame, 14, 0, "Async Gradient Reduce", - tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 14, 2, "Buffer size (MB)", - tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") - - components.label(frame, 15, 0, "Temp Device", - tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 15, 1, self.ui_state, "temp_device") - + self.build_general_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) + return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) def create_data_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -352,77 +283,35 @@ def create_data_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", - tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") - - # latent caching - components.label(frame, 1, 0, "Latent Caching", - tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") - - # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", - tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") - + self.build_data_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) + return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) + def create_training_tab(self, master) -> CtkTrainingTabView: + return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) + def create_cloud_tab(self, master) -> CtkCloudTabView: + return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) def create_sampling_tab(self, master): master.grid_rowconfigure(0, weight=0) master.grid_rowconfigure(1, weight=1) master.grid_columnconfigure(0, weight=1) - # sample after top_frame = ctk.CTkFrame(master=master, corner_radius=0) top_frame.grid(row=0, column=0, sticky="nsew") sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - components.label(top_frame, 0, 0, "Sample After", - tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") - - components.label(top_frame, 0, 2, "Skip First", - tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") - - components.label(top_frame, 0, 4, "Format", - tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ - ("PNG", ImageFormat.PNG), - ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") - - components.button(top_frame, 0, 6, "sample now", self.sample_now) - - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) - - components.label(sub_frame, 0, 0, "Non-EMA Sampling", - tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") - - components.label(sub_frame, 0, 2, "Samples to Tensorboard", - tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") + self.build_sampling_tab_header(top_frame, sub_frame, self.controller, self.ui_state) - # table frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=1, column=0, sticky="nsew") - return SamplingTab(frame, self.train_config, self.ui_state) + return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) def create_backup_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -431,48 +320,7 @@ def create_backup_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # backup after - components.label(frame, 0, 0, "Backup After", - tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") - - # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) - - # rolling backup - components.label(frame, 1, 0, "Rolling Backup", - tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") - - # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", - tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") - - # backup before save - components.label(frame, 2, 0, "Backup Before Save", - tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") - - # save after - components.label(frame, 3, 0, "Save Every", - tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") - - # save now - components.button(frame, 3, 3, "save now", self.save_now) - - # skip save - components.label(frame, 4, 0, "Skip First", - tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") - - # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", - tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - + self.build_backup_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame @@ -483,48 +331,12 @@ def embedding_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # embedding model name - components.label(frame, 0, 0, "Base embedding", - tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier - ) - - # token count - components.label(frame, 1, 0, "Token count", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") - - # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") - - # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", - tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") - - # placeholder - components.label(frame, 4, 0, "Placeholder", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") - - # output embedding - components.label(frame, 5, 0, "Output embedding", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - + self.build_embedding_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) def create_tools_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -533,34 +345,13 @@ def create_tools_tab(self, master): frame.grid_columnconfigure(2, minsize=50) frame.grid_columnconfigure(3, weight=0) frame.grid_columnconfigure(4, weight=1) - - # dataset - components.label(frame, 0, 0, "Dataset Tools", - tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) - - # video tools - components.label(frame, 1, 0, "Video Tools", - tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) - - # convert model - components.label(frame, 2, 0, "Convert Model Tools", - tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) - - # sample - components.label(frame, 3, 0, "Sampling Tool", - tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) - - components.label(frame, 4, 0, "Profiling Tool", - tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - + self.build_tools_tab_content(frame, self.controller, self.ui_state) frame.pack(fill="both", expand=1) return frame + def open_profiling_tool(self): + self.profiling_window.deiconify() + def change_model_type(self, model_type: ModelType): if self.model_tab: self.model_tab.refresh_ui() @@ -585,7 +376,7 @@ def change_training_method(self, training_method: TrainingMethod): self.tabview.delete("embedding") if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: self.embedding_tab(self.tabview.add("embedding")) @@ -596,294 +387,27 @@ def load_preset(self): if self.additional_embeddings_tab: self.additional_embeddings_tab.refresh_ui() - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) - - def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: - spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step - remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - - if train_progress.global_step <= 30: - return "Estimating ..." - - td = datetime.timedelta(seconds=total_eta) - days = td.days - hours, remainder = divmod(td.seconds, 3600) - minutes, seconds = divmod(remainder, 60) - if days > 0: - return f"{days}d {hours}h" - elif hours > 0: - return f"{hours}h {minutes}m" - elif minutes > 0: - return f"{minutes}m {seconds}s" - else: - return f"{seconds}s" - - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") - - def delete_eta_label(self): - self.eta_label.configure(text="") - - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) - - def on_update_status(self, status: str): - self.status_label.configure(text=status) - - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) - - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) - - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) - - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() - - def open_profiling_tool(self): - self.profiling_window.deiconify() - - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) - - if not zip_path: + def _set_training_button_style(self, mode: str): + if not self.training_button: return - - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" - - self.on_update_status("Generating debug package...") - - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") - - - def open_manual_sample_window (self): - training_callbacks = self.training_callbacks - training_commands = self.training_commands - - if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, - use_external_model=True, - callbacks=training_callbacks, - commands=training_commands, - ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() - - def __training_thread_function(self): - error_caught = False - - self.training_callbacks = TrainCallbacks( - on_update_train_progress=self.on_update_train_progress, - on_update_status=self.on_update_status, - ) - - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) - try: - trainer.start() - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - - self.start_time = time.monotonic() - trainer.train() - except Exception: - if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) - error_caught = True - traceback.print_exc() - - trainer.end() - - # clear gpu memory - del trainer - - self.training_thread = None - self.training_commands = None - torch.clear_autocast_cache() - torch_gc() - - if error_caught: - self.on_update_status("Error: check the console for details") - else: - self.on_update_status("Stopped") - self.delete_eta_label() - - # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash - self.after(0, self._set_training_button_idle) - - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self.after(0, self._start_always_on_tensorboard) - - def start_training(self): - if self.training_thread is None: - self.save_default() - - # --- pre-training validation gate --- - errors = flush_and_validate_all() - - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) - return - - self._set_training_button_running() - - if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - self.training_commands = TrainCommands() - torch_gc() - - self.training_thread = threading.Thread(target=self.__training_thread_function) - self.training_thread.start() - else: - self._set_training_button_stopping() - self.on_update_status("Stopping ...") - self.training_commands.stop() - - def save_default(self): - self.top_bar_component.save_default() - self.concepts_tab.save_current_config() - self.sampling_tab.save_current_config() - self.additional_embeddings_tab.save_current_config() + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) def export_training(self): file_path = filedialog.asksaveasfilename(filetypes=[ ("All Files", "*.*"), ("json", "*.json"), ], initialdir=".", initialfile="config.json") - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.save() - - def _check_start_always_on_tensorboard(self): - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _start_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") - - os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - - tensorboard_args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.train_config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - try: - self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) - except Exception: - self.always_on_tensorboard_subprocess = None - - def _stop_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - try: - self.always_on_tensorboard_subprocess.terminate() - self.always_on_tensorboard_subprocess.wait(timeout=5) - except subprocess.TimeoutExpired: - self.always_on_tensorboard_subprocess.kill() - except Exception: - pass - finally: - self.always_on_tensorboard_subprocess = None - - def _on_workspace_dir_change(self, new_workspace_dir: str): - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_workspace_dir_change_trace(self, *args): - new_workspace_dir = self.train_config.workspace_dir - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_always_on_tensorboard_toggle(self): - if self.train_config.tensorboard_always_on: - if not (self.training_thread and self.train_config.tensorboard): - self._start_always_on_tensorboard() - else: - if not (self.training_thread and self.train_config.tensorboard): - self._stop_always_on_tensorboard() + self.controller.export_training(file_path) - def _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + if not zip_path: return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") + self.controller.generate_debug_package(Path(zip_path) / "OneTrainer_debug_report.zip") diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bcca11ae9..bc29488dd 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -1,40 +1,27 @@ -from modules.ui.OffloadingWindow import OffloadingWindow -from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow -from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow -from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow -from modules.util import create -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod -from modules.util.enum.LearningRateScaler import LearningRateScaler -from modules.util.enum.LearningRateScheduler import LearningRateScheduler -from modules.util.enum.LossScaler import LossScaler -from modules.util.enum.LossWeight import LossWeight -from modules.util.enum.Optimizer import Optimizer -from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.optimizer_util import change_optimizer -from modules.util.ui import components -from modules.util.ui.UIState import UIState -from modules.util.ui.validation_helpers import check_range, validate_resolution -import customtkinter as ctk +from modules.ui.BaseTrainingTabView import BaseTrainingTabView +from modules.ui.CtkOffloadingWindowView import CtkOffloadingWindowView +from modules.ui.CtkOptimizerParamsWindowView import CtkOptimizerParamsWindowView +from modules.ui.CtkSchedulerParamsWindowView import CtkSchedulerParamsWindowView +from modules.ui.CtkTimestepDistributionWindowView import CtkTimestepDistributionWindowView +from modules.ui.TrainingTabController import TrainingTabController +from modules.util.ui import ctk_components +import customtkinter as ctk -class TrainingTab: - def __init__(self, master, train_config: TrainConfig, ui_state: UIState): - super().__init__() +class CtkTrainingTabView(BaseTrainingTabView): + def __init__(self, master, controller: TrainingTabController, ui_state): + BaseTrainingTabView.__init__(self, ctk_components) self.master = master - self.train_config = train_config + self.controller = controller self.ui_state = ui_state + self.scroll_frame = None master.grid_rowconfigure(0, weight=1) master.grid_columnconfigure(0, weight=1) - self.scroll_frame = None - self.refresh_ui() def refresh_ui(self): @@ -60,797 +47,31 @@ def refresh_ui(self): column_2.grid(row=0, column=2, sticky="nsew") column_2.grid_columnconfigure(0, weight=1) - if self.train_config.model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2) - - - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1) - self.__create_text_encoder_n_frame(column_0, 2, i=2) - self.__create_embedding_frame(column_0, 3) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_unet_frame(column_1, 1) - self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_wuerstchen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0, supports_circular_padding=True) - self.__create_prior_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 0) - self.__create_loss_frame(column_2, 1) - self.__create_layer_frame(column_2, 2) - - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2, supports_vb_loss=True) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True, supports_sequence_length=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_flux_2_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_chroma_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_qwen_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_z_image_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_ernie_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_sana_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_frame(column_0, 1) - self.__create_embedding_frame(column_0, 2) - - self.__create_base2_frame(column_1, 0) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_embedding_frame(column_0, 4) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __setup_hi_dream_ui(self, column_0, column_1, column_2): - self.__create_base_frame(column_0, 0) - self.__create_text_encoder_n_frame(column_0, 1, i=1, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 2, i=2, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 3, i=3, supports_include=True) - self.__create_text_encoder_n_frame(column_0, 4, i=4, supports_include=True, supports_layer_skip=False) - self.__create_embedding_frame(column_0, 5) - - self.__create_base2_frame(column_1, 0, video_training_enabled=True) - self.__create_transformer_frame(column_1, 1) - self.__create_noise_frame(column_1, 2) - - self.__create_masked_frame(column_2, 1) - self.__create_loss_frame(column_2, 2) - self.__create_layer_frame(column_2, 3) - - def __create_base_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # optimizer - components.label(frame, 0, 0, "Optimizer", - tooltip="The type of optimizer") - components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], self.ui_state, "optimizer.optimizer", - command=self.__restore_optimizer_config, adv_command=self.__open_optimizer_params_window) - - # learning rate scheduler - # Wackiness will ensue when reloading configs if we don't check and clear this first. - if hasattr(self, "lr_scheduler_comp"): - delattr(self, "lr_scheduler_comp") - delattr(self, "lr_scheduler_adv_comp") - components.label(frame, 1, 0, "Learning Rate Scheduler", - tooltip="Learning rate scheduler that automatically changes the learning rate during training") - _, d = components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], self.ui_state, - "learning_rate_scheduler", command=self.__restore_scheduler_config, - adv_command=self.__open_scheduler_params_window) - self.lr_scheduler_comp = d['component'] - self.lr_scheduler_adv_comp = d['button_component'] - # Initial call requires the presence of self.lr_scheduler_adv_comp. - self.__restore_scheduler_config(self.ui_state.get_var("learning_rate_scheduler").get()) - - # learning rate - components.label(frame, 2, 0, "Learning Rate", - tooltip="The base learning rate") - components.entry(frame, 2, 1, self.ui_state, "learning_rate", required=True) - - # learning rate warmup steps - components.label(frame, 3, 0, "Learning Rate Warmup Steps", - tooltip="The number of steps it takes to gradually increase the learning rate from 0 to the specified learning rate. Values >1 are interpeted as a fixed number of steps, values <=1 are intepreted as a percentage of the total training steps (ex. 0.2 = 20% of the total step count)") - components.entry(frame, 3, 1, self.ui_state, "learning_rate_warmup_steps") - - # learning rate min factor - components.label(frame, 4, 0, "Learning Rate Min Factor", - tooltip="Unit = float. Method = percentage. For a factor of 0.1, the final LR will be 10% of the initial LR. If the initial LR is 1e-4, the final LR will be 1e-5.") - components.entry(frame, 4, 1, self.ui_state, "learning_rate_min_factor", - extra_validate=check_range(lower=0, upper=0.99, message="Learning rate min factor must be between 0 and 0.99")) - - # learning rate cycles - components.label(frame, 5, 0, "Learning Rate Cycles", - tooltip="The number of learning rate cycles. This is only applicable if the learning rate scheduler supports cycles") - components.entry(frame, 5, 1, self.ui_state, "learning_rate_cycles") - - # epochs - components.label(frame, 6, 0, "Epochs", - tooltip="The number of epochs for a full training run") - components.entry(frame, 6, 1, self.ui_state, "epochs", required=True) - - # batch size - components.label(frame, 7, 0, "Local Batch Size", - tooltip="The batch size of one training step. If you use multiple GPUs, this is the batch size of each GPU (local batch size).") - components.entry(frame, 7, 1, self.ui_state, "batch_size", required=True) - - # accumulation steps - components.label(frame, 8, 0, "Accumulation Steps", - tooltip="Number of accumulation steps. Increase this number to trade batch size for training speed") - components.entry(frame, 8, 1, self.ui_state, "gradient_accumulation_steps", required=True) - - # Learning Rate Scaler - components.label(frame, 9, 0, "Learning Rate Scaler", - tooltip="Selects the type of learning rate scaling to use during training. Functionally equated as: LR * SQRT(selection)") - components.options(frame, 9, 1, [str(x) for x in list(LearningRateScaler)], self.ui_state, - "learning_rate_scaler") - - # clip grad norm - components.label(frame, 10, 0, "Clip Grad Norm", - tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") - components.entry(frame, 10, 1, self.ui_state, "clip_grad_norm") - - def __create_base2_frame(self, master, row, video_training_enabled: bool=False, supports_circular_padding: bool=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - row = 0 - - # ema - components.label(frame, row, 0, "EMA", - tooltip="EMA averages the training progress over many steps, better preserving different concepts in big datasets") - components.options(frame, row, 1, [str(x) for x in list(EMAMode)], self.ui_state, "ema") - row += 1 - - # ema decay - components.label(frame, row, 0, "EMA Decay", - tooltip="Decay parameter of the EMA model. Higher numbers will average more steps. For datasets of hundreds or thousands of images, set this to 0.9999. For smaller datasets, set it to 0.999 or even 0.998") - components.entry(frame, row, 1, self.ui_state, "ema_decay", - extra_validate=check_range(lower=0.5, upper=1, - message="EMA decay must be between 0.5 and 1")) - row += 1 - - # ema update step interval - components.label(frame, row, 0, "EMA Update Step Interval", - tooltip="Number of steps between EMA update steps") - components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") - row += 1 - - # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) - row += 1 - - # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") - row += 1 - - # train dtype - components.label(frame, row, 0, "Train Data Type", - tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ - ("float32", DataType.FLOAT_32), - ("float16", DataType.FLOAT_16), - ("bfloat16", DataType.BFLOAT_16), - ("tfloat32", DataType.TFLOAT_32), - ], self.ui_state, "train_dtype") - row += 1 - - # fallback train dtype - components.label(frame, row, 0, "Fallback Train Data Type", - tooltip="The mixed precision data type used for training stages that don't support float16 data types. This can increase training speed, but reduces precision") - components.options_kv(frame, row, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "fallback_train_dtype") - row += 1 - - # autocast cache - components.label(frame, row, 0, "Autocast Cache", - tooltip="Enables the autocast cache. Disabling this reduces memory usage, but increases training time") - components.switch(frame, row, 1, self.ui_state, "enable_autocast_cache") - row += 1 - - # resolution - components.label(frame, row, 0, "Resolution", - tooltip="The resolution used for training. Optionally specify multiple resolutions separated by a comma, or a single exact resolution in the format x") - components.entry(frame, row, 1, self.ui_state, "resolution", required=True, - extra_validate=validate_resolution()) - row += 1 - - # frames - if video_training_enabled: - components.label(frame, row, 0, "Frames", - tooltip="The number of frames used for training.") - components.entry(frame, row, 1, self.ui_state, "frames", required=True) - row += 1 - - # force circular padding - if supports_circular_padding: - components.label(frame, row, 0, "Force Circular Padding", - tooltip="Enables circular padding for all conv layers to better train seamless images") - components.switch(frame, row, 1, self.ui_state, "force_circular_padding") - - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", - tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + callbacks = { + 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), + 'open_optimizer_params': self._open_optimizer_params_window, + 'restore_scheduler': self._restore_scheduler_config, + 'open_scheduler_params': self._open_scheduler_params_window, + 'open_offloading': self._open_offloading_window, + 'open_timestep_distribution': self._open_timestep_distribution_window, + } - # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) - if supports_training: - # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", - tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", - "text_encoder.stop_training_after_unit", supports_time_units=False) - - # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", - tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") - - if supports_clip_skip: - # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") - - if supports_sequence_length: - # text encoder sequence length - components.label(frame, row, 0, "Text Encoder Sequence Length", - tooltip="Number of tokens for captions") - components.entry(frame, row, 1, self.ui_state, "text_encoder_sequence_length") - row += 1 - - def __create_text_encoder_n_frame( - self, - master, - row: int, - i: int, - supports_include: bool = False, - supports_layer_skip: bool = True, - supports_sequence_length: bool = False, - ): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - row = 0 - - suffix = f"_{i}" if i > 1 else "" - - if supports_include: - # include text encoder - components.label(frame, row, 0, f"Include Text Encoder {i}", - tooltip=f"Includes text encoder {i} in the training run") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.include") - row += 1 - - # train text encoder - components.label(frame, row, 0, f"Train Text Encoder {i}", - tooltip=f"Enables training the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") - row += 1 - - # train text encoder embedding - components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", - tooltip=f"Enables training embeddings for the text encoder {i} model") - components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train_embedding") - row += 1 - - # dropout - components.label(frame, row, 0, "Dropout Probability", - tooltip=f"The Probability for dropping the text encoder {i} conditioning") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.dropout_probability") - row += 1 - - # train text encoder epochs - components.label(frame, row, 0, "Stop Training After", - tooltip=f"When to stop training the text encoder {i}") - components.time_entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.stop_training_after", - f"text_encoder{suffix}.stop_training_after_unit", supports_time_units=False) - row += 1 - - # text encoder learning rate - components.label(frame, row, 0, f"Text Encoder {i} Learning Rate", - tooltip=f"The learning rate of the text encoder {i}. Overrides the base learning rate") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}.learning_rate") - row += 1 - - if supports_layer_skip: - # text encoder layer skip (clip skip) - components.label(frame, row, 0, f"Text Encoder {i} Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_layer_skip") - row += 1 - - if supports_sequence_length: - # text encoder sequence length - components.label(frame, row, 0, f"Text Encoder {i} Sequence Length", - tooltip="Overrides the number of tokens used for captions. If empty, the model default is used, which is 512 on Flux. Comfy samples with 256 tokens though. 77 is the default only for backwards compatibility.") - components.entry(frame, row, 1, self.ui_state, f"text_encoder{suffix}_sequence_length") - row += 1 - - def __create_embedding_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - - # embedding learning rate - components.label(frame, 0, 0, "Embeddings Learning Rate", - tooltip="The learning rate of embeddings. Overrides the base learning rate") - components.entry(frame, 0, 1, self.ui_state, "embedding_learning_rate") - - # preserve embedding norm - components.label(frame, 1, 0, "Preserve Embedding Norm", - tooltip="Rescales each trained embedding to the median embedding norm") - components.switch(frame, 1, 1, self.ui_state, "preserve_embedding_norm") - - def __create_unet_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train unet - components.label(frame, 0, 0, "Train UNet", - tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") - - # train unet epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", - supports_time_units=False) - - # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", - tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") - - # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", - tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") - rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") - - def __create_prior_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train prior - components.label(frame, 0, 0, "Train Prior", - tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") - - # train prior epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", - supports_time_units=False) - - # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", - tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") - - def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # train transformer - components.label(frame, 0, 0, "Train Transformer", - tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") - - # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", - tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", - supports_time_units=False) - - # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", - tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") - - if supports_force_attention_mask: - # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", - tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") - - if supports_guidance_scale: - # guidance scale - components.label(frame, 4, 0, "Guidance Scale", - tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") - - def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # offset noise weight - components.label(frame, 0, 0, "Offset Noise Weight", - tooltip="The weight of offset noise added to each training step") - components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight") - - if supports_generalized_offset_noise: - # generalized offset noise weight - generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise", - tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.") - generalised_offset_label.configure(wraplength=130, justify="left") - components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise") - - # perturbation noise weight - components.label(frame, 2, 0, "Perturbation Noise Weight", - tooltip="The weight of perturbation noise added to each training step") - components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight") - - # timestep distribution - components.label(frame, 3, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution", - adv_command=self.__open_timestep_distribution_window) - - # min noising strength - components.label(frame, 4, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True) - - # max noising strength - components.label(frame, 5, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True) - - # noising weight - components.label(frame, 6, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True) - - # noising bias - components.label(frame, 7, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True) - - # timestep shift - components.label(frame, 8, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) - - if supports_dynamic_timestep_shifting: - # dynamic timestep shifting - components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - - - - def __create_masked_frame(self, master, row): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # Masked Training - components.label(frame, 0, 0, "Masked Training", - tooltip="Masks the training samples to let the model focus on certain parts of the image. When enabled, one mask image is loaded for each training sample.") - components.switch(frame, 0, 1, self.ui_state, "masked_training") - - # unmasked probability - components.label(frame, 1, 0, "Unmasked Probability", - tooltip="When masked training is enabled, specifies the number of training steps done on unmasked samples") - components.entry(frame, 1, 1, self.ui_state, "unmasked_probability", - extra_validate=check_range(lower=0, upper=1, message="Unmasked probability must be between 0 and 1")) - - # unmasked weight - components.label(frame, 2, 0, "Unmasked Weight", - tooltip="When masked training is enabled, specifies the loss weight of areas outside the masked region") - components.entry(frame, 2, 1, self.ui_state, "unmasked_weight", - extra_validate=check_range(lower=0, upper=1, message="Unmasked weight must be between 0 and 1")) - - # normalize masked area loss - components.label(frame, 3, 0, "Normalize Masked Area Loss", - tooltip="When masked training is enabled, normalizes the loss for each sample based on the sizes of the masked region") - components.switch(frame, 3, 1, self.ui_state, "normalize_masked_area_loss") - - # masked prior preservation - components.label(frame, 4, 0, "Masked Prior Preservation Weight", - tooltip="Preserves regions outside the mask using the original untrained model output as a target. Only available for LoRA training. If enabled, use a low unmasked weight.") - components.entry(frame, 4, 1, self.ui_state, "masked_prior_preservation_weight", - extra_validate=check_range(lower=0, upper=1, message="Masked prior preservation weight must be between 0 and 1")) - - # use custom conditioning image - components.label(frame, 5, 0, "Custom Conditioning Image", - tooltip="When custom conditioning image is enabled, will use png postfix with -condlabel instead of automatically generated.It's suitable for special scenarios, such as object removal, allowing the model to learn a certain behavior concept") - components.switch(frame, 5, 1, self.ui_state, "custom_conditioning_image") - - def __create_loss_frame(self, master, row, supports_vb_loss: bool = False): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - # MSE Strength - components.label(frame, 0, 0, "MSE Strength", - tooltip="Mean Squared Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 0, 1, self.ui_state, "mse_strength", required=True) - - # MAE Strength - components.label(frame, 1, 0, "MAE Strength", - tooltip="Mean Absolute Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 1, 1, self.ui_state, "mae_strength", required=True) - - # log-cosh Strength - components.label(frame, 2, 0, "log-cosh Strength", - tooltip="Log - Hyperbolic cosine Error strength for custom loss settings. Strengths should generally sum to 1.") - components.entry(frame, 2, 1, self.ui_state, "log_cosh_strength", required=True) - - # Huber Strength - components.label(frame, 3, 0, "Huber Strength", - tooltip="Huber loss strength for custom loss settings. Less sensitive to outliers than MSE. Strengths should generally sum to 1.") - components.entry(frame, 3, 1, self.ui_state, "huber_strength", required=True) - - # Huber Delta - components.label(frame, 4, 0, "Huber Delta", - tooltip="Delta parameter for huber loss") - components.entry(frame, 4, 1, self.ui_state, "huber_delta", required=True) - - if supports_vb_loss: - # VB Strength - components.label(frame, 5, 0, "VB Strength", - tooltip="Variational lower-bound strength for custom loss settings. Should be set to 1 for variational diffusion models") - components.entry(frame, 5, 1, self.ui_state, "vb_loss_strength", required=True) - - # Loss Weight function - components.label(frame, 6, 0, "Loss Weight Function", - tooltip="Choice of loss weight function. Can help the model learn details more accurately.") - components.options(frame, 6, 1, [str(x) for x in list(LossWeight) - if x.supports_flow_matching() == self.train_config.model_type.is_flow_matching() - or x == LossWeight.CONSTANT - ], - self.ui_state, "loss_weight_fn") - - row = 7 - - # Loss weight strength - if not self.train_config.model_type.is_flow_matching(): - components.label(frame, row, 0, "Gamma", - tooltip="Inverse strength of loss weighting. Range: 1-20, only applies to Min SNR and P2.") - components.entry(frame, row, 1, self.ui_state, "loss_weight_strength", - extra_validate=check_range(lower=1, upper=20, message="Gamma must be between 1 and 20")) - row += 1 - - # Loss Scaler - components.label(frame, row, 0, "Loss Scaler", - tooltip="Selects the type of loss scaling to use during training. Functionally equated as: Loss * selection") - components.options(frame, row, 1, [str(x) for x in list(LossScaler)], self.ui_state, "loss_scaler") - row += 1 - - def __create_layer_frame(self, master, row): - cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) - presets = cls.LAYER_PRESETS if cls is not None else {"full": []} - components.layer_filter_entry(master, row, 0, self.ui_state, - preset_var_name="layer_filter_preset", presets=presets, - preset_label="Layer Filter", - preset_tooltip="Select a preset defining which layers to train, or select 'Custom' to define your own.\nA blank 'custom' field or 'Full' will train all layers.", - entry_var_name="layer_filter", - entry_tooltip="Comma-separated list of diffusion layers to train. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained", - regex_var_name="layer_filter_regex", - regex_tooltip="If enabled, layer filter patterns are interpreted as regular expressions. Otherwise, simple substring matching is used.", - ) - - - def __on_layer_filter_preset_change(self): - if not self.layer_selector: + def _restore_scheduler_config(self, variable): + if not hasattr(self, 'lr_scheduler_adv_comp'): return - selected = self.ui_state.get_var("layer_filter_preset").get() - self.__preset_set_layer_choice(selected) - - def __hide_layer_entry(self): - if self.layer_entry and self.layer_entry.winfo_manager(): - self.layer_entry.grid_remove() - - def __show_layer_entry(self): - if self.layer_entry and not self.layer_entry.winfo_manager(): - self.layer_entry.grid() - - def __open_optimizer_params_window(self): - window = OptimizerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" + self.lr_scheduler_adv_comp.configure(state=state) - def __open_scheduler_params_window(self): - window = SchedulerParamsWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + def _open_optimizer_params_window(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) - def __open_timestep_distribution_window(self): - window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) + def _open_scheduler_params_window(self): + self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - - def __restore_optimizer_config(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def __restore_scheduler_config(self, variable): - if not hasattr(self, 'lr_scheduler_adv_comp'): - return + def _open_timestep_distribution_window(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) - if variable == "CUSTOM": - self.lr_scheduler_adv_comp.configure(state="normal") - else: - self.lr_scheduler_adv_comp.configure(state="disabled") + def _open_offloading_window(self): + self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py index c3291e6ea..c272c891c 100644 --- a/modules/ui/CtkVideoToolUIView.py +++ b/modules/ui/CtkVideoToolUIView.py @@ -1,33 +1,23 @@ -import concurrent.futures -import math -import os -import pathlib -import random -import shlex -import subprocess -import threading -import webbrowser -from fractions import Fraction from tkinter import filedialog +from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView +from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image -from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState -import av import customtkinter as ctk -import cv2 -import scenedetect -from PIL import Image +PAD = ctk_components.PAD -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): + +class CtkVideoToolUIView(BaseVideoToolUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: VideoToolUIController, *args, **kwargs): ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseVideoToolUIView.__init__(self, ctk_components) + + self.controller = controller + ui_state = CtkUIState(self, controller.args) self.title("Video Tools") self.geometry("600x720") @@ -42,836 +32,97 @@ def __init__( tabview = ctk.CTkTabview(self) tabview.grid(row=0, column=0, sticky="nsew") - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): + clip_frame = ctk.CTkScrollableFrame(tabview.add("extract clips"), fg_color="transparent") + clip_frame.grid_columnconfigure(0, weight=0, minsize=120) + clip_frame.grid_columnconfigure(1, weight=0, minsize=200) + clip_frame.grid_columnconfigure(2, weight=0) + clip_frame.grid_columnconfigure(3, weight=1) + self.build_clip_extract_tab(clip_frame, controller, ui_state) + clip_frame.pack(fill="both", expand=1) + + image_frame = ctk.CTkScrollableFrame(tabview.add("extract images"), fg_color="transparent") + image_frame.grid_columnconfigure(0, weight=0, minsize=120) + image_frame.grid_columnconfigure(1, weight=0, minsize=200) + image_frame.grid_columnconfigure(2, weight=0) + image_frame.grid_columnconfigure(3, weight=1) + self.build_image_extract_tab(image_frame, controller, ui_state) + image_frame.pack(fill="both", expand=1) + + download_frame = ctk.CTkScrollableFrame(tabview.add("download"), fg_color="transparent") + download_frame.grid_columnconfigure(0, weight=0, minsize=120) + download_frame.grid_columnconfigure(1, weight=0, minsize=200) + download_frame.grid_columnconfigure(2, weight=0) + download_frame.grid_columnconfigure(3, weight=1) + self.build_video_download_tab(download_frame, controller, ui_state) + download_frame.pack(fill="both", expand=1) + + self._build_status_bar(self) + + def _build_status_bar(self, master): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=1, column=0) frame.grid_columnconfigure(0, weight=0, minsize=160) frame.grid_columnconfigure(1, weight=0, minsize=300) frame.grid_columnconfigure(2, weight=1) - #create preview image preview_path = "resources/icons/icon.png" preview = load_image(preview_path, 'RGB') preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image = ctk.CTkImage(light_image=preview, size=preview.size) self.preview_image_label = ctk.CTkLabel( master=frame, text="Preview image", image=self.preview_image, height=150, width=150, compound="top") self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - #displays progress and messages that also go to terminal self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) self.status_label.insert(index="1.0", text="Current status") self.status_label.configure(state="disabled") self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", - tooltip="If enabled, detect cuts in the input video and split at those points. \ - Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - # maximum length - components.label(frame, 6, 0, "Max Length (s)", - tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") - - # Set FPS - components.label(frame, 7, 0, "Set FPS", - tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") - - # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", - tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 9, 0, "Crop Variation", - tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # image capture rate - components.label(frame, 5, 0, "Images/sec", - tooltip="Number of images to capture per second of video. \ - Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") - - # blur removal - components.label(frame, 6, 0, "Blur Removal", - tooltip="Threshold for removal of blurry images, relative to all others. \ - For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") - - # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", - tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 8, 0, "Crop Variation", - tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # link - components.label(frame, 0, 0, "Single Link", - tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) - - # link list - components.label(frame, 1, 0, "Link List", - tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) - - # output directory - components.label(frame, 2, 0, "Output", - tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - - # additional args - components.label(frame, 3, 0, "Additional Args", - tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ - Default args will hide most terminal outputs.") - self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) - self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) - self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress --format mp4") - components.button(frame, 3, 2, "yt-dlp info", - command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): - input_videos = [] - if not batch_mode: - path = pathlib.Path(input_path_single) - if path.is_file(): - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - return [path] - else: - self.__update_status("Invalid video file!") - return [] - else: - self.__update_status("No file specified, or invalid file path!") - return [] - else: - input_videos = [] - if not pathlib.Path(input_path_dir).is_dir() or input_path_dir == "": - self.__update_status("Invalid input directory!") - return [] - # Only traverse supported extensions to avoid opening every file. - lower_exts = {e.lower() for e in SUPPORTED_VIDEO_EXTENSIONS} - for path in pathlib.Path(input_path_dir).rglob("*"): - if path.is_file() and path.suffix.lower() in lower_exts: - vid = cv2.VideoCapture(str(path)) - ok = False - try: - if vid.isOpened(): - ok, _ = vid.read() - finally: - vid.release() - if ok: - input_videos.append(path) - self.__update_status(f'Found {len(input_videos)} videos to process') - return input_videos - - def __run_in_thread(self, target, *args): - """Clear status box and run target function in a daemon thread.""" + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + var = ui_state.get_var(var_name) + textbox = ctk.CTkTextbox(master, width=width, height=height, border_width=2) + textbox.insert("1.0", var.get()) + textbox.grid(row=row, column=col, rowspan=2, sticky="w", padx=PAD, pady=PAD) + + def on_text_change(event=None): + var.set(textbox.get("1.0", "end-1c")) + + textbox.bind("", on_text_change) + return textbox + + def _create_browse_dir_button(self, master, row, ui_state, var_name): + def browse(): + path = filedialog.askdirectory() + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): + def browse(): + path = filedialog.askopenfilename(filetypes=filetypes) + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def update_status(self, status_text: str): self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") + self.status_label.insert(index="end", text=status_text + "\n") self.status_label.configure(state="disabled") - t = threading.Thread(target=target, args=args) - t.daemon = True - t.start() - - @staticmethod - def __parse_timestamp_to_frames(timestamp: str, fps: float) -> int: - return int(sum(int(x) * 60 ** i for i, x in enumerate(reversed(timestamp.split(':')))) * fps) - - def __get_safe_fps(self, video: cv2.VideoCapture, video_path: str) -> float: - fps = video.get(cv2.CAP_PROP_FPS) or 0.0 - if fps <= 0: - self.__update_status(f'Warning: Could not read FPS for "{os.path.basename(video_path)}". Falling back to 30 FPS.') - return 30.0 - return fps - - @staticmethod - def __get_output_dir(use_subdir: bool, batch_mode: bool, output_entry: str, - video_path, input_dir: str) -> str: - if use_subdir and batch_mode: - return os.path.join(output_entry, - os.path.splitext(os.path.relpath(video_path, input_dir))[0]) - elif use_subdir: - return os.path.join(output_entry, - os.path.splitext(os.path.basename(video_path))[0]) - return output_entry - - def __get_random_aspect(self, height: int, width: int, variation: float) -> tuple[int, int, int, int]: - # Return original dimensions and no offset if variation is zero - if variation == 0: - return 0, height, 0, width - - old_aspect = height/width - variation_scaled = old_aspect*variation - if old_aspect > 1.2: #tall image - new_aspect = min(4.0, max(1.0, random.triangular(old_aspect-(variation_scaled*1.5), old_aspect+(variation_scaled/2), old_aspect))) - elif old_aspect < 0.85: #wide image - new_aspect = max(0.25, min(1.0, random.triangular(old_aspect-(variation_scaled/2), old_aspect+(variation_scaled*1.5), old_aspect))) - else: #square image - new_aspect = random.triangular(old_aspect-variation_scaled, old_aspect+variation_scaled) - - new_aspect = round(new_aspect, 2) - #keep the height the same if reducing width, and vice versa - if new_aspect > old_aspect: - new_height = int(height) - new_width = int(width*(old_aspect/new_aspect)) - elif new_aspect < old_aspect: - new_height = int(height*(new_aspect/old_aspect)) - new_width = int(width) - else: - new_height = int(height) - new_width = int(width) - - #random offset in dimension that was cropped - position_x = random.randint(0, width-new_width) - position_y = random.randint(0, height-new_height) - return position_y, new_height, position_x, new_width - - def find_main_contour(self, frame): - #outline image to find main content and exclude black bars often present on letterboxed videos - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - _, frame_thresh = cv2.threshold(frame_grayscale, 15, 255, cv2.THRESH_BINARY) - frame_contours, _ = cv2.findContours(frame_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - if frame_contours: - #select largest contour by area - frame_maincontour = max(frame_contours, key=lambda c: cv2.contourArea(c)) - x1, y1, w1, h1 = cv2.boundingRect(frame_maincontour) - else: #fallback if no contours detected - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - #if bounding box did not detect the correct area, likely due to all-black frame - if not frame_contours or h1 < 10 or w1 < 10: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - return x1, y1, w1, h1 - - def __extract_clips_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_clips_multi, batch_mode) - - def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") - return - if max_length <= 0.25: - self.__update_status("Max Length of clips must be > 0.25 seconds.") - return - if target_fps < 0: - self.__update_status("Target FPS must be a positive number (or 0 to skip fps re-encoding).") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if len(input_videos) == 0: # exit if no paths found - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) - executor.submit(self.__extract_clips, - str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), - crop_variation, target_fps, output_directory) - - if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') - else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') - - def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, - split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - max_length_frames = int(max_length * vid_fps) - min_length_frames = max(int(0.25 * vid_fps), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - - if split_at_cuts: - #use scenedetect to find cuts, based on start/end frame number - self.__update_status(f'Detecting scenes in "{os.path.basename(video_path)}"') - timecode_list = scenedetect.detect( - video_path=str(video_path), - detector=scenedetect.AdaptiveDetector(), - start_time=int(timestamp_min_frame), - end_time=int(timestamp_max_frame)) - scene_list = [(x[0].get_frames(), x[1].get_frames()) for x in timecode_list] - if not scene_list: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - else: - scene_list = [(timestamp_min_frame, timestamp_max_frame)] - - scene_list_split = [] - for scene in scene_list: - length = scene[1]-scene[0] - if length > max_length_frames: #check for any scenes longer than max length - n = math.ceil(length/max_length_frames) #divide into n new scenes - new_length = int(length/n) - new_splits = range(scene[0], scene[1]+min_length_frames, new_length) #divide clip into closest chunks to max_length - for i, _n in enumerate(new_splits[:-1]): - if new_splits[i + 1] - new_splits[i] > min_length_frames: - scene_list_split.append((new_splits[i], new_splits[i + 1])) - elif length > (min_length_frames + 2): - # Trim first/last frame to avoid transition artifacts - scene_list_split.append((scene[0] + 1, scene[1] - 1)) - - self.__update_status(f'Video "{os.path.basename(video_path)}" being split into {len(scene_list_split)} clips in "{output_dir}"') - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - futures = [ - executor.submit(self.__save_clip, scene, video_path, target_fps, - remove_borders, crop_variation, output_dir) - for scene in scene_list_split - ] - for future in concurrent.futures.as_completed(futures): - exc = future.exception() - if exc is not None: - self.__update_status(f'Error saving clip: {exc}') - - video.release() - - def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float, - remove_borders: bool, crop_variation: float, output_dir: str): - basename, ext = os.path.splitext(os.path.basename(video_path)) - video = cv2.VideoCapture(str(video_path)) - fps = self.__get_safe_fps(video, video_path) - os.makedirs(output_dir, exist_ok=True) - output_name = f'{output_dir}{os.sep}{basename}_{scene[0]}-{scene[1]}' - output_ext = ".mp4" - - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) #set to middle of scene - frame_number = int(video.get(cv2.CAP_PROP_POS_FRAMES)) - success, frame = video.read() - if not success or frame is None: - self.__update_status(f'Failed to read frame from "{os.path.basename(video_path)}" at {int(frame_number)}. Skipping clip.') - video.release() - return - - # Blend random frames to detect borders, avoiding incorrect crop from black frames - if remove_borders: - frame_blend = frame - for i in range(5): - random_frame = random.randint(scene[0], scene[1]) - video.set(cv2.CAP_PROP_POS_FRAMES, random_frame) - success, frame = video.read() - if not success or frame is None: - continue - a = 1/(i+1) - b = 1-a - frame_blend = cv2.addWeighted(frame, a, frame_blend, b, 0) - x1, y1, w1, h1 = self.find_main_contour(frame_blend) - else: - x1 = 0 - y1 = 0 - h1, w1, _ = frame.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - # Ensure dimensions are even, required - h2 -= h2 % 2 - w2 -= w2 % 2 - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frames {scene[0]}-{scene[1]} at size {w2}x{h2}', end="\r") - video.set(cv2.CAP_PROP_POS_FRAMES, (scene[1] + scene[0])//2) - success, frame = video.read() - if success: - try: - preview = Image.fromarray( - cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) - #truncate filename of long files so UI doesn't shift around - filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') - except Exception: - pass - video.release() - - if target_fps <= 0: - target_fps = fps - - output_path = f'{output_name}{output_ext}' - self.__write_clip_av(video_path, output_path, scene, fps, target_fps, - x1 + x2, y1 + y2, w2, h2) - - @staticmethod - def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], - src_fps: float, target_fps: float, - crop_x: int, crop_y: int, crop_w: int, crop_h: int): - start_sec = scene[0] / src_fps - end_sec = scene[1] / src_fps - rate_frac = Fraction(target_fps).limit_denominator(10000) - stream_time_base = Fraction(rate_frac.denominator, rate_frac.numerator) - - with av.open(video_path) as input_container: - in_video = input_container.streams.video[0] - in_video.thread_type = 'AUTO' - in_audio = input_container.streams.audio[0] if input_container.streams.audio else None - - with av.open(output_path, mode='w') as output_container: - out_video = output_container.add_stream('libx264', rate=rate_frac) - out_video.width = crop_w - out_video.height = crop_h - out_video.pix_fmt = 'yuv420p' - out_video.time_base = stream_time_base - - out_audio = output_container.add_stream_from_template(in_audio) if in_audio else None - - input_container.seek(int(start_sec * 1_000_000)) - - out_frame_idx = 0 - out_time_step = 1.0 / target_fps - video_done = False - decode_streams = [s for s in (in_video, in_audio) if s is not None] - - for packet in input_container.demux(decode_streams): - if packet.stream == in_video: - if video_done: - continue - for frame in packet.decode(): - if frame.time is None or frame.time < start_sec: - continue - if frame.time >= end_sec: - video_done = True - break - - # FPS conversion: skip frames when source fps > target fps - if frame.time < start_sec + out_frame_idx * out_time_step: - continue - - img = frame.to_ndarray(format='bgr24') - cropped = img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w] - out_frame = av.VideoFrame.from_ndarray(cropped, format='bgr24') - out_frame.pts = out_frame_idx - out_frame.time_base = stream_time_base - - for out_pkt in out_video.encode(out_frame): - output_container.mux(out_pkt) - out_frame_idx += 1 - - elif packet.stream == in_audio and out_audio is not None: - if packet.dts is None: - continue - pkt_time = float(packet.pts * packet.time_base) - if pkt_time < start_sec or pkt_time >= end_sec: - continue - # Re-timestamp audio relative to clip start - packet.pts = int((pkt_time - start_sec) / packet.time_base) - packet.dts = packet.pts - packet.stream = out_audio - output_container.mux(packet) - - # Flush video encoder - for pkt in out_video.encode(): - output_container.mux(pkt) - - def __extract_images_button(self, batch_mode: bool): - self.__run_in_thread(self.__extract_images_multi, batch_mode) - - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - # validate numeric inputs - try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() - except ValueError: - self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") - return - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - return - if not (0.0 <= blur_threshold < 1.0): - self.__update_status("Blur Removal must be between 0.0 and 1.0.") - return - if not (0.0 <= crop_variation < 1.0): - self.__update_status("Crop Variation must be between 0.0 and 1.0.") - return - - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) - if not input_videos: - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for video_path in input_videos: - output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) - executor.submit(self.__save_frames, - str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), - crop_variation, output_directory) - if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') - else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') - - def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, - blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): - video = cv2.VideoCapture(video_path) - vid_fps = self.__get_safe_fps(video, video_path) - if capture_rate <= 0: - self.__update_status("Images/sec must be > 0.") - video.release() - return - image_rate = max(int(vid_fps / capture_rate), 1) - total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) or 0 - timestamp_max_frame = min(self.__parse_timestamp_to_frames(timestamp_max, vid_fps), max(total_frames - 1, 0)) - timestamp_min_frame = min(self.__parse_timestamp_to_frames(timestamp_min, vid_fps), timestamp_max_frame) - frame_range = range(timestamp_min_frame, timestamp_max_frame, image_rate) - frame_list = [] - - for n in frame_range: - #pick frame from random triangular distribution around center of each "chunk" of the video - frame = abs(int(random.triangular(n-(image_rate/2), n+(image_rate/2)))) - frame = max(0, min(frame, max(total_frames - 1, 0))) - frame_list.append(frame) - - self.__update_status(f'Video "{os.path.basename(video_path)}" will be split into {len(frame_list)} images in "{output_dir}"') - - output_list = [] - for f in frame_list: - video.set(cv2.CAP_PROP_POS_FRAMES, f) - success, frame = video.read() - if success and frame is not None: - frame_grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - frame_sharpness = cv2.Laplacian(frame_grayscale, cv2.CV_64F).var() - output_list.append((f, frame_sharpness)) - - if not output_list: - self.__update_status(f'No frames extracted from "{os.path.basename(video_path)}" in the selected range.') - video.release() - return - - output_list_sorted = sorted(output_list, key=lambda x: x[1]) - cutoff = int(blur_threshold * len(output_list_sorted)) - output_list_cut = output_list_sorted[cutoff:] - self.__update_status(f'{cutoff} blurriest images have been dropped from "{os.path.basename(video_path)}"') - - basename, ext = os.path.splitext(os.path.basename(video_path)) - os.makedirs(output_dir, exist_ok=True) - - for f in output_list_cut: - filename = f'{output_dir}{os.sep}{basename}_{f[0]}.jpg' - video.set(cv2.CAP_PROP_POS_FRAMES, f[0]) - success, frame = video.read() - - #crop out borders of frame - if remove_borders and success and frame is not None: - x1, y1, w1, h1 = self.find_main_contour(frame) - frame_cropped = frame[y1:y1+h1, x1:x1+w1] - else: - frame_cropped = frame if success and frame is not None else None - if frame_cropped is not None: - x1 = 0 - y1 = 0 - h1, w1, _ = frame_cropped.shape - - y2, h2, x2, w2 = self.__get_random_aspect(h1, w1, crop_variation) - - if success and frame is not None and frame_cropped is not None: - print(end='\x1b[2K') #clear terminal so next line can overwrite it - print(f'Saving frame {f[0]} at size {w2}x{h2}', end="\r") - try: - preview = Image.fromarray( - cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) - preview.thumbnail((150, 150)) - filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') - except Exception: - pass # preview update is non-critical - - cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) - video.release() - - def __download_button(self, batch_mode: bool): - self.__run_in_thread(self.__download_multi, batch_mode) - - def __update_status(self, status_text: str): - print(status_text) + def clear_status(self): self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.delete(index1="1.0", index2="end") self.status_label.configure(state="disabled") - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": - self.__update_status("Invalid output directory!") - return - - if not batch_mode: - ydl_urls = [self.download_link_entry.get()] - elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) - if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": - with open(ydl_path) as file: - ydl_urls = file.readlines() - else: - self.__update_status("Invalid link list!") - return - - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: - for url in ydl_urls: - executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) - - self.__update_status(f'Completed {len(ydl_urls)} downloads.') - - def __download_video(self, url: str, output_dir: str, output_args: str): - url = (url or "").strip() - if not url: - self.__update_status("Empty URL, skipping download.") - return - - #Respect quotes and split into list to run as yt-dlp command - additional_args = shlex.split(output_args.strip()) if output_args and output_args.strip() else [] - cmd = ["yt-dlp", "-o", "%(title)s.%(ext)s", "-P", output_dir] + additional_args + [url] - - self.__update_status(f'Downloading {url}') - subprocess.run(cmd) - self.__update_status(f'Download {url} done!') + def update_preview(self, preview_image, label_text: str): + self.preview_image.configure(light_image=preview_image, size=preview_image.size) + self.preview_image_label.configure(text=label_text) diff --git a/modules/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py index 1690879f1..2b2411b28 100644 --- a/modules/ui/GenerateCaptionsWindowController.py +++ b/modules/ui/GenerateCaptionsWindowController.py @@ -1,133 +1,28 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog - -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateCaptionsWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating captions for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) +class GenerateCaptionsWindowController: + def __init__(self, parent): self.parent = parent + self.view = None - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all captions", "Create if absent", "Add as new line"] - self.model_var = ctk.StringVar(self, "Blip") - self.models = ["Blip", "Blip2", "WD14 VIT v2"] - - self.title("Batch generate captions") - self.geometry("360x360") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) - self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) - self.caption_entry = ctk.CTkEntry(self.frame, width=200) - self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) - self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.prefix_entry = ctk.CTkEntry(self.frame, width=200) - self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + def create_window(self, parent_window, path, parent_include_subdirectories, view_cls): + self.view = view_cls(parent_window, self, path, parent_include_subdirectories) + return self.view - self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) - self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.postfix_entry = ctk.CTkEntry(self.frame, width=200) - self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self.create_captions) - self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() - - def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_captions(self): - self.parent.load_captioning_model(self.model_var.get()) + def create_captions(self, model_name, path, initial_caption, caption_prefix, caption_postfix, mode_str, include_subdirectories): + self.parent.load_captioning_model(model_name) mode = { "Replace all captions": "replace", "Create if absent": "fill", "Add as new line": "add", - }[self.mode_var.get()] + }[mode_str] self.parent.captioning_model.caption_folder( - sample_dir=self.path_entry.get(), - initial_caption=self.caption_entry.get(), - caption_prefix=self.prefix_entry.get(), - caption_postfix=self.postfix_entry.get(), + sample_dir=path, + initial_caption=initial_caption, + caption_prefix=caption_prefix, + caption_postfix=caption_postfix, mode=mode, - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), + progress_callback=self.view.set_progress, + include_subdirectories=include_subdirectories, ) self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() diff --git a/modules/ui/GenerateMasksWindowController.py b/modules/ui/GenerateMasksWindowController.py index daff0d3d5..6c154ef46 100644 --- a/modules/ui/GenerateMasksWindowController.py +++ b/modules/ui/GenerateMasksWindowController.py @@ -1,127 +1,14 @@ -import contextlib -import tkinter as tk -from tkinter import filedialog - -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class GenerateMasksWindow(ctk.CTkToplevel): - def __init__(self, parent, path, parent_include_subdirectories, *args, **kwargs): - """ - Window for generating masks for a folder of images - - Parameters: - parent (`Tk`): the parent window - path (`str`): the path to the folder - parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox - """ - super().__init__(parent, *args, **kwargs) - +class GenerateMasksWindowController: + def __init__(self, parent): self.parent = parent - if path is None: - path = "" - - self.mode_var = ctk.StringVar(self, "Create if absent") - self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] - self.model_var = ctk.StringVar(self, "ClipSeg") - self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] - - self.title("Batch generate masks") - self.geometry("360x430") - self.resizable(True, True) - - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) - self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) - self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) - self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) - - self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) - self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) - self.path_entry = ctk.CTkEntry(self.frame, width=150) - self.path_entry.insert(0, path) - self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) - self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - - self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) - self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) - self.prompt_entry = ctk.CTkEntry(self.frame, width=200) - self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - - self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) - self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) - self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) - self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + self.view = None - self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) - self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) - self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") - self.threshold_entry.insert(0, "0.3") - self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + def create_window(self, parent_window, path, parent_include_subdirectories, view_cls): + self.view = view_cls(parent_window, self, path, parent_include_subdirectories) + return self.view - self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) - self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) - self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") - self.smooth_entry.insert(0, 5) - self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) - self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) - self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") - self.expand_entry.insert(0, 10) - self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - - self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) - self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) - self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") - self.alpha_entry.insert(0, 1) - self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) - self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) - self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) - self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) - self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) - self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) - self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) - self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) - - self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self.create_masks) - self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) - - self.frame.pack(fill="both", expand=True) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def browse_for_path(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, filedialog.END) - entry_box.insert(0, path) - self.focus_set() - - def set_progress(self, value, max_value): - progress = value / max_value - self.progress.set(progress) - self.progress_label.configure(text=f"{value}/{max_value}") - self.progress.update() - - def create_masks(self): - self.parent.load_masking_model(self.model_var.get()) + def create_masks(self, model_name, path, prompt, mode_str, alpha_str, threshold_str, smooth_str, expand_str, include_subdirectories): + self.parent.load_masking_model(model_name) mode = { "Replace all masks": "replace", @@ -129,23 +16,17 @@ def create_masks(self): "Add to existing": "add", "Subtract from existing": "subtract", "Blend with existing": "blend", - }[self.mode_var.get()] + }[mode_str] self.parent.masking_model.mask_folder( - sample_dir=self.path_entry.get(), - prompts=[self.prompt_entry.get()], + sample_dir=path, + prompts=[prompt], mode=mode, - alpha=float(self.alpha_entry.get()), - threshold=float(self.threshold_entry.get()), - smooth_pixels=int(self.smooth_entry.get()), - expand_pixels=int(self.expand_entry.get()), - progress_callback=self.set_progress, - include_subdirectories=self.include_subdirectories_var.get(), + alpha=float(alpha_str), + threshold=float(threshold_str), + smooth_pixels=int(smooth_str), + expand_pixels=int(expand_str), + progress_callback=self.view.set_progress, + include_subdirectories=include_subdirectories, ) self.parent.load_image() - - def destroy(self): - with contextlib.suppress(tk.TclError): - self.grab_release() - - super().destroy() diff --git a/modules/ui/LoraTabController.py b/modules/ui/LoraTabController.py new file mode 100644 index 000000000..46aa4aab3 --- /dev/null +++ b/modules/ui/LoraTabController.py @@ -0,0 +1,22 @@ + +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelType import PeftType + + +class LoraTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_peft_types(self) -> list[tuple[str, PeftType]]: + return [ + ("LoRA", PeftType.LORA), + ("LoHa", PeftType.LOHA), + ("OFT v2", PeftType.OFT_2), + ] + + def get_lora_weight_dtypes(self) -> list[tuple[str, DataType]]: + return [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ] diff --git a/modules/ui/ModelTabController.py b/modules/ui/ModelTabController.py new file mode 100644 index 000000000..69f603f78 --- /dev/null +++ b/modules/ui/ModelTabController.py @@ -0,0 +1,13 @@ + + +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig + + +class ModelTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_presets(self) -> dict: + cls = create.get_model_setup_class(self.train_config.model_type, self.train_config.training_method) + return cls.LAYER_PRESETS if cls is not None else {"full": []} diff --git a/modules/ui/MuonAdamWindowController.py b/modules/ui/MuonAdamWindowController.py new file mode 100644 index 000000000..834cae400 --- /dev/null +++ b/modules/ui/MuonAdamWindowController.py @@ -0,0 +1,28 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.Optimizer import Optimizer +from modules.util.optimizer_util import OPTIMIZER_DEFAULT_PARAMETERS + +MUON_AUX_ADAM_DEFAULTS = { + "beta1": 0.9, + "beta2": 0.999, + "eps": 1e-8, + "weight_decay": 0.0, +} + + + + +class MuonAdamWindowController: + def __init__(self, config: TrainConfig, parent_optimizer_type: Optimizer): + self.config = config + self.parent_optimizer_type = parent_optimizer_type + + def get_title(self) -> str: + if self.parent_optimizer_type == Optimizer.MUON: + return "Muon's Auxiliary AdamW Settings" + return "Muon_adv's Auxiliary AdamW_adv Settings" + + def get_adam_params_def(self) -> dict: + if self.parent_optimizer_type == Optimizer.MUON: + return MUON_AUX_ADAM_DEFAULTS + return OPTIMIZER_DEFAULT_PARAMETERS[Optimizer.ADAMW_ADV] diff --git a/modules/ui/OffloadingWindowController.py b/modules/ui/OffloadingWindowController.py new file mode 100644 index 000000000..cc9bad142 --- /dev/null +++ b/modules/ui/OffloadingWindowController.py @@ -0,0 +1,6 @@ +from modules.util.config.TrainConfig import TrainConfig + + +class OffloadingWindowController: + def __init__(self, config: TrainConfig): + self.config = config diff --git a/modules/ui/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py index 16063c26c..37f838907 100644 --- a/modules/ui/OptimizerParamsWindowController.py +++ b/modules/ui/OptimizerParamsWindowController.py @@ -1,7 +1,5 @@ -import contextlib -from tkinter import TclError -from modules.ui.MuonAdamWindow import MUON_AUX_ADAM_DEFAULTS, MuonAdamWindow +from modules.ui.MuonAdamWindowController import MUON_AUX_ADAM_DEFAULTS from modules.util.config.TrainConfig import TrainConfig, TrainOptimizerConfig from modules.util.enum.Optimizer import Optimizer from modules.util.optimizer_util import ( @@ -10,261 +8,27 @@ load_optimizer_defaults, update_optimizer_config, ) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState -import customtkinter as ctk +class OptimizerParamsWindowController: + def __init__(self, config: TrainConfig): + self.config = config -class OptimizerParamsWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - train_config: TrainConfig, - ui_state, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) + def restore_optimizer_config(self, ui_state): + optimizer_config = change_optimizer(self.config) + ui_state.get_var("optimizer").update(optimizer_config) - self.parent = parent - self.train_config = train_config - self.ui_state = ui_state - self.optimizer_ui_state = ui_state.get_var("optimizer") - self.protocol("WM_DELETE_WINDOW", self.on_window_close) - self.muon_adam_button = None + def load_defaults(self, ui_state): + optimizer_config = load_optimizer_defaults(self.config) + ui_state.get_var("optimizer").update(optimizer_config) - self.title("Optimizer Settings") - self.geometry("800x500") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) - - components.button(self, 1, 0, "ok", command=self.on_window_close) - self.main_frame(self.frame) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def main_frame(self, master): - # Optimizer - components.label(master, 0, 0, "Optimizer", - tooltip="The type of optimizer") - - # Create the optimizer dropdown menu and set the command - components.options(master, 0, 1, [str(x) for x in list(Optimizer)], self.optimizer_ui_state, "optimizer", - command=self.on_optimizer_change) - - # Defaults Button - components.label(master, 0, 3, "Optimizer Defaults", - tooltip="Load default settings for the selected optimizer") - components.button(self.frame, 0, 4, "Load Defaults", self.load_defaults, - tooltip="Load default settings for the selected optimizer") - - self.create_dynamic_ui(master) - - def clear_dynamic_ui(self, master): - with contextlib.suppress(TclError): - for widget in master.winfo_children(): - grid_info = widget.grid_info() - if int(grid_info["row"]) >= 1: - widget.destroy() - - def create_dynamic_ui( - self, - master, - ): - - # Lookup for the title and tooltip for a key - # @formatter:off - KEY_DETAIL_MAP = { - 'adam_w_mode': {'title': 'Adam W Mode', 'tooltip': 'Whether to use weight decay correction for Adam optimizer.', 'type': 'bool'}, - 'alpha': {'title': 'Alpha', 'tooltip': 'Smoothing parameter for RMSprop and others.', 'type': 'float'}, - 'amsgrad': {'title': 'AMSGrad', 'tooltip': 'Whether to use the AMSGrad variant for Adam.', 'type': 'bool'}, - 'beta1': {'title': 'Beta1', 'tooltip': 'optimizer_momentum term.', 'type': 'float'}, - 'beta2': {'title': 'Beta2', 'tooltip': 'Coefficients for computing running averages of gradient.', 'type': 'float'}, - 'beta3': {'title': 'Beta3', 'tooltip': 'Coefficient for computing the Prodigy stepsize.', 'type': 'float'}, - 'bias_correction': {'title': 'Bias Correction', 'tooltip': 'Whether to use bias correction in optimization algorithms like Adam.', 'type': 'bool'}, - 'block_wise': {'title': 'Block Wise', 'tooltip': 'Whether to perform block-wise model update.', 'type': 'bool'}, - 'capturable': {'title': 'Capturable', 'tooltip': 'Whether some property of the optimizer can be captured.', 'type': 'bool'}, - 'centered': {'title': 'Centered', 'tooltip': 'Whether to center the gradient before scaling. Great for stabilizing the training process.', 'type': 'bool'}, - 'clip_threshold': {'title': 'Clip Threshold', 'tooltip': 'Clipping value for gradients.', 'type': 'float'}, - 'd0': {'title': 'Initial D', 'tooltip': 'Initial D estimate for D-adaptation.', 'type': 'float'}, - 'd_coef': {'title': 'D Coefficient', 'tooltip': 'Coefficient in the expression for the estimate of d.', 'type': 'float'}, - 'dampening': {'title': 'Dampening', 'tooltip': 'Dampening for optimizer_momentum.', 'type': 'float'}, - 'decay_rate': {'title': 'Decay Rate', 'tooltip': 'Rate of decay for moment estimation.', 'type': 'float'}, - 'decouple': {'title': 'Decouple', 'tooltip': 'Use AdamW style optimizer_decoupled weight decay.', 'type': 'bool'}, - 'differentiable': {'title': 'Differentiable', 'tooltip': 'Whether the optimization function is optimizer_differentiable.', 'type': 'bool'}, - 'eps': {'title': 'EPS', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'eps2': {'title': 'EPS 2', 'tooltip': 'A small value to prevent division by zero.', 'type': 'float'}, - 'foreach': {'title': 'ForEach', 'tooltip': 'Whether to use a foreach implementation if available. This implementation is usually faster.', 'type': 'bool'}, - 'fsdp_in_use': {'title': 'FSDP in Use', 'tooltip': 'Flag for using sharded parameters.', 'type': 'bool'}, - 'fused': {'title': 'Fused', 'tooltip': 'Whether to use a fused implementation if available. This implementation is usually faster and requires less memory.', 'type': 'bool'}, - 'fused_back_pass': {'title': 'Fused Back Pass', 'tooltip': 'Whether to fuse the back propagation pass with the optimizer step. This reduces VRAM usage, but is not compatible with gradient accumulation.', 'type': 'bool'}, - 'growth_rate': {'title': 'Growth Rate', 'tooltip': 'Limit for D estimate growth rate.', 'type': 'float'}, - 'initial_accumulator_value': {'title': 'Initial Accumulator Value', 'tooltip': 'Initial value for Adagrad optimizer.', 'type': 'float'}, - 'initial_accumulator': {'title': 'Initial Accumulator', 'tooltip': 'Sets the starting value for both moment estimates to ensure numerical stability and balanced adaptive updates early in training.', 'type': 'float'}, - 'is_paged': {'title': 'Is Paged', 'tooltip': 'Whether the optimizer\'s internal state should be paged to CPU.', 'type': 'bool'}, - 'log_every': {'title': 'Log Every', 'tooltip': 'Intervals at which logging should occur.', 'type': 'int'}, - 'lr_decay': {'title': 'LR Decay', 'tooltip': 'Rate at which learning rate decreases.', 'type': 'float'}, - 'max_unorm': {'title': 'Max Unorm', 'tooltip': 'Maximum value for gradient clipping by norms.', 'type': 'float'}, - 'maximize': {'title': 'Maximize', 'tooltip': 'Whether to optimizer_maximize the optimization function.', 'type': 'bool'}, - 'min_8bit_size': {'title': 'Min 8bit Size', 'tooltip': 'Minimum tensor size for 8-bit quantization.', 'type': 'int'}, - 'quant_block_size': {'title': 'Quant Block Size', 'tooltip': 'Size of a block of normalized 8-bit quantization data. Larger values increase memory efficiency at the cost of data precision.', 'type': 'int'}, - 'momentum': {'title': 'optimizer_momentum', 'tooltip': 'Factor to accelerate SGD in relevant direction.', 'type': 'float'}, - 'nesterov': {'title': 'Nesterov', 'tooltip': 'Whether to enable Nesterov optimizer_momentum.', 'type': 'bool'}, - 'no_prox': {'title': 'No Prox', 'tooltip': 'Whether to use proximity updates or not.', 'type': 'bool'}, - 'optim_bits': {'title': 'Optim Bits', 'tooltip': 'Number of bits used for optimization.', 'type': 'int'}, - 'percentile_clipping': {'title': 'Percentile Clipping', 'tooltip': 'Gradient clipping based on percentile values.', 'type': 'int'}, - 'relative_step': {'title': 'Relative Step', 'tooltip': 'Whether to use a relative step size.', 'type': 'bool'}, - 'safeguard_warmup': {'title': 'Safeguard Warmup', 'tooltip': 'Avoid issues during warm-up stage.', 'type': 'bool'}, - 'scale_parameter': {'title': 'Scale Parameter', 'tooltip': 'Whether to scale the parameter or not.', 'type': 'bool'}, - 'stochastic_rounding': {'title': 'Stochastic Rounding', 'tooltip': 'Stochastic rounding for weight updates. Improves quality when using bfloat16 weights.', 'type': 'bool'}, - 'use_bias_correction': {'title': 'Bias Correction', 'tooltip': 'Turn on Adam\'s bias correction.', 'type': 'bool'}, - 'use_triton': {'title': 'Use Triton', 'tooltip': 'Whether Triton optimization should be used.', 'type': 'bool'}, - 'warmup_init': {'title': 'Warmup Initialization', 'tooltip': 'Whether to warm-up the optimizer initialization.', 'type': 'bool'}, - 'weight_decay': {'title': 'Weight Decay', 'tooltip': 'Regularization to prevent overfitting.', 'type': 'float'}, - 'weight_lr_power': {'title': 'Weight LR Power', 'tooltip': 'During warmup, the weights in the average will be equal to lr raised to this power. Set to 0 for no weighting.', 'type': 'float'}, - 'decoupled_decay': {'title': 'Decoupled Decay', 'tooltip': 'If set as True, then the optimizer uses decoupled weight decay as in AdamW.', 'type': 'bool'}, - 'fixed_decay': {'title': 'Fixed Decay', 'tooltip': '(When Decoupled Decay is True:) Applies fixed weight decay when True; scales decay with learning rate when False.', 'type': 'bool'}, - 'rectify': {'title': 'Rectify', 'tooltip': 'Perform the rectified update similar to RAdam.', 'type': 'bool'}, - 'degenerated_to_sgd': {'title': 'Degenerated to SGD', 'tooltip': 'Performs SGD update when gradient variance is high.', 'type': 'bool'}, - 'k': {'title': 'K', 'tooltip': 'Number of vector projected per iteration.', 'type': 'int'}, - 'xi': {'title': 'Xi', 'tooltip': 'Term used in vector projections to avoid division by zero.', 'type': 'float'}, - 'n_sma_threshold': {'title': 'N SMA Threshold', 'tooltip': 'Number of SMA threshold.', 'type': 'int'}, - 'ams_bound': {'title': 'AMS Bound', 'tooltip': 'Whether to use the AMSBound variant.', 'type': 'bool'}, - 'r': {'title': 'R', 'tooltip': 'EMA factor.', 'type': 'float'}, - 'adanorm': {'title': 'AdaNorm', 'tooltip': 'Whether to use the AdaNorm variant', 'type': 'bool'}, - 'adam_debias': {'title': 'Adam Debias', 'tooltip': 'Only correct the denominator to avoid inflating step sizes early in training.', 'type': 'bool'}, - 'slice_p': {'title': 'Slice parameters', 'tooltip': 'Reduce memory usage by calculating LR adaptation statistics on only every pth entry of each tensor. For values greater than 1 this is an approximation to standard Prodigy. Values ~11 are reasonable.', 'type': 'int'}, - 'cautious': {'title': 'Cautious', 'tooltip': 'Whether to use the Cautious variant', 'type': 'bool'}, - 'weight_decay_by_lr': {'title': 'weight_decay_by_lr', 'tooltip': 'Automatically adjust weight decay based on lr', 'type': 'bool'}, - 'prodigy_steps': {'title': 'prodigy_steps', 'tooltip': 'Turn off Prodigy after N steps', 'type': 'int'}, - 'use_speed': {'title': 'use_speed', 'tooltip': 'use_speed method', 'type': 'bool'}, - 'split_groups': {'title': 'split_groups', 'tooltip': 'Use split groups when training multiple params(uNet,TE..)', 'type': 'bool'}, - 'split_groups_mean': {'title': 'split_groups_mean', 'tooltip': 'Use mean for split groups', 'type': 'bool'}, - 'factored': {'title': 'factored', 'tooltip': 'Use factored', 'type': 'bool'}, - 'factored_fp32': {'title': 'factored_fp32', 'tooltip': 'Use factored_fp32', 'type': 'bool'}, - 'use_stableadamw': {'title': 'use_stableadamw', 'tooltip': 'Use use_stableadamw for gradient scaling', 'type': 'bool'}, - 'use_cautious': {'title': 'use_cautious', 'tooltip': 'Use cautious method', 'type': 'bool'}, - 'use_grams': {'title': 'use_grams', 'tooltip': 'Use grams method', 'type': 'bool'}, - 'use_adopt': {'title': 'use_adopt', 'tooltip': 'Use adopt method', 'type': 'bool'}, - 'd_limiter': {'title': 'd_limiter', 'tooltip': 'Prevent over-estimated LRs when gradients and EMA are still stabilizing', 'type': 'bool'}, - 'use_schedulefree': {'title': 'use_schedulefree', 'tooltip': 'Use Schedulefree method', 'type': 'bool'}, - 'use_orthograd': {'title': 'use_orthograd', 'tooltip': 'Use orthograd method', 'type': 'bool'}, - 'nnmf_factor': {'title': 'Factored Optimizer', 'tooltip': 'Enables a memory-efficient mode by applying fast low-rank factorization to the optimizers states. It combines factorization for magnitudes with 1-bit compression for signs, drastically reducing VRAM usage and allowing for larger models or batch sizes. This is an approximation which may slightly alter training dynamics.', 'type': 'bool'}, - 'orthogonal_gradient': {'title': 'OrthoGrad', 'tooltip': 'Reduces overfitting by removing the gradient component parallel to the weight, thus improving generalization.', 'type': 'bool'}, - 'use_atan2': {'title': 'Atan2 Scaling', 'tooltip': 'A robust replacement for eps, which also incorporates gradient clipping, bounding and stabilizing the optimizer updates.', 'type': 'bool'}, - 'use_AdEMAMix': {'title': 'AdEMAMix EMA', 'tooltip': 'Adds a second, slow-moving EMA, which is combined with the primary momentum to stabilize updates, and accelerate the training.', 'type': 'bool'}, - 'beta3_ema': {'title': 'Beta3 EMA', 'tooltip': 'Coefficient for slow-moving EMA of AdEMAMix.', 'type': 'float'}, - 'beta1_warmup': {'title': 'Beta1 Warmup Steps', 'tooltip': 'Number of warmup steps to gradually increase beta1 from Minimum Beta1 Value to its final value. During warmup, beta1 increases linearly. leave it empty to disable warmup and use constant beta1.', 'type': 'int'}, - 'min_beta1': {'title': 'Minimum Beta1', 'tooltip': 'Starting beta1 value for warmup scheduling. Used only when beta1 warmup is enabled. Lower values allow faster initial adaptation, while higher values provide more smoothing. The final beta1 value is specified in the beta1 parameter.', 'type': 'float'}, - 'Simplified_AdEMAMix': {'title': 'Simplified AdEMAMix', 'tooltip': "Enables a simplified, single-EMA variant of AdEMAMix. Instead of blending two moving averages (fast and slow momentum), this version combines the raw current gradient (controlled by 'Grad α') directly with a single theory-based momentum. This makes the optimizer highly responsive to recent gradient information, which can accelerate training in all batch size scenarios when tuned correctly.", 'type': 'bool'}, - 'alpha_grad': {'title': 'Grad α', 'tooltip': 'Controls the mixing coefficient between raw gradients and momentum gradients in Simplified AdEMAMix. Higher values (e.g., 10-100) emphasize recent gradients, suitable for small batch sizes to reduce noise. Lower values (e.g., 0-1) emphasize historical gradients, suitable for large batch sizes for stability. Setting to 0 uses only momentum gradients without raw gradient contribution.', 'type': 'float'}, - 'kourkoutas_beta': {'title': 'Kourkoutas Beta', 'tooltip': 'Enables a layer-wise dynamic β₂ adaptation. This feature makes the optimizer more responsive to "spiky" gradients by lowering β₂ during periods of high variance, and more stable during calm periods by raising β₂ towards its maximum. It can significantly improve training stability and final loss.', 'type': 'bool'}, - 'schedulefree_c': {'title': 'Schedule free averaging strength', 'tooltip': 'Larger values = more responsive (shorter averaging window); smaller values = smoother (longer window). Set to 0 to disable and use the original Schedule-Free rule. Short small batches (≈6-12); long/large-batch (≈50-200).', 'type': 'float'}, - 'ns_steps': {'title': 'Newton-Schulz Iterations', 'tooltip': 'Controls the number of iterations for update orthogonalization. Higher values improve the updates quality but make each step slower. Lower values are faster per step but may be less effective.', 'type': 'int'}, - 'MuonWithAuxAdam': {'title': 'MuonWithAuxAdam', 'tooltip': 'Whether to use the standard way of Muon. Non-hidden layers fallback to ADAMW, and MUON takes the rest. Note: The auxiliary Adam (ADAMW) is typically only relevant for training "full" LoRA (LoRA for all layers) or full finetune and is irrelevant for most common LoRA use cases.', 'type': 'bool'}, - 'muon_hidden_layers': {'title': 'Hidden Layers', 'tooltip': 'Comma-separated list of hidden layers to train using Muon. Regular expressions (if toggled) are supported. Any model layer with a matching name will be trained using Muon. If None is provided it will default to using automatic way of finding hidden layers.', 'type': 'str'}, - 'muon_adam_regex': {'title': 'Use Regex', 'tooltip': 'Whether to use regular expressions for hidden layers.', 'type': 'bool'}, - 'muon_adam_lr': {'title': 'Auxiliary Adam LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer. If empty, it will use the main learning rate.', 'type': 'float'}, - 'muon_te1_adam_lr': {'title': 'AuxAdam TE1 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the first text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'muon_te2_adam_lr': {'title': 'AuxAdam TE2 LR', 'tooltip': 'Learning rate for the auxiliary AdamW optimizer for the second text encoder. If empty, it will use the Auxiliary Adam LR.', 'type': 'float'}, - 'rms_rescaling': {'title': 'RMS Rescaling', 'tooltip': 'Muon already scales its updates to approximate and use the same learning rate (LR) as Adam. This option integrates a more accurate method to match the Adam LR, but it is slower.', 'type': 'bool'}, - 'normuon_variant': {'title': 'NorMuon Variant', 'tooltip': 'Enables the NorMuon optimizer variant, which combines Muon orthogonalization with per-neuron adaptive learning rates for better convergence and balanced parameter updates. Costs only one scalar state buffer per parameter group, size few KBs, maintaining high memory efficiency.', 'type': 'bool'}, - 'beta2_normuon': {'title': 'NorMuon Beta2', 'tooltip': 'Exponential decay rate for the neuron-wise second-moment estimator in NorMuon (analogous to Adams beta2). Controls how past squared updates influence current normalization.', 'type': 'float'}, - 'low_rank_ortho': {'title': 'Low-rank Orthogonalization', 'tooltip': 'Use low-rank orthogonalization to accelerate Muon by orthogonalizing only in a low-dimensional subspace, improving speed and noise robustness.', 'type': 'bool'}, - 'ortho_rank': {'title': 'Ortho Rank', 'tooltip': 'Target rank for low-rank orthogonalization. Controls the dimensionality of the subspace used for efficient and noise-robust orthogonalization.', 'type': 'int'}, - 'accelerated_ns': {'title': 'Accelerated Newton-Schulz', 'tooltip': 'Applies an enhanced Newton-Schulz variant that replaces heuristic coefficients with optimal coefficients derived at each step. This improves performance and convergence by reducing the number of required operations.', 'type': 'bool'}, - 'cautious_wd': {'title': 'Cautious Weight Decay', 'tooltip': 'Applies weight decay only to parameter coordinates whose signs align with the optimizer update direction. This preserves the original optimization objective while still benefiting from regularization effects, leading to improved convergence and better final performance.', 'type': 'bool'}, - 'approx_mars': {'title': 'Approx MARS-M', 'tooltip': 'Enables Approximated MARS-M, a variance reduction technique. It uses the previous step\'s gradient to correct the current update, leading to lower losses and improved convergence stability. This requires additional state to store the previous gradient.', 'type': 'bool'}, - 'auto_kappa_p': {'title': 'Auto Lion-K', 'tooltip': 'Automatically determines the optimal P-value based on layer dimensions. Uses p=2.0 (Spherical) for 4D (Conv) tensors for stability and rotational invariance, and p=1.0 (Sign) for 2D (Linear) tensors for sparsity. Overrides the manual P-value. Recommend for unet models.', 'type': 'bool'}, - 'compile': {'title': 'Compiled Optimizer', 'tooltip': 'Enables PyTorch compilation for the optimizer internal step logic. This is intended to improve performance by allowing PyTorch to fuse operations and optimize the computational graph.', 'type': 'bool'}, - } - # @formatter:on - - if not self.winfo_exists(): # check if this window isn't open - return - - selected_optimizer = self.train_config.optimizer.optimizer - - # Extract the keys for the selected optimizer - for index, key in enumerate(OPTIMIZER_DEFAULT_PARAMETERS[selected_optimizer].keys()): - if key not in KEY_DETAIL_MAP: - continue - arg_info = KEY_DETAIL_MAP[key] - - title = arg_info['title'] - tooltip = arg_info['tooltip'] - type = arg_info['type'] - - row = (index // 2) + 1 - col = 3 * (index % 2) - - components.label(master, row, col, title, tooltip=tooltip) - - if key == 'MuonWithAuxAdam': - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col + 1, columnspan=2, sticky="ew", padx=0, pady=0) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - - components.switch(frame, 0, 0, self.optimizer_ui_state, key, command=self.update_user_pref) - - self.muon_adam_button = components.button( - frame, 0, 1, "...", self.open_muon_adam_window, - tooltip="Configure the auxiliary AdamW_adv optimizer", - width=20, padx=5 ) - self.toggle_muon_adam_button() - elif type != 'bool': - components.entry(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - else: - components.switch(master, row, col + 1, self.optimizer_ui_state, key, - command=self.update_user_pref) - - def update_user_pref(self, *args): - update_optimizer_config(self.train_config) - self.toggle_muon_adam_button() - - def on_optimizer_change(self, *args): - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.clear_dynamic_ui(self.frame) - self.create_dynamic_ui(self.frame) - - def load_defaults(self, *args): - optimizer_config = load_optimizer_defaults(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - def on_window_close(self): - self.destroy() - - def toggle_muon_adam_button(self): - if self.muon_adam_button and self.muon_adam_button.winfo_exists(): - muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() - self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") - - def open_muon_adam_window(self): - current_optimizer = self.train_config.optimizer.optimizer + def on_close(self): + update_optimizer_config(self.config) + def prepare_muon_adam_config(self) -> tuple['TrainOptimizerConfig', Optimizer]: + current_optimizer = self.config.optimizer.optimizer adam_config = TrainOptimizerConfig.default_values() - current_state = self.train_config.optimizer.muon_adam_config + current_state = self.config.optimizer.muon_adam_config if current_optimizer == Optimizer.MUON: defaults = MUON_AUX_ADAM_DEFAULTS @@ -281,8 +45,7 @@ def open_muon_adam_window(self): # Should not happen if TrainConfig defines it as dict, but for safety adam_config = current_state - temp_adam_ui_state = UIState(self, adam_config) - window = MuonAdamWindow(self, self.train_config, temp_adam_ui_state, current_optimizer) - self.wait_window(window) + return adam_config, current_optimizer - self.train_config.optimizer.muon_adam_config = adam_config.to_dict() + def save_muon_adam_config(self, adam_config: 'TrainOptimizerConfig'): + self.config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/ProfilingWindowController.py b/modules/ui/ProfilingWindowController.py new file mode 100644 index 000000000..4e1f2d980 --- /dev/null +++ b/modules/ui/ProfilingWindowController.py @@ -0,0 +1,25 @@ +import faulthandler + + +class ProfilingWindowController: + def __init__(self): + self.view = None + + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + + def dump_stack(self): + with open('stacks.txt', 'w') as f: + faulthandler.dump_traceback(f) + self.view.set_message('Stack dumped to stacks.txt') + + def start_profiler(self): + from scalene import scalene_profiler + scalene_profiler.start() + self.view.set_profiling_active(True) + + def end_profiler(self): + from scalene import scalene_profiler + scalene_profiler.stop() + self.view.set_profiling_active(False) diff --git a/modules/ui/SampleFrameController.py b/modules/ui/SampleFrameController.py new file mode 100644 index 000000000..474c52ab8 --- /dev/null +++ b/modules/ui/SampleFrameController.py @@ -0,0 +1,17 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType + + +class SampleFrameController: + def __init__(self, sample: SampleConfig, model_type: ModelType): + self.sample = sample + self.model_type = model_type + + def is_flow_matching(self) -> bool: + return self.model_type.is_flow_matching() + + def is_inpainting_model(self) -> bool: + return self.model_type.has_conditioning_image_input() + + def is_video_model(self) -> bool: + return self.model_type.is_video_model() diff --git a/modules/ui/SampleParamsWindowController.py b/modules/ui/SampleParamsWindowController.py new file mode 100644 index 000000000..abe7c8b33 --- /dev/null +++ b/modules/ui/SampleParamsWindowController.py @@ -0,0 +1,8 @@ +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.ModelType import ModelType + + +class SampleParamsWindowController: + def __init__(self, sample: SampleConfig, model_type: ModelType | None = None): + self.sample = sample + self.model_type = model_type diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py index 0f91ad2fa..d1d24d643 100644 --- a/modules/ui/SampleWindowController.py +++ b/modules/ui/SampleWindowController.py @@ -1,49 +1,34 @@ -import contextlib import copy import os -import tkinter as tk -import traceback from modules.model.BaseModel import BaseModel from modules.modelSampler.BaseModelSampler import ( BaseModelSampler, - ModelSamplerOutput, ) -from modules.ui.SampleFrame import SampleFrame from modules.util import create from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.SampleConfig import SampleConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.FileType import FileType from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.time_util import get_string_timestamp -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import torch -import customtkinter as ctk -from PIL import Image - -class SampleWindow(ctk.CTkToplevel): +class SampleWindowController: def __init__( self, - parent, train_config: TrainConfig, use_external_model: bool, callbacks: TrainCallbacks | None = None, commands: TrainCommands | None = None, - *args, **kwargs ): - super().__init__(parent, *args, **kwargs) - - self.title("Sample") - self.geometry("1200x800") - self.resizable(True, True) + self.current_train_config = train_config + self.use_external_model = use_external_model + self.callbacks = callbacks + self.commands = commands if not use_external_model: self.initial_train_config = TrainConfig.default_values().from_dict(train_config.to_dict()) @@ -55,52 +40,19 @@ def __init__( #TODO why is there a current_train_config and an initial_train_config? #current_train_config doesn't seem to ever change - self.current_train_config = train_config - self.callbacks = callbacks - self.commands = commands # get model specific defaults model_type = train_config.model_type self.sample = SampleConfig.default_values(model_type) - self.ui_state = UIState(self, self.sample) - if use_external_model: - self.callbacks.set_on_sample_custom(self.__update_preview) - self.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - else: + if not use_external_model: self.model = None self.model_sampler = None - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_rowconfigure(3, weight=0) - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1) - - prompt_frame = SampleFrame(self, self.sample, self.ui_state, include_settings=False, model_type=model_type) - prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + def get_model_type(self): + return self.current_train_config.model_type - settings_frame = SampleFrame(self, self.sample, self.ui_state, include_prompt=False, model_type=model_type) - settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - # image - self.image = ctk.CTkImage( - light_image=self.__dummy_image(), - size=(512, 512) - ) - - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) - image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") - - self.progress = components.progress(self, 2, 0) - components.button(self, 3, 0, "sample", self.__sample) - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def __load_model(self) -> BaseModel: + def load_model(self) -> BaseModel: model_loader = create.create_model_loader( model_type=self.initial_train_config.model_type, training_method=self.initial_train_config.training_method, @@ -148,7 +100,7 @@ def __load_model(self) -> BaseModel: return model - def __create_sampler(self, model: BaseModel) -> BaseModelSampler: + def create_sampler(self, model: BaseModel) -> BaseModelSampler: return create.create_model_sampler( train_device=torch.device(self.initial_train_config.train_device), temp_device=torch.device(self.initial_train_config.temp_device), @@ -157,22 +109,7 @@ def __create_sampler(self, model: BaseModel) -> BaseModelSampler: training_method=self.initial_train_config.training_method, ) - def __update_preview(self, sampler_output: ModelSamplerOutput): - if sampler_output.file_type == FileType.IMAGE: - image = sampler_output.data - self.image.configure( - light_image=image, - size=(image.width, image.height), - ) - - def __update_progress(self, progress: int, max_progress: int): - self.progress.set(progress / max_progress) - self.update() - - def __dummy_image(self) -> Image: - return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - - def __sample(self): + def do_sample(self, on_sample, on_update_progress): sample = copy.copy(self.sample) if self.commands: @@ -180,8 +117,8 @@ def __sample(self): else: if self.model is None: # lazy initialization - self.model = self.__load_model() - self.model_sampler = self.__create_sampler(self.model) + self.model = self.load_model() + self.model_sampler = self.create_sampler(self.model) sample.from_train_config(self.current_train_config) @@ -205,23 +142,6 @@ def __sample(self): image_format=self.current_train_config.sample_image_format, video_format=self.current_train_config.sample_video_format, audio_format=self.current_train_config.sample_audio_format, - on_sample=self.__update_preview, - on_update_progress=self.__update_progress, + on_sample=on_sample, + on_update_progress=on_update_progress, ) - - def destroy(self): - try: - if hasattr(self, "_icon_image_ref"): - del self._icon_image_ref - - # Remove any pending after callbacks - for after_id in self.tk.call('after', 'info'): - with contextlib.suppress(tk.TclError, RuntimeError): - self.after_cancel(after_id) - - super().destroy() - except (tk.TclError, RuntimeError) as e: - print(f"Error destroying window: {e}") - except Exception as e: - print(f"Unexpected error destroying window: {e}") - traceback.print_exc() diff --git a/modules/ui/SamplingTabController.py b/modules/ui/SamplingTabController.py new file mode 100644 index 000000000..ef95c77fa --- /dev/null +++ b/modules/ui/SamplingTabController.py @@ -0,0 +1,14 @@ +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.config.SampleConfig import SampleConfig +from modules.util.config.TrainConfig import TrainConfig + + +class SamplingTabController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def create_new_element(self) -> SampleConfig: + return SampleConfig.default_values(self.train_config.model_type) + + def open_element_window(self, parent, sample_config, ui_state, view_cls): + return view_cls(parent, SampleParamsWindowController(sample_config, model_type=self.train_config.model_type), ui_state) diff --git a/modules/ui/SchedulerParamsWindowController.py b/modules/ui/SchedulerParamsWindowController.py new file mode 100644 index 000000000..363391b04 --- /dev/null +++ b/modules/ui/SchedulerParamsWindowController.py @@ -0,0 +1,17 @@ +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.LearningRateScheduler import LearningRateScheduler + + +class SchedulerParamsWindowController: + def __init__(self, config: TrainConfig): + self.config = config + + def is_custom_scheduler(self) -> bool: + return self.config.learning_rate_scheduler is LearningRateScheduler.CUSTOM + +class KvParamsController: + def __init__(self, train_config: TrainConfig): + self.train_config = train_config + + def create_new_element(self) -> dict[str, str]: + return {"key": "", "value": ""} diff --git a/modules/ui/TimestepDistributionWindowController.py b/modules/ui/TimestepDistributionWindowController.py index 21e41ce3e..682508599 100644 --- a/modules/ui/TimestepDistributionWindowController.py +++ b/modules/ui/TimestepDistributionWindowController.py @@ -1,21 +1,11 @@ -from modules.modelSetup.mixin.ModelSetupNoiseMixin import ( - ModelSetupNoiseMixin, -) +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.TimestepDistribution import TimestepDistribution -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState import torch from torch import Tensor -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager -from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg - class TimestepGenerator(ModelSetupNoiseMixin): @@ -59,128 +49,20 @@ def generate(self) -> Tensor: ) -class TimestepDistributionWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.title("Timestep Distribution") - self.geometry("900x600") - self.resizable(True, True) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - frame.grid_rowconfigure(7, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Timestep Distribution", - tooltip="Selects the function to sample timesteps during training", - wide_tooltip=True) - components.options(frame, 0, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, - "timestep_distribution") - - # min noising strength - components.label(frame, 1, 0, "Min Noising Strength", - tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained") - components.entry(frame, 1, 1, self.ui_state, "min_noising_strength") - - # max noising strength - components.label(frame, 2, 0, "Max Noising Strength", - tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition") - components.entry(frame, 2, 1, self.ui_state, "max_noising_strength") +class TimestepDistributionWindowController: + def __init__(self, config: TrainConfig): + self.train_config = config - # noising weight - components.label(frame, 3, 0, "Noising Weight", - tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 3, 1, self.ui_state, "noising_weight") + def get_distribution_options(self) -> list[str]: + return [str(x) for x in list(TimestepDistribution)] - # noising bias - components.label(frame, 4, 0, "Noising Bias", - tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.") - components.entry(frame, 4, 1, self.ui_state, "noising_bias") - - # timestep shift - components.label(frame, 5, 0, "Timestep Shift", - tooltip="Shift the timestep distribution. Use the preview to see more details.") - components.entry(frame, 5, 1, self.ui_state, "timestep_shift") - - # dynamic timestep shifting - components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 6, 1, self.ui_state, "dynamic_timestep_shifting") - - - # plot - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - fig, ax = plt.subplots() - self.ax = ax - self.canvas = FigureCanvasTkAgg(fig, master=frame) - self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) - - fig.set_facecolor(background_color) - ax.set_facecolor(background_color) - ax.spines['bottom'].set_color(text_color) - ax.spines['left'].set_color(text_color) - ax.spines['top'].set_color(text_color) - ax.spines['right'].set_color(text_color) - ax.tick_params(axis='x', colors=text_color, which="both") - ax.tick_params(axis='y', colors=text_color, which="both") - ax.xaxis.label.set_color(text_color) - ax.yaxis.label.set_color(text_color) - - self.__update_preview() - - # update button - components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) - - frame.pack(fill="both", expand=1) - return frame - - def __update_preview(self): + def generate_preview_data(self) -> Tensor: generator = TimestepGenerator( - timestep_distribution=self.config.timestep_distribution, - min_noising_strength=self.config.min_noising_strength, - max_noising_strength=self.config.max_noising_strength, - noising_weight=self.config.noising_weight, - noising_bias=self.config.noising_bias, - timestep_shift=self.config.timestep_shift, + timestep_distribution=self.train_config.timestep_distribution, + min_noising_strength=self.train_config.min_noising_strength, + max_noising_strength=self.train_config.max_noising_strength, + noising_weight=self.train_config.noising_weight, + noising_bias=self.train_config.noising_bias, + timestep_shift=self.train_config.timestep_shift, ) - - self.ax.cla() - self.ax.hist(generator.generate(), bins=1000, range=(0, 999)) - self.canvas.draw() - - def __ok(self): - self.destroy() + return generator.generate() diff --git a/modules/ui/TopBarController.py b/modules/ui/TopBarController.py index 820fdb71a..31d3fc1be 100644 --- a/modules/ui/TopBarController.py +++ b/modules/ui/TopBarController.py @@ -1,260 +1,103 @@ -import json import os -import traceback import webbrowser -from collections.abc import Callable -from contextlib import suppress from modules.util import path_util -from modules.util.config.SecretsConfig import SecretsConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.optimizer_util import change_optimizer from modules.util.path_util import write_json_atomic -from modules.util.ui import components, dialogs -from modules.util.ui.UIState import UIState -import customtkinter as ctk - -class TopBar: - def __init__( - self, - master, - train_config: TrainConfig, - ui_state: UIState, - change_model_type_callback: Callable[[ModelType], None], - change_training_method_callback: Callable[[TrainingMethod], None], - load_preset_callback: Callable[[], None], - ): - self.master = master - self.train_config = train_config - self.ui_state = ui_state - self.change_model_type_callback = change_model_type_callback - self.change_training_method_callback = change_training_method_callback - self.load_preset_callback = load_preset_callback - - self.dir = "training_presets" - - self.config_ui_data = { - "config_name": path_util.canonical_join(self.dir, "#.json") - } - self.config_ui_state = UIState(master, self.config_ui_data) - - self.configs = [("", path_util.canonical_join(self.dir, "#.json"))] - self.__load_available_config_names() - - self.current_config = [] - - self.frame = ctk.CTkFrame(master=master, corner_radius=0) - self.frame.grid(row=0, column=0, sticky="nsew") - - self.training_method = None - - # title - components.app_title(self.frame, 0, 0) - - # dropdown - self.configs_dropdown = None - self.__create_configs_dropdown() - - # remove button - # TODO - # components.icon_button(self.frame, 0, 2, "-", self.__remove_config) - - # Wiki button - components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) - - # save button - components.button(self.frame, 0, 3, "Save config", self.__save_config, - tooltip="Save the current configuration in a custom preset", width=90) - - # padding - self.frame.grid_columnconfigure(5, weight=1) - - # model type - components.options_kv( - master=self.frame, - row=0, - column=6, - values=[ #TODO simplify - ("SD1.5", ModelType.STABLE_DIFFUSION_15), - ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), - ("SD2.0", ModelType.STABLE_DIFFUSION_20), - ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), - ("SD2.1", ModelType.STABLE_DIFFUSION_21), - ("SD3", ModelType.STABLE_DIFFUSION_3), - ("SD3.5", ModelType.STABLE_DIFFUSION_35), - ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), - ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), - ("Wuerstchen v2", ModelType.WUERSTCHEN_2), - ("Stable Cascade", ModelType.STABLE_CASCADE_1), - ("PixArt Alpha", ModelType.PIXART_ALPHA), - ("PixArt Sigma", ModelType.PIXART_SIGMA), - ("Flux Dev.1", ModelType.FLUX_DEV_1), - ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), - ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), - ("Sana", ModelType.SANA), - ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), - ("HiDream Full", ModelType.HI_DREAM_FULL), - ("Chroma1", ModelType.CHROMA_1), - ("QwenImage", ModelType.QWEN), - ("Z-Image", ModelType.Z_IMAGE), - ("Ernie Image", ModelType.ERNIE), - ], - ui_state=self.ui_state, - var_name="model_type", - command=self.__change_model_type, - ) - - def __create_training_method(self): - if self.training_method: - self.training_method.destroy() - - values = [] +class TopBarController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_model_types(self) -> list[tuple[str, ModelType]]: + return [ #TODO simplify + ("SD1.5", ModelType.STABLE_DIFFUSION_15), + ("SD1.5 Inpainting", ModelType.STABLE_DIFFUSION_15_INPAINTING), + ("SD2.0", ModelType.STABLE_DIFFUSION_20), + ("SD2.0 Inpainting", ModelType.STABLE_DIFFUSION_20_INPAINTING), + ("SD2.1", ModelType.STABLE_DIFFUSION_21), + ("SD3", ModelType.STABLE_DIFFUSION_3), + ("SD3.5", ModelType.STABLE_DIFFUSION_35), + ("SDXL", ModelType.STABLE_DIFFUSION_XL_10_BASE), + ("SDXL Inpainting", ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING), + ("Wuerstchen v2", ModelType.WUERSTCHEN_2), + ("Stable Cascade", ModelType.STABLE_CASCADE_1), + ("PixArt Alpha", ModelType.PIXART_ALPHA), + ("PixArt Sigma", ModelType.PIXART_SIGMA), + ("Flux Dev.1", ModelType.FLUX_DEV_1), + ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), + ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Sana", ModelType.SANA), + ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), + ("HiDream Full", ModelType.HI_DREAM_FULL), + ("Chroma1", ModelType.CHROMA_1), + ("QwenImage", ModelType.QWEN), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ] + + def get_training_methods(self, model_type: ModelType) -> list[tuple[str, TrainingMethod]]: #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ + if model_type.is_stable_diffusion(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ + elif model_type.is_stable_diffusion_3() \ + or model_type.is_stable_diffusion_xl() \ + or model_type.is_wuerstchen() \ + or model_type.is_pixart() \ + or model_type.is_flux_1() \ + or model_type.is_sana() \ + or model_type.is_hunyuan_video() \ + or model_type.is_hi_dream() \ + or model_type.is_chroma(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ("Embedding", TrainingMethod.EMBEDDING), ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ + elif model_type.is_qwen() \ + or model_type.is_z_image() \ + or model_type.is_flux_2() \ + or model_type.is_ernie(): + return [ ("Fine Tune", TrainingMethod.FINE_TUNE), ("LoRA", TrainingMethod.LORA), ] + return [] - # training method - self.training_method = components.options_kv( - master=self.frame, - row=0, - column=7, - values=values, - ui_state=self.ui_state, - var_name="training_method", - command=self.change_training_method_callback, - ) - - def __change_model_type(self, model_type: ModelType): - self.change_model_type_callback(model_type) - self.__create_training_method() - - def __create_configs_dropdown(self): - if self.configs_dropdown is not None: - self.configs_dropdown.grid_forget() - - self.configs_dropdown = components.options_kv( - self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config - ) - - def __load_available_config_names(self): - if os.path.isdir(self.dir): - for path in os.listdir(self.dir): + def load_available_config_names(self, dir: str) -> list[tuple[str, str]]: + configs = [("", path_util.canonical_join(dir, "#.json"))] + if os.path.isdir(dir): + for path in os.listdir(dir): if path != "#.json": - path = path_util.canonical_join(self.dir, path) + path = path_util.canonical_join(dir, path) if path.endswith(".json") and os.path.isfile(path): name = os.path.basename(path) name = os.path.splitext(name)[0] - self.configs.append((name, path)) - self.configs.sort() + configs.append((name, path)) + configs.sort() + return configs - def __save_to_file(self, name) -> str: + def save_to_file(self, name) -> str: name = path_util.safe_filename(name) path = path_util.canonical_join("training_presets", f"{name}.json") - write_json_atomic(path, self.train_config.to_settings_dict(secrets=False)) - return path - def __save_secrets(self, path) -> str: + def save_secrets(self, path) -> str: write_json_atomic(path, self.train_config.secrets.to_dict()) return path def open_wiki(self): webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False) - def __save_new_config(self, name): - path = self.__save_to_file(name) - - is_new_config = name not in [x[0] for x in self.configs] - - if is_new_config: - self.configs.append((name, path)) - self.configs.sort() - - if self.config_ui_data["config_name"] != path_util.canonical_join(self.dir, f"{name}.json"): - self.config_ui_state.get_var("config_name").set(path_util.canonical_join(self.dir, f"{name}.json")) - - if is_new_config: - self.__create_configs_dropdown() - - def __save_config(self): - default_value = self.configs_dropdown.get() - while default_value.startswith('#'): - default_value = default_value[1:] - - dialogs.StringInputDialog( - parent=self.master, - title="name", - question="Config Name", - callback=self.__save_new_config, - default_value=default_value, - validate_callback=lambda x: not x.startswith("#") - ) - - def __load_current_config(self, filename): - try: - basename = os.path.basename(filename) - is_built_in_preset = basename.startswith("#") and basename != "#.json" - - with open(filename, "r") as f: - loaded_dict = json.load(f) - default_config = TrainConfig.default_values() - if is_built_in_preset: - # always assume built-in configs are saved in the most recent version - loaded_dict["__version"] = default_config.config_version - loaded_config = default_config.from_dict(loaded_dict).to_unpacked_config() - - with suppress(FileNotFoundError), open("secrets.json", "r") as f: - secrets_dict=json.load(f) - loaded_config.secrets = SecretsConfig.default_values().from_dict(secrets_dict) - - self.train_config.from_dict(loaded_config.to_dict()) - self.ui_state.update(loaded_config) - - optimizer_config = change_optimizer(self.train_config) - self.ui_state.get_var("optimizer").update(optimizer_config) - - self.load_preset_callback() - except FileNotFoundError: - pass - except Exception: - print(traceback.format_exc()) - - def __remove_config(self): - # TODO - pass - def save_default(self): - self.__save_to_file("#") - self.__save_secrets("secrets.json") + self.save_to_file("#") + self.save_secrets("secrets.json") diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index ba90d2e64..15d1eaff6 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -9,47 +9,25 @@ import time import traceback import webbrowser -from collections.abc import Callable from contextlib import suppress from pathlib import Path -from tkinter import filedialog, messagebox import scripts.generate_debug_report -from modules.ui.AdditionalEmbeddingsTab import AdditionalEmbeddingsTab -from modules.ui.CaptionUI import CaptionUI -from modules.ui.CloudTab import CloudTab -from modules.ui.ConceptTab import ConceptTab -from modules.ui.ConvertModelUI import ConvertModelUI -from modules.ui.LoraTab import LoraTab -from modules.ui.ModelTab import ModelTab -from modules.ui.ProfilingWindow import ProfilingWindow -from modules.ui.SampleWindow import SampleWindow -from modules.ui.SamplingTab import SamplingTab -from modules.ui.TopBar import TopBar -from modules.ui.TrainingTab import TrainingTab -from modules.ui.VideoToolUI import VideoToolUI +from modules.ui.BaseTrainUIView import BaseTrainUIView +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.ui.SampleWindowController import SampleWindowController +from modules.ui.VideoToolUIController import VideoToolUIController from modules.util import create from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.DataType import DataType -from modules.util.enum.GradientReducePrecision import GradientReducePrecision -from modules.util.enum.ImageFormat import ImageFormat -from modules.util.enum.ModelType import ModelType -from modules.util.enum.PathIOType import PathIOType -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState from modules.util.ui.validation import flush_and_validate_all import torch -import customtkinter as ctk -from customtkinter import AppearanceModeTracker - # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 if platform.system() == "Windows": @@ -57,547 +35,25 @@ # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE -class TrainUI(ctk.CTk): - set_step_progress: Callable[[int, int], None] - set_epoch_progress: Callable[[int, int], None] - - status_label: ctk.CTkLabel | None - training_button: ctk.CTkButton | None - training_callbacks: TrainCallbacks | None - training_commands: TrainCommands | None - - _TRAIN_BUTTON_STYLES = { - "idle": { - "text": "Start Training", - "state": "normal", - "fg_color": "#198754", - "hover_color": "#146c43", - "text_color": "white", - "text_color_disabled": "white", - }, - "running": { - "text": "Stop Training", - "state": "normal", - "fg_color": "#dc3545", - "hover_color": "#bb2d3b", - "text_color": "white", - }, - "stopping": { - "text": "Stopping...", - "state": "disabled", - "fg_color": "#dc3545", - "hover_color": "#dc3545", - "text_color": "white", - "text_color_disabled": "white", - }, - } - - def __init__(self): - super().__init__() - - self.title("OneTrainer") - self.geometry("1100x740") - - self.after(100, lambda: self._set_icon()) - - # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration - ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") - ctk.set_default_color_theme("blue") - - self.train_config = TrainConfig.default_values() - self.ui_state = UIState(self, self.train_config) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_columnconfigure(0, weight=1) - - self.status_label = None - self.eta_label = None - self.training_button = None - self.export_button = None - self.tabview = None - - self.model_tab = None - self.training_tab = None - self.lora_tab = None - self.cloud_tab = None - self.additional_embeddings_tab = None - - self.top_bar_component = self.top_bar(self) - self.content_frame(self) - self.bottom_bar(self) - self.training_thread = None - self.training_callbacks = None - self.training_commands = None +class TrainUIController: + def __init__(self, config: TrainConfig): + self.train_config = config + self.view: BaseTrainUIView | None = None + self.training_thread = None + self.training_callbacks: TrainCallbacks | None = None + self.training_commands: TrainCommands | None = None self.always_on_tensorboard_subprocess = None - self.current_workspace_dir = self.train_config.workspace_dir - self._check_start_always_on_tensorboard() - - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self._on_workspace_dir_change_trace) - - # Persistent profiling window. - self.profiling_window = ProfilingWindow(self) - - self.protocol("WM_DELETE_WINDOW", self.__close) - - def __close(self): - self.top_bar_component.save_default() - self._stop_always_on_tensorboard() - if hasattr(self, 'workspace_dir_trace_id'): - self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) - self.quit() - - def top_bar(self, master): - return TopBar( - master, - self.train_config, - self.ui_state, - self.change_model_type, - self.change_training_method, - self.load_preset, - ) - - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) - - def bottom_bar(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=2, column=0, sticky="nsew") - - self.set_step_progress, self.set_epoch_progress = components.double_progress(frame, 0, 0, "step", "epoch") - - # status + ETA container - self.status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - self.status_frame.grid(row=0, column=1, sticky="w") - self.status_frame.grid_rowconfigure(0, weight=0) - self.status_frame.grid_rowconfigure(1, weight=0) - self.status_frame.grid_columnconfigure(0, weight=1) - - self.status_label = components.label(self.status_frame, 0, 0, "", pad=0, - tooltip="Current status of the training run") - self.eta_label = components.label(self.status_frame, 1, 0, "", pad=0) - - # padding - frame.grid_columnconfigure(2, weight=1) - - - # export button - self.export_button = components.button(frame, 0, 3, "Export", self.export_training, - width=60, padx=5, pady=(15, 0), - tooltip="Export the current configuration as a script to run without a UI") - - # debug button - components.button(frame, 0, 4, "Debug", self.generate_debug_package, - width=60, padx=(5, 25), pady=(15, 0), - tooltip="Generate a zip file with config.json, debug_report.log and settings diff, use this to report bugs or issues") - - # tensorboard button - components.button(frame, 0, 5, "Tensorboard", self.open_tensorboard, - width=100, padx=(0, 5), pady=(15, 0)) - - # training button - self.training_button = components.button(frame, 0, 6, "Start Training", self.start_training, - padx=(5, 20), pady=(15, 0)) - self._set_training_button_style("idle") # centralized styling - - return frame - - def content_frame(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - frame.grid_rowconfigure(0, weight=1) - frame.grid_columnconfigure(0, weight=1) - - self.tabview = ctk.CTkTabview(frame) - self.tabview.grid(row=0, column=0, sticky="nsew") - - self.general_tab = self.create_general_tab(self.tabview.add("general")) - self.model_tab = self.create_model_tab(self.tabview.add("model")) - self.data_tab = self.create_data_tab(self.tabview.add("data")) - self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) - self.training_tab = self.create_training_tab(self.tabview.add("training")) - self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) - self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) - self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) - self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) - self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) - - self.change_training_method(self.train_config.training_method) - - return frame - - def create_general_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # workspace dir - components.label(frame, 0, 0, "Workspace Directory", - tooltip="The directory where all files of this training run are saved") - components.path_entry(frame, 0, 1, self.ui_state, "workspace_dir", mode="dir", command=self._on_workspace_dir_change) - - # cache dir - components.label(frame, 0, 2, "Cache Directory", - tooltip="The directory where cached data is saved") - components.path_entry(frame, 0, 3, self.ui_state, "cache_dir", mode="dir") - - # continue from previous backup - components.label(frame, 2, 0, "Continue from last backup", - tooltip="Automatically continues training from the last backup saved in /backup") - components.switch(frame, 2, 1, self.ui_state, "continue_last_backup") - - # only cache - components.label(frame, 2, 2, "Only Cache", - tooltip="Only populate the cache, without any training") - components.switch(frame, 2, 3, self.ui_state, "only_cache") - - # TODO: In Phase 4 rework the general tab. - # prevent overwrites - components.label(frame, 3, 0, "Prevent Overwrites", - tooltip="When enabled, output paths that already exist on disk will be flagged as invalid to avoid accidental overwrites") - components.switch(frame, 3, 1, self.ui_state, "prevent_overwrites") - - # debug - components.label(frame, 4, 0, "Debug mode", - tooltip="Save debug information during the training into the debug directory") - components.switch(frame, 4, 1, self.ui_state, "debug_mode") - - components.label(frame, 4, 2, "Debug Directory", - tooltip="The directory where debug data is saved") - components.path_entry(frame, 4, 3, self.ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) - - # tensorboard - components.label(frame, 6, 0, "Tensorboard", - tooltip="Starts the Tensorboard Web UI during training") - components.switch(frame, 6, 1, self.ui_state, "tensorboard") - - components.label(frame, 6, 2, "Always-On Tensorboard", - tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") - components.switch(frame, 6, 3, self.ui_state, "tensorboard_always_on", command=self._on_always_on_tensorboard_toggle) - - components.label(frame, 7, 0, "Expose Tensorboard", - tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") - components.switch(frame, 7, 1, self.ui_state, "tensorboard_expose") - components.label(frame, 7, 2, "Tensorboard Port", - tooltip="Port to use for Tensorboard link") - components.entry(frame, 7, 3, self.ui_state, "tensorboard_port") - - - # validation - components.label(frame, 8, 0, "Validation", - tooltip="Enable validation steps and add new graph in tensorboard") - components.switch(frame, 8, 1, self.ui_state, "validation") - - components.label(frame, 8, 2, "Validate after", - tooltip="The interval used when validate training") - components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") - - # device - components.label(frame, 10, 0, "Dataloader Threads", - tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) - - components.label(frame, 11, 0, "Train Device", - tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") - components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) - - components.label(frame, 12, 0, "Multi-GPU", - tooltip="Enable multi-GPU training") - components.switch(frame, 12, 1, self.ui_state, "multi_gpu") - components.label(frame, 12, 2, "Device Indexes", - tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") - components.entry(frame, 12, 3, self.ui_state, "device_indexes") - - components.label(frame, 13, 0, "Gradient Reduce Precision", - tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" - "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" - "FLOAT_32: Reduce gradients in float32\n" - "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", - wide_tooltip=True) - components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, - "gradient_reduce_precision") - - components.label(frame, 13, 2, "Fused Gradient Reduce", - tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") - - components.label(frame, 14, 0, "Async Gradient Reduce", - tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 14, 2, "Buffer size (MB)", - tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") - - components.label(frame, 15, 0, "Temp Device", - tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 15, 1, self.ui_state, "temp_device") - - frame.pack(fill="both", expand=1) - return frame - - def create_model_tab(self, master): - return ModelTab(master, self.train_config, self.ui_state) - - def create_data_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # aspect ratio bucketing - components.label(frame, 0, 0, "Aspect Ratio Bucketing", - tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") - components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") - - # latent caching - components.label(frame, 1, 0, "Latent Caching", - tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") - - # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", - tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") - - frame.pack(fill="both", expand=1) - return frame - - def create_concepts_tab(self, master): - return ConceptTab(master, self.train_config, self.ui_state) - - def create_training_tab(self, master) -> TrainingTab: - return TrainingTab(master, self.train_config, self.ui_state) - - def create_cloud_tab(self, master) -> CloudTab: - return CloudTab(master, self.train_config, self.ui_state,parent=self) - - def create_sampling_tab(self, master): - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) - - # sample after - top_frame = ctk.CTkFrame(master=master, corner_radius=0) - top_frame.grid(row=0, column=0, sticky="nsew") - sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") - sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) - - components.label(top_frame, 0, 0, "Sample After", - tooltip="The interval used when automatically sampling from the model during training") - components.time_entry(top_frame, 0, 1, self.ui_state, "sample_after", "sample_after_unit") - - components.label(top_frame, 0, 2, "Skip First", - tooltip="Start sampling automatically after this interval has elapsed.") - components.entry(top_frame, 0, 3, self.ui_state, "sample_skip_first", width=50, sticky="nw") - - components.label(top_frame, 0, 4, "Format", - tooltip="File Format used when saving samples") - components.options_kv(top_frame, 0, 5, [ - ("PNG", ImageFormat.PNG), - ("JPG", ImageFormat.JPG), - ], self.ui_state, "sample_image_format") - - components.button(top_frame, 0, 6, "sample now", self.sample_now) - - components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window ) - - components.label(sub_frame, 0, 0, "Non-EMA Sampling", - tooltip="Whether to include non-ema sampling when using ema.") - components.switch(sub_frame, 0, 1, self.ui_state, "non_ema_sampling") - - components.label(sub_frame, 0, 2, "Samples to Tensorboard", - tooltip="Whether to include sample images in the Tensorboard output.") - components.switch(sub_frame, 0, 3, self.ui_state, "samples_to_tensorboard") - - # table - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - return SamplingTab(frame, self.train_config, self.ui_state) - - def create_backup_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # backup after - components.label(frame, 0, 0, "Backup After", - tooltip="The interval used when automatically creating model backups during training") - components.time_entry(frame, 0, 1, self.ui_state, "backup_after", "backup_after_unit") - - # backup now - components.button(frame, 0, 3, "backup now", self.backup_now) - - # rolling backup - components.label(frame, 1, 0, "Rolling Backup", - tooltip="If rolling backups are enabled, older backups are deleted automatically") - components.switch(frame, 1, 1, self.ui_state, "rolling_backup") - - # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", - tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") - - # backup before save - components.label(frame, 2, 0, "Backup Before Save", - tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") - - # save after - components.label(frame, 3, 0, "Save Every", - tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") - - # save now - components.button(frame, 3, 3, "save now", self.save_now) - - # skip save - components.label(frame, 4, 0, "Skip First", - tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") - - # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", - tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") - - frame.pack(fill="both", expand=1) - return frame - - def embedding_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # embedding model name - components.label(frame, 0, 0, "Base embedding", - tooltip="The base embedding to train on. Leave empty to create a new embedding") - components.path_entry( - frame, 0, 1, self.ui_state, "embedding.model_name", - mode="file", path_modifier=components.json_path_modifier - ) + self.current_workspace_dir = config.workspace_dir + self.start_time: float = 0.0 - # token count - components.label(frame, 1, 0, "Token count", - tooltip="The token count used when creating a new embedding. Leave empty to auto detect from the initial embedding text.") - components.entry(frame, 1, 1, self.ui_state, "embedding.token_count") - - # initial embedding text - components.label(frame, 2, 0, "Initial embedding text", - tooltip="The initial embedding text used when creating a new embedding") - components.entry(frame, 2, 1, self.ui_state, "embedding.initial_embedding_text") - - # embedding weight dtype - components.label(frame, 3, 0, "Embedding Weight Data Type", - tooltip="The Embedding weight data type used for training. This can reduce memory consumption, but reduces precision") - components.options_kv(frame, 3, 1, [ - ("float32", DataType.FLOAT_32), - ("bfloat16", DataType.BFLOAT_16), - ], self.ui_state, "embedding_weight_dtype") - - # placeholder - components.label(frame, 4, 0, "Placeholder", - tooltip="The placeholder used when using the embedding in a prompt") - components.entry(frame, 4, 1, self.ui_state, "embedding.placeholder") - - # output embedding - components.label(frame, 5, 0, "Output embedding", - tooltip="Output embeddings are calculated at the output of the text encoder, not the input. This can improve results for larger text encoders and lower VRAM usage.") - components.switch(frame, 5, 1, self.ui_state, "embedding.is_output_embedding") - - frame.pack(fill="both", expand=1) - return frame - - def create_additional_embeddings_tab(self, master): - return AdditionalEmbeddingsTab(master, self.train_config, self.ui_state) - - def create_tools_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - # dataset - components.label(frame, 0, 0, "Dataset Tools", - tooltip="Open the captioning tool") - components.button(frame, 0, 1, "Open", self.open_dataset_tool) - - # video tools - components.label(frame, 1, 0, "Video Tools", - tooltip="Open the video tools") - components.button(frame, 1, 1, "Open", self.open_video_tool) - - # convert model - components.label(frame, 2, 0, "Convert Model Tools", - tooltip="Open the model conversion tool") - components.button(frame, 2, 1, "Open", self.open_convert_model_tool) - - # sample - components.label(frame, 3, 0, "Sampling Tool", - tooltip="Open the model sampling tool") - components.button(frame, 3, 1, "Open", self.open_sampling_tool) - - components.label(frame, 4, 0, "Profiling Tool", - tooltip="Open the profiling tools.") - components.button(frame, 4, 1, "Open", self.open_profiling_tool) - - frame.pack(fill="both", expand=1) - return frame - - def change_model_type(self, model_type: ModelType): - if self.model_tab: - self.model_tab.refresh_ui() - - if self.training_tab: - self.training_tab.refresh_ui() - - if self.lora_tab: - self.lora_tab.refresh_ui() - - def change_training_method(self, training_method: TrainingMethod): - if not self.tabview: - return - - if self.model_tab: - self.model_tab.refresh_ui() - - if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: - self.tabview.delete("LoRA") - self.lora_tab = None - if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: - self.tabview.delete("embedding") - - if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = LoraTab(self.tabview.add("LoRA"), self.train_config, self.ui_state) - if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: - self.embedding_tab(self.tabview.add("embedding")) - - def load_preset(self): - if not self.tabview: - return - - if self.additional_embeddings_tab: - self.additional_embeddings_tab.refresh_ui() + def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) + self.view.on_update_progress(train_progress.epoch_step, max_step, train_progress.epoch, max_epoch, eta_str) - def open_tensorboard(self): - webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + def on_update_status(self, status: str): + self.view.on_update_status(status) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: spent_total = time.monotonic() - self.start_time @@ -621,85 +77,136 @@ def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, ma else: return f"{seconds}s" - def set_eta_label(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - eta_str = self._calculate_eta_string(train_progress, max_step, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") + def _check_start_always_on_tensorboard(self): + if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() - def delete_eta_label(self): - self.eta_label.configure(text="") + def _start_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + self._stop_always_on_tensorboard() - def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): - self.set_step_progress(train_progress.epoch_step, max_step) - self.set_epoch_progress(train_progress.epoch, max_epoch) - self.set_eta_label(train_progress, max_step, max_epoch) + tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") + tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") - def on_update_status(self, status: str): - self.status_label.configure(text=status) + os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - def open_dataset_tool(self): - window = CaptionUI(self, None, False) - self.wait_window(window) + tensorboard_args = [ + tensorboard_executable, + "--logdir", + tensorboard_log_dir, + "--port", + str(self.train_config.tensorboard_port), + "--samples_per_plugin=images=100,scalars=10000", + ] - def open_video_tool(self): - window = VideoToolUI(self) - self.wait_window(window) + if self.train_config.tensorboard_expose: + tensorboard_args.append("--bind_all") - def open_convert_model_tool(self): - window = ConvertModelUI(self) - self.wait_window(window) + try: + self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) + except Exception: + self.always_on_tensorboard_subprocess = None - def open_sampling_tool(self): - if not self.training_callbacks and not self.training_commands: - window = SampleWindow( - self, - use_external_model=False, - train_config=self.train_config, - ) - self.wait_window(window) - torch_gc() + def _stop_always_on_tensorboard(self): + if self.always_on_tensorboard_subprocess: + try: + self.always_on_tensorboard_subprocess.terminate() + self.always_on_tensorboard_subprocess.wait(timeout=5) + except subprocess.TimeoutExpired: + self.always_on_tensorboard_subprocess.kill() + except Exception: + pass + finally: + self.always_on_tensorboard_subprocess = None - def open_profiling_tool(self): - self.profiling_window.deiconify() + def _on_workspace_dir_change(self, new_workspace_dir: str): + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir - def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() - if not zip_path: - return + def _on_workspace_dir_change_trace(self, *args): + new_workspace_dir = self.train_config.workspace_dir + if new_workspace_dir != self.current_workspace_dir: + self.current_workspace_dir = new_workspace_dir - zip_path = Path(zip_path) / "OneTrainer_debug_report.zip" + if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: + self._start_always_on_tensorboard() - self.on_update_status("Generating debug package...") + def _on_always_on_tensorboard_toggle(self): + if self.train_config.tensorboard_always_on: + if not (self.training_thread and self.train_config.tensorboard): + self._start_always_on_tensorboard() + else: + if not (self.training_thread and self.train_config.tensorboard): + self._stop_always_on_tensorboard() - try: - config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) - scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) - self.on_update_status(f"Debug package saved to {zip_path.name}") - except Exception as e: - traceback.print_exc() - self.on_update_status(f"Error generating debug package: {e}") + def open_tensorboard(self): + webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) + + def open_dataset_tool(self, parent, view_cls): + return CaptionUIController(None, False).create_window(parent, view_cls) + def open_video_tool(self, parent, view_cls): + return VideoToolUIController().create_window(parent, view_cls) + + def open_convert_model_tool(self, parent, view_cls): + return ConvertModelUIController().create_window(parent, view_cls) + + def open_sampling_tool(self, parent, view_cls): + if not self.training_callbacks and not self.training_commands: + controller = SampleWindowController( + self.train_config, + use_external_model=False, + ) + window = view_cls(parent, controller) + parent.show_window(window) + torch_gc() - def open_manual_sample_window (self): + def open_manual_sample_window(self, parent, view_cls): training_callbacks = self.training_callbacks training_commands = self.training_commands if training_callbacks and training_commands: - window = SampleWindow( - self, - train_config=self.train_config, + controller = SampleWindowController( + self.train_config, use_external_model=True, callbacks=training_callbacks, commands=training_commands, ) - self.wait_window(window) - training_callbacks.set_on_sample_custom() + window = view_cls(parent, controller) + parent.show_window(window) + parent.connect_window_closed(window, lambda: training_callbacks.set_on_sample_custom()) + + def sample_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.sample_default() + + def backup_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.backup() + + def save_now(self): + train_commands = self.training_commands + if train_commands: + train_commands.save() + + def export_training(self, file_path: str): + with open(file_path, "w") as f: + json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) + + def generate_debug_package(self, zip_path: Path): + self.view.on_update_status("Generating debug package...") + try: + config_json_string = json.dumps(self.train_config.to_pack_dict(secrets=False)) + scripts.generate_debug_report.create_debug_package(str(zip_path), config_json_string) + self.view.on_update_status(f"Debug package saved to {zip_path.name}") + except Exception as e: + traceback.print_exc() + self.view.on_update_status(f"Error generating debug package: {e}") def __training_thread_function(self): error_caught = False @@ -709,17 +216,17 @@ def __training_thread_function(self): on_update_status=self.on_update_status, ) - trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.cloud_tab.reattach) + trainer = create.create_trainer(self.train_config, self.training_callbacks, self.training_commands, reattach=self.view.get_cloud_reattach()) try: trainer.start() if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.view.sync_cloud_secrets() self.start_time = time.monotonic() trainer.train() except Exception: if self.train_config.cloud.enabled: - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.view.sync_cloud_secrets() error_caught = True traceback.print_exc() @@ -737,30 +244,23 @@ def __training_thread_function(self): self.on_update_status("Error: check the console for details") else: self.on_update_status("Stopped") - self.delete_eta_label() - # queue UI update on Tk main thread; _set_training_button_idle applies shared styles, avoid potential race/crash - self.after(0, self._set_training_button_idle) + # queue UI update on Tk main thread; on_training_stopped applies shared styles, avoid potential race/crash + self.view.schedule_on_main_thread(lambda: self.view.on_training_stopped(error_caught)) if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self.after(0, self._start_always_on_tensorboard) + self.view.schedule_on_main_thread(self._start_always_on_tensorboard) def start_training(self): if self.training_thread is None: - self.save_default() + self.view.save_default() - # --- pre-training validation gate --- errors = flush_and_validate_all() - if errors: - bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) + self.view.show_validation_errors(errors) return - self._set_training_button_running() + self.view.on_training_started() if self.train_config.tensorboard and not self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: self._stop_always_on_tensorboard() @@ -771,119 +271,6 @@ def start_training(self): self.training_thread = threading.Thread(target=self.__training_thread_function) self.training_thread.start() else: - self._set_training_button_stopping() + self.view.on_training_stopping() self.on_update_status("Stopping ...") self.training_commands.stop() - - def save_default(self): - self.top_bar_component.save_default() - self.concepts_tab.save_current_config() - self.sampling_tab.save_current_config() - self.additional_embeddings_tab.save_current_config() - - def export_training(self): - file_path = filedialog.asksaveasfilename(filetypes=[ - ("All Files", "*.*"), - ("json", "*.json"), - ], initialdir=".", initialfile="config.json") - - if file_path: - with open(file_path, "w") as f: - json.dump(self.train_config.to_pack_dict(secrets=False), f, indent=4) - - def sample_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.sample_default() - - def backup_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.backup() - - def save_now(self): - train_commands = self.training_commands - if train_commands: - train_commands.save() - - def _check_start_always_on_tensorboard(self): - if self.train_config.tensorboard_always_on and not self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _start_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - self._stop_always_on_tensorboard() - - tensorboard_executable = os.path.join(os.path.dirname(sys.executable), "tensorboard") - tensorboard_log_dir = os.path.join(self.train_config.workspace_dir, "tensorboard") - - os.makedirs(Path(tensorboard_log_dir).absolute(), exist_ok=True) - - tensorboard_args = [ - tensorboard_executable, - "--logdir", - tensorboard_log_dir, - "--port", - str(self.train_config.tensorboard_port), - "--samples_per_plugin=images=100,scalars=10000", - ] - - if self.train_config.tensorboard_expose: - tensorboard_args.append("--bind_all") - - try: - self.always_on_tensorboard_subprocess = subprocess.Popen(tensorboard_args) - except Exception: - self.always_on_tensorboard_subprocess = None - - def _stop_always_on_tensorboard(self): - if self.always_on_tensorboard_subprocess: - try: - self.always_on_tensorboard_subprocess.terminate() - self.always_on_tensorboard_subprocess.wait(timeout=5) - except subprocess.TimeoutExpired: - self.always_on_tensorboard_subprocess.kill() - except Exception: - pass - finally: - self.always_on_tensorboard_subprocess = None - - def _on_workspace_dir_change(self, new_workspace_dir: str): - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_workspace_dir_change_trace(self, *args): - new_workspace_dir = self.train_config.workspace_dir - if new_workspace_dir != self.current_workspace_dir: - self.current_workspace_dir = new_workspace_dir - - if self.train_config.tensorboard_always_on and self.always_on_tensorboard_subprocess: - self._start_always_on_tensorboard() - - def _on_always_on_tensorboard_toggle(self): - if self.train_config.tensorboard_always_on: - if not (self.training_thread and self.train_config.tensorboard): - self._start_always_on_tensorboard() - else: - if not (self.training_thread and self.train_config.tensorboard): - self._stop_always_on_tensorboard() - - def _set_training_button_style(self, mode: str): - if not self.training_button: - return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: - return - self.training_button.configure(**style) - - def _set_training_button_idle(self): - self._set_training_button_style("idle") - - def _set_training_button_running(self): - self._set_training_button_style("running") - - def _set_training_button_stopping(self): - self._set_training_button_style("stopping") diff --git a/modules/ui/TrainingTabController.py b/modules/ui/TrainingTabController.py new file mode 100644 index 000000000..858ca0701 --- /dev/null +++ b/modules/ui/TrainingTabController.py @@ -0,0 +1,39 @@ + +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.ui.SchedulerParamsWindowController import SchedulerParamsWindowController +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util import create +from modules.util.config.TrainConfig import TrainConfig +from modules.util.optimizer_util import change_optimizer + + +class TrainingTabController: + def __init__(self, config: TrainConfig): + self.config = config + + def restore_optimizer_config(self, ui_state): + optimizer_config = change_optimizer(self.config) + ui_state.get_var("optimizer").update(optimizer_config) + + def get_layer_presets(self) -> dict: + cls = create.get_model_setup_class(self.config.model_type, self.config.training_method) + return cls.LAYER_PRESETS if cls is not None else {"full": []} + + def is_flow_matching(self) -> bool: + return self.config.model_type.is_flow_matching() + + def is_custom_scheduler_value(self, value: str) -> bool: + return value == "CUSTOM" + + def open_optimizer_params_window(self, parent, ui_state, view_cls): + return view_cls(parent, OptimizerParamsWindowController(self.config), ui_state) + + def open_scheduler_params_window(self, parent, ui_state, view_cls): + return view_cls(parent, SchedulerParamsWindowController(self.config), ui_state) + + def open_timestep_distribution_window(self, parent, ui_state, view_cls): + return view_cls(parent, TimestepDistributionWindowController(self.config), ui_state) + + def open_offloading_window(self, parent, ui_state, view_cls): + return view_cls(parent, OffloadingWindowController(self.config), ui_state) diff --git a/modules/ui/VideoToolUIController.py b/modules/ui/VideoToolUIController.py index c3291e6ea..294f458f4 100644 --- a/modules/ui/VideoToolUIController.py +++ b/modules/ui/VideoToolUIController.py @@ -6,338 +6,54 @@ import shlex import subprocess import threading -import webbrowser from fractions import Fraction -from tkinter import filedialog -from modules.util.image_util import load_image from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS -from modules.util.ui import components import av -import customtkinter as ctk import cv2 import scenedetect from PIL import Image -class VideoToolUI(ctk.CTkToplevel): - def __init__( - self, - parent, - *args, **kwargs, - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - - self.title("Video Tools") - self.geometry("600x720") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - self.clip_extract_tab = self.__clip_extract_tab(tabview.add("extract clips")) - self.image_extract_tab = self.__image_extract_tab(tabview.add("extract images")) - self.video_download_tab = self.__video_download_tab(tabview.add("download")) - self.status_bar(self) - - def status_bar(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=1, column=0) - frame.grid_columnconfigure(0, weight=0, minsize=160) - frame.grid_columnconfigure(1, weight=0, minsize=300) - frame.grid_columnconfigure(2, weight=1) - - #create preview image - preview_path = "resources/icons/icon.png" - preview = load_image(preview_path, 'RGB') - preview.thumbnail((150, 150)) - self.preview_image= ctk.CTkImage(light_image=preview, size=preview.size) - self.preview_image_label = ctk.CTkLabel( - master=frame, text="Preview image", image=self.preview_image, height=150, width=150, - compound="top") - self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - - #displays progress and messages that also go to terminal - self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) - self.status_label.insert(index="1.0", text="Current status") - self.status_label.configure(state="disabled") - self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) - - def __clip_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.clip_single_entry = ctk.CTkEntry(frame, width=190) - self.clip_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.clip_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.clip_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.clip_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_clips_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.clip_time_start_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.clip_time_start_entry.insert(0, "00:00:00") - self.clip_time_end_entry = ctk.CTkEntry(frame, width=100) - self.clip_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.clip_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.clip_list_entry = ctk.CTkEntry(frame, width=190) - self.clip_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.clip_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_list_entry)) - self.clip_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_clips_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted clips will be saved.") - self.clip_output_entry = ctk.CTkEntry(frame, width=190) - self.clip_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.clip_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.clip_output_entry)) - self.clip_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_clip = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_clip_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_clip, text="") - self.output_subdir_clip_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # split at cuts - self.split_at_cuts = ctk.BooleanVar(self, False) - components.label(frame, 5, 0, "Split at Cuts", - tooltip="If enabled, detect cuts in the input video and split at those points. \ - Otherwise will split at any point, and clips may contain cuts.") - self.split_cuts_entry = ctk.CTkSwitch(frame, variable=self.split_at_cuts, text="") - self.split_cuts_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - - # maximum length - components.label(frame, 6, 0, "Max Length (s)", - tooltip="Maximum length in seconds for saved clips, larger clips will be broken into multiple small clips.") - self.clip_length_entry = ctk.CTkEntry(frame, width=220) - self.clip_length_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.clip_length_entry.insert(0, "3") - - # Set FPS - components.label(frame, 7, 0, "Set FPS", - tooltip="FPS to convert output videos to, set to 0 to keep original rate.") - self.clip_fps_entry = ctk.CTkEntry(frame, width=220) - self.clip_fps_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - self.clip_fps_entry.insert(0, "24.0") - - # Remove borders - self.clip_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 8, 0, "Remove Borders", - tooltip="Remove black borders from output clip") - self.clip_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.clip_bordercrop, text="") - self.clip_bordercrop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 9, 0, "Crop Variation", - tooltip="Output clips will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square videos. Set to 0 to use only base aspect.") - self.clip_crop_entry = ctk.CTkEntry(frame, width=220) - self.clip_crop_entry.grid(row=9, column=1, sticky="w", padx=5, pady=5) - self.clip_crop_entry.insert(0, "0.2") - - # object filter - currently unused, may implement in future - # components.label(frame, 9, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 9, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 9, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __image_extract_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # single video - components.label(frame, 0, 0, "Single Video", - tooltip="Link to single video file to process.") - self.image_single_entry = ctk.CTkEntry(frame, width=190) - self.image_single_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - self.image_single_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.image_single_entry, - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))] - )) - self.image_single_button.grid(row=0, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 0, 2, "Extract Single", - command=lambda: self.__extract_images_button(False)) - - # time range - components.label(frame, 1, 0, " Time Range", - tooltip="Time range to limit selection for single video, \ - format as hour:minute:second, minute:second, or seconds.") - self.image_time_start_entry = ctk.CTkEntry(frame, width=100) - self.image_time_start_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.image_time_start_entry.insert(0, "00:00:00") - self.image_time_end_entry = ctk.CTkEntry(frame, width=100) - self.image_time_end_entry.grid(row=1, column=1, sticky="e", padx=5, pady=5) - self.image_time_end_entry.insert(0, "99:99:99") - - # directory of videos - components.label(frame, 2, 0, "Directory", - tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.image_list_entry = ctk.CTkEntry(frame, width=190) - self.image_list_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.image_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_list_entry)) - self.image_list_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 2, 2, "Extract Directory", - command=lambda: self.__extract_images_button(True)) - - # output directory - components.label(frame, 3, 0, "Output", - tooltip="Path to folder where extracted images will be saved.") - self.image_output_entry = ctk.CTkEntry(frame, width=190) - self.image_output_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) - self.image_output_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_dir(self.image_output_entry)) - self.image_output_button.grid(row=3, column=1, sticky="e", padx=5, pady=5) - - # output to subdirectories - self.output_subdir_img = ctk.BooleanVar(self, False) - components.label(frame, 4, 0, "Output to\nSubdirectories", - tooltip="If enabled, files are saved to subfolders based on filename and input directory. \ - Otherwise will all be saved to the top level of the output directory.") - self.output_subdir_img_entry = ctk.CTkSwitch(frame, variable=self.output_subdir_img, text="") - self.output_subdir_img_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) - - # image capture rate - components.label(frame, 5, 0, "Images/sec", - tooltip="Number of images to capture per second of video. \ - Images will be taken at semi-random frames around the specified frequency.") - self.capture_rate_entry = ctk.CTkEntry(frame, width=220) - self.capture_rate_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) - self.capture_rate_entry.insert(0, "0.5") - - # blur removal - components.label(frame, 6, 0, "Blur Removal", - tooltip="Threshold for removal of blurry images, relative to all others. \ - For example at 0.2, the blurriest 20%% of the final selected frames will not be saved.") - self.blur_threshold_entry = ctk.CTkEntry(frame, width=220) - self.blur_threshold_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) - self.blur_threshold_entry.insert(0, "0.2") - - # Remove borders - self.image_bordercrop = ctk.BooleanVar(self, False) - components.label(frame, 7, 0, "Remove Borders", - tooltip="Remove black borders from output image") - self.image_bordercrop_entry = ctk.CTkSwitch(frame, variable=self.image_bordercrop, text="") - self.image_bordercrop_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) - - # Crop Variation - components.label(frame, 8, 0, "Crop Variation", - tooltip="Output images will be randomly cropped to +- the base aspect ratio, \ - somewhat biased towards making square images. Set to 0 to use only base sapect.") - self.image_crop_entry = ctk.CTkEntry(frame, width=220) - self.image_crop_entry.grid(row=8, column=1, sticky="w", padx=5, pady=5) - self.image_crop_entry.insert(0, "0.2") - - # # object filter - currently unused, may implement in future - # components.label(frame, 5, 0, "Object Filter", - # tooltip="Detect general features using Haar-Cascade classifier, and choose how to deal with clips where it is detected") - # components.options(frame, 5, 1, ["NONE", "FACE", "EYE", "BODY"], self.video_ui_state, "filter_object") - # components.options(frame, 5, 2, ["INCLUDE", "EXCLUDE", "SUBFOLDER"], self.video_ui_state, "filter_behavior") - - frame.pack(fill="both", expand=1) - return frame - - def __video_download_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0, minsize=120) - frame.grid_columnconfigure(1, weight=0, minsize=200) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - - # link - components.label(frame, 0, 0, "Single Link", - tooltip="Link to video/playlist to download. Uses yt-dlp, supports youtube, twitch, instagram, and many other sites.") - self.download_link_entry = ctk.CTkEntry(frame, width=220) - self.download_link_entry.grid(row=0, column=1, sticky="w", padx=5, pady=5) - components.button(frame, 0, 2, "Download Link", command=lambda: self.__download_button(False)) - - # link list - components.label(frame, 1, 0, "Link List", - tooltip="Path to txt file with list of links separated by newlines.") - self.download_list_entry = ctk.CTkEntry(frame, width=190) - self.download_list_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) - self.download_list_button = ctk.CTkButton(frame, width=30, text="...", - command=lambda: self.__browse_for_file(self.download_list_entry, [("Text file", ".txt")])) - self.download_list_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) - components.button(frame, 1, 2, "Download List", command=lambda: self.__download_button(True)) - - # output directory - components.label(frame, 2, 0, "Output", - tooltip="Path to folder where downloaded videos will be saved.") - self.download_output_entry = ctk.CTkEntry(frame, width=190) - self.download_output_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) - self.download_output_button = ctk.CTkButton(frame, width=30, text="...", command=lambda: self.__browse_for_dir(self.download_output_entry)) - self.download_output_button.grid(row=2, column=1, sticky="e", padx=5, pady=5) - - # additional args - components.label(frame, 3, 0, "Additional Args", - tooltip="Any additional arguments to pass to yt-dlp, for example '--restrict-filenames --force-overwrite'. \ - Default args will hide most terminal outputs.") - self.download_args_entry = ctk.CTkTextbox(frame, width=220, height=90, border_width=2) - self.download_args_entry.grid(row=3, column=1, rowspan=2, sticky="w", padx=5, pady=5) - self.download_args_entry.insert(index="1.0", text="--quiet --no-warnings --progress --format mp4") - components.button(frame, 3, 2, "yt-dlp info", - command=lambda: webbrowser.open("https://github.com/yt-dlp/yt-dlp?tab=readme-ov-file#usage-and-options", new=0, autoraise=False)) - - frame.pack(fill="both", expand=1) - return frame - - def __browse_for_dir(self, entry_box): - # get the path from the user - path = filedialog.askdirectory() - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() - - def __browse_for_file(self, entry_box, filetypes): - # get the path from the user - path = filedialog.askopenfilename(filetypes=filetypes) - # set the path to the entry box - # delete entry box text - entry_box.focus_set() - entry_box.delete(0, ctk.END) - entry_box.insert(0, path) - self.focus_set() +class VideoToolUIController: + def __init__(self): + self.view = None + self.args = { + "clip_single": "", + "clip_list": "", + "clip_time_start": "00:00:00", + "clip_time_end": "99:99:99", + "clip_output": "", + "output_subdir_clip": False, + "split_cuts": False, + "clip_length": "3", + "clip_fps": "24.0", + "clip_bordercrop": False, + "clip_crop": "0.2", + "image_single": "", + "image_list": "", + "image_time_start": "00:00:00", + "image_time_end": "99:99:99", + "image_output": "", + "output_subdir_img": False, + "capture_rate": "0.5", + "blur_threshold": "0.2", + "image_bordercrop": False, + "image_crop": "0.2", + "download_link": "", + "download_list": "", + "download_output": "", + "download_args": "--quiet --no-warnings --progress --format mp4", + } + + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + + def __update_status(self, status_text: str): + print(status_text) + self.view.update_status(status_text) def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_dir: str): input_videos = [] @@ -382,9 +98,7 @@ def __get_vid_paths(self, batch_mode: bool, input_path_single: str, input_path_d def __run_in_thread(self, target, *args): """Clear status box and run target function in a daemon thread.""" - self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") - self.status_label.configure(state="disabled") + self.view.clear_status() t = threading.Thread(target=target, args=args) t.daemon = True t.start() @@ -463,22 +177,20 @@ def find_main_contour(self, frame): h1, w1, _ = frame.shape return x1, y1, w1, h1 - def __extract_clips_button(self, batch_mode: bool): + def extract_clips_button(self, batch_mode: bool): self.__run_in_thread(self.__extract_clips_multi, batch_mode) def __extract_clips_multi(self, batch_mode: bool): - if not pathlib.Path(self.clip_output_entry.get()).is_dir() or self.clip_output_entry.get() == "": + p = self.args + if not pathlib.Path(p['clip_output']).is_dir() or p['clip_output'] == "": self.__update_status("Invalid output directory!") return # validate numeric inputs try: - max_length = float(self.clip_length_entry.get()) - crop_variation = float(self.clip_crop_entry.get()) - target_fps = float(self.clip_fps_entry.get()) - input_single_entry = self.clip_single_entry.get() - input_multiple_entry = self.clip_list_entry.get() - output_entry = self.clip_output_entry.get() + max_length = float(p['clip_length']) + crop_variation = float(p['clip_crop']) + target_fps = float(p['clip_fps']) except ValueError: self.__update_status("Invalid numeric input for Max Length, Crop Variation, or FPS.") return @@ -492,26 +204,26 @@ def __extract_clips_multi(self, batch_mode: bool): self.__update_status("Crop Variation must be between 0.0 and 1.0.") return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + input_videos = self.__get_vid_paths(batch_mode, p['clip_single'], p['clip_list']) if len(input_videos) == 0: # exit if no paths found return with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for video_path in input_videos: output_directory = self.__get_output_dir( - self.output_subdir_clip_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.clip_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.clip_time_end_entry.get()) + p['output_subdir_clip'], batch_mode, + p['clip_output'], video_path, p['clip_list']) + time_start = "00:00:00" if batch_mode else p['clip_time_start'] + time_end = "99:99:99" if batch_mode else p['clip_time_end'] executor.submit(self.__extract_clips, str(video_path), time_start, time_end, max_length, - self.split_at_cuts.get(), bool(self.clip_bordercrop_entry.get()), + p['split_cuts'], p['clip_bordercrop'], crop_variation, target_fps, output_directory) if batch_mode: - self.__update_status(f'Clip extraction from all videos in "{input_multiple_entry}" complete') + self.__update_status(f'Clip extraction from all videos in "{p["clip_list"]}" complete') else: - self.__update_status(f'Clip extraction from "{input_single_entry}" complete') + self.__update_status(f'Clip extraction from "{p["clip_single"]}" complete') def __extract_clips(self, video_path: str, timestamp_min: str, timestamp_max: str, max_length: float, split_at_cuts: bool, remove_borders: bool, crop_variation: float, target_fps: float, output_dir: str): @@ -614,11 +326,9 @@ def __save_clip(self, scene: tuple[int, int], video_path: str, target_fps: float preview = Image.fromarray( cv2.cvtColor(frame[y1+y2:y1+y2+h2, x1+x2:x1+x2+w2], cv2.COLOR_BGR2RGB)) preview.thumbnail((150, 150)) - self.preview_image.configure(light_image=preview, size=preview.size) #truncate filename of long files so UI doesn't shift around filename_truncated = basename + ext if len(basename) < 20 else basename[:18] + ".." + ext - self.preview_image_label.configure( - text=f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') + self.view.update_preview(preview, f'{filename_truncated}\nFrames: {scene[0]}-{scene[1]}\nSize: {w2}x{h2}') except Exception: pass video.release() @@ -701,22 +411,20 @@ def __write_clip_av(video_path: str, output_path: str, scene: tuple[int, int], for pkt in out_video.encode(): output_container.mux(pkt) - def __extract_images_button(self, batch_mode: bool): + def extract_images_button(self, batch_mode: bool): self.__run_in_thread(self.__extract_images_multi, batch_mode) - def __extract_images_multi(self, batch_mode : bool): - if not pathlib.Path(self.image_output_entry.get()).is_dir() or self.image_output_entry.get() == "": + def __extract_images_multi(self, batch_mode: bool): + p = self.args + if not pathlib.Path(p['image_output']).is_dir() or p['image_output'] == "": self.__update_status("Invalid output directory!") return # validate numeric inputs try: - capture_rate = float(self.capture_rate_entry.get()) - blur_threshold = float(self.blur_threshold_entry.get()) - crop_variation = float(self.image_crop_entry.get()) - input_single_entry = self.image_single_entry.get() - input_multiple_entry = self.image_list_entry.get() - output_entry = self.image_output_entry.get() + capture_rate = float(p['capture_rate']) + blur_threshold = float(p['blur_threshold']) + crop_variation = float(p['image_crop']) except ValueError: self.__update_status("Invalid numeric input for Images/sec, Blur Removal, or Crop Variation.") return @@ -730,25 +438,25 @@ def __extract_images_multi(self, batch_mode : bool): self.__update_status("Crop Variation must be between 0.0 and 1.0.") return - input_videos = self.__get_vid_paths(batch_mode, input_single_entry, input_multiple_entry) + input_videos = self.__get_vid_paths(batch_mode, p['image_single'], p['image_list']) if not input_videos: return with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for video_path in input_videos: output_directory = self.__get_output_dir( - self.output_subdir_img_entry.get(), batch_mode, - output_entry, video_path, input_multiple_entry) - time_start = "00:00:00" if batch_mode else str(self.image_time_start_entry.get()) - time_end = "99:99:99" if batch_mode else str(self.image_time_end_entry.get()) + p['output_subdir_img'], batch_mode, + p['image_output'], video_path, p['image_list']) + time_start = "00:00:00" if batch_mode else p['image_time_start'] + time_end = "99:99:99" if batch_mode else p['image_time_end'] executor.submit(self.__save_frames, str(video_path), time_start, time_end, capture_rate, - blur_threshold, self.image_bordercrop.get(), + blur_threshold, p['image_bordercrop'], crop_variation, output_directory) if batch_mode: - self.__update_status(f'Image extraction from all videos in {input_multiple_entry} complete') + self.__update_status(f'Image extraction from all videos in {p["image_list"]} complete') else: - self.__update_status(f'Image extraction from "{input_single_entry}" complete') + self.__update_status(f'Image extraction from "{p["image_single"]}" complete') def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, capture_rate: float, blur_threshold: float, remove_borders: bool, crop_variation: float, output_dir: str): @@ -821,32 +529,26 @@ def __save_frames(self, video_path: str, timestamp_min: str, timestamp_max: str, cv2.cvtColor(frame_cropped[y2:y2+h2, x2:x2+w2], cv2.COLOR_BGR2RGB)) preview.thumbnail((150, 150)) filename_truncated = basename + ext if len(basename) < 20 else basename[:17] + "..." + ext - self.preview_image.configure(light_image=preview, size=preview.size) - self.preview_image_label.configure(text=f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') + self.view.update_preview(preview, f'{filename_truncated}\nFrame: {f[0]}\nSize: {w2}x{h2}') except Exception: pass # preview update is non-critical cv2.imwrite(filename, frame_cropped[y2:y2+h2, x2:x2+w2]) video.release() - def __download_button(self, batch_mode: bool): + def download_button(self, batch_mode: bool): self.__run_in_thread(self.__download_multi, batch_mode) - def __update_status(self, status_text: str): - print(status_text) - self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") - self.status_label.configure(state="disabled") - def __download_multi(self, batch_mode: bool): - if not pathlib.Path(self.download_output_entry.get()).is_dir() or self.download_output_entry.get() == "": + p = self.args + if not pathlib.Path(p['download_output']).is_dir() or p['download_output'] == "": self.__update_status("Invalid output directory!") return if not batch_mode: - ydl_urls = [self.download_link_entry.get()] + ydl_urls = [p['download_link']] elif batch_mode: - ydl_path = pathlib.Path(self.download_list_entry.get()) + ydl_path = pathlib.Path(p['download_list']) if ydl_path.is_file() and ydl_path.suffix.lower() == ".txt": with open(ydl_path) as file: ydl_urls = file.readlines() @@ -857,8 +559,8 @@ def __download_multi(self, batch_mode: bool): with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: for url in ydl_urls: executor.submit(self.__download_video, - url.strip(), self.download_output_entry.get(), - self.download_args_entry.get("0.0", ctk.END)) + url.strip(), p['download_output'], + p['download_args']) self.__update_status(f'Completed {len(ydl_urls)} downloads.') diff --git a/modules/util/path_util.py b/modules/util/path_util.py index 94f802423..7d25c8e8c 100644 --- a/modules/util/path_util.py +++ b/modules/util/path_util.py @@ -1,5 +1,6 @@ import json import os.path +from pathlib import Path from typing import Any @@ -52,3 +53,8 @@ def supported_video_extensions() -> set[str]: def is_supported_video_extension(extension: str) -> bool: return extension.lower() in SUPPORTED_VIDEO_EXTENSIONS + + +def json_path_modifier(x: str | Path) -> Path: + x = Path(x).absolute() + return x.parent if x.suffix == ".json" else x diff --git a/modules/util/ui/CtkUIState.py b/modules/util/ui/CtkUIState.py new file mode 100644 index 000000000..0096ca3e6 --- /dev/null +++ b/modules/util/ui/CtkUIState.py @@ -0,0 +1,23 @@ +import tkinter as tk +from typing import Any + +from modules.util.ui.UIState import BaseUIState + + +class CtkUIState(BaseUIState): + def __init__(self, master, obj): + self.master = master + super().__init__(obj) + + def _make_str_var(self, initial_value: Any): + var = tk.StringVar(master=self.master) + var.set(initial_value) + return var + + def _make_bool_var(self, initial_value: Any): + var = tk.BooleanVar(master=self.master) + var.set(initial_value) + return var + + def _make_nested_state(self, obj: Any) -> "CtkUIState": + return CtkUIState(self.master, obj) diff --git a/modules/util/ui/UIState.py b/modules/util/ui/UIState.py index 8b13d23f7..73d73ed34 100644 --- a/modules/util/ui/UIState.py +++ b/modules/util/ui/UIState.py @@ -1,4 +1,4 @@ -import tkinter as tk +from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass from enum import Enum @@ -8,13 +8,12 @@ from modules.util.type_util import issubclass_safe -class UIState: +class BaseUIState(ABC): __vars: dict[str, Any] __var_traces: dict[str, dict[int, Callable[[], None]]] __latest_var_trace_id: int - def __init__(self, master, obj): - self.master = master + def __init__(self, obj): self.obj = obj self.__var_types: dict[str, type] = {} @@ -25,13 +24,24 @@ def __init__(self, master, obj): self.__var_traces = {name: {} for name in self.__vars} self.__latest_var_trace_id = 0 + @abstractmethod + def _make_str_var(self, initial_value: Any): + pass + + @abstractmethod + def _make_bool_var(self, initial_value: Any): + pass + + @abstractmethod + def _make_nested_state(self, obj: Any) -> "BaseUIState": + pass + def update(self, obj): self.obj = obj self.__set_vars(obj) def get_var(self, name): split_name = name.split('.') - if len(split_name) == 1: return self.__vars[split_name[0]] else: @@ -72,7 +82,6 @@ def update(_0, _1, _2): else: setattr(obj, name, string_var) self.__call_var_traces(name) - return update def __set_enum_var(self, obj, is_dict, name, var, var_type, nullable): @@ -92,7 +101,6 @@ def update(_0, _1, _2): else: setattr(obj, name, var_type[string_var]) self.__call_var_traces(name) - return update def __set_bool_var(self, obj, is_dict, name, var): @@ -104,7 +112,6 @@ def update(_0, _1, _2): def update(_0, _1, _2): setattr(obj, name, var.get()) self.__call_var_traces(name) - return update def __set_int_var(self, obj, is_dict, name, var, nullable): @@ -138,7 +145,6 @@ def update(_0, _1, _2): except ValueError: setattr(obj, name, None) self.__call_var_traces(name) - return update def __set_float_var(self, obj, is_dict, name, var, nullable): @@ -172,12 +178,10 @@ def update(_0, _1, _2): except ValueError: setattr(obj, name, None) self.__call_var_traces(name) - return update def __create_vars(self, obj): new_vars = {} - is_dict = isinstance(obj, dict) is_config = isinstance(obj, BaseConfig) @@ -190,61 +194,48 @@ def __create_vars(self, obj): obj_var = getattr(obj, name) if issubclass_safe(var_type, BaseConfig): - var = UIState(self.master, obj_var) - new_vars[name] = var + new_vars[name] = self._make_nested_state(obj_var) elif var_type is str: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else obj_var) + var = self._make_str_var("" if obj_var is None else obj_var) var.trace_add("write", self.__set_str_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var elif issubclass_safe(var_type, Enum): - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_enum_var(obj, is_dict, name, var, var_type, obj.nullables[name])) new_vars[name] = var elif var_type is bool: - var = tk.BooleanVar(master=self.master) - var.set(obj_var or False) + var = self._make_bool_var(obj_var or False) var.trace_add("write", self.__set_bool_var(obj, is_dict, name, var)) new_vars[name] = var elif var_type is int: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_int_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var elif var_type is float: - var = tk.StringVar(master=self.master) - var.set("" if obj_var is None else str(obj_var)) + var = self._make_str_var("" if obj_var is None else str(obj_var)) var.trace_add("write", self.__set_float_var(obj, is_dict, name, var, obj.nullables[name])) new_vars[name] = var else: iterable = obj.items() if is_dict else vars(obj).items() - for name, obj_var in iterable: - if isinstance(obj_var, str): - var = tk.StringVar(master=self.master) - var.set(obj_var) + var = self._make_str_var(obj_var) var.trace_add("write", self.__set_str_var(obj, is_dict, name, var, False)) new_vars[name] = var elif isinstance(obj_var, Enum): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_enum_var(obj, is_dict, name, var, type(obj_var), False)) new_vars[name] = var elif isinstance(obj_var, bool): - var = tk.BooleanVar(master=self.master) - var.set(obj_var) + var = self._make_bool_var(obj_var) var.trace_add("write", self.__set_bool_var(obj, is_dict, name, var)) new_vars[name] = var elif isinstance(obj_var, int): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_int_var(obj, is_dict, name, var, False)) new_vars[name] = var elif isinstance(obj_var, float): - var = tk.StringVar(master=self.master) - var.set(str(obj_var)) + var = self._make_str_var(str(obj_var)) var.trace_add("write", self.__set_float_var(obj, is_dict, name, var, False)) new_vars[name] = var @@ -253,7 +244,6 @@ def __create_vars(self, obj): def __set_vars(self, obj): is_dict = isinstance(obj, dict) is_config = isinstance(obj, BaseConfig) - iterable = obj.items() if is_dict else vars(obj).items() if is_config: for name, var_type in obj.types.items(): @@ -274,6 +264,7 @@ def __set_vars(self, obj): var = self.__vars[name] var.set("" if obj_var is None else str(obj_var)) else: + iterable = obj.items() if is_dict else vars(obj).items() for name, obj_var in iterable: if isinstance(obj_var, str): var = self.__vars[name] @@ -288,27 +279,26 @@ def __set_vars(self, obj): var = self.__vars[name] var.set(str(obj_var)) - # metadata api + @dataclass(frozen=True) + class VarMeta: + type: type | None + nullable: bool + default: Any + def _resolve_state_and_leaf(self, name: str): parts = name.split('.') - state: UIState = self + state: BaseUIState = self for part in parts[:-1]: state = state.get_var(part) - if not isinstance(state, UIState): + if not isinstance(state, BaseUIState): return None, None return state, parts[-1] - @dataclass(frozen=True) - class VarMeta: - type: type | None - nullable: bool - default: Any - - def get_field_metadata(self, name: str) -> "UIState.VarMeta": + def get_field_metadata(self, name: str) -> "BaseUIState.VarMeta": state, leaf = self._resolve_state_and_leaf(name) if state is None: - return UIState.VarMeta(None, False, None) - return UIState.VarMeta( + return BaseUIState.VarMeta(None, False, None) + return BaseUIState.VarMeta( state.__var_types.get(leaf), state.__var_nullables.get(leaf, False), state.__var_defaults.get(leaf, None), diff --git a/modules/util/ui/components.py b/modules/util/ui/ctk_components.py similarity index 88% rename from modules/util/ui/components.py rename to modules/util/ui/ctk_components.py index dd7b89719..e462f72a1 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/ctk_components.py @@ -8,9 +8,9 @@ from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit from modules.util.path_util import supported_image_extensions +from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator +from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ToolTip import ToolTip -from modules.util.ui.UIState import UIState -from modules.util.ui.validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator import customtkinter as ctk from customtkinter.windows.widgets.scaling import CTkScalingBaseClass @@ -34,11 +34,13 @@ def app_title(master, row, column): label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) -def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0): +def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0, underline=False): component = ctk.CTkLabel(master, text=text, wraplength=wraplength) component.grid(row=row, column=column, padx=pad, pady=pad, sticky="nw") if tooltip: ToolTip(component, tooltip, wide=wide_tooltip) + if underline: + component.configure(font=ctk.CTkFont(underline=True)) return component @@ -46,7 +48,7 @@ def entry( master, row, column, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, command: Callable[[], None] | None = None, tooltip: str = "", @@ -108,13 +110,8 @@ def new_destroy(): return component -def json_path_modifier(x: str | Path) -> Path: - x = Path(x).absolute() - return x.parent if x.suffix == ".json" else x - - def path_entry( - master, row, column, ui_state: UIState, var_name: str, + master, row, column, ui_state: CtkUIState, var_name: str, *, mode: Literal["file", "dir"] = "file", io_type: PathIOType = PathIOType.INPUT, @@ -124,9 +121,10 @@ def path_entry( command: Callable[[str], None] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, + columnspan: int = 1, ): frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="new", columnspan=columnspan) frame.grid_columnconfigure(0, weight=1) @@ -216,7 +214,7 @@ def _frame_destroy(): return frame -def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_name, supports_time_units: bool = True): +def time_entry(master, row, column, ui_state: CtkUIState, var_name: str, unit_var_name, supports_time_units: bool = True): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") @@ -239,7 +237,7 @@ def time_entry(master, row, column, ui_state: UIState, var_name: str, unit_var_n return frame -def layer_filter_entry(master, row, column, ui_state: UIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): +def layer_filter_entry(master, row, column, ui_state: CtkUIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): frame = ctk.CTkFrame(master=master, corner_radius=5, fg_color=frame_color) frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) @@ -353,6 +351,15 @@ def icon_button(master, row, column, text, command): return component +def colored_icon_button(master, row, column, text, fg_color, command, padx=0): + component = ctk.CTkButton( + master=master, width=20, height=20, text=text, + corner_radius=2, fg_color=fg_color, command=command, + ) + component.grid(row=row, column=column, padx=padx) + return component + + def button(master, row, column, text, command, tooltip=None, **kwargs): # Pop grid-specific parameters from kwargs, using PAD as the default if not provided. padx = kwargs.pop('padx', PAD) @@ -365,7 +372,7 @@ def button(master, row, column, text, command, tooltip=None, **kwargs): return component -def options(master, row, column, values, ui_state: UIState, var_name: str, command: Callable[[str], None] | None = None): +def options(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None): component = ctk.CTkOptionMenu(master, values=values, variable=ui_state.get_var(var_name), command=command) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") @@ -385,7 +392,7 @@ def destroy(self): return component -def options_adv(master, row, column, values, ui_state: UIState, var_name: str, +def options_adv(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None, adv_command: Callable[[], None] | None = None): frame = ctk.CTkFrame(master, fg_color="transparent") frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") @@ -417,7 +424,7 @@ def destroy(self): return frame, {'component': component, 'button_component': button_component} -def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: UIState, var_name: str, +def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: CtkUIState, var_name: str, command: Callable[[Any], None] | None = None): var = ui_state.get_var(var_name) keys = [key for key, value in values] @@ -475,16 +482,19 @@ def switch( master, row, column, - ui_state: UIState, + ui_state: CtkUIState, var_name: str, command: Callable[[], None] | None = None, text: str = "", + width: int | None = None, ): var = ui_state.get_var(var_name) if command: trace_id = ui_state.add_var_trace(var_name, command) component = ctk.CTkSwitch(master, variable=var, text=text, command=command) + if width is not None: + component.configure(width=width) component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") def create_destroy(component): @@ -545,3 +555,34 @@ def set_2(value, max_value): description_2_component.configure(text=f"{value}/{max_value}") return set_1, set_2 + + +def section_frame(master, row: int, col: int = 0): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=col, padx=PAD // 2, pady=PAD // 2, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + return frame + + +def inline_frame(master, row: int, col: int, columnspan: int = 1): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col, columnspan=columnspan, sticky="ew", padx=0, pady=0) + return frame + + +def set_widget_enabled(widget, enabled: bool) -> None: + state = "normal" if enabled else "disabled" + if isinstance(widget, ctk.CTkFrame): + for child in widget.children.values(): + with contextlib.suppress(Exception): + child.configure(state=state) + else: + widget.configure(state=state) + + +def set_label_text(label, text: str) -> None: + label.configure(text=str(text)) + + +def call_after(widget, delay_ms: int, func) -> None: + widget.after(delay_ms, func) diff --git a/modules/util/ui/ctk_validation.py b/modules/util/ui/ctk_validation.py index 9f44de8eb..5347ba6d4 100644 --- a/modules/util/ui/ctk_validation.py +++ b/modules/util/ui/ctk_validation.py @@ -1,18 +1,21 @@ from __future__ import annotations import contextlib -import os -import re -import sys import tkinter as tk -from collections import deque from collections.abc import Callable -from pathlib import PurePosixPath, PureWindowsPath from typing import TYPE_CHECKING, Any -from urllib.parse import urlparse -from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType +from modules.util.ui.validation import ( + DEBOUNCE_TYPING_MS, + DEFAULT_MAX_UNDO, + ERROR_BORDER_COLOR, + UNDO_DEBOUNCE_MS, + BaseFieldValidator, + UndoHistory, + _active_validators, + _validate_path_field, +) if TYPE_CHECKING: from modules.util.ui.UIState import UIState @@ -20,159 +23,6 @@ import customtkinter as ctk -DEBOUNCE_TYPING_MS = 250 -UNDO_DEBOUNCE_MS = 500 -ERROR_BORDER_COLOR = "#dc3545" - -_active_validators: set[FieldValidator] = set() - -TRAILING_SLASH_RE = re.compile(r"[\\/]$") -ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") -HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") - -_INVALID_CHARS = {chr(c) for c in range(32)} -_IS_WINDOWS = sys.platform == "win32" -if _IS_WINDOWS: - _INVALID_CHARS |= set('<>"|?*') - - -def _is_huggingface_repo_or_file(value: str) -> bool: - trimmed = value.strip() - - if trimmed.startswith("https://"): - parsed = urlparse(trimmed) - if parsed.netloc not in {"huggingface.co", "huggingface.com"}: - return False - parts = parsed.path.strip("/").split("/") - if len(parts) >= 5 and parts[2] in {"resolve", "blob"}: - return bool(ENDS_WITH_EXT.search(parts[-1])) - return False - - if len(trimmed) > 96: - return False - if " " in trimmed or "\t" in trimmed: - return False - if "—" in trimmed or ".." in trimmed: - return False - if trimmed.startswith(("\\\\", "//", "/")): - return False - if len(trimmed) >= 2 and trimmed[1] == ":" and trimmed[0].isalpha(): - return False - if trimmed.count("/") != 1: - return False - - return bool(HUGGINGFACE_REPO_RE.match(trimmed)) - - -def _has_invalid_chars(value: str) -> bool: - return bool(_INVALID_CHARS.intersection(value)) - - -def _check_overwrite(path: str, *, is_dir: bool, prevent: bool) -> str | None: - if not prevent: - return None - abs_path = os.path.abspath(path) - if is_dir and os.path.isdir(abs_path): - return "Output folder already exists (overwrite prevented)" - if not is_dir and os.path.isfile(abs_path): - return "Output file already exists (overwrite prevented)" - return None - - -def validate_path( - value: str, - io_type: PathIOType = PathIOType.INPUT, - *, - prevent_overwrites: bool = False, - output_format: str | None = None, -) -> str | None: - """Return an error string if *value* is an invalid path, else ``None``.""" - trimmed = value.strip() - - if not trimmed: - return "Path is empty" - if TRAILING_SLASH_RE.search(trimmed): - return "Path must not end with a slash" - if _has_invalid_chars(trimmed): - return "Path contains invalid characters" - - if trimmed.startswith("cloud:"): - cloud_path = trimmed[6:] - if not cloud_path: - return "Cloud path is empty" - if cloud_path.startswith(("http://", "https://")): - return "Cloud path cannot be a URL" - if "\\" in cloud_path: - return "Cloud path must use forward slashes (/)" - return None - - if io_type == PathIOType.INPUT and _is_huggingface_repo_or_file(trimmed): - return None - - if io_type == PathIOType.INPUT: - if not os.path.exists(os.path.abspath(trimmed)): - return "Input path does not exist" - - if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): - if not os.path.isdir(os.path.dirname(os.path.abspath(trimmed))): - return "Parent folder does not exist" - - if io_type == PathIOType.MODEL and output_format is not None: - if output_format == "DIFFUSERS": - if ENDS_WITH_EXT.search(trimmed): - return "Diffusers output must be a directory path, not a file" - return _check_overwrite(trimmed, is_dir=True, prevent=prevent_overwrites) - - try: - expected_ext = ModelFormat[output_format].file_extension() - except KeyError: - expected_ext = "" - - if expected_ext: - suffix = (PureWindowsPath(trimmed) if _IS_WINDOWS else PurePosixPath(trimmed)).suffix.lower() - if suffix != expected_ext: - return f"Extension must be '{expected_ext}' for {output_format} format" - return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) - - if io_type == PathIOType.OUTPUT: - return _check_overwrite(trimmed, is_dir=False, prevent=prevent_overwrites) - - return None - -DEFAULT_MAX_UNDO = 20 - - -class UndoHistory: - def __init__(self, max_size: int = DEFAULT_MAX_UNDO): - self._stack: deque[str] = deque(maxlen=max_size) - self._redo_stack: list[str] = [] - - def push(self, value: str): - if self._stack and self._stack[-1] == value: - return - self._stack.append(value) - self._redo_stack.clear() - - def undo(self, current: str) -> str | None: - if not self._stack: - return None - top = self._stack[-1] - if top == current and len(self._stack) > 1: - self._redo_stack.append(self._stack.pop()) - return self._stack[-1] - elif top != current: - self._redo_stack.append(current) - return top - return None - - def redo(self) -> str | None: - if not self._redo_stack: - return None - value = self._redo_stack.pop() - self._stack.append(value) - return value - - class DebounceTimer: def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): self.widget = widget @@ -199,7 +49,7 @@ def cancel(self): self._after_id = None -class FieldValidator: +class FieldValidator(BaseFieldValidator): def __init__( self, component: ctk.CTkEntry, @@ -210,12 +60,9 @@ def __init__( extra_validate: Callable[[str], str | None] | None = None, required: bool = False, ): + super().__init__(ui_state, var_name, extra_validate, required) self.component = component self.var = var - self.ui_state = ui_state - self.var_name = var_name - self._extra_validate = extra_validate - self._required = required try: self._original_border_color = component.cget("border_color") @@ -227,7 +74,6 @@ def __init__( self._real_var_trace_name: str | None = None self._syncing = False self._touched = False - self._bound = False self._debounce: DebounceTimer | None = None self._undo_debounce: DebounceTimer | None = None @@ -309,59 +155,12 @@ def _commit(self) -> None: self.var.set(shadow_val) self._syncing = False - def validate(self, value: str) -> str | None: - """Return an error string if *value* is invalid, else None.""" - meta = self.ui_state.get_field_metadata(self.var_name) - declared_type = meta.type - nullable = meta.nullable - default_val = meta.default - - if value == "": - if self._required: - return "Value required" - if nullable: - return None - if declared_type is str: - if default_val == "": - return None - return "Value required" - return None - - try: - if declared_type is int: - v = int(value) - if v < 0: - return "Value must be non-negative" - elif declared_type is float: - v = float(value) - if v < 0: - return "Value must be non-negative" - elif declared_type is bool: - if value.lower() not in ("true", "false", "0", "1"): - return "Invalid bool" - except ValueError: - return "Invalid value" - - if self._extra_validate is not None: - return self._extra_validate(value) - - return None - def _apply_error(self) -> None: self.component.configure(border_color=ERROR_BORDER_COLOR) def _clear_error(self) -> None: self.component.configure(border_color=self._original_border_color) - def _validate_and_style(self, value: str) -> bool: - error = self.validate(value) - if error is None: - self._clear_error() - return True - else: - self._apply_error() - return False - def _on_shadow_write(self, *_args) -> None: if self._syncing: return @@ -436,6 +235,21 @@ def _on_redo(self, _e=None) -> str: self._set_value(next_val) return "break" + def flush(self) -> str | None: + if self._debounce: + self._debounce.cancel() + + value = self._shadow_var.get() + error = self.validate(value) + + if error is not None: + self._apply_error() + else: + self._clear_error() + self._commit() + + return error + class PathValidator(FieldValidator): """FieldValidator with additional path-specific checks.""" @@ -454,48 +268,14 @@ def __init__( super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) self.io_type = io_type - def _get_var_safe(self, name: str) -> tk.Variable | None: - try: - return self.ui_state.get_var(name) - except (KeyError, AttributeError): - return None - def validate(self, value: str) -> str | None: base_err = super().validate(value) if base_err is not None: return base_err if value == "": return None - - prevent_var = self._get_var_safe("prevent_overwrites") - format_var = self._get_var_safe("output_model_format") - return validate_path( - value, - io_type=self.io_type, - prevent_overwrites=prevent_var.get() if prevent_var is not None else False, - output_format=format_var.get() if format_var is not None else None, - ) + return _validate_path_field(self.ui_state, self.io_type, value) def revalidate(self) -> None: if self.component.winfo_exists(): self._validate_and_style(self._shadow_var.get()) - - -def flush_and_validate_all() -> list[str]: - invalid: list[str] = [] - - for v in list(_active_validators): - if v._debounce: - v._debounce.cancel() - - value = v._shadow_var.get() - error = v.validate(value) - - if error is not None: - v._apply_error() - invalid.append(f"{v.var_name}: {error}") - else: - v._clear_error() - v._commit() - - return invalid diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 9f44de8eb..abb131be3 100644 --- a/modules/util/ui/validation.py +++ b/modules/util/ui/validation.py @@ -1,31 +1,26 @@ from __future__ import annotations -import contextlib import os import re import sys -import tkinter as tk +from abc import ABC, abstractmethod from collections import deque from collections.abc import Callable from pathlib import PurePosixPath, PureWindowsPath -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from urllib.parse import urlparse from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.PathIOType import PathIOType if TYPE_CHECKING: - from modules.util.ui.UIState import UIState - - import customtkinter as ctk + from modules.util.ui.UIState import BaseUIState DEBOUNCE_TYPING_MS = 250 UNDO_DEBOUNCE_MS = 500 ERROR_BORDER_COLOR = "#dc3545" -_active_validators: set[FieldValidator] = set() - TRAILING_SLASH_RE = re.compile(r"[\\/]$") ENDS_WITH_EXT = re.compile(r"\.[A-Za-z0-9]+$") HUGGINGFACE_REPO_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") @@ -173,141 +168,37 @@ def redo(self) -> str | None: return value -class DebounceTimer: - def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): - self.widget = widget - self.delay_ms = delay_ms - self.callback = callback - self._after_id: str | None = None - - def call(self, *args, **kwargs): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) - - def fire(): - self._after_id = None - self.callback(*args, **kwargs) - - with contextlib.suppress(tk.TclError): - self._after_id = self.widget.after(self.delay_ms, fire) - - def cancel(self): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) - self._after_id = None - - -class FieldValidator: +class BaseFieldValidator(ABC): def __init__( self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, + ui_state: BaseUIState, var_name: str, - max_undo: int = DEFAULT_MAX_UNDO, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, ): - self.component = component - self.var = var self.ui_state = ui_state self.var_name = var_name self._extra_validate = extra_validate self._required = required - - try: - self._original_border_color = component.cget("border_color") - except Exception: - self._original_border_color = "gray50" - - self._shadow_var = tk.StringVar(master=component) - self._shadow_trace_name: str | None = None - self._real_var_trace_name: str | None = None - self._syncing = False - self._touched = False self._bound = False - self._debounce: DebounceTimer | None = None - self._undo_debounce: DebounceTimer | None = None - self._undo = UndoHistory(max_undo) - - def attach(self) -> None: - self._shadow_var.set(self.var.get()) - self._swap_textvariable(self._shadow_var) - - self._debounce = DebounceTimer( - self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire - ) - self._undo_debounce = DebounceTimer( - self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot - ) - - self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) - self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) - - self.component.bind("", self._on_focus_in) - self.component.bind("", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("", self._on_focus_out) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_enter) - - self._bound = True - _active_validators.add(self) - - def detach(self) -> None: - if not self._bound: - return - self._bound = False - _active_validators.discard(self) - - self._commit() - - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() - - if self._shadow_trace_name: - with contextlib.suppress(Exception): - self._shadow_var.trace_remove("write", self._shadow_trace_name) - self._shadow_trace_name = None - - if self._real_var_trace_name: - with contextlib.suppress(Exception): - self.var.trace_remove("write", self._real_var_trace_name) - self._real_var_trace_name = None - - self._swap_textvariable(self.var) - - def _swap_textvariable(self, new_var: tk.Variable) -> None: - comp = self.component - if comp._textvariable_callback_name: - with contextlib.suppress(Exception): - comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] - comp._textvariable_callback_name = "" + @abstractmethod + def _apply_error(self) -> None: + pass - comp.configure(textvariable=new_var) + @abstractmethod + def _clear_error(self) -> None: + pass - if new_var is not None: - comp._textvariable_callback_name = new_var.trace_add( - "write", comp._textvariable_callback - ) + @abstractmethod + def flush(self) -> str | None: + pass - def _commit(self) -> None: - shadow_val = self._shadow_var.get() - if shadow_val != self.var.get(): - self._syncing = True - self.var.set(shadow_val) - self._syncing = False + def _get_var_safe(self, name: str): + try: + return self.ui_state.get_var(name) + except (KeyError, AttributeError): + return None def validate(self, value: str) -> str | None: """Return an error string if *value* is invalid, else None.""" @@ -347,12 +238,6 @@ def validate(self, value: str) -> str | None: return None - def _apply_error(self) -> None: - self.component.configure(border_color=ERROR_BORDER_COLOR) - - def _clear_error(self) -> None: - self.component.configure(border_color=self._original_border_color) - def _validate_and_style(self, value: str) -> bool: error = self.validate(value) if error is None: @@ -362,140 +247,31 @@ def _validate_and_style(self, value: str) -> bool: self._apply_error() return False - def _on_shadow_write(self, *_args) -> None: - if self._syncing: - return - if not self._touched: - # external sync or initial set — commit immediately - self._commit() - if self._debounce: - self._debounce.cancel() - return - if self._debounce: - self._debounce.call() - if self._undo_debounce: - self._undo_debounce.call() - - def _on_real_var_write(self, *_args) -> None: - if self._syncing: - return - # external change (preset load, file dialog, etc) — sync to shadow var - self._syncing = True - self._shadow_var.set(self.var.get()) - self._syncing = False - self._validate_and_style(self._shadow_var.get()) - - def _push_undo_snapshot(self) -> None: - self._undo.push(self._shadow_var.get()) - - def _on_debounce_fire(self) -> None: - val = self._shadow_var.get() - if self._validate_and_style(val): - self._commit() - - def _on_focus_in(self, _e=None) -> None: - self._touched = False - self._undo.push(self._shadow_var.get()) - - def _on_user_input(self, _e=None) -> None: - self._touched = True - - def _on_focus_out(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() - if self._touched: - if self._validate_and_style(self._shadow_var.get()): - self._commit() - self._undo.push(self._shadow_var.get()) - - def _on_enter(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._touched: - if self._validate_and_style(self._shadow_var.get()): - self._commit() - - def _set_value(self, value: str) -> None: - self._syncing = True - self._shadow_var.set(value) - self._syncing = False - if self._validate_and_style(value): - self._commit() - - def _on_undo(self, _e=None) -> str: - previous = self._undo.undo(self._shadow_var.get()) - if previous is not None: - self._set_value(previous) - return "break" - - def _on_redo(self, _e=None) -> str: - next_val = self._undo.redo() - if next_val is not None: - self._set_value(next_val) - return "break" - - -class PathValidator(FieldValidator): - """FieldValidator with additional path-specific checks.""" - - def __init__( - self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, - var_name: str, - io_type: PathIOType = PathIOType.INPUT, - max_undo: int = DEFAULT_MAX_UNDO, - extra_validate: Callable[[str], str | None] | None = None, - required: bool = False, - ): - super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) - self.io_type = io_type - - def _get_var_safe(self, name: str) -> tk.Variable | None: - try: - return self.ui_state.get_var(name) - except (KeyError, AttributeError): - return None - - def validate(self, value: str) -> str | None: - base_err = super().validate(value) - if base_err is not None: - return base_err - if value == "": - return None - - prevent_var = self._get_var_safe("prevent_overwrites") - format_var = self._get_var_safe("output_model_format") - return validate_path( - value, - io_type=self.io_type, - prevent_overwrites=prevent_var.get() if prevent_var is not None else False, - output_format=format_var.get() if format_var is not None else None, - ) - def revalidate(self) -> None: - if self.component.winfo_exists(): - self._validate_and_style(self._shadow_var.get()) +_active_validators: set[BaseFieldValidator] = set() def flush_and_validate_all() -> list[str]: invalid: list[str] = [] - for v in list(_active_validators): - if v._debounce: - v._debounce.cancel() - - value = v._shadow_var.get() - error = v.validate(value) - + error = v.flush() if error is not None: - v._apply_error() invalid.append(f"{v.var_name}: {error}") - else: - v._clear_error() - v._commit() - return invalid + + +def _validate_path_field(ui_state: BaseUIState, io_type: PathIOType, value: str) -> str | None: + try: + prevent_var = ui_state.get_var("prevent_overwrites") + except (KeyError, AttributeError): + prevent_var = None + try: + format_var = ui_state.get_var("output_model_format") + except (KeyError, AttributeError): + format_var = None + return validate_path( + value, + io_type=io_type, + prevent_overwrites=prevent_var.get() if prevent_var is not None else False, + output_format=format_var.get() if format_var is not None else None, + ) diff --git a/scripts/train_ui.py b/scripts/train_ui.py index 46ee8f1e6..562c73feb 100644 --- a/scripts/train_ui.py +++ b/scripts/train_ui.py @@ -2,11 +2,11 @@ script_imports() -from modules.ui.TrainUI import TrainUI +from modules.ui.CtkTrainUIView import CtkTrainUIView def main(): - ui = TrainUI() + ui = CtkTrainUIView() ui.mainloop() From f7e4159212fe1908cecc0966f764d2e535392863 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:09:00 +0200 Subject: [PATCH 11/67] copy: add PySide6 view copies mirroring CTK structure Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 1 + .../ui/PySide6AdditionalEmbeddingsTabView.py | 51 ++ modules/ui/PySide6CaptionUIView.py | 228 +++++++ modules/ui/PySide6CloudTabView.py | 45 ++ modules/ui/PySide6ConceptTabView.py | 176 ++++++ modules/ui/PySide6ConceptWindowView.py | 173 ++++++ modules/ui/PySide6ConfigListView.py | 71 +++ modules/ui/PySide6ConvertModelUIView.py | 34 + .../ui/PySide6GenerateCaptionsWindowView.py | 118 ++++ modules/ui/PySide6GenerateMasksWindowView.py | 141 +++++ modules/ui/PySide6LoraTabView.py | 41 ++ modules/ui/PySide6ModelTabView.py | 48 ++ modules/ui/PySide6MuonAdamWindowView.py | 37 ++ modules/ui/PySide6OffloadingWindowView.py | 31 + .../ui/PySide6OptimizerParamsWindowView.py | 91 +++ modules/ui/PySide6ProfilingWindowView.py | 49 ++ modules/ui/PySide6SampleFrameView.py | 43 ++ modules/ui/PySide6SampleParamsWindowView.py | 32 + modules/ui/PySide6SampleWindowView.py | 102 +++ modules/ui/PySide6SamplingTabView.py | 50 ++ .../ui/PySide6SchedulerParamsWindowView.py | 96 +++ .../PySide6TimestepDistributionWindowView.py | 84 +++ modules/ui/PySide6TopBarView.py | 50 ++ modules/ui/PySide6TrainUIView.py | 413 ++++++++++++ modules/ui/PySide6TrainingTabView.py | 77 +++ modules/ui/PySide6VideoToolUIView.py | 128 ++++ modules/util/ui/pyside6_components.py | 588 ++++++++++++++++++ modules/util/ui/pyside6_validation.py | 281 +++++++++ 28 files changed, 3279 insertions(+) create mode 100644 modules/ui/PySide6AdditionalEmbeddingsTabView.py create mode 100644 modules/ui/PySide6CaptionUIView.py create mode 100644 modules/ui/PySide6CloudTabView.py create mode 100644 modules/ui/PySide6ConceptTabView.py create mode 100644 modules/ui/PySide6ConceptWindowView.py create mode 100644 modules/ui/PySide6ConfigListView.py create mode 100644 modules/ui/PySide6ConvertModelUIView.py create mode 100644 modules/ui/PySide6GenerateCaptionsWindowView.py create mode 100644 modules/ui/PySide6GenerateMasksWindowView.py create mode 100644 modules/ui/PySide6LoraTabView.py create mode 100644 modules/ui/PySide6ModelTabView.py create mode 100644 modules/ui/PySide6MuonAdamWindowView.py create mode 100644 modules/ui/PySide6OffloadingWindowView.py create mode 100644 modules/ui/PySide6OptimizerParamsWindowView.py create mode 100644 modules/ui/PySide6ProfilingWindowView.py create mode 100644 modules/ui/PySide6SampleFrameView.py create mode 100644 modules/ui/PySide6SampleParamsWindowView.py create mode 100644 modules/ui/PySide6SampleWindowView.py create mode 100644 modules/ui/PySide6SamplingTabView.py create mode 100644 modules/ui/PySide6SchedulerParamsWindowView.py create mode 100644 modules/ui/PySide6TimestepDistributionWindowView.py create mode 100644 modules/ui/PySide6TopBarView.py create mode 100644 modules/ui/PySide6TrainUIView.py create mode 100644 modules/ui/PySide6TrainingTabView.py create mode 100644 modules/ui/PySide6VideoToolUIView.py create mode 100644 modules/util/ui/pyside6_components.py create mode 100644 modules/util/ui/pyside6_validation.py diff --git a/.gitignore b/.gitignore index da19fd74e..92e3d5612 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,4 @@ pixi.toml train.bat debug_report.log config_diff.txt +PLAN.md diff --git a/modules/ui/PySide6AdditionalEmbeddingsTabView.py b/modules/ui/PySide6AdditionalEmbeddingsTabView.py new file mode 100644 index 000000000..fc24c61d1 --- /dev/null +++ b/modules/ui/PySide6AdditionalEmbeddingsTabView.py @@ -0,0 +1,51 @@ + +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseAdditionalEmbeddingsTabView import BaseAdditionalEmbeddingsTabView, BaseEmbeddingWidgetView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState + +import customtkinter as ctk + + +class CtkAdditionalEmbeddingsTabView(CtkConfigListView, BaseAdditionalEmbeddingsTabView): + + def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, + attr_name="additional_embeddings", + enable_key="train", + from_external_file=False, + add_button_text="add embedding", + is_full_width=True, + show_toggle_button=True, + ) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkEmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + + +class CtkEmbeddingWidgetView(BaseEmbeddingWidgetView, ctk.CTkFrame): + + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseEmbeddingWidgetView.__init__(self, ctk_components) + + self.element = element + ui_state = CtkUIState(self, element) + + self.grid_columnconfigure(0, weight=1) + + top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + top_frame.grid(row=0, column=0, sticky="nsew") + top_frame.grid_columnconfigure(3, weight=1) + top_frame.grid_columnconfigure(5, weight=1) + + bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") + bottom_frame.grid(row=1, column=0, sticky="nsew") + bottom_frame.grid_columnconfigure(7, weight=1) + + self.build_content(top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller) + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/PySide6CaptionUIView.py b/modules/ui/PySide6CaptionUIView.py new file mode 100644 index 000000000..281036912 --- /dev/null +++ b/modules/ui/PySide6CaptionUIView.py @@ -0,0 +1,228 @@ +from tkinter import filedialog + +from modules.ui.BaseCaptionUIView import BaseCaptionUIView +from modules.ui.CaptionUIController import CaptionUIController +from modules.ui.CtkGenerateCaptionsWindowView import CtkGenerateCaptionsWindowView +from modules.ui.CtkGenerateMasksWindowView import CtkGenerateMasksWindowView +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon + +import customtkinter as ctk +from customtkinter import ScalingTracker, ThemeManager +from PIL import Image + + +class CtkCaptionUIView(BaseCaptionUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: CaptionUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseCaptionUIView.__init__(self, ctk_components) + self.protocol("WM_DELETE_WINDOW", controller.on_close) + + self.controller = controller + controller.view = self + self.config_ui_state = CtkUIState(self, controller.config_ui_data) + self.enable_mask_editing_var = ctk.BooleanVar() + self.mask_editing_alpha = None + self.prompt_var = None + self.prompt_component = None + self.image = None + self.image_label = None + self.file_list = None + self.image_labels = [] + + self.title("OneTrainer") + self.geometry("1280x980") + self.resizable(False, False) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + top_frame = ctk.CTkFrame(self) + top_frame.grid(row=0, column=0, sticky="nsew") + self.build_top_bar(top_frame, controller, self.config_ui_state) + + self.bottom_frame = ctk.CTkFrame(self) + self.bottom_frame.grid(row=1, column=0, sticky="nsew") + self.bottom_frame.grid_rowconfigure(0, weight=1) + self.bottom_frame.grid_columnconfigure(0, weight=0) + self.bottom_frame.grid_columnconfigure(1, weight=1) + + self.file_list_column(self.bottom_frame) + self.content_column(self.bottom_frame) + self.controller.load_directory() + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def file_list_column(self, master): + if self.file_list is not None: + self.image_labels = [] + self.file_list.destroy() + + self.file_list = ctk.CTkScrollableFrame(master, width=300) + self.file_list.grid(row=0, column=0, sticky="nsew") + + for i, filename in enumerate(self.controller.image_rel_paths): + def __create_switch_image(index): + def __switch_image(event): + self.controller.switch_image(index) + + return __switch_image + + label = ctk.CTkLabel(self.file_list, text=filename) + label.bind("", __create_switch_image(i)) + + self.image_labels.append(label) + label.grid(row=i, column=0, padx=5, sticky="nsw") + + def content_column(self, master): + image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) + + right_frame = ctk.CTkFrame(master, fg_color="transparent") + right_frame.grid(row=0, column=1, sticky="nsew") + + right_frame.grid_columnconfigure(4, weight=1) + right_frame.grid_rowconfigure(1, weight=1) + + self.build_mask_buttons(right_frame) + + # checkbox to enable mask editing + self.enable_mask_editing_var = ctk.BooleanVar() + self.enable_mask_editing_var.set(False) + enable_mask_editing_checkbox = ctk.CTkCheckBox( + right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) + enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") + + # mask alpha textbox + self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") + self.mask_editing_alpha.insert(0, "1.0") + self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) + self.bind_key_events(self.mask_editing_alpha) + + mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) + mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") + + # image + self.image = ctk.CTkImage( + light_image=image, + size=(self.controller.image_size, self.controller.image_size) + ) + self.image_label = ctk.CTkLabel( + master=right_frame, text="", image=self.image, + height=self.controller.image_size, width=self.controller.image_size + ) + self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") + + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + self.image_label.bind("", self.edit_mask) + bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) + + # prompt + self.prompt_var = ctk.StringVar() + self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) + self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") + self.bind_key_events(self.prompt_component) + self.prompt_component.focus_set() + + def bind_key_events(self, component): + component.bind("", lambda e: self.controller.next_image()) + component.bind("", lambda e: self.controller.previous_image()) + component.bind("", self.save) + component.bind("", self.toggle_mask) + component.bind("", self.draw_mask_editing_mode) + component.bind("", self.fill_mask_editing_mode) + + def refresh_file_list(self): + self.file_list_column(self.bottom_frame) + + def focus_prompt(self): + self.prompt_component.focus_set() + + def on_image_switched(self, old_index, new_index, prompt): + if len(self.image_labels) > 0 and old_index < len(self.image_labels): + self.image_labels[old_index].configure( + text_color=ThemeManager.theme["CTkLabel"]["text_color"]) + self.image_labels[new_index].configure(text_color="#FF0000") + self.refresh_image() + self.prompt_var.set(prompt) + + def on_image_cleared(self): + image = Image.new("RGB", (512, 512), (0, 0, 0)) + self.image.configure(light_image=image) + + def refresh_image(self): + pil_image, size = self.controller.get_display_image() + self.image.configure(light_image=pil_image, size=size) + + def draw_mask_radius(self, delta, raw_event): + self.controller.update_mask_draw_radius(delta) + + def edit_mask(self, event): + if not self.enable_mask_editing_var.get(): + return + + if event.widget != self.image_label.children["!label"]: + return + + display_scaling = ScalingTracker.get_window_scaling(self) + + event_x = event.x / display_scaling + event_y = event.y / display_scaling + + is_right = False + is_left = False + if event.state & 0x0100 or event.num == 1: # left mouse button + is_left = True + elif event.state & 0x0400 or event.num == 3: # right mouse button + is_right = True + + try: + alpha = float(self.mask_editing_alpha.get()) + except Exception: + alpha = 1.0 + + self.controller.handle_edit_mask(event_x, event_y, is_left, is_right, alpha) + + def save(self, event): + self.controller.save(self.prompt_var.get()) + + def draw_mask_editing_mode(self, *args): + self.controller.set_mask_editing_mode('draw') + + if args: + # disable default event + return "break" + return None + + def fill_mask_editing_mode(self, *args): + self.controller.set_mask_editing_mode('fill') + + def toggle_mask(self, *args): + self.controller.toggle_mask() + self.refresh_image() + + def open_directory(self): + new_dir = filedialog.askdirectory() + + if new_dir: + self.controller.dir = new_dir + self.controller.load_directory(include_subdirectories=self.controller.config_ui_data["include_subdirectories"]) + + def open_mask_window(self): + self.wait_window(self.controller.open_mask_window(self, CtkGenerateMasksWindowView)) + self.controller.switch_image(self.controller.current_image_index) + + def open_caption_window(self): + self.wait_window(self.controller.open_caption_window(self, CtkGenerateCaptionsWindowView)) + self.controller.switch_image(self.controller.current_image_index) + + def open_in_explorer(self): + self.controller.open_in_explorer() + + def destroy(self): + self.controller._release_models() + super().destroy() diff --git a/modules/ui/PySide6CloudTabView.py b/modules/ui/PySide6CloudTabView.py new file mode 100644 index 000000000..0a5249069 --- /dev/null +++ b/modules/ui/PySide6CloudTabView.py @@ -0,0 +1,45 @@ + + +from modules.ui.BaseCloudTabView import BaseCloudTabView +from modules.ui.CloudTabController import CloudTabController +from modules.util.ui import ctk_components + +import customtkinter as ctk + + +class CtkCloudTabView(BaseCloudTabView): + def __init__(self, master, controller: CloudTabController, ui_state): + BaseCloudTabView.__init__(self, ctk_components) + self.master = master + self.controller = controller + self.ui_state = ui_state + + self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, weight=0) + self.frame.grid_columnconfigure(3, weight=1) + self.frame.grid_columnconfigure(4, weight=0) + self.frame.grid_columnconfigure(5, weight=1) + + self.build_content(self.frame, controller, ui_state) + + self.frame.pack(fill="both", expand=1) + + + def _on_set_gpu_types(self): + self.gpu_types_menu.configure(values=self.controller.get_gpu_types()) + + def _make_reattach_frame(self, frame): + reattach_frame = ctk.CTkFrame(frame, fg_color="transparent") + reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") + reattach_frame.grid_columnconfigure(0, weight=1) + reattach_frame.grid_columnconfigure(1, weight=1) + return reattach_frame + + def _make_create_frame(self, frame): + create_frame = ctk.CTkFrame(frame, fg_color="transparent") + create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") + create_frame.grid_columnconfigure(0, weight=0) + create_frame.grid_columnconfigure(1, weight=1) + return create_frame diff --git a/modules/ui/PySide6ConceptTabView.py b/modules/ui/PySide6ConceptTabView.py new file mode 100644 index 000000000..5b3e86ac9 --- /dev/null +++ b/modules/ui/PySide6ConceptTabView.py @@ -0,0 +1,176 @@ +from tkinter import BooleanVar, StringVar + +from modules.ui.BaseConceptTabView import BaseConceptTabView, BaseConceptWidgetView +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkConceptWindowView import CtkConceptWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.util.ui import ctk_components +from modules.util.ui.ctk_validation import DebounceTimer +from modules.util.ui.CtkUIState import CtkUIState + +import customtkinter as ctk + + +class CtkConceptTabView(CtkConfigListView, BaseConceptTabView): + + def __init__(self, master, controller: ConceptTabController, ui_state): + # Pre-initialize before CtkConfigListView.__init__ because _reset_filters is + # called during build() via options_kv's immediate update_var() call. + self.search_var = StringVar() + self.filter_var = StringVar(value="ALL") + self.show_disabled_var = BooleanVar(value=True) + + CtkConfigListView.__init__( + self, master, controller, ui_state, + from_external_file=True, + attr_name="concept_file_name", + config_dir="training_concepts", + default_config_name="concepts.json", + add_button_text="Add Concept", + add_button_tooltip="Adds a new concept to the current config.", + is_full_width=False, + show_toggle_button=True, + ) + self._toolbar = None + self._toolbar_is_wrapped = False + self._add_search_bar() + self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return self.controller.open_element_window(self.master, self.current_config[i], ui_state[0], ui_state[1], ui_state[2], CtkConceptWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + + def _add_search_bar(self): + toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") + toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") + toolbar.grid_columnconfigure(2, weight=1) + self._toolbar = toolbar + + ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0, 5)) + self.search_var = StringVar() + self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, + placeholder_text="Filter...", width=200) + self.search_entry.grid(row=0, column=1) + self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) + self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) + + ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) + + ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0, 5)) + self.filter_var = StringVar(value="ALL") + ctk.CTkOptionMenu(toolbar, values=self._FILTER_TYPES, + variable=self.filter_var, command=lambda x: self._update_filters(), + width=150).grid(row=0, column=4) + + self.show_disabled_var = BooleanVar(value=True) + self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, + command=self._update_filters, width=100) + self.show_disabled_checkbox.grid(row=0, column=5, padx=(10, 0)) + self._refresh_show_disabled_text() + + ctk.CTkButton(toolbar, text="Clear", width=50, + command=self._reset_filters).grid(row=0, column=6, padx=(10, 0)) + + def _maybe_reposition_toolbar(self, width): + if not self._toolbar: + return + threshold = 1070 + want_wrapped = width < threshold + if want_wrapped == self._toolbar_is_wrapped: + return + self._toolbar_is_wrapped = want_wrapped + if want_wrapped: + self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) + else: + self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + + def _reset_filters(self): + self.search_var.set("") + self.filter_var.set("ALL") + self.show_disabled_var.set(True) + self._update_filters() + + def _refresh_show_disabled_text(self): + try: + disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) + except (AttributeError, TypeError): + disabled_count = 0 + text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" + try: + if getattr(self, 'show_disabled_checkbox', None): + self.show_disabled_checkbox.configure(text=text) + except (AttributeError, RuntimeError): + pass + + +class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): + + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): + ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") + BaseConceptWidgetView.__init__(self, ctk_components) + + self.concept = concept + self.ui_state = CtkUIState(self, concept) + self.image_ui_state = CtkUIState(self, concept.image) + self.text_ui_state = CtkUIState(self, concept.text) + self.i = i + + self.grid_rowconfigure(1, weight=1) + + self.image = ctk.CTkImage( + light_image=self._get_preview_image(), + size=(150, 150) + ) + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150) + image_label.grid(row=0, column=0) + + self.name_label = self.components.label(self, 1, 0, self._get_display_name(), pad=5, wraplength=140) + + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i), + ) + close_button.place(x=0, y=0) + + clone_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="+", + corner_radius=2, + fg_color="#00C000", + command=lambda: clone_command(self.i, controller.randomize_seed), + ) + clone_button.place(x=25, y=0) + + enabled_switch = ctk.CTkSwitch( + master=self, + width=40, + variable=self.ui_state.get_var("enabled"), + text="", + command=save_command, + ) + enabled_switch.place(x=110, y=0) + + image_label.bind( + "", + lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) + ) + + def configure_element(self): + self.name_label.configure(text=self._get_display_name()) + self.image.configure(light_image=self._get_preview_image()) + self._clear_search_cache() + + def place_in_list(self): + index = getattr(self, 'visible_index', self.i) + x = index % 6 + y = index // 6 + self.grid(row=y, column=x, pady=5, padx=5) diff --git a/modules/ui/PySide6ConceptWindowView.py b/modules/ui/PySide6ConceptWindowView.py new file mode 100644 index 000000000..60c0f57fe --- /dev/null +++ b/modules/ui/PySide6ConceptWindowView.py @@ -0,0 +1,173 @@ +import threading + +from modules.ui.BaseConceptWindowView import BaseConceptWindowView +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class CtkConceptWindowView(BaseConceptWindowView, ctk.CTkToplevel): + def __init__( + self, + parent, + controller: ConceptWindowController, + ui_state, + image_ui_state, + text_ui_state, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConceptWindowView.__init__(self, ctk_components) + + self.controller = controller + self.image_preview_file_index = 0 + self.preview_augmentations = ctk.BooleanVar(self, True) + self.bucket_fig = None + + self.title("Concept") + self.geometry("800x700") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + # general tab + general_frame = ctk.CTkScrollableFrame(tabview.add("general"), fg_color="transparent") + general_frame.grid_columnconfigure(1, weight=1) + general_frame.grid_columnconfigure(2, weight=1) + self.build_general_tab(general_frame, controller, ui_state, text_ui_state) + general_frame.pack(fill="both", expand=1) + + # image augmentation tab + image_aug_master = tabview.add("image augmentation") + image_aug_frame = ctk.CTkScrollableFrame(image_aug_master, fg_color="transparent") + image_aug_frame.grid_columnconfigure(0, weight=0) + image_aug_frame.grid_columnconfigure(1, weight=0) + image_aug_frame.grid_columnconfigure(2, weight=0) + image_aug_frame.grid_columnconfigure(3, weight=1) + self.build_image_augmentation_tab(image_aug_frame, controller, image_ui_state) + + # image + image_preview, filename_preview, caption_preview = controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) + self.image = ctk.CTkImage( + light_image=image_preview, + size=image_preview.size, + ) + image_label = ctk.CTkLabel(master=image_aug_frame, text="", image=self.image, height=300, width=300) + image_label.grid(row=0, column=4, rowspan=6) + + # refresh preview + update_button_frame = ctk.CTkFrame(master=image_aug_frame, corner_radius=0, fg_color="transparent") + update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") + update_button_frame.grid_columnconfigure(1, weight=1) + + prev_preview_button = self.components.button(update_button_frame, 0, 0, "<", command=self._prev_image_preview) + self.components.button(update_button_frame, 0, 1, "Update Preview", command=self._update_image_preview) + next_preview_button = self.components.button(update_button_frame, 0, 2, ">", command=self._next_image_preview) + preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self._update_image_preview) + preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) + + prev_preview_button.configure(width=40) + next_preview_button.configure(width=40) + + #caption and filename preview + self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) + self.filename_preview.grid(row=2, column=0, columnspan=3) + self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) + + image_aug_frame.pack(fill="both", expand=1) + + # text augmentation tab + text_aug_frame = ctk.CTkScrollableFrame(tabview.add("text augmentation"), fg_color="transparent") + text_aug_frame.grid_columnconfigure(0, weight=0) + text_aug_frame.grid_columnconfigure(1, weight=0) + text_aug_frame.grid_columnconfigure(2, weight=0) + text_aug_frame.grid_columnconfigure(3, weight=1) + self.build_text_augmentation_tab(text_aug_frame, controller, text_ui_state) + text_aug_frame.pack(fill="both", expand=1) + + # statistics tab + stats_frame = ctk.CTkScrollableFrame(tabview.add("statistics"), fg_color="transparent") + stats_frame.grid_columnconfigure(0, weight=0, minsize=150) + stats_frame.grid_columnconfigure(1, weight=0, minsize=150) + stats_frame.grid_columnconfigure(2, weight=0, minsize=150) + stats_frame.grid_columnconfigure(3, weight=0, minsize=150) + self.build_concept_stats_tab(stats_frame, controller) + + #aspect bucketing plot, mostly copied from timestep preview graph + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + plt.set_loglevel('WARNING') #suppress errors about data type in bar chart + + assert self.bucket_fig is None + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) + self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=stats_frame) + self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + self.bucket_fig.tight_layout() + self.bucket_fig.subplots_adjust(bottom=0.15) + + self.bucket_fig.set_facecolor(background_color) + self.bucket_ax.set_facecolor(background_color) + self.bucket_ax.spines['bottom'].set_color(self.text_color) + self.bucket_ax.spines['left'].set_color(self.text_color) + self.bucket_ax.spines['top'].set_visible(False) + self.bucket_ax.spines['right'].set_color(self.text_color) + self.bucket_ax.tick_params(axis='x', colors=self.text_color, which="both") + self.bucket_ax.tick_params(axis='y', colors=self.text_color, which="both") + self.bucket_ax.xaxis.label.set_color(self.text_color) + self.bucket_ax.yaxis.label.set_color(self.text_color) + + stats_frame.pack(fill="both", expand=1) + + #automatic concept scan + self.scan_thread = threading.Thread(target=controller.auto_update_concept_stats, args=[self], daemon=True) + self.scan_thread.start() + + self.components.button(self, 1, 0, "ok", self._ok) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def _prev_image_preview(self): + self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) + self._update_image_preview() + + def _next_image_preview(self): + self.image_preview_file_index += 1 + self._update_image_preview() + + def _update_image_preview(self): + image_preview, filename_preview, caption_preview = self.controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) + self.image.configure(light_image=image_preview, size=image_preview.size) + self.filename_preview.configure(text=filename_preview) + self.caption_preview.configure(state="normal") + self.caption_preview.delete(index1="1.0", index2="end") + self.caption_preview.insert(index="1.0", text=caption_preview) + self.caption_preview.configure(state="disabled") + + def destroy(self): + if self.bucket_fig is not None: + plt.close(self.bucket_fig) + self.bucket_fig = None + + super().destroy() + + def _ok(self): + self.destroy() diff --git a/modules/ui/PySide6ConfigListView.py b/modules/ui/PySide6ConfigListView.py new file mode 100644 index 000000000..72995bfcc --- /dev/null +++ b/modules/ui/PySide6ConfigListView.py @@ -0,0 +1,71 @@ +import contextlib +from abc import ABC + +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util.ui import ctk_components, dialogs + +import customtkinter as ctk + + +class CtkConfigListView(BaseConfigListView, ABC): + + def __init__( + self, + master, + controller, + ui_state, + from_external_file: bool, + attr_name: str = "", + enable_key: str = "enabled", + config_dir: str = "", + default_config_name: str = "", + add_button_text: str = "", + add_button_tooltip: str = "", + is_full_width: bool = False, + show_toggle_button: bool = False, + ): + BaseConfigListView.__init__(self, ctk_components) + + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.build( + master, controller, ui_state, from_external_file, + attr_name=attr_name, + enable_key=enable_key, + config_dir=config_dir, + default_config_name=default_config_name, + add_button_text=add_button_text, + add_button_tooltip=add_button_tooltip, + is_full_width=is_full_width, + show_toggle_button=show_toggle_button, + ) + + def _create_top_frame(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew") + return frame + + def _create_element_list_frame(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid(row=1, column=0, sticky="nsew") + if self.is_full_width: + frame.grid_columnconfigure(0, weight=1) + return frame + + def _wait_for_window(self, window): + self.master.wait_window(window) + + def _remove_widget_from_layout(self, widget): + widget.grid_remove() + + def _destroy_widget(self, widget): + with contextlib.suppress(AttributeError): + widget.destroy() + + def _destroy_frame(self, frame): + frame.destroy() + + def _show_name_dialog(self, callback): + dialogs.StringInputDialog(self.master, "name", "Name", callback) diff --git a/modules/ui/PySide6ConvertModelUIView.py b/modules/ui/PySide6ConvertModelUIView.py new file mode 100644 index 000000000..782637348 --- /dev/null +++ b/modules/ui/PySide6ConvertModelUIView.py @@ -0,0 +1,34 @@ +from modules.ui.BaseConvertModelUIView import BaseConvertModelUIView +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkConvertModelUIView(BaseConvertModelUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: ConvertModelUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseConvertModelUIView.__init__(self, ctk_components) + + ui_state = CtkUIState(self, controller.convert_model_args) + + self.title("Convert models") + self.geometry("550x350") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + + self.build_content(self.frame, controller, ui_state) + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def set_converting(self, active): + self.button.configure(state="disabled" if active else "normal") diff --git a/modules/ui/PySide6GenerateCaptionsWindowView.py b/modules/ui/PySide6GenerateCaptionsWindowView.py new file mode 100644 index 000000000..09d82f74b --- /dev/null +++ b/modules/ui/PySide6GenerateCaptionsWindowView.py @@ -0,0 +1,118 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.ui.BaseGenerateCaptionsWindowView import BaseGenerateCaptionsWindowView +from modules.ui.GenerateCaptionsWindowController import GenerateCaptionsWindowController +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkGenerateCaptionsWindowView(BaseGenerateCaptionsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateCaptionsWindowController, path, parent_include_subdirectories, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + if path is None: + path = "" + + self.controller = controller + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all captions", "Create if absent", "Add as new line"] + self.model_var = ctk.StringVar(self, "Blip") + self.models = ["Blip", "Blip2", "WD14 VIT v2"] + + self.title("Batch generate captions") + self.geometry("360x360") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.caption_label = ctk.CTkLabel(self.frame, text="Initial Caption", width=100) + self.caption_label.grid(row=2, column=0, sticky="w", padx=5, pady=5) + self.caption_entry = ctk.CTkEntry(self.frame, width=200) + self.caption_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.prefix_label = ctk.CTkLabel(self.frame, text="Caption Prefix", width=100) + self.prefix_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.prefix_entry = ctk.CTkEntry(self.frame, width=200) + self.prefix_entry.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.postfix_label = ctk.CTkLabel(self.frame, text="Caption Postfix", width=100) + self.postfix_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.postfix_entry = ctk.CTkEntry(self.frame, width=200) + self.postfix_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.create_captions_button = ctk.CTkButton(self.frame, text="Create Captions", width=310, command=self._on_create_captions) + self.create_captions_button.grid(row=8, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def _on_create_captions(self): + self.controller.create_captions( + model_name=self.model_var.get(), + path=self.path_entry.get(), + initial_caption=self.caption_entry.get(), + caption_prefix=self.prefix_entry.get(), + caption_postfix=self.postfix_entry.get(), + mode_str=self.mode_var.get(), + include_subdirectories=self.include_subdirectories_var.get(), + ) + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/PySide6GenerateMasksWindowView.py b/modules/ui/PySide6GenerateMasksWindowView.py new file mode 100644 index 000000000..631179fac --- /dev/null +++ b/modules/ui/PySide6GenerateMasksWindowView.py @@ -0,0 +1,141 @@ +import contextlib +import tkinter as tk +from tkinter import filedialog + +from modules.ui.BaseGenerateMasksWindowView import BaseGenerateMasksWindowView +from modules.ui.GenerateMasksWindowController import GenerateMasksWindowController +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkGenerateMasksWindowView(BaseGenerateMasksWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: GenerateMasksWindowController, path, parent_include_subdirectories, *args, **kwargs): + """ + Window for generating masks for a folder of images + + Parameters: + parent (`Tk`): the parent window + path (`str`): the path to the folder + parent_include_subdirectories (`bool`): whether to include subdirectories. used to set the default value of the include subdirectories checkbox + """ + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + + self.controller = controller + if path is None: + path = "" + + self.mode_var = ctk.StringVar(self, "Create if absent") + self.modes = ["Replace all masks", "Create if absent", "Add to existing", "Subtract from existing", "Blend with existing"] + self.model_var = ctk.StringVar(self, "ClipSeg") + self.models = ["ClipSeg", "Rembg", "Rembg-Human", "Hex Color"] + + self.title("Batch generate masks") + self.geometry("360x430") + self.resizable(True, True) + + self.frame = ctk.CTkFrame(self, width=600, height=300) + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.model_label = ctk.CTkLabel(self.frame, text="Model", width=100) + self.model_label.grid(row=0, column=0, sticky="w", padx=5, pady=5) + self.model_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.model_var, values=self.models, dynamic_resizing=False, width=200) + self.model_dropdown.grid(row=0, column=1, sticky="w", padx=5, pady=5) + + self.path_label = ctk.CTkLabel(self.frame, text="Folder", width=100) + self.path_label.grid(row=1, column=0, sticky="w",padx=5, pady=5) + self.path_entry = ctk.CTkEntry(self.frame, width=150) + self.path_entry.insert(0, path) + self.path_entry.grid(row=1, column=1, sticky="w", padx=5, pady=5) + self.path_button = ctk.CTkButton(self.frame, width=30, text="...", command=lambda: self.browse_for_path(self.path_entry)) + self.path_button.grid(row=1, column=1, sticky="e", padx=5, pady=5) + + self.prompt_label = ctk.CTkLabel(self.frame, text="Prompt", width=100) + self.prompt_label.grid(row=2, column=0, sticky="w",padx=5, pady=5) + self.prompt_entry = ctk.CTkEntry(self.frame, width=200) + self.prompt_entry.grid(row=2, column=1, sticky="w", padx=5, pady=5) + + self.mode_label = ctk.CTkLabel(self.frame, text="Mode", width=100) + self.mode_label.grid(row=3, column=0, sticky="w", padx=5, pady=5) + self.mode_dropdown = ctk.CTkOptionMenu(self.frame, variable=self.mode_var, values=self.modes, dynamic_resizing=False, width=200) + self.mode_dropdown.grid(row=3, column=1, sticky="w", padx=5, pady=5) + + self.threshold_label = ctk.CTkLabel(self.frame, text="Threshold", width=100) + self.threshold_label.grid(row=4, column=0, sticky="w", padx=5, pady=5) + self.threshold_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="0.0 - 1.0") + self.threshold_entry.insert(0, "0.3") + self.threshold_entry.grid(row=4, column=1, sticky="w", padx=5, pady=5) + + self.smooth_label = ctk.CTkLabel(self.frame, text="Smooth", width=100) + self.smooth_label.grid(row=5, column=0, sticky="w", padx=5, pady=5) + self.smooth_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="5") + self.smooth_entry.insert(0, 5) + self.smooth_entry.grid(row=5, column=1, sticky="w", padx=5, pady=5) + + self.expand_label = ctk.CTkLabel(self.frame, text="Expand", width=100) + self.expand_label.grid(row=6, column=0, sticky="w", padx=5, pady=5) + self.expand_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="10") + self.expand_entry.insert(0, 10) + self.expand_entry.grid(row=6, column=1, sticky="w", padx=5, pady=5) + + self.alpha_label = ctk.CTkLabel(self.frame, text="Alpha", width=100) + self.alpha_label.grid(row=7, column=0, sticky="w", padx=5, pady=5) + self.alpha_entry = ctk.CTkEntry(self.frame, width=200, placeholder_text="1") + self.alpha_entry.insert(0, 1) + self.alpha_entry.grid(row=7, column=1, sticky="w", padx=5, pady=5) + + self.include_subdirectories_label = ctk.CTkLabel(self.frame, text="Include subfolders", width=100) + self.include_subdirectories_label.grid(row=8, column=0, sticky="w", padx=5, pady=5) + self.include_subdirectories_var = ctk.BooleanVar(self, parent_include_subdirectories) + self.include_subdirectories_switch = ctk.CTkSwitch(self.frame, text="", variable=self.include_subdirectories_var) + self.include_subdirectories_switch.grid(row=8, column=1, sticky="w", padx=5, pady=5) + + self.progress_label = ctk.CTkLabel(self.frame, text="Progress: 0/0", width=100) + self.progress_label.grid(row=9, column=0, sticky="w", padx=5, pady=5) + self.progress = ctk.CTkProgressBar(self.frame, orientation="horizontal", mode="determinate", width=200) + self.progress.grid(row=9, column=1, sticky="w", padx=5, pady=5) + + self.create_masks_button = ctk.CTkButton(self.frame, text="Create Masks", width=310, command=self._on_create_masks) + self.create_masks_button.grid(row=10, column=0, columnspan=2, sticky="w", padx=5, pady=5) + + self.frame.pack(fill="both", expand=True) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def browse_for_path(self, entry_box): + # get the path from the user + path = filedialog.askdirectory() + # set the path to the entry box + # delete entry box text + entry_box.focus_set() + entry_box.delete(0, filedialog.END) + entry_box.insert(0, path) + self.focus_set() + + def set_progress(self, value, max_value): + progress = value / max_value + self.progress.set(progress) + self.progress_label.configure(text=f"{value}/{max_value}") + self.progress.update() + + def _on_create_masks(self): + self.controller.create_masks( + model_name=self.model_var.get(), + path=self.path_entry.get(), + prompt=self.prompt_entry.get(), + mode_str=self.mode_var.get(), + alpha_str=self.alpha_entry.get(), + threshold_str=self.threshold_entry.get(), + smooth_str=self.smooth_entry.get(), + expand_str=self.expand_entry.get(), + include_subdirectories=self.include_subdirectories_var.get(), + ) + + def destroy(self): + with contextlib.suppress(tk.TclError): + self.grab_release() + + super().destroy() diff --git a/modules/ui/PySide6LoraTabView.py b/modules/ui/PySide6LoraTabView.py new file mode 100644 index 000000000..8caa1f171 --- /dev/null +++ b/modules/ui/PySide6LoraTabView.py @@ -0,0 +1,41 @@ + +from modules.ui.BaseLoraTabView import BaseLoraTabView +from modules.ui.LoraTabController import LoraTabController +from modules.util.enum.ModelType import PeftType +from modules.util.ui import ctk_components + +import customtkinter as ctk + + +class CtkLoraTabView(BaseLoraTabView): + def __init__(self, master, controller: LoraTabController, ui_state): + BaseLoraTabView.__init__(self, ctk_components) + self.master = master + self.controller = controller + self.ui_state = ui_state + self.scroll_frame = None + self.options_frame = None + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + self.scroll_frame.grid_columnconfigure(0, weight=0) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=2) + self.build(self.scroll_frame, self.controller, self.ui_state, self.setup_lora) + + def setup_lora(self, peft_type: PeftType): + if self.options_frame: + self.options_frame.destroy() + self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") + self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") + master = self.options_frame + master.grid_columnconfigure(0, weight=0, uniform="a") + master.grid_columnconfigure(1, weight=1, uniform="a") + master.grid_columnconfigure(2, minsize=50, uniform="a") + master.grid_columnconfigure(3, weight=0, uniform="a") + master.grid_columnconfigure(4, weight=1, uniform="a") + self.build_lora_options(master, self.controller, self.ui_state, peft_type) diff --git a/modules/ui/PySide6ModelTabView.py b/modules/ui/PySide6ModelTabView.py new file mode 100644 index 000000000..5b43b7dca --- /dev/null +++ b/modules/ui/PySide6ModelTabView.py @@ -0,0 +1,48 @@ + + +from modules.ui.BaseModelTabView import BaseModelTabView +from modules.ui.ModelTabController import ModelTabController +from modules.util.ui import ctk_components + +import customtkinter as ctk + + +class CtkModelTabView(BaseModelTabView): + def __init__(self, master, controller: ModelTabController, ui_state): + BaseModelTabView.__init__(self, ctk_components) + self.master = master + self.controller = controller + self.ui_state = ui_state + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.scroll_frame = None + + self.refresh_ui() + + def _make_svd_frames(self, parent, row: int): + svd_label_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_label_frame.grid(row=row, column=3, sticky="nsew") + svd_entry_frame = ctk.CTkFrame(parent, fg_color="transparent") + svd_entry_frame.grid(row=row, column=4, sticky="nsew") + return svd_label_frame, svd_entry_frame + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + self.scroll_frame.grid_columnconfigure(0, weight=1) + + base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) + base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") + + base_frame.grid_columnconfigure(0, weight=0) + base_frame.grid_columnconfigure(1, weight=10) # , minsize=500) + base_frame.grid_columnconfigure(2, minsize=50) + base_frame.grid_columnconfigure(3, weight=0) + base_frame.grid_columnconfigure(4, weight=1) + + self.build_content(base_frame, self.controller, self.ui_state) diff --git a/modules/ui/PySide6MuonAdamWindowView.py b/modules/ui/PySide6MuonAdamWindowView.py new file mode 100644 index 000000000..3dc48ab74 --- /dev/null +++ b/modules/ui/PySide6MuonAdamWindowView.py @@ -0,0 +1,37 @@ +from modules.ui.BaseMuonAdamWindowView import BaseMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkMuonAdamWindowView(BaseMuonAdamWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: MuonAdamWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseMuonAdamWindowView.__init__(self, ctk_components) + + self.title(controller.get_title()) + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) diff --git a/modules/ui/PySide6OffloadingWindowView.py b/modules/ui/PySide6OffloadingWindowView.py new file mode 100644 index 000000000..b752f1602 --- /dev/null +++ b/modules/ui/PySide6OffloadingWindowView.py @@ -0,0 +1,31 @@ +from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkOffloadingWindowView(BaseOffloadingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OffloadingWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOffloadingWindowView.__init__(self, ctk_components) + + self.title("Offloading") + self.geometry("800x400") + self.resizable(True, True) + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid_columnconfigure(0, weight=1) + frame.grid_columnconfigure(1, weight=1) + self.build_content(frame, controller, ui_state) + frame.pack(fill="both", expand=1) + frame.grid(row=0, column=0, sticky='nsew') + self.components.button(self, 1, 0, "ok", self.destroy) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) diff --git a/modules/ui/PySide6OptimizerParamsWindowView.py b/modules/ui/PySide6OptimizerParamsWindowView.py new file mode 100644 index 000000000..ecc1f6a38 --- /dev/null +++ b/modules/ui/PySide6OptimizerParamsWindowView.py @@ -0,0 +1,91 @@ +import contextlib +from tkinter import TclError + +from modules.ui.BaseOptimizerParamsWindowView import BaseOptimizerParamsWindowView +from modules.ui.CtkMuonAdamWindowView import CtkMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkOptimizerParamsWindowView(BaseOptimizerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: OptimizerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseOptimizerParamsWindowView.__init__(self, ctk_components) + + self.controller = controller + self.ui_state = ui_state + self.optimizer_ui_state = ui_state.get_var("optimizer") + self.muon_adam_button = None + self.protocol("WM_DELETE_WINDOW", self.on_window_close) + + self.title("Optimizer Settings") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + + self.frame.grid_columnconfigure(0, weight=0) + self.frame.grid_columnconfigure(1, weight=1) + self.frame.grid_columnconfigure(2, minsize=50) + self.frame.grid_columnconfigure(3, weight=0) + self.frame.grid_columnconfigure(4, weight=1) + + self.components.button(self, 1, 0, "ok", command=self.on_window_close) + self.build_content(self.frame, controller, ui_state, self.optimizer_ui_state, + self.on_optimizer_change, self._load_defaults) + self._rebuild_dynamic_ui() + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def _rebuild_dynamic_ui(self): + with contextlib.suppress(TclError): + for widget in self.frame.winfo_children(): + grid_info = widget.grid_info() + if int(grid_info["row"]) >= 1: + widget.destroy() + + if not self.winfo_exists(): + return + + self.build_dynamic_content(self.frame, self.controller, self.optimizer_ui_state, + self.update_user_pref, self.open_muon_adam_window) + self.toggle_muon_adam_button() + + def update_user_pref(self, *args): + self.controller.on_close() + self.toggle_muon_adam_button() + + def on_optimizer_change(self, *args): + self.controller.restore_optimizer_config(self.ui_state) + self._rebuild_dynamic_ui() + + def _load_defaults(self, *args): + self.controller.load_defaults(self.ui_state) + + def on_window_close(self): + self.destroy() + + def toggle_muon_adam_button(self): + if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + + def open_muon_adam_window(self): + adam_config, current_optimizer = self.controller.prepare_muon_adam_config() + temp_adam_ui_state = CtkUIState(self, adam_config) + window = CtkMuonAdamWindowView(self, MuonAdamWindowController(self.controller.config, current_optimizer), temp_adam_ui_state) + self.wait_window(window) + self.controller.save_muon_adam_config(adam_config) diff --git a/modules/ui/PySide6ProfilingWindowView.py b/modules/ui/PySide6ProfilingWindowView.py new file mode 100644 index 000000000..15d5055a0 --- /dev/null +++ b/modules/ui/PySide6ProfilingWindowView.py @@ -0,0 +1,49 @@ + +from modules.ui.BaseProfilingWindowView import BaseProfilingWindowView +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkProfilingWindowView(BaseProfilingWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: ProfilingWindowController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseProfilingWindowView.__init__(self, ctk_components) + + self._controller = controller + + self.title("Profiling") + self.geometry("512x512") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=0) + self.grid_rowconfigure(2, weight=1) + self.grid_columnconfigure(0, weight=1) + + # Bottom bar + self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) + self._bottom_bar.grid(row=2, column=0, sticky="sew") + + self.build_content(self, self._bottom_bar, controller) + + self.protocol("WM_DELETE_WINDOW", self.withdraw) + self.withdraw() + self.after(200, lambda: set_window_icon(self)) + + def set_message(self, text): + self._message_label.configure(text=text) + + def set_profiling_active(self, active): + if active: + self._message_label.configure(text='Profiling active...') + self._profile_button.configure(text='End Profiling') + self._profile_button.configure(command=self._controller.end_profiler) + else: + self._message_label.configure(text='Inactive') + self._profile_button.configure(text='Start Profiling') + self._profile_button.configure(command=self._controller.start_profiler) diff --git a/modules/ui/PySide6SampleFrameView.py b/modules/ui/PySide6SampleFrameView.py new file mode 100644 index 000000000..167b25692 --- /dev/null +++ b/modules/ui/PySide6SampleFrameView.py @@ -0,0 +1,43 @@ +from modules.ui.BaseSampleFrameView import BaseSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.util.ui import ctk_components + +import customtkinter as ctk + + +class CtkSampleFrameView(BaseSampleFrameView, ctk.CTkFrame): + def __init__( + self, + parent, + controller: SampleFrameController, + ui_state, + include_prompt: bool = True, + include_settings: bool = True, + ): + ctk.CTkFrame.__init__(self, parent, fg_color="transparent") + BaseSampleFrameView.__init__(self, ctk_components) + + if include_prompt and include_prompt: + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_columnconfigure(0, weight=1) + + top_frame = None + if include_prompt: + top_frame = ctk.CTkFrame(self, fg_color="transparent") + top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + top_frame.grid_columnconfigure(0, weight=0) + top_frame.grid_columnconfigure(1, weight=1) + + bottom_frame = None + if include_settings: + bottom_frame = ctk.CTkFrame(self, fg_color="transparent") + bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + bottom_frame.grid_columnconfigure(0, weight=0) + bottom_frame.grid_columnconfigure(1, weight=1) + bottom_frame.grid_columnconfigure(2, weight=0) + bottom_frame.grid_columnconfigure(3, weight=1) + + self.build_content(top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings) diff --git a/modules/ui/PySide6SampleParamsWindowView.py b/modules/ui/PySide6SampleParamsWindowView.py new file mode 100644 index 000000000..8229a19f6 --- /dev/null +++ b/modules/ui/PySide6SampleParamsWindowView.py @@ -0,0 +1,32 @@ +from modules.ui.BaseSampleParamsWindowView import BaseSampleParamsWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkSampleParamsWindowView(BaseSampleParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SampleParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleParamsWindowView.__init__(self, ctk_components) + + self.title("Sample") + self.geometry("800x500") + self.resizable(True, True) + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) + frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + + self.components.button(self, 1, 0, "ok", self.destroy) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) diff --git a/modules/ui/PySide6SampleWindowView.py b/modules/ui/PySide6SampleWindowView.py new file mode 100644 index 000000000..3b67f03d5 --- /dev/null +++ b/modules/ui/PySide6SampleWindowView.py @@ -0,0 +1,102 @@ +import contextlib +import tkinter as tk +import traceback + +from modules.modelSampler.BaseModelSampler import ( + ModelSamplerOutput, +) +from modules.ui.BaseSampleWindowView import BaseSampleWindowView +from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleWindowController import SampleWindowController +from modules.util.enum.FileType import FileType +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from PIL import Image + + +class CtkSampleWindowView(BaseSampleWindowView, ctk.CTkToplevel): + def __init__( + self, + parent, + controller: SampleWindowController, + *args, **kwargs + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSampleWindowView.__init__(self, ctk_components) + + self.title("Sample") + self.geometry("1200x800") + self.resizable(True, True) + + model_type = controller.get_model_type() + self.ui_state = CtkUIState(self, controller.sample) + + if controller.use_external_model: + controller.callbacks.set_on_sample_custom(self.__update_preview) + controller.callbacks.set_on_update_sample_custom_progress(self.__update_progress) + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_rowconfigure(3, weight=0) + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1) + + prompt_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_settings=False) + prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") + + settings_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_prompt=False) + settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") + + # image + self.image = ctk.CTkImage( + light_image=self.__dummy_image(), + size=(512, 512) + ) + + image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) + image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + + self.progress = self.components.progress(self, 2, 0) + self.components.button(self, 3, 0, "sample", + lambda: controller.do_sample(self.__update_preview, self.__update_progress)) + + self.wait_visibility() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) + + def __update_preview(self, sampler_output: ModelSamplerOutput): + if sampler_output.file_type == FileType.IMAGE: + image = sampler_output.data + self.image.configure( + light_image=image, + size=(image.width, image.height), + ) + + def __update_progress(self, progress: int, max_progress: int): + self.progress.set(progress / max_progress) + self.update() + + def __dummy_image(self) -> Image: + return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) + + def destroy(self): + try: + if hasattr(self, "_icon_image_ref"): + del self._icon_image_ref + + # Remove any pending after callbacks + for after_id in self.tk.call('after', 'info'): + with contextlib.suppress(tk.TclError, RuntimeError): + self.after_cancel(after_id) + + super().destroy() + except (tk.TclError, RuntimeError) as e: + print(f"Error destroying window: {e}") + except Exception as e: + print(f"Unexpected error destroying window: {e}") + traceback.print_exc() diff --git a/modules/ui/PySide6SamplingTabView.py b/modules/ui/PySide6SamplingTabView.py new file mode 100644 index 000000000..dfe1a704a --- /dev/null +++ b/modules/ui/PySide6SamplingTabView.py @@ -0,0 +1,50 @@ +from modules.ui.BaseSamplingTabView import BaseSampleWidgetView, BaseSamplingTabView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.CtkSampleParamsWindowView import CtkSampleParamsWindowView +from modules.ui.SamplingTabController import SamplingTabController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState + +import customtkinter as ctk + + +class CtkSamplingTabView(CtkConfigListView, BaseSamplingTabView): + def __init__(self, master, controller: SamplingTabController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, + from_external_file=True, + attr_name="sample_definition_file_name", + config_dir="training_samples", + default_config_name="samples.json", + add_button_text="Add Sample", + add_button_tooltip="Add a new sample configuration.", + is_full_width=True, + show_toggle_button=True, + ) + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + return self.controller.open_element_window(self.master, self.current_config[i], ui_state, CtkSampleParamsWindowView) + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return CtkSampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) + + +class CtkSampleWidgetView(BaseSampleWidgetView, ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") + BaseSampleWidgetView.__init__(self, ctk_components) + + self.ui_state = CtkUIState(self, element) + + self.grid_columnconfigure(10, weight=1) + + self.build_content(self, element, self.ui_state, i, open_command, remove_command, clone_command, save_command) + + def _bind_save(self, save_command): + self.width_entry.bind('', lambda _: save_command()) + self.height_entry.bind('', lambda _: save_command()) + self.seed_entry.bind('', lambda _: save_command()) + self.prompt_entry.bind('', lambda _: save_command()) + + def place_in_list(self): + self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") diff --git a/modules/ui/PySide6SchedulerParamsWindowView.py b/modules/ui/PySide6SchedulerParamsWindowView.py new file mode 100644 index 000000000..9a7c41b96 --- /dev/null +++ b/modules/ui/PySide6SchedulerParamsWindowView.py @@ -0,0 +1,96 @@ +from modules.ui.BaseSchedulerParamsWindowView import BaseKvParamsView, BaseSchedulerParamsWindowView +from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.SchedulerParamsWindowController import KvParamsController, SchedulerParamsWindowController +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk + + +class CtkKvParamsView(CtkConfigListView, BaseKvParamsView): + def __init__(self, master, controller: KvParamsController, ui_state): + CtkConfigListView.__init__( + self, master, controller, ui_state, + attr_name="scheduler_params", + from_external_file=False, + add_button_text="add parameter", + is_full_width=True, + ) + BaseKvParamsView.__init__(self, ctk_components) + + def refresh_ui(self): + self._create_element_list() + + def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): + return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + +class KvWidget(ctk.CTkFrame): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__(master=master, bg_color="transparent") + self.element = element + self.ui_state = CtkUIState(self, element) + self.i = i + self.save_command = save_command + + self.grid_columnconfigure(0, weight=0) + self.grid_columnconfigure(1, weight=1, uniform=1) + self.grid_columnconfigure(2, weight=1, uniform=1) + + close_button = ctk.CTkButton( + master=self, + width=20, + height=20, + text="X", + corner_radius=2, + fg_color="#C00000", + command=lambda: remove_command(self.i)) + close_button.grid(row=0, column=0) + + # Key + tooltip_key = "Key name for an argument in your scheduler" + self.key = ctk_components.entry(self, 0, 1, self.ui_state, "key", + tooltip=tooltip_key, wide_tooltip=True) + self.key.bind("", lambda _: save_command()) + self.key.configure(width=50) + + # Value + tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." + self.value = ctk_components.entry(self, 0, 2, self.ui_state, "value", + tooltip=tooltip_val, wide_tooltip=True) + self.value.bind("", lambda _: save_command()) + self.value.configure(width=50) + + def place_in_list(self): + self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") + + +class CtkSchedulerParamsWindowView(BaseSchedulerParamsWindowView, ctk.CTkToplevel): + def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseSchedulerParamsWindowView.__init__(self, ctk_components) + + self.title("Learning Rate Scheduler Settings") + self.geometry("800x400") + self.resizable(True, True) + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + frame = ctk.CTkFrame(self) + frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + + expand_frame = ctk.CTkFrame(frame, bg_color="transparent") + expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + + self.components.button(self, 1, 0, "ok", command=self.destroy) + self.build_content(frame, controller, ui_state) + CtkKvParamsView(expand_frame, KvParamsController(controller.config), ui_state) + + self.wait_visibility() + self.grab_set() + self.focus_set() + self.after(200, lambda: set_window_icon(self)) diff --git a/modules/ui/PySide6TimestepDistributionWindowView.py b/modules/ui/PySide6TimestepDistributionWindowView.py new file mode 100644 index 000000000..69f2ae7d3 --- /dev/null +++ b/modules/ui/PySide6TimestepDistributionWindowView.py @@ -0,0 +1,84 @@ + +from modules.ui.BaseTimestepDistributionWindowView import BaseTimestepDistributionWindowView +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util.ui import ctk_components +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker, ThemeManager +from matplotlib import pyplot as plt +from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg + + +class CtkTimestepDistributionWindowView(BaseTimestepDistributionWindowView, ctk.CTkToplevel): + def __init__( + self, + parent, + controller: TimestepDistributionWindowController, + ui_state, + *args, **kwargs, + ): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseTimestepDistributionWindowView.__init__(self, ctk_components) + + self.title("Timestep Distribution") + self.geometry("900x600") + self.resizable(True, True) + + self.controller = controller + self.ax = None + self.canvas = None + + self.grid_rowconfigure(0, weight=1) + self.grid_columnconfigure(0, weight=1) + + frame = ctk.CTkScrollableFrame(self, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=0) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + frame.grid_rowconfigure(7, weight=1) + + self.build_content(frame, controller, ui_state) + + # matplotlib chart (CTK-only: needs winfo_rgb from the toplevel) + appearance_mode = AppearanceModeTracker.get_mode() + background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) + text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) + background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" + text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" + + fig, ax = plt.subplots() + self.ax = ax + self.canvas = FigureCanvasTkAgg(fig, master=frame) + self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) + + fig.set_facecolor(background_color) + ax.set_facecolor(background_color) + ax.spines['bottom'].set_color(text_color) + ax.spines['left'].set_color(text_color) + ax.spines['top'].set_color(text_color) + ax.spines['right'].set_color(text_color) + ax.tick_params(axis='x', colors=text_color, which="both") + ax.tick_params(axis='y', colors=text_color, which="both") + ax.xaxis.label.set_color(text_color) + ax.yaxis.label.set_color(text_color) + + self.__update_preview() + + # update button + ctk_components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + + frame.pack(fill="both", expand=1) + frame.grid(row=0, column=0, sticky='nsew') + ctk_components.button(self, 1, 0, "ok", self.destroy) + + self.wait_visibility() + self.after(200, lambda: set_window_icon(self)) + self.grab_set() + self.focus_set() + + def __update_preview(self): + self.ax.cla() + self.ax.hist(self.controller.generate_preview_data(), bins=1000, range=(0, 999)) + self.canvas.draw() diff --git a/modules/ui/PySide6TopBarView.py b/modules/ui/PySide6TopBarView.py new file mode 100644 index 000000000..ee1bcfa75 --- /dev/null +++ b/modules/ui/PySide6TopBarView.py @@ -0,0 +1,50 @@ +from collections.abc import Callable + +from modules.ui.BaseTopBarView import BaseTopBarView +from modules.ui.TopBarController import TopBarController +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ui import ctk_components, dialogs +from modules.util.ui.CtkUIState import CtkUIState + +import customtkinter as ctk + + +class CtkTopBarView(BaseTopBarView): + def __init__( + self, + master, + controller: TopBarController, + ui_state, + change_model_type_callback: Callable[[ModelType], None], + change_training_method_callback: Callable[[TrainingMethod], None], + load_preset_callback: Callable[[], None], + ): + BaseTopBarView.__init__(self, ctk_components) + + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=0, column=0, sticky="nsew") + + self.build(frame, master, controller, ui_state, change_model_type_callback, change_training_method_callback, load_preset_callback) + + def _make_config_ui_state(self, master, data): + return CtkUIState(master, data) + + def _get_dropdown_text(self, widget) -> str: + return widget.get() + + def _setup_frame_column_weight(self): + self.frame.grid_columnconfigure(5, weight=1) + + def _forget_dropdown(self): + self.configs_dropdown.grid_forget() + + def _show_save_dialog(self, default_value: str, callback): + dialogs.StringInputDialog( + parent=self.master, + title="name", + question="Config Name", + callback=callback, + default_value=default_value, + validate_callback=lambda x: not x.startswith("#"), + ) diff --git a/modules/ui/PySide6TrainUIView.py b/modules/ui/PySide6TrainUIView.py new file mode 100644 index 000000000..d3c1a70b5 --- /dev/null +++ b/modules/ui/PySide6TrainUIView.py @@ -0,0 +1,413 @@ +import ctypes +import platform +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from tkinter import filedialog, messagebox + +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseTrainUIView import BaseTrainUIView +from modules.ui.CloudTabController import CloudTabController +from modules.ui.ConceptTabController import ConceptTabController +from modules.ui.CtkAdditionalEmbeddingsTabView import CtkAdditionalEmbeddingsTabView +from modules.ui.CtkCaptionUIView import CtkCaptionUIView +from modules.ui.CtkCloudTabView import CtkCloudTabView +from modules.ui.CtkConceptTabView import CtkConceptTabView +from modules.ui.CtkConvertModelUIView import CtkConvertModelUIView +from modules.ui.CtkLoraTabView import CtkLoraTabView +from modules.ui.CtkModelTabView import CtkModelTabView +from modules.ui.CtkProfilingWindowView import CtkProfilingWindowView +from modules.ui.CtkSampleWindowView import CtkSampleWindowView +from modules.ui.CtkSamplingTabView import CtkSamplingTabView +from modules.ui.CtkTopBarView import CtkTopBarView +from modules.ui.CtkTrainingTabView import CtkTrainingTabView +from modules.ui.CtkVideoToolUIView import CtkVideoToolUIView +from modules.ui.LoraTabController import LoraTabController +from modules.ui.ModelTabController import ModelTabController +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.ui.SamplingTabController import SamplingTabController +from modules.ui.TopBarController import TopBarController +from modules.ui.TrainingTabController import TrainingTabController +from modules.ui.TrainUIController import TrainUIController +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ui_utils import set_window_icon + +import customtkinter as ctk +from customtkinter import AppearanceModeTracker + +# chunk for forcing Windows to ignore DPI scaling when moving between monitors +# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 +if platform.system() == "Windows": + with suppress(Exception): + # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically + ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE + + +class CtkTrainUIView(BaseTrainUIView, ctk.CTk): + set_step_progress: Callable[[int, int], None] + set_epoch_progress: Callable[[int, int], None] + + status_label: ctk.CTkLabel | None + training_button: ctk.CTkButton | None + + _TRAIN_BUTTON_STYLES = { + "idle": { + "text": "Start Training", + "state": "normal", + "fg_color": "#198754", + "hover_color": "#146c43", + "text_color": "white", + "text_color_disabled": "white", + }, + "running": { + "text": "Stop Training", + "state": "normal", + "fg_color": "#dc3545", + "hover_color": "#bb2d3b", + "text_color": "white", + }, + "stopping": { + "text": "Stopping...", + "state": "disabled", + "fg_color": "#dc3545", + "hover_color": "#dc3545", + "text_color": "white", + "text_color_disabled": "white", + }, + } + + def __init__(self): + ctk.CTk.__init__(self) + BaseTrainUIView.__init__(self, ctk_components) + + self.title("OneTrainer") + self.geometry("1100x740") + + self.after(100, lambda: self._set_icon()) + + # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration + ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") + ctk.set_default_color_theme("blue") + + self.train_config = TrainConfig.default_values() + self.ui_state = CtkUIState(self, self.train_config) + + self.controller = TrainUIController(self.train_config) + self.controller.view = self + + self.grid_rowconfigure(0, weight=0) + self.grid_rowconfigure(1, weight=1) + self.grid_rowconfigure(2, weight=0) + self.grid_columnconfigure(0, weight=1) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = None + self.tabview = None + + self.model_tab = None + self.training_tab = None + self.lora_tab = None + self.cloud_tab = None + self.additional_embeddings_tab = None + + self.top_bar_component = self.top_bar(self) + self.content_frame(self) + self.bottom_bar(self) + + self.controller._check_start_always_on_tensorboard() + + self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self.controller._on_workspace_dir_change_trace) + + # Persistent profiling window. + self._profiling_controller = ProfilingWindowController() + self.profiling_window = self._profiling_controller.create_window(self, CtkProfilingWindowView) + + self.protocol("WM_DELETE_WINDOW", self.__close) + + def __close(self): + self.top_bar_component.save_default() + self.controller._stop_always_on_tensorboard() + if hasattr(self, 'workspace_dir_trace_id'): + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + self.quit() + + # --- BaseTrainUIView abstract method implementations --- + + def on_update_status(self, status: str): + self.status_label.configure(text=status) + + def on_training_started(self): + self._set_training_button_style("running") + + def on_training_stopped(self, error_caught: bool): + self.eta_label.configure(text="") + self._set_training_button_style("idle") + + def on_training_stopping(self): + self._set_training_button_style("stopping") + + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): + self.set_step_progress(epoch_step, max_step) + self.set_epoch_progress(epoch, max_epoch) + if eta_str is not None: + self.eta_label.configure(text=f"ETA: {eta_str}") + else: + self.eta_label.configure(text="") + + def schedule_on_main_thread(self, fn: Callable): + self.after(0, fn) + + def get_cloud_reattach(self) -> bool: + return self.cloud_tab.reattach + + def save_default(self): + self.top_bar_component.save_default() + self.concepts_tab.save_current_config() + self.sampling_tab.save_current_config() + self.additional_embeddings_tab.save_current_config() + + def show_validation_errors(self, errors: list[str]): + bullet_list = "\n".join(f"• {e}" for e in errors) + messagebox.showerror( + "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}", + ) + + def open_dataset_tool(self): + self.wait_window(self.controller.open_dataset_tool(self, CtkCaptionUIView)) + + def open_video_tool(self): + self.wait_window(self.controller.open_video_tool(self, CtkVideoToolUIView)) + + def open_convert_model_tool(self): + self.wait_window(self.controller.open_convert_model_tool(self, CtkConvertModelUIView)) + + def open_sampling_tool(self): + self.controller.open_sampling_tool(self, CtkSampleWindowView) + + def open_manual_sample_window(self): + self.controller.open_manual_sample_window(self, CtkSampleWindowView) + + def wait_window(self, window): + ctk.CTk.wait_window(self, window) + + def show_window(self, window): + window.focus_set() + + def connect_window_closed(self, window, callback): + window.bind("", lambda _: callback()) + + # --- CTK layout and frame builders --- + + def _set_icon(self): + """Set the window icon safely after window is ready""" + set_window_icon(self) + + def top_bar(self, master): + return CtkTopBarView( + master, + TopBarController(self.train_config), + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def bottom_bar(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=2, column=0, sticky="nsew") + + # status + ETA container + status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") + status_frame.grid(row=0, column=1, sticky="w") + status_frame.grid_rowconfigure(0, weight=0) + status_frame.grid_rowconfigure(1, weight=0) + status_frame.grid_columnconfigure(0, weight=1) + + # padding + frame.grid_columnconfigure(2, weight=1) + + self.build_bottom_bar_content(frame, status_frame, self.controller, self.ui_state) + self._set_training_button_style("idle") # centralized styling + + return frame + + def content_frame(self, master): + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_columnconfigure(0, weight=1) + + self.tabview = ctk.CTkTabview(frame) + self.tabview.grid(row=0, column=0, sticky="nsew") + + self.general_tab = self.create_general_tab(self.tabview.add("general")) + self.model_tab = self.create_model_tab(self.tabview.add("model")) + self.data_tab = self.create_data_tab(self.tabview.add("data")) + self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) + self.training_tab = self.create_training_tab(self.tabview.add("training")) + self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) + self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) + self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) + self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) + self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) + + self.change_training_method(self.train_config.training_method) + + return frame + + def create_general_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, weight=0) + frame.grid_columnconfigure(3, weight=1) + self.build_general_tab_content(frame, self.controller, self.ui_state) + frame.pack(fill="both", expand=1) + return frame + + def create_model_tab(self, master): + return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) + + def create_data_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + self.build_data_tab_content(frame, self.controller, self.ui_state) + frame.pack(fill="both", expand=1) + return frame + + def create_concepts_tab(self, master): + return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) + + def create_training_tab(self, master) -> CtkTrainingTabView: + return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) + + def create_cloud_tab(self, master) -> CtkCloudTabView: + return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) + + def create_sampling_tab(self, master): + master.grid_rowconfigure(0, weight=0) + master.grid_rowconfigure(1, weight=1) + master.grid_columnconfigure(0, weight=1) + + top_frame = ctk.CTkFrame(master=master, corner_radius=0) + top_frame.grid(row=0, column=0, sticky="nsew") + sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") + sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + + self.build_sampling_tab_header(top_frame, sub_frame, self.controller, self.ui_state) + + frame = ctk.CTkFrame(master=master, corner_radius=0) + frame.grid(row=1, column=0, sticky="nsew") + + return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) + + def create_backup_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + self.build_backup_tab_content(frame, self.controller, self.ui_state) + frame.pack(fill="both", expand=1) + return frame + + def embedding_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + self.build_embedding_tab_content(frame, self.controller, self.ui_state) + frame.pack(fill="both", expand=1) + return frame + + def create_additional_embeddings_tab(self, master): + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) + + def create_tools_tab(self, master): + frame = ctk.CTkScrollableFrame(master, fg_color="transparent") + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + frame.grid_columnconfigure(2, minsize=50) + frame.grid_columnconfigure(3, weight=0) + frame.grid_columnconfigure(4, weight=1) + self.build_tools_tab_content(frame, self.controller, self.ui_state) + frame.pack(fill="both", expand=1) + return frame + + def open_profiling_tool(self): + self.profiling_window.deiconify() + + def change_model_type(self, model_type: ModelType): + if self.model_tab: + self.model_tab.refresh_ui() + + if self.training_tab: + self.training_tab.refresh_ui() + + if self.lora_tab: + self.lora_tab.refresh_ui() + + def change_training_method(self, training_method: TrainingMethod): + if not self.tabview: + return + + if self.model_tab: + self.model_tab.refresh_ui() + + if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: + self.tabview.delete("LoRA") + self.lora_tab = None + if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: + self.tabview.delete("embedding") + + if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) + if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: + self.embedding_tab(self.tabview.add("embedding")) + + def load_preset(self): + if not self.tabview: + return + + if self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def _set_training_button_style(self, mode: str): + if not self.training_button: + return + style = self._TRAIN_BUTTON_STYLES.get(mode) + if not style: + return + self.training_button.configure(**style) + + def export_training(self): + file_path = filedialog.asksaveasfilename(filetypes=[ + ("All Files", "*.*"), + ("json", "*.json"), + ], initialdir=".", initialfile="config.json") + if file_path: + self.controller.export_training(file_path) + + def generate_debug_package(self): + zip_path = filedialog.askdirectory( + initialdir=".", + title="Select Directory to Save Debug Package" + ) + if not zip_path: + return + self.controller.generate_debug_package(Path(zip_path) / "OneTrainer_debug_report.zip") diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py new file mode 100644 index 000000000..bc29488dd --- /dev/null +++ b/modules/ui/PySide6TrainingTabView.py @@ -0,0 +1,77 @@ + +from modules.ui.BaseTrainingTabView import BaseTrainingTabView +from modules.ui.CtkOffloadingWindowView import CtkOffloadingWindowView +from modules.ui.CtkOptimizerParamsWindowView import CtkOptimizerParamsWindowView +from modules.ui.CtkSchedulerParamsWindowView import CtkSchedulerParamsWindowView +from modules.ui.CtkTimestepDistributionWindowView import CtkTimestepDistributionWindowView +from modules.ui.TrainingTabController import TrainingTabController +from modules.util.ui import ctk_components + +import customtkinter as ctk + + +class CtkTrainingTabView(BaseTrainingTabView): + def __init__(self, master, controller: TrainingTabController, ui_state): + BaseTrainingTabView.__init__(self, ctk_components) + + self.master = master + self.controller = controller + self.ui_state = ui_state + self.scroll_frame = None + + master.grid_rowconfigure(0, weight=1) + master.grid_columnconfigure(0, weight=1) + + self.refresh_ui() + + def refresh_ui(self): + if self.scroll_frame: + self.scroll_frame.destroy() + + self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") + self.scroll_frame.grid(row=0, column=0, sticky="nsew") + + self.scroll_frame.grid_columnconfigure(0, weight=1) + self.scroll_frame.grid_columnconfigure(1, weight=1) + self.scroll_frame.grid_columnconfigure(2, weight=1) + + column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_0.grid(row=0, column=0, sticky="nsew") + column_0.grid_columnconfigure(0, weight=1) + + column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_1.grid(row=0, column=1, sticky="nsew") + column_1.grid_columnconfigure(0, weight=1) + + column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") + column_2.grid(row=0, column=2, sticky="nsew") + column_2.grid_columnconfigure(0, weight=1) + + callbacks = { + 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), + 'open_optimizer_params': self._open_optimizer_params_window, + 'restore_scheduler': self._restore_scheduler_config, + 'open_scheduler_params': self._open_scheduler_params_window, + 'open_offloading': self._open_offloading_window, + 'open_timestep_distribution': self._open_timestep_distribution_window, + } + + self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) + + def _restore_scheduler_config(self, variable): + if not hasattr(self, 'lr_scheduler_adv_comp'): + return + state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" + self.lr_scheduler_adv_comp.configure(state=state) + + def _open_optimizer_params_window(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) + + def _open_scheduler_params_window(self): + self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) + + def _open_timestep_distribution_window(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) + + def _open_offloading_window(self): + self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) diff --git a/modules/ui/PySide6VideoToolUIView.py b/modules/ui/PySide6VideoToolUIView.py new file mode 100644 index 000000000..c272c891c --- /dev/null +++ b/modules/ui/PySide6VideoToolUIView.py @@ -0,0 +1,128 @@ +from tkinter import filedialog + +from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView +from modules.ui.VideoToolUIController import VideoToolUIController +from modules.util.image_util import load_image +from modules.util.ui import ctk_components +from modules.util.ui.CtkUIState import CtkUIState + +import customtkinter as ctk + +PAD = ctk_components.PAD + + +class CtkVideoToolUIView(BaseVideoToolUIView, ctk.CTkToplevel): + def __init__(self, parent, controller: VideoToolUIController, *args, **kwargs): + ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) + BaseVideoToolUIView.__init__(self, ctk_components) + + self.controller = controller + ui_state = CtkUIState(self, controller.args) + + self.title("Video Tools") + self.geometry("600x720") + self.resizable(True, True) + self.wait_visibility() + self.focus_set() + + self.grid_rowconfigure(0, weight=1) + self.grid_rowconfigure(1, weight=0) + self.grid_columnconfigure(0, weight=1) + + tabview = ctk.CTkTabview(self) + tabview.grid(row=0, column=0, sticky="nsew") + + clip_frame = ctk.CTkScrollableFrame(tabview.add("extract clips"), fg_color="transparent") + clip_frame.grid_columnconfigure(0, weight=0, minsize=120) + clip_frame.grid_columnconfigure(1, weight=0, minsize=200) + clip_frame.grid_columnconfigure(2, weight=0) + clip_frame.grid_columnconfigure(3, weight=1) + self.build_clip_extract_tab(clip_frame, controller, ui_state) + clip_frame.pack(fill="both", expand=1) + + image_frame = ctk.CTkScrollableFrame(tabview.add("extract images"), fg_color="transparent") + image_frame.grid_columnconfigure(0, weight=0, minsize=120) + image_frame.grid_columnconfigure(1, weight=0, minsize=200) + image_frame.grid_columnconfigure(2, weight=0) + image_frame.grid_columnconfigure(3, weight=1) + self.build_image_extract_tab(image_frame, controller, ui_state) + image_frame.pack(fill="both", expand=1) + + download_frame = ctk.CTkScrollableFrame(tabview.add("download"), fg_color="transparent") + download_frame.grid_columnconfigure(0, weight=0, minsize=120) + download_frame.grid_columnconfigure(1, weight=0, minsize=200) + download_frame.grid_columnconfigure(2, weight=0) + download_frame.grid_columnconfigure(3, weight=1) + self.build_video_download_tab(download_frame, controller, ui_state) + download_frame.pack(fill="both", expand=1) + + self._build_status_bar(self) + + def _build_status_bar(self, master): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=1, column=0) + frame.grid_columnconfigure(0, weight=0, minsize=160) + frame.grid_columnconfigure(1, weight=0, minsize=300) + frame.grid_columnconfigure(2, weight=1) + + preview_path = "resources/icons/icon.png" + preview = load_image(preview_path, 'RGB') + preview.thumbnail((150, 150)) + self.preview_image = ctk.CTkImage(light_image=preview, size=preview.size) + self.preview_image_label = ctk.CTkLabel( + master=frame, text="Preview image", image=self.preview_image, height=150, width=150, + compound="top") + self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) + + self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) + self.status_label.insert(index="1.0", text="Current status") + self.status_label.configure(state="disabled") + self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + var = ui_state.get_var(var_name) + textbox = ctk.CTkTextbox(master, width=width, height=height, border_width=2) + textbox.insert("1.0", var.get()) + textbox.grid(row=row, column=col, rowspan=2, sticky="w", padx=PAD, pady=PAD) + + def on_text_change(event=None): + var.set(textbox.get("1.0", "end-1c")) + + textbox.bind("", on_text_change) + return textbox + + def _create_browse_dir_button(self, master, row, ui_state, var_name): + def browse(): + path = filedialog.askdirectory() + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): + def browse(): + path = filedialog.askopenfilename(filetypes=filetypes) + if path: + ui_state.get_var(var_name).set(path) + self.focus_set() + + button = ctk.CTkButton(master, width=30, text="...", command=browse) + button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + return button + + def update_status(self, status_text: str): + self.status_label.configure(state="normal") + self.status_label.insert(index="end", text=status_text + "\n") + self.status_label.configure(state="disabled") + + def clear_status(self): + self.status_label.configure(state="normal") + self.status_label.delete(index1="1.0", index2="end") + self.status_label.configure(state="disabled") + + def update_preview(self, preview_image, label_text: str): + self.preview_image.configure(light_image=preview_image, size=preview_image.size) + self.preview_image_label.configure(text=label_text) diff --git a/modules/util/ui/pyside6_components.py b/modules/util/ui/pyside6_components.py new file mode 100644 index 000000000..e462f72a1 --- /dev/null +++ b/modules/util/ui/pyside6_components.py @@ -0,0 +1,588 @@ +import contextlib +import tkinter as tk +from collections.abc import Callable +from pathlib import Path +from tkinter import filedialog +from typing import Any, Literal + +from modules.util.enum.PathIOType import PathIOType +from modules.util.enum.TimeUnit import TimeUnit +from modules.util.path_util import supported_image_extensions +from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator +from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui.ToolTip import ToolTip + +import customtkinter as ctk +from customtkinter.windows.widgets.scaling import CTkScalingBaseClass +from PIL import Image + +PAD = 10 + + +def app_title(master, row, column): + frame = ctk.CTkFrame(master) + frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") + + image_component = ctk.CTkImage( + Image.open("resources/icons/icon.png").resize((40, 40), Image.Resampling.LANCZOS), + size=(40, 40) + ) + image_label_component = ctk.CTkLabel(frame, image=image_component, text="") + image_label_component.grid(row=0, column=0, padx=PAD, pady=PAD) + + label_component = ctk.CTkLabel(frame, text="OneTrainer", font=ctk.CTkFont(size=20, weight="bold")) + label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) + + +def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0, underline=False): + component = ctk.CTkLabel(master, text=text, wraplength=wraplength) + component.grid(row=row, column=column, padx=pad, pady=pad, sticky="nw") + if tooltip: + ToolTip(component, tooltip, wide=wide_tooltip) + if underline: + component.configure(font=ctk.CTkFont(underline=True)) + return component + + +def entry( + master, + row, + column, + ui_state: CtkUIState, + var_name: str, + command: Callable[[], None] | None = None, + tooltip: str = "", + wide_tooltip: bool = False, + width: int = 140, + sticky: str = "new", + max_undo: int | None = None, + validator_factory: Callable[..., FieldValidator] | None = None, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, +): + var = ui_state.get_var(var_name) + trace_id = None + if command: + trace_id = ui_state.add_var_trace(var_name, command) + + component = ctk.CTkEntry(master, textvariable=var, width=width) + component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) + + if validator_factory is not None: + validator = validator_factory( + component, var, ui_state, var_name, + max_undo=max_undo or DEFAULT_MAX_UNDO, + extra_validate=extra_validate, + required=required, + ) + else: + validator = FieldValidator( + component, var, ui_state, var_name, + max_undo=max_undo or DEFAULT_MAX_UNDO, + extra_validate=extra_validate, + required=required, + ) + validator.attach() + component._validator = validator # type: ignore[attr-defined] + + original_destroy = component.destroy + + def new_destroy(): + validator.detach() + + # 'temporary' fix until https://github.com/TomSchimansky/CustomTkinter/pull/2077 is merged + # unfortunately Tom has admitted to forgetting about how to maintain CTK so this likely will never be merged + if component._textvariable_callback_name: + with contextlib.suppress(tk.TclError): + component._textvariable.trace_remove("write", component._textvariable_callback_name) # type: ignore[union-attr] + component._textvariable_callback_name = "" + + if command is not None and trace_id is not None: + ui_state.remove_var_trace(var_name, trace_id) + + original_destroy() + + component.destroy = new_destroy # type: ignore[assignment] + + if tooltip: + ToolTip(component, tooltip, wide=wide_tooltip) + + return component + + +def path_entry( + master, row, column, ui_state: CtkUIState, var_name: str, + *, + mode: Literal["file", "dir"] = "file", + io_type: PathIOType = PathIOType.INPUT, + path_modifier: Callable[[str], str | Path] | None = None, + allow_model_files: bool = True, + allow_image_files: bool = False, + command: Callable[[str], None] | None = None, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + columnspan: int = 1, +): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="new", columnspan=columnspan) + + frame.grid_columnconfigure(0, weight=1) + + def _path_validator_factory(comp, var, state, name, **kw): + return PathValidator(comp, var, state, name, io_type=io_type, **kw) + + entry_component = entry( + frame, row=0, column=0, ui_state=ui_state, var_name=var_name, + validator_factory=_path_validator_factory, + extra_validate=extra_validate, + required=required, + ) + + trace_ids = [] + if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): + validator = getattr(entry_component, '_validator', None) + if validator is not None: + for dep_var_name in ("prevent_overwrites", "output_model_format"): + with contextlib.suppress(KeyError, AttributeError): + dep_var = ui_state.get_var(dep_var_name) + tid = dep_var.trace_add("write", lambda *_a: validator.revalidate()) + trace_ids.append((dep_var, tid)) + + use_save_dialog = io_type in (PathIOType.OUTPUT, PathIOType.MODEL) + + def __open_dialog(): + # Determine currently selected filename and/or directory + current_dir, current_filename = None, None + current_path_str = ui_state.get_var(var_name).get() or None + + if current_path_str is not None: + current_path = Path(current_path_str) + if mode == "file": + current_dir = str(current_path.parent) + current_filename = str(current_path.name) + elif mode == "dir": + current_dir = str(current_path.parent) + current_filename = None + + if mode == "dir": + chosen = filedialog.askdirectory(initialdir=current_dir) + else: + filetypes = [ + ("All Files", "*.*"), + ] + + if allow_model_files: + filetypes.extend([ + ("Diffusers", "model_index.json"), + ("Checkpoint", "*.ckpt *.pt *.bin"), + ("Safetensors", "*.safetensors"), + ]) + if allow_image_files: + filetypes.extend([ + ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), + ]) + + if use_save_dialog: + chosen = filedialog.asksaveasfilename(filetypes=filetypes, initialdir=current_dir, + initialfile=current_filename) + else: + chosen = filedialog.askopenfilename(filetypes=filetypes, initialdir=current_dir, + initialfile=current_filename) + + if chosen: + if path_modifier: + chosen = path_modifier(chosen) + + chosen_str = str(chosen) + ui_state.get_var(var_name).set(chosen_str) + + if command: + command(chosen_str) + + button_component = ctk.CTkButton(frame, text="...", width=40, command=__open_dialog) + button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") + + if trace_ids: + original_frame_destroy = frame.destroy + def _frame_destroy(): + for dep_var, tid in trace_ids: + with contextlib.suppress(tk.TclError, ValueError): + dep_var.trace_remove("write", tid) + original_frame_destroy() + frame.destroy = _frame_destroy # type: ignore[assignment] + + return frame + + +def time_entry(master, row, column, ui_state: CtkUIState, var_name: str, unit_var_name, supports_time_units: bool = True): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") + + frame.grid_columnconfigure(0, weight=0) + frame.grid_columnconfigure(1, weight=1) + + entry(frame, row=0, column=0, ui_state=ui_state, var_name=var_name, width=50) + + values = [str(x) for x in list(TimeUnit)] + if not supports_time_units: + values = [str(x) for x in list(TimeUnit) if not x.is_time_unit()] + + unit_component = ctk.CTkOptionMenu( + frame, + values=values, + variable=ui_state.get_var(unit_var_name), + width=100, + ) + unit_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="new") + + return frame + +def layer_filter_entry(master, row, column, ui_state: CtkUIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): + frame = ctk.CTkFrame(master=master, corner_radius=5, fg_color=frame_color) + frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + + layer_entry = entry( + frame, 1, 0, ui_state, entry_var_name, + tooltip=entry_tooltip + ) + layer_entry_fg_color = layer_entry.cget("fg_color") + layer_entry_text_color = layer_entry.cget("text_color") + + regex_label = label( + frame, 2, 0, "Use Regex", + tooltip=regex_tooltip, + ) + regex_switch = switch( + frame, 2, 1, ui_state, regex_var_name + ) + + # Let the user set their own layer filter + # TODO + #if self.train_config.layer_filter and self.train_config.layer_filter_preset == "custom": + # self.prior_custom = self.train_config.layer_filter + #else: + # self.prior_custom = "" + + layer_entry.grid_configure(columnspan=2, sticky="ew") + + presets_list = list(presets.keys()) + ["custom"] + + + def hide_layer_entry(): + if layer_entry and layer_entry.winfo_manager(): + layer_entry.grid_remove() + + def show_layer_entry(): + if layer_entry and not layer_entry.winfo_manager(): + layer_entry.grid() + + + def preset_set_layer_choice(selected: str): + if not selected or selected not in presets_list: + selected = presets_list[0] + + if selected == "custom": + # Allow editing + regex toggle + show_layer_entry() + layer_entry.configure(state="normal", fg_color=layer_entry_fg_color, text_color=layer_entry_text_color) + #layer_entry.cget('textvariable').set("") + regex_label.grid() + regex_switch.grid() + else: + # Preserve custom text before overwriting + #if self.prior_selected == "custom": + # self.prior_custom = self.layer_entry.get() + + # Resolve preset definition (list[str] OR {'patterns': [...], 'regex': bool}) + preset_def = presets.get(selected, []) + if isinstance(preset_def, dict): + patterns = preset_def.get("patterns", []) + preset_uses_regex = bool(preset_def.get("regex", False)) + else: + patterns = preset_def + preset_uses_regex = False + + disabled_color = ("gray85", "gray17") + disabled_text_color = ("gray30", "gray70") + layer_entry.configure(state="disabled", fg_color=disabled_color, text_color=disabled_text_color) + layer_entry.cget('textvariable').set(",".join(patterns)) + + ui_state.get_var(entry_var_name).set(",".join(patterns)) + ui_state.get_var(regex_var_name).set(preset_uses_regex) + + regex_label.grid_remove() + regex_switch.grid_remove() + + if selected == "full" and not patterns: + hide_layer_entry() + else: + show_layer_entry() + +# self.prior_selected = selected + + label(frame, 0, 0, preset_label, + tooltip=preset_tooltip) + + + ui_state.remove_all_var_traces(preset_var_name) + + layer_selector = options( + frame, 0, 1, presets_list, ui_state, preset_var_name, + command=preset_set_layer_choice + ) + + def on_layer_filter_preset_change(): + if not layer_selector: + return + selected = ui_state.get_var(preset_var_name).get() + preset_set_layer_choice(selected) + + ui_state.add_var_trace( + preset_var_name, + on_layer_filter_preset_change, + ) + + preset_set_layer_choice(layer_selector.get()) + +def icon_button(master, row, column, text, command): + component = ctk.CTkButton(master, text=text, width=40, command=command) + component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky="new") + return component + + +def colored_icon_button(master, row, column, text, fg_color, command, padx=0): + component = ctk.CTkButton( + master=master, width=20, height=20, text=text, + corner_radius=2, fg_color=fg_color, command=command, + ) + component.grid(row=row, column=column, padx=padx) + return component + + +def button(master, row, column, text, command, tooltip=None, **kwargs): + # Pop grid-specific parameters from kwargs, using PAD as the default if not provided. + padx = kwargs.pop('padx', PAD) + pady = kwargs.pop('pady', PAD) + + component = ctk.CTkButton(master, text=text, command=command, **kwargs) + component.grid(row=row, column=column, padx=padx, pady=pady, sticky="new") + if tooltip: + ToolTip(component, tooltip, x_position=25) + return component + + +def options(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None): + component = ctk.CTkOptionMenu(master, values=values, variable=ui_state.get_var(var_name), command=command) + component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") + + # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged + def create_destroy(component): + orig_destroy = component.destroy + + def destroy(self): + orig_destroy() + CTkScalingBaseClass.destroy(self) + + return destroy + + destroy = create_destroy(component._dropdown_menu) + component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] + + return component + + +def options_adv(master, row, column, values, ui_state: CtkUIState, var_name: str, + command: Callable[[str], None] | None = None, adv_command: Callable[[], None] | None = None): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") + + frame.grid_columnconfigure(0, weight=1) + + component = ctk.CTkOptionMenu(frame, values=values, variable=ui_state.get_var(var_name), command=command) + component.grid(row=0, column=0, padx=PAD, pady=(PAD, PAD), sticky="new") + + button_component = ctk.CTkButton(frame, text="…", width=20, command=adv_command) + button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") + + if command: + command(ui_state.get_var(var_name).get()) # call command once to set the initial value + + # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged + def create_destroy(component): + orig_destroy = component.destroy + + def destroy(self): + orig_destroy() + CTkScalingBaseClass.destroy(self) + + return destroy + + destroy = create_destroy(component._dropdown_menu) + component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] + + return frame, {'component': component, 'button_component': button_component} + + +def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: CtkUIState, var_name: str, + command: Callable[[Any], None] | None = None): + var = ui_state.get_var(var_name) + keys = [key for key, value in values] + + # if the current value is not valid, select the first option + if var.get() not in [str(value) for key, value in values] and len(keys) > 0: + var.set(values[0][1]) + + deactivate_update_var = False + + def update_component(text): + for key, value in values: + if text == key: + nonlocal deactivate_update_var + deactivate_update_var = True + var.set(value) + if command: + command(value) + deactivate_update_var = False + break + + component = ctk.CTkOptionMenu(master, values=keys, command=update_component) + component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") + + def update_var(): + if not deactivate_update_var: + for key, value in values: + if var.get() == str(value): + if component.winfo_exists(): # the component could already be destroyed + component.set(key) + if command: + command(value) + break + + var.trace_add("write", lambda _0, _1, _2: update_var()) + update_var() # call update_var once to set the initial value + + # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged + def create_destroy(component): + orig_destroy = component.destroy + + def destroy(self): + orig_destroy() + CTkScalingBaseClass.destroy(self) + + return destroy + + destroy = create_destroy(component._dropdown_menu) + component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] + + return component + + +def switch( + master, + row, + column, + ui_state: CtkUIState, + var_name: str, + command: Callable[[], None] | None = None, + text: str = "", + width: int | None = None, +): + var = ui_state.get_var(var_name) + if command: + trace_id = ui_state.add_var_trace(var_name, command) + + component = ctk.CTkSwitch(master, variable=var, text=text, command=command) + if width is not None: + component.configure(width=width) + component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") + + def create_destroy(component): + orig_destroy = component.destroy + + def destroy(self): + if command is not None: + ui_state.remove_var_trace(var_name, trace_id) + + orig_destroy() + + return destroy + + destroy = create_destroy(component) + component.destroy = lambda: destroy(component) # type: ignore[assignment] + + return component + + +def progress(master, row, column): + component = ctk.CTkProgressBar(master) + component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="ew") + return component + + +def double_progress(master, row, column, label_1, label_2): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=column, padx=0, pady=0, sticky="nsew") + + frame.grid_rowconfigure(0, weight=1) + frame.grid_rowconfigure(1, weight=1) + frame.grid_columnconfigure(0, weight=1) + + label_1_component = ctk.CTkLabel(frame, text=label_1) + label_1_component.grid(row=0, column=0, padx=(PAD, PAD), pady=(0, 0), sticky="new") + + label_2_component = ctk.CTkLabel(frame, text=label_2) + label_2_component.grid(row=1, column=0, padx=(PAD, PAD), pady=(0, 0), sticky="sew") + + progress_1_component = ctk.CTkProgressBar(frame) + progress_1_component.grid(row=0, column=1, padx=(PAD, PAD), pady=(PAD, 0), sticky="new") + + progress_2_component = ctk.CTkProgressBar(frame) + progress_2_component.grid(row=1, column=1, padx=(PAD, PAD), pady=(0, PAD), sticky="sew") + + description_1_component = ctk.CTkLabel(frame, text="") + description_1_component.grid(row=0, column=2, padx=(PAD, PAD), pady=(0, 0), sticky="new") + + description_2_component = ctk.CTkLabel(frame, text="") + description_2_component.grid(row=1, column=2, padx=(PAD, PAD), pady=(0, 0), sticky="sew") + + def set_1(value, max_value): + progress_1_component.set(value / max_value) + description_1_component.configure(text=f"{value}/{max_value}") + + def set_2(value, max_value): + progress_2_component.set(value / max_value) + description_2_component.configure(text=f"{value}/{max_value}") + + return set_1, set_2 + + +def section_frame(master, row: int, col: int = 0): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=col, padx=PAD // 2, pady=PAD // 2, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + return frame + + +def inline_frame(master, row: int, col: int, columnspan: int = 1): + frame = ctk.CTkFrame(master, fg_color="transparent") + frame.grid(row=row, column=col, columnspan=columnspan, sticky="ew", padx=0, pady=0) + return frame + + +def set_widget_enabled(widget, enabled: bool) -> None: + state = "normal" if enabled else "disabled" + if isinstance(widget, ctk.CTkFrame): + for child in widget.children.values(): + with contextlib.suppress(Exception): + child.configure(state=state) + else: + widget.configure(state=state) + + +def set_label_text(label, text: str) -> None: + label.configure(text=str(text)) + + +def call_after(widget, delay_ms: int, func) -> None: + widget.after(delay_ms, func) diff --git a/modules/util/ui/pyside6_validation.py b/modules/util/ui/pyside6_validation.py new file mode 100644 index 000000000..5347ba6d4 --- /dev/null +++ b/modules/util/ui/pyside6_validation.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import contextlib +import tkinter as tk +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from modules.util.enum.PathIOType import PathIOType +from modules.util.ui.validation import ( + DEBOUNCE_TYPING_MS, + DEFAULT_MAX_UNDO, + ERROR_BORDER_COLOR, + UNDO_DEBOUNCE_MS, + BaseFieldValidator, + UndoHistory, + _active_validators, + _validate_path_field, +) + +if TYPE_CHECKING: + from modules.util.ui.UIState import UIState + + import customtkinter as ctk + + +class DebounceTimer: + def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): + self.widget = widget + self.delay_ms = delay_ms + self.callback = callback + self._after_id: str | None = None + + def call(self, *args, **kwargs): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + + def fire(): + self._after_id = None + self.callback(*args, **kwargs) + + with contextlib.suppress(tk.TclError): + self._after_id = self.widget.after(self.delay_ms, fire) + + def cancel(self): + if self._after_id: + with contextlib.suppress(tk.TclError): + self.widget.after_cancel(self._after_id) + self._after_id = None + + +class FieldValidator(BaseFieldValidator): + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + super().__init__(ui_state, var_name, extra_validate, required) + self.component = component + self.var = var + + try: + self._original_border_color = component.cget("border_color") + except Exception: + self._original_border_color = "gray50" + + self._shadow_var = tk.StringVar(master=component) + self._shadow_trace_name: str | None = None + self._real_var_trace_name: str | None = None + self._syncing = False + self._touched = False + + self._debounce: DebounceTimer | None = None + self._undo_debounce: DebounceTimer | None = None + self._undo = UndoHistory(max_undo) + + def attach(self) -> None: + self._shadow_var.set(self.var.get()) + self._swap_textvariable(self._shadow_var) + + self._debounce = DebounceTimer( + self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire + ) + self._undo_debounce = DebounceTimer( + self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot + ) + + self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) + self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) + + self.component.bind("", self._on_focus_in) + self.component.bind("", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("<>", self._on_user_input) + self.component.bind("", self._on_focus_out) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_undo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_redo) + self.component.bind("", self._on_enter) + + self._bound = True + _active_validators.add(self) + + def detach(self) -> None: + if not self._bound: + return + self._bound = False + _active_validators.discard(self) + + self._commit() + + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + + if self._shadow_trace_name: + with contextlib.suppress(Exception): + self._shadow_var.trace_remove("write", self._shadow_trace_name) + self._shadow_trace_name = None + + if self._real_var_trace_name: + with contextlib.suppress(Exception): + self.var.trace_remove("write", self._real_var_trace_name) + self._real_var_trace_name = None + + self._swap_textvariable(self.var) + + def _swap_textvariable(self, new_var: tk.Variable) -> None: + comp = self.component + if comp._textvariable_callback_name: + with contextlib.suppress(Exception): + comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] + comp._textvariable_callback_name = "" + + comp.configure(textvariable=new_var) + + if new_var is not None: + comp._textvariable_callback_name = new_var.trace_add( + "write", comp._textvariable_callback + ) + + def _commit(self) -> None: + shadow_val = self._shadow_var.get() + if shadow_val != self.var.get(): + self._syncing = True + self.var.set(shadow_val) + self._syncing = False + + def _apply_error(self) -> None: + self.component.configure(border_color=ERROR_BORDER_COLOR) + + def _clear_error(self) -> None: + self.component.configure(border_color=self._original_border_color) + + def _on_shadow_write(self, *_args) -> None: + if self._syncing: + return + if not self._touched: + # external sync or initial set — commit immediately + self._commit() + if self._debounce: + self._debounce.cancel() + return + if self._debounce: + self._debounce.call() + if self._undo_debounce: + self._undo_debounce.call() + + def _on_real_var_write(self, *_args) -> None: + if self._syncing: + return + # external change (preset load, file dialog, etc) — sync to shadow var + self._syncing = True + self._shadow_var.set(self.var.get()) + self._syncing = False + self._validate_and_style(self._shadow_var.get()) + + def _push_undo_snapshot(self) -> None: + self._undo.push(self._shadow_var.get()) + + def _on_debounce_fire(self) -> None: + val = self._shadow_var.get() + if self._validate_and_style(val): + self._commit() + + def _on_focus_in(self, _e=None) -> None: + self._touched = False + self._undo.push(self._shadow_var.get()) + + def _on_user_input(self, _e=None) -> None: + self._touched = True + + def _on_focus_out(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._undo_debounce: + self._undo_debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + self._undo.push(self._shadow_var.get()) + + def _on_enter(self, _e=None) -> None: + if self._debounce: + self._debounce.cancel() + if self._touched: + if self._validate_and_style(self._shadow_var.get()): + self._commit() + + def _set_value(self, value: str) -> None: + self._syncing = True + self._shadow_var.set(value) + self._syncing = False + if self._validate_and_style(value): + self._commit() + + def _on_undo(self, _e=None) -> str: + previous = self._undo.undo(self._shadow_var.get()) + if previous is not None: + self._set_value(previous) + return "break" + + def _on_redo(self, _e=None) -> str: + next_val = self._undo.redo() + if next_val is not None: + self._set_value(next_val) + return "break" + + def flush(self) -> str | None: + if self._debounce: + self._debounce.cancel() + + value = self._shadow_var.get() + error = self.validate(value) + + if error is not None: + self._apply_error() + else: + self._clear_error() + self._commit() + + return error + + +class PathValidator(FieldValidator): + """FieldValidator with additional path-specific checks.""" + + def __init__( + self, + component: ctk.CTkEntry, + var: tk.Variable, + ui_state: UIState, + var_name: str, + io_type: PathIOType = PathIOType.INPUT, + max_undo: int = DEFAULT_MAX_UNDO, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + ): + super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) + self.io_type = io_type + + def validate(self, value: str) -> str | None: + base_err = super().validate(value) + if base_err is not None: + return base_err + if value == "": + return None + return _validate_path_field(self.ui_state, self.io_type, value) + + def revalidate(self) -> None: + if self.component.winfo_exists(): + self._validate_and_style(self._shadow_var.get()) From a36aa74239cd8d34b3adde868f7e00c85515c689 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 10 May 2026 14:09:43 +0200 Subject: [PATCH 12/67] feat: implement PySide6 views replacing CTK placeholder copies Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 1 + .../ui/PySide6AdditionalEmbeddingsTabView.py | 45 +- modules/ui/PySide6CaptionUIView.py | 240 +---- modules/ui/PySide6CloudTabView.py | 51 +- modules/ui/PySide6ConceptTabView.py | 256 +++-- modules/ui/PySide6ConceptWindowView.py | 270 +++--- modules/ui/PySide6ConfigListView.py | 65 +- modules/ui/PySide6ConvertModelUIView.py | 42 +- modules/ui/PySide6LoraTabView.py | 52 +- modules/ui/PySide6ModelTabView.py | 51 +- modules/ui/PySide6MuonAdamWindowView.py | 46 +- modules/ui/PySide6OffloadingWindowView.py | 40 +- .../ui/PySide6OptimizerParamsWindowView.py | 82 +- modules/ui/PySide6ProfilingWindowView.py | 64 +- modules/ui/PySide6SampleFrameView.py | 47 +- modules/ui/PySide6SampleParamsWindowView.py | 39 +- modules/ui/PySide6SampleWindowView.py | 145 ++- modules/ui/PySide6SamplingTabView.py | 56 +- .../ui/PySide6SchedulerParamsWindowView.py | 113 ++- .../PySide6TimestepDistributionWindowView.py | 98 +- modules/ui/PySide6TopBarView.py | 46 +- modules/ui/PySide6TrainUIView.py | 463 +++++---- modules/ui/PySide6TrainingTabView.py | 100 +- modules/ui/PySide6VideoToolUIView.py | 225 +++-- modules/util/ui/PySide6UIState.py | 18 + modules/util/ui/QtVar.py | 41 + modules/util/ui/pyside6_abc.py | 7 + modules/util/ui/pyside6_components.py | 887 ++++++++++-------- modules/util/ui/pyside6_validation.py | 287 ++---- requirements-global.txt | 1 + scripts/train_ui_pyside6.py | 43 + scripts/video_tool_ui.py | 2 +- 32 files changed, 1935 insertions(+), 1988 deletions(-) create mode 100644 modules/util/ui/PySide6UIState.py create mode 100644 modules/util/ui/QtVar.py create mode 100644 modules/util/ui/pyside6_abc.py create mode 100644 scripts/train_ui_pyside6.py diff --git a/.gitignore b/.gitignore index 92e3d5612..efe8dc001 100644 --- a/.gitignore +++ b/.gitignore @@ -38,4 +38,5 @@ pixi.toml train.bat debug_report.log config_diff.txt +CLAUDE.md PLAN.md diff --git a/modules/ui/PySide6AdditionalEmbeddingsTabView.py b/modules/ui/PySide6AdditionalEmbeddingsTabView.py index fc24c61d1..196526138 100644 --- a/modules/ui/PySide6AdditionalEmbeddingsTabView.py +++ b/modules/ui/PySide6AdditionalEmbeddingsTabView.py @@ -1,17 +1,16 @@ - from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController from modules.ui.BaseAdditionalEmbeddingsTabView import BaseAdditionalEmbeddingsTabView, BaseEmbeddingWidgetView -from modules.ui.CtkConfigListView import CtkConfigListView -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk +from PySide6.QtWidgets import QWidget -class CtkAdditionalEmbeddingsTabView(CtkConfigListView, BaseAdditionalEmbeddingsTabView): +class PySide6AdditionalEmbeddingsTabView(PySide6ConfigListView, BaseAdditionalEmbeddingsTabView): def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state): - CtkConfigListView.__init__( + PySide6ConfigListView.__init__( self, master, controller, ui_state, attr_name="additional_embeddings", enable_key="train", @@ -22,30 +21,34 @@ def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_sta ) def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return CtkEmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + return PySide6EmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) -class CtkEmbeddingWidgetView(BaseEmbeddingWidgetView, ctk.CTkFrame): +class PySide6EmbeddingWidgetView(BaseEmbeddingWidgetView, QWidget): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command, controller): - ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") - BaseEmbeddingWidgetView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseEmbeddingWidgetView.__init__(self, pyside6_components) self.element = element - ui_state = CtkUIState(self, element) + ui_state = PySide6UIState(element) - self.grid_columnconfigure(0, weight=1) + pyside6_components._layout(self).setColumnStretch(0, 1) - top_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - top_frame.grid(row=0, column=0, sticky="nsew") - top_frame.grid_columnconfigure(3, weight=1) - top_frame.grid_columnconfigure(5, weight=1) + top_frame = QWidget(self) + pyside6_components._layout(top_frame).setColumnStretch(3, 1) + pyside6_components._layout(top_frame).setColumnStretch(5, 1) + pyside6_components._layout(self).addWidget(top_frame, 0, 0) - bottom_frame = ctk.CTkFrame(master=self, corner_radius=0, fg_color="transparent") - bottom_frame.grid(row=1, column=0, sticky="nsew") - bottom_frame.grid_columnconfigure(7, weight=1) + bottom_frame = QWidget(self) + pyside6_components._layout(bottom_frame).setColumnStretch(7, 1) + pyside6_components._layout(self).addWidget(bottom_frame, 1, 0) self.build_content(top_frame, bottom_frame, ui_state, i, save_command, remove_command, clone_command, controller) def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") + pyside6_components._layout(self.parent()).addWidget(self, getattr(self, 'visible_index', self.i), 0) + self.show() + + def destroy(self): + self.deleteLater() diff --git a/modules/ui/PySide6CaptionUIView.py b/modules/ui/PySide6CaptionUIView.py index 281036912..69e2b5286 100644 --- a/modules/ui/PySide6CaptionUIView.py +++ b/modules/ui/PySide6CaptionUIView.py @@ -1,228 +1,12 @@ -from tkinter import filedialog - -from modules.ui.BaseCaptionUIView import BaseCaptionUIView -from modules.ui.CaptionUIController import CaptionUIController -from modules.ui.CtkGenerateCaptionsWindowView import CtkGenerateCaptionsWindowView -from modules.ui.CtkGenerateMasksWindowView import CtkGenerateMasksWindowView -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import bind_mousewheel, set_window_icon - -import customtkinter as ctk -from customtkinter import ScalingTracker, ThemeManager -from PIL import Image - - -class CtkCaptionUIView(BaseCaptionUIView, ctk.CTkToplevel): - def __init__(self, parent, controller: CaptionUIController, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseCaptionUIView.__init__(self, ctk_components) - self.protocol("WM_DELETE_WINDOW", controller.on_close) - - self.controller = controller - controller.view = self - self.config_ui_state = CtkUIState(self, controller.config_ui_data) - self.enable_mask_editing_var = ctk.BooleanVar() - self.mask_editing_alpha = None - self.prompt_var = None - self.prompt_component = None - self.image = None - self.image_label = None - self.file_list = None - self.image_labels = [] - - self.title("OneTrainer") - self.geometry("1280x980") - self.resizable(False, False) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) - - top_frame = ctk.CTkFrame(self) - top_frame.grid(row=0, column=0, sticky="nsew") - self.build_top_bar(top_frame, controller, self.config_ui_state) - - self.bottom_frame = ctk.CTkFrame(self) - self.bottom_frame.grid(row=1, column=0, sticky="nsew") - self.bottom_frame.grid_rowconfigure(0, weight=1) - self.bottom_frame.grid_columnconfigure(0, weight=0) - self.bottom_frame.grid_columnconfigure(1, weight=1) - - self.file_list_column(self.bottom_frame) - self.content_column(self.bottom_frame) - self.controller.load_directory() - - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - def file_list_column(self, master): - if self.file_list is not None: - self.image_labels = [] - self.file_list.destroy() - - self.file_list = ctk.CTkScrollableFrame(master, width=300) - self.file_list.grid(row=0, column=0, sticky="nsew") - - for i, filename in enumerate(self.controller.image_rel_paths): - def __create_switch_image(index): - def __switch_image(event): - self.controller.switch_image(index) - - return __switch_image - - label = ctk.CTkLabel(self.file_list, text=filename) - label.bind("", __create_switch_image(i)) - - self.image_labels.append(label) - label.grid(row=i, column=0, padx=5, sticky="nsw") - - def content_column(self, master): - image = Image.new("RGBA", (512, 512), (0, 0, 0, 0)) - - right_frame = ctk.CTkFrame(master, fg_color="transparent") - right_frame.grid(row=0, column=1, sticky="nsew") - - right_frame.grid_columnconfigure(4, weight=1) - right_frame.grid_rowconfigure(1, weight=1) - - self.build_mask_buttons(right_frame) - - # checkbox to enable mask editing - self.enable_mask_editing_var = ctk.BooleanVar() - self.enable_mask_editing_var.set(False) - enable_mask_editing_checkbox = ctk.CTkCheckBox( - right_frame, text="Enable Mask Editing", variable=self.enable_mask_editing_var, width=50) - enable_mask_editing_checkbox.grid(row=0, column=2, padx=25, pady=5, sticky="w") - - # mask alpha textbox - self.mask_editing_alpha = ctk.CTkEntry(master=right_frame, width=40, placeholder_text="1.0") - self.mask_editing_alpha.insert(0, "1.0") - self.mask_editing_alpha.grid(row=0, column=3, sticky="e", padx=5, pady=5) - self.bind_key_events(self.mask_editing_alpha) - - mask_editing_alpha_label = ctk.CTkLabel(right_frame, text="Brush Alpha", width=75) - mask_editing_alpha_label.grid(row=0, column=4, padx=0, pady=5, sticky="w") - - # image - self.image = ctk.CTkImage( - light_image=image, - size=(self.controller.image_size, self.controller.image_size) - ) - self.image_label = ctk.CTkLabel( - master=right_frame, text="", image=self.image, - height=self.controller.image_size, width=self.controller.image_size - ) - self.image_label.grid(row=1, column=0, columnspan=5, sticky="nsew") - - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - self.image_label.bind("", self.edit_mask) - bind_mousewheel(self.image_label, {self.image_label.children["!label"]}, self.draw_mask_radius) - - # prompt - self.prompt_var = ctk.StringVar() - self.prompt_component = ctk.CTkEntry(right_frame, textvariable=self.prompt_var) - self.prompt_component.grid(row=2, column=0, columnspan=5, pady=5, sticky="new") - self.bind_key_events(self.prompt_component) - self.prompt_component.focus_set() - - def bind_key_events(self, component): - component.bind("", lambda e: self.controller.next_image()) - component.bind("", lambda e: self.controller.previous_image()) - component.bind("", self.save) - component.bind("", self.toggle_mask) - component.bind("", self.draw_mask_editing_mode) - component.bind("", self.fill_mask_editing_mode) - - def refresh_file_list(self): - self.file_list_column(self.bottom_frame) - - def focus_prompt(self): - self.prompt_component.focus_set() - - def on_image_switched(self, old_index, new_index, prompt): - if len(self.image_labels) > 0 and old_index < len(self.image_labels): - self.image_labels[old_index].configure( - text_color=ThemeManager.theme["CTkLabel"]["text_color"]) - self.image_labels[new_index].configure(text_color="#FF0000") - self.refresh_image() - self.prompt_var.set(prompt) - - def on_image_cleared(self): - image = Image.new("RGB", (512, 512), (0, 0, 0)) - self.image.configure(light_image=image) - - def refresh_image(self): - pil_image, size = self.controller.get_display_image() - self.image.configure(light_image=pil_image, size=size) - - def draw_mask_radius(self, delta, raw_event): - self.controller.update_mask_draw_radius(delta) - - def edit_mask(self, event): - if not self.enable_mask_editing_var.get(): - return - - if event.widget != self.image_label.children["!label"]: - return - - display_scaling = ScalingTracker.get_window_scaling(self) - - event_x = event.x / display_scaling - event_y = event.y / display_scaling - - is_right = False - is_left = False - if event.state & 0x0100 or event.num == 1: # left mouse button - is_left = True - elif event.state & 0x0400 or event.num == 3: # right mouse button - is_right = True - - try: - alpha = float(self.mask_editing_alpha.get()) - except Exception: - alpha = 1.0 - - self.controller.handle_edit_mask(event_x, event_y, is_left, is_right, alpha) - - def save(self, event): - self.controller.save(self.prompt_var.get()) - - def draw_mask_editing_mode(self, *args): - self.controller.set_mask_editing_mode('draw') - - if args: - # disable default event - return "break" - return None - - def fill_mask_editing_mode(self, *args): - self.controller.set_mask_editing_mode('fill') - - def toggle_mask(self, *args): - self.controller.toggle_mask() - self.refresh_image() - - def open_directory(self): - new_dir = filedialog.askdirectory() - - if new_dir: - self.controller.dir = new_dir - self.controller.load_directory(include_subdirectories=self.controller.config_ui_data["include_subdirectories"]) - - def open_mask_window(self): - self.wait_window(self.controller.open_mask_window(self, CtkGenerateMasksWindowView)) - self.controller.switch_image(self.controller.current_image_index) - - def open_caption_window(self): - self.wait_window(self.controller.open_caption_window(self, CtkGenerateCaptionsWindowView)) - self.controller.switch_image(self.controller.current_image_index) - - def open_in_explorer(self): - self.controller.open_in_explorer() - - def destroy(self): - self.controller._release_models() - super().destroy() +from PySide6.QtWidgets import QDialog, QLabel, QPushButton, QVBoxLayout + + +class PySide6CaptionUIView(QDialog): + def __init__(self, parent, controller): + super().__init__(parent) + self.setWindowTitle("Dataset Tool") + lo = QVBoxLayout(self) + lo.addWidget(QLabel("The dataset tool has not been ported to Qt6 yet.")) + ok = QPushButton("OK") + ok.clicked.connect(self.accept) + lo.addWidget(ok) diff --git a/modules/ui/PySide6CloudTabView.py b/modules/ui/PySide6CloudTabView.py index 0a5249069..1fa1e2791 100644 --- a/modules/ui/PySide6CloudTabView.py +++ b/modules/ui/PySide6CloudTabView.py @@ -1,45 +1,42 @@ - - from modules.ui.BaseCloudTabView import BaseCloudTabView from modules.ui.CloudTabController import CloudTabController -from modules.util.ui import ctk_components +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QWidget -import customtkinter as ctk +class PySide6CloudTabView(BaseCloudTabView, QWidget, metaclass=QtABCMeta): -class CtkCloudTabView(BaseCloudTabView): def __init__(self, master, controller: CloudTabController, ui_state): - BaseCloudTabView.__init__(self, ctk_components) - self.master = master + QWidget.__init__(self, master) + BaseCloudTabView.__init__(self, pyside6_components) + self.controller = controller self.ui_state = ui_state - self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, weight=0) - self.frame.grid_columnconfigure(3, weight=1) - self.frame.grid_columnconfigure(4, weight=0) - self.frame.grid_columnconfigure(5, weight=1) - - self.build_content(self.frame, controller, ui_state) - - self.frame.pack(fill="both", expand=1) + scroll, frame = pyside6_components.scrollable_frame(self) + pyside6_components._layout(self).addWidget(scroll, 0, 0) + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(3, 1) + lo.setColumnStretch(5, 1) + self.frame = frame + self.build_content(frame, controller, ui_state) def _on_set_gpu_types(self): - self.gpu_types_menu.configure(values=self.controller.get_gpu_types()) + self.gpu_types_menu.clear() + self.gpu_types_menu.addItems(self.controller.get_gpu_types()) def _make_reattach_frame(self, frame): - reattach_frame = ctk.CTkFrame(frame, fg_color="transparent") - reattach_frame.grid(row=9, column=3, padx=0, pady=0, sticky="new") - reattach_frame.grid_columnconfigure(0, weight=1) - reattach_frame.grid_columnconfigure(1, weight=1) + reattach_frame = QWidget(frame) + pyside6_components._layout(frame).addWidget(reattach_frame, 9, 3) + pyside6_components._layout(reattach_frame).setColumnStretch(0, 1) return reattach_frame def _make_create_frame(self, frame): - create_frame = ctk.CTkFrame(frame, fg_color="transparent") - create_frame.grid(row=1, column=5, padx=0, pady=0, sticky="new") - create_frame.grid_columnconfigure(0, weight=0) - create_frame.grid_columnconfigure(1, weight=1) + create_frame = QWidget(frame) + pyside6_components._layout(frame).addWidget(create_frame, 1, 5) + pyside6_components._layout(create_frame).setColumnStretch(1, 1) return create_frame diff --git a/modules/ui/PySide6ConceptTabView.py b/modules/ui/PySide6ConceptTabView.py index 5b3e86ac9..f55e6f5f7 100644 --- a/modules/ui/PySide6ConceptTabView.py +++ b/modules/ui/PySide6ConceptTabView.py @@ -1,26 +1,26 @@ -from tkinter import BooleanVar, StringVar - from modules.ui.BaseConceptTabView import BaseConceptTabView, BaseConceptWidgetView from modules.ui.ConceptTabController import ConceptTabController -from modules.ui.CtkConceptWindowView import CtkConceptWindowView -from modules.ui.CtkConfigListView import CtkConfigListView -from modules.util.ui import ctk_components -from modules.util.ui.ctk_validation import DebounceTimer -from modules.util.ui.CtkUIState import CtkUIState +from modules.ui.PySide6ConceptWindowView import PySide6ConceptWindowView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState +from modules.util.ui.QtVar import QtVar -import customtkinter as ctk +from PIL.ImageQt import ImageQt +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import QCheckBox, QComboBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QWidget -class CtkConceptTabView(CtkConfigListView, BaseConceptTabView): +class PySide6ConceptTabView(PySide6ConfigListView, BaseConceptTabView): def __init__(self, master, controller: ConceptTabController, ui_state): - # Pre-initialize before CtkConfigListView.__init__ because _reset_filters is - # called during build() via options_kv's immediate update_var() call. - self.search_var = StringVar() - self.filter_var = StringVar(value="ALL") - self.show_disabled_var = BooleanVar(value=True) + # Pre-initialize before PySide6ConfigListView.__init__ because _reset_filters is + # called during build() via options_kv's initial command fire. + self.search_var = QtVar("") + self.filter_var = QtVar("ALL") + self.show_disabled_var = QtVar(True) - CtkConfigListView.__init__( + PySide6ConfigListView.__init__( self, master, controller, ui_state, from_external_file=True, attr_name="concept_file_name", @@ -31,146 +31,142 @@ def __init__(self, master, controller: ConceptTabController, ui_state): is_full_width=False, show_toggle_button=True, ) - self._toolbar = None - self._toolbar_is_wrapped = False self._add_search_bar() - self.top_frame.bind('', lambda e: self._maybe_reposition_toolbar(e.width)) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return self.controller.open_element_window(self.master, self.current_config[i], ui_state[0], ui_state[1], ui_state[2], CtkConceptWindowView) + def open_element_window(self, i, ui_state): + return self.controller.open_element_window(self.master, self.current_config[i], ui_state[0], ui_state[1], ui_state[2], PySide6ConceptWindowView) def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return CtkConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + return PySide6ConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) def _add_search_bar(self): - toolbar = ctk.CTkFrame(self.top_frame, fg_color="transparent") - toolbar.grid(row=0, column=4, columnspan=2, padx=10, sticky="ew") - toolbar.grid_columnconfigure(2, weight=1) - self._toolbar = toolbar - - ctk.CTkLabel(toolbar, text="Search:").grid(row=0, column=0, padx=(0, 5)) - self.search_var = StringVar() - self.search_entry = ctk.CTkEntry(toolbar, textvariable=self.search_var, - placeholder_text="Filter...", width=200) - self.search_entry.grid(row=0, column=1) - self._search_debouncer = DebounceTimer(self.search_entry, 300, lambda: self._update_filters()) - self.search_var.trace_add("write", lambda *_: self._search_debouncer.call()) - - ctk.CTkLabel(toolbar, text="").grid(row=0, column=2, padx=5) - - ctk.CTkLabel(toolbar, text="Type:").grid(row=0, column=3, padx=(0, 5)) - self.filter_var = StringVar(value="ALL") - ctk.CTkOptionMenu(toolbar, values=self._FILTER_TYPES, - variable=self.filter_var, command=lambda x: self._update_filters(), - width=150).grid(row=0, column=4) - - self.show_disabled_var = BooleanVar(value=True) - self.show_disabled_checkbox = ctk.CTkCheckBox(toolbar, text="Show Disabled", variable=self.show_disabled_var, - command=self._update_filters, width=100) - self.show_disabled_checkbox.grid(row=0, column=5, padx=(10, 0)) - self._refresh_show_disabled_text() - - ctk.CTkButton(toolbar, text="Clear", width=50, - command=self._reset_filters).grid(row=0, column=6, padx=(10, 0)) - - def _maybe_reposition_toolbar(self, width): - if not self._toolbar: - return - threshold = 1070 - want_wrapped = width < threshold - if want_wrapped == self._toolbar_is_wrapped: - return - self._toolbar_is_wrapped = want_wrapped - if want_wrapped: - self._toolbar.grid_configure(row=1, column=0, columnspan=8, sticky="ew", padx=10) - else: - self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + toolbar = QWidget(self.top_frame) + row_lo = QHBoxLayout(toolbar) + row_lo.setContentsMargins(0, 0, 0, 0) + pyside6_components._layout(self.top_frame).addWidget(toolbar, 0, 4) + + self.search_var = QtVar("") + search_entry = QLineEdit(toolbar) + search_entry.setPlaceholderText("Filter...") + search_entry.setFixedWidth(200) + row_lo.addWidget(QLabel("Search:", toolbar)) + row_lo.addWidget(search_entry) + + def _on_search(text): + self.search_var.set(text) + self._update_filters() + search_entry.textChanged.connect(_on_search) + + self.filter_var = QtVar("ALL") + filter_combo = QComboBox(toolbar) + filter_combo.addItems(self._FILTER_TYPES) + filter_combo.setFixedWidth(150) + row_lo.addWidget(QLabel("Type:", toolbar)) + row_lo.addWidget(filter_combo) + + def _on_filter(text): + self.filter_var.set(text) + self._update_filters() + filter_combo.currentTextChanged.connect(_on_filter) + + self.show_disabled_var = QtVar(True) + show_disabled_cb = QCheckBox("Show Disabled", toolbar) + show_disabled_cb.setChecked(True) + row_lo.addWidget(show_disabled_cb) + + def _on_show_disabled(state): + self.show_disabled_var.set(bool(state)) + self._update_filters() + show_disabled_cb.stateChanged.connect(_on_show_disabled) + + clear_btn = QPushButton("Clear", toolbar) + clear_btn.setFixedWidth(50) + clear_btn.clicked.connect(self._reset_filters) + row_lo.addWidget(clear_btn) def _reset_filters(self): - self.search_var.set("") - self.filter_var.set("ALL") - self.show_disabled_var.set(True) + if self.search_var is not None: + self.search_var.set("") + if self.filter_var is not None: + self.filter_var.set("ALL") + if self.show_disabled_var is not None: + self.show_disabled_var.set(True) self._update_filters() - def _refresh_show_disabled_text(self): - try: - disabled_count = sum(1 for c in getattr(self, 'current_config', []) if getattr(c, 'enabled', True) is False) - except (AttributeError, TypeError): - disabled_count = 0 - text = f"Show Disabled ({disabled_count})" if disabled_count > 0 else "Show Disabled" - try: - if getattr(self, 'show_disabled_checkbox', None): - self.show_disabled_checkbox.configure(text=text) - except (AttributeError, RuntimeError): - pass - -class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): +class PySide6ConceptWidgetView(BaseConceptWidgetView, QWidget): def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): - ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") - BaseConceptWidgetView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseConceptWidgetView.__init__(self, pyside6_components) self.concept = concept - self.ui_state = CtkUIState(self, concept) - self.image_ui_state = CtkUIState(self, concept.image) - self.text_ui_state = CtkUIState(self, concept.text) + self.ui_state = PySide6UIState(concept) + self.image_ui_state = PySide6UIState(concept.image) + self.text_ui_state = PySide6UIState(concept.text) self.i = i - self.grid_rowconfigure(1, weight=1) + self.setFixedSize(160, 180) - self.image = ctk.CTkImage( - light_image=self._get_preview_image(), - size=(150, 150) - ) - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=150, width=150) - image_label.grid(row=0, column=0) - - self.name_label = self.components.label(self, 1, 0, self._get_display_name(), pad=5, wraplength=140) - - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i), + image = self._get_preview_image() + pixmap = QPixmap.fromImage(ImageQt(image.convert("RGBA"))) + self.image_label = QLabel(self) + self.image_label.setPixmap(pixmap) + self.image_label.setFixedSize(150, 150) + self.image_label.move(5, 0) + self.image_label.mousePressEvent = lambda _: open_command( + self.i, (self.ui_state, self.image_ui_state, self.text_ui_state) ) - close_button.place(x=0, y=0) - - clone_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="+", - corner_radius=2, - fg_color="#00C000", - command=lambda: clone_command(self.i, controller.randomize_seed), - ) - clone_button.place(x=25, y=0) - - enabled_switch = ctk.CTkSwitch( - master=self, - width=40, - variable=self.ui_state.get_var("enabled"), - text="", - command=save_command, - ) - enabled_switch.place(x=110, y=0) - image_label.bind( - "", - lambda event: open_command(self.i, (self.ui_state, self.image_ui_state, self.text_ui_state)) + self.name_label = QLabel(self._get_display_name(), self) + self.name_label.setWordWrap(True) + self.name_label.setFixedWidth(140) + self.name_label.move(5, 153) + + close_btn = QPushButton("X", self) + close_btn.setFixedSize(20, 20) + close_btn.setStyleSheet("background-color: #C00000; color: white;") + close_btn.move(5, 0) + close_btn.clicked.connect(lambda: remove_command(self.i)) + + clone_btn = QPushButton("+", self) + clone_btn.setFixedSize(20, 20) + clone_btn.setStyleSheet("background-color: #00C000; color: white;") + clone_btn.move(30, 0) + clone_btn.clicked.connect(lambda: clone_command(self.i, controller.randomize_seed)) + + enabled_cb = QCheckBox(self) + enabled_cb.setChecked(concept.enabled) + enabled_cb.setFixedSize(20, 20) + enabled_cb.setStyleSheet("QCheckBox::indicator { width: 20px; height: 20px; }") + enabled_cb.move(135, 0) + enabled_cb.stateChanged.connect(lambda state: ( + setattr(concept, 'enabled', bool(state)), + save_command(), + )) + self.ui_state.get_var("enabled")._bind_widget( + lambda v: enabled_cb.setChecked(bool(v)) ) def configure_element(self): - self.name_label.configure(text=self._get_display_name()) - self.image.configure(light_image=self._get_preview_image()) - self._clear_search_cache() + self.name_label.setText(self._get_display_name()) + image = self._get_preview_image() + pixmap = QPixmap.fromImage(ImageQt(image.convert("RGBA"))) + self.image_label.setPixmap(pixmap) + try: + if hasattr(self.concept, '_search_cache'): + delattr(self.concept, '_search_cache') + except AttributeError: + pass def place_in_list(self): index = getattr(self, 'visible_index', self.i) x = index % 6 y = index // 6 - self.grid(row=y, column=x, pady=5, padx=5) + lo = pyside6_components._layout(self.parent()) + lo.addWidget(self, y, x) + lo.setColumnStretch(6, 1) + self.show() + + def destroy(self): + self.deleteLater() diff --git a/modules/ui/PySide6ConceptWindowView.py b/modules/ui/PySide6ConceptWindowView.py index 60c0f57fe..87355a9a6 100644 --- a/modules/ui/PySide6ConceptWindowView.py +++ b/modules/ui/PySide6ConceptWindowView.py @@ -2,16 +2,27 @@ from modules.ui.BaseConceptWindowView import BaseConceptWindowView from modules.ui.ConceptWindowController import ConceptWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg - - -class CtkConceptWindowView(BaseConceptWindowView, ctk.CTkToplevel): +from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg +from PIL.ImageQt import ImageQt +from PySide6.QtCore import Qt +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import ( + QCheckBox, + QDialog, + QGridLayout, + QLabel, + QPushButton, + QScrollArea, + QTabWidget, + QTextEdit, + QWidget, +) + + +class PySide6ConceptWindowView(BaseConceptWindowView, QDialog): def __init__( self, parent, @@ -19,108 +30,130 @@ def __init__( ui_state, image_ui_state, text_ui_state, - *args, **kwargs, ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseConceptWindowView.__init__(self, ctk_components) + QDialog.__init__(self, parent) + BaseConceptWindowView.__init__(self, pyside6_components) self.controller = controller self.image_preview_file_index = 0 - self.preview_augmentations = ctk.BooleanVar(self, True) + self._preview_augmentations = True self.bucket_fig = None - self.title("Concept") - self.geometry("800x700") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - # general tab - general_frame = ctk.CTkScrollableFrame(tabview.add("general"), fg_color="transparent") - general_frame.grid_columnconfigure(1, weight=1) - general_frame.grid_columnconfigure(2, weight=1) - self.build_general_tab(general_frame, controller, ui_state, text_ui_state) - general_frame.pack(fill="both", expand=1) - - # image augmentation tab - image_aug_master = tabview.add("image augmentation") - image_aug_frame = ctk.CTkScrollableFrame(image_aug_master, fg_color="transparent") - image_aug_frame.grid_columnconfigure(0, weight=0) - image_aug_frame.grid_columnconfigure(1, weight=0) - image_aug_frame.grid_columnconfigure(2, weight=0) - image_aug_frame.grid_columnconfigure(3, weight=1) - self.build_image_augmentation_tab(image_aug_frame, controller, image_ui_state) - - # image - image_preview, filename_preview, caption_preview = controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) - self.image = ctk.CTkImage( - light_image=image_preview, - size=image_preview.size, + self.setWindowTitle("Concept") + self.resize(800, 700) + + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + + tabs = QTabWidget(self) + outer.addWidget(tabs, 0, 0) + + _pad = pyside6_components.PAD + + # --- general tab --- + gen_scroll = QScrollArea() + gen_scroll.setWidgetResizable(True) + gen_frame = QWidget() + gen_scroll.setWidget(gen_frame) + pyside6_components._layout(gen_frame).setContentsMargins(_pad, _pad, _pad, _pad) + pyside6_components._layout(gen_frame).setColumnStretch(1, 1) + pyside6_components._layout(gen_frame).setColumnStretch(2, 1) + self.build_general_tab(gen_frame, controller, ui_state, text_ui_state) + pyside6_components._pack_form(gen_frame) + tabs.addTab(gen_scroll, "general") + + # --- image augmentation tab --- + img_scroll = QScrollArea() + img_scroll.setWidgetResizable(True) + img_outer = QWidget() + img_scroll.setWidget(img_outer) + lo_img_outer = pyside6_components._layout(img_outer) + lo_img_outer.setContentsMargins(_pad, _pad, _pad, _pad) + lo_img_outer.setColumnStretch(0, 1) + + # form in its own widget so the preview panel can't affect row heights + img_form = QWidget(img_outer) + img_form_lo = pyside6_components._layout(img_form) + img_form_lo.setColumnStretch(3, 1) + self.build_image_augmentation_tab(img_form, controller, image_ui_state) + pyside6_components._pack_form(img_form) + lo_img_outer.addWidget(img_form, 0, 0, Qt.AlignTop) + + # preview panel alongside the form + image_preview, filename_preview, caption_preview = controller.get_preview_image( + self.image_preview_file_index, self._preview_augmentations ) - image_label = ctk.CTkLabel(master=image_aug_frame, text="", image=self.image, height=300, width=300) - image_label.grid(row=0, column=4, rowspan=6) - - # refresh preview - update_button_frame = ctk.CTkFrame(master=image_aug_frame, corner_radius=0, fg_color="transparent") - update_button_frame.grid(row=6, column=4, rowspan=6, sticky="nsew") - update_button_frame.grid_columnconfigure(1, weight=1) - - prev_preview_button = self.components.button(update_button_frame, 0, 0, "<", command=self._prev_image_preview) - self.components.button(update_button_frame, 0, 1, "Update Preview", command=self._update_image_preview) - next_preview_button = self.components.button(update_button_frame, 0, 2, ">", command=self._next_image_preview) - preview_augmentations_switch = ctk.CTkSwitch(update_button_frame, text="Show Augmentations", variable=self.preview_augmentations, command=self._update_image_preview) - preview_augmentations_switch.grid(row=1, column=0, columnspan=3, padx=5, pady=5) - - prev_preview_button.configure(width=40) - next_preview_button.configure(width=40) - - #caption and filename preview - self.filename_preview = ctk.CTkLabel(master=update_button_frame, text=filename_preview, width=300, anchor="nw", justify="left", padx=10, wraplength=280) - self.filename_preview.grid(row=2, column=0, columnspan=3) - self.caption_preview = ctk.CTkTextbox(master=update_button_frame, width = 300, height = 150, wrap="word", border_width=2) - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - self.caption_preview.grid(row=3, column=0, columnspan=3, rowspan=3) - - image_aug_frame.pack(fill="both", expand=1) - - # text augmentation tab - text_aug_frame = ctk.CTkScrollableFrame(tabview.add("text augmentation"), fg_color="transparent") - text_aug_frame.grid_columnconfigure(0, weight=0) - text_aug_frame.grid_columnconfigure(1, weight=0) - text_aug_frame.grid_columnconfigure(2, weight=0) - text_aug_frame.grid_columnconfigure(3, weight=1) - self.build_text_augmentation_tab(text_aug_frame, controller, text_ui_state) - text_aug_frame.pack(fill="both", expand=1) - - # statistics tab - stats_frame = ctk.CTkScrollableFrame(tabview.add("statistics"), fg_color="transparent") - stats_frame.grid_columnconfigure(0, weight=0, minsize=150) - stats_frame.grid_columnconfigure(1, weight=0, minsize=150) - stats_frame.grid_columnconfigure(2, weight=0, minsize=150) - stats_frame.grid_columnconfigure(3, weight=0, minsize=150) + preview_panel = QWidget(img_outer) + pb_lo = QGridLayout(preview_panel) + + self._image_label = QLabel(preview_panel) + self._image_label.setFixedSize(300, 300) + self._image_label.setPixmap(QPixmap.fromImage(ImageQt(image_preview.convert("RGBA"))).scaled(300, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation)) + pb_lo.addWidget(self._image_label, 0, 0, 1, 3) + + prev_btn = QPushButton("<", preview_panel) + prev_btn.setFixedWidth(40) + prev_btn.clicked.connect(self._prev_image_preview) + update_btn = QPushButton("Update Preview", preview_panel) + update_btn.clicked.connect(self._update_image_preview) + next_btn = QPushButton(">", preview_panel) + next_btn.setFixedWidth(40) + next_btn.clicked.connect(self._next_image_preview) + self._aug_checkbox = QCheckBox("Show Augmentations", preview_panel) + self._aug_checkbox.setChecked(True) + self._aug_checkbox.toggled.connect(lambda checked: self._on_aug_toggle(checked)) + pb_lo.addWidget(prev_btn, 1, 0) + pb_lo.addWidget(update_btn, 1, 1) + pb_lo.addWidget(next_btn, 1, 2) + pb_lo.addWidget(self._aug_checkbox, 2, 0, 1, 3) + + self._filename_label = QLabel(filename_preview, preview_panel) + self._filename_label.setWordWrap(True) + self._filename_label.setFixedWidth(300) + pb_lo.addWidget(self._filename_label, 3, 0, 1, 3) + + self._caption_box = QTextEdit(preview_panel) + self._caption_box.setReadOnly(True) + self._caption_box.setPlainText(caption_preview) + self._caption_box.setFixedSize(300, 150) + pb_lo.addWidget(self._caption_box, 4, 0, 1, 3) + + lo_img_outer.addWidget(preview_panel, 0, 1, Qt.AlignTop) + tabs.addTab(img_scroll, "image augmentation") + + # --- text augmentation tab --- + text_scroll = QScrollArea() + text_scroll.setWidgetResizable(True) + text_frame = QWidget() + text_scroll.setWidget(text_frame) + pyside6_components._layout(text_frame).setContentsMargins(_pad, _pad, _pad, _pad) + pyside6_components._layout(text_frame).setColumnStretch(3, 1) + self.build_text_augmentation_tab(text_frame, controller, text_ui_state) + pyside6_components._pack_form(text_frame) + tabs.addTab(text_scroll, "text augmentation") + + # --- statistics tab --- + stats_scroll = QScrollArea() + stats_scroll.setWidgetResizable(True) + stats_frame = QWidget() + stats_scroll.setWidget(stats_frame) + stats_lo = pyside6_components._layout(stats_frame) + stats_lo.setContentsMargins(_pad, _pad, _pad, _pad) + stats_lo.setColumnMinimumWidth(0, 150) + stats_lo.setColumnMinimumWidth(1, 150) + stats_lo.setColumnMinimumWidth(2, 150) + stats_lo.setColumnMinimumWidth(3, 150) self.build_concept_stats_tab(stats_frame, controller) - #aspect bucketing plot, mostly copied from timestep preview graph - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - self.text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - plt.set_loglevel('WARNING') #suppress errors about data type in bar chart - - assert self.bucket_fig is None - self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7,3)) - self.canvas = FigureCanvasTkAgg(self.bucket_fig, master=stats_frame) - self.canvas.get_tk_widget().grid(row=19, column=0, columnspan=4, rowspan=2) + plt.set_loglevel('WARNING') + self.bucket_fig, self.bucket_ax = plt.subplots(figsize=(7, 3)) + self.canvas = FigureCanvasQTAgg(self.bucket_fig) self.bucket_fig.tight_layout() self.bucket_fig.subplots_adjust(bottom=0.15) + palette = self.palette() + self.text_color = palette.text().color().name() + background_color = palette.window().color().name() self.bucket_fig.set_facecolor(background_color) self.bucket_ax.set_facecolor(background_color) self.bucket_ax.spines['bottom'].set_color(self.text_color) @@ -132,18 +165,22 @@ def __init__( self.bucket_ax.xaxis.label.set_color(self.text_color) self.bucket_ax.yaxis.label.set_color(self.text_color) - stats_frame.pack(fill="both", expand=1) + stats_lo.addWidget(self.canvas, 19, 0, 2, 4) + + tabs.addTab(stats_scroll, "statistics") + + ok = QPushButton("ok", self) + ok.clicked.connect(self._ok) + outer.addWidget(ok, 1, 0) #automatic concept scan self.scan_thread = threading.Thread(target=controller.auto_update_concept_stats, args=[self], daemon=True) self.scan_thread.start() - self.components.button(self, 1, 0, "ok", self._ok) - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + def _on_aug_toggle(self, checked: bool): + self._preview_augmentations = checked + self._update_image_preview() def _prev_image_preview(self): self.image_preview_file_index = max(self.image_preview_file_index - 1, 0) @@ -154,20 +191,17 @@ def _next_image_preview(self): self._update_image_preview() def _update_image_preview(self): - image_preview, filename_preview, caption_preview = self.controller.get_preview_image(self.image_preview_file_index, self.preview_augmentations.get()) - self.image.configure(light_image=image_preview, size=image_preview.size) - self.filename_preview.configure(text=filename_preview) - self.caption_preview.configure(state="normal") - self.caption_preview.delete(index1="1.0", index2="end") - self.caption_preview.insert(index="1.0", text=caption_preview) - self.caption_preview.configure(state="disabled") - - def destroy(self): + image_preview, filename_preview, caption_preview = self.controller.get_preview_image( + self.image_preview_file_index, self._preview_augmentations + ) + self._image_label.setPixmap( + QPixmap.fromImage(ImageQt(image_preview.convert("RGBA"))).scaled(300, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation) + ) + self._filename_label.setText(filename_preview) + self._caption_box.setPlainText(caption_preview) + + def _ok(self): if self.bucket_fig is not None: plt.close(self.bucket_fig) self.bucket_fig = None - - super().destroy() - - def _ok(self): - self.destroy() + self.accept() diff --git a/modules/ui/PySide6ConfigListView.py b/modules/ui/PySide6ConfigListView.py index 72995bfcc..6a369d6c7 100644 --- a/modules/ui/PySide6ConfigListView.py +++ b/modules/ui/PySide6ConfigListView.py @@ -2,12 +2,12 @@ from abc import ABC from modules.ui.BaseConfigListView import BaseConfigListView -from modules.util.ui import ctk_components, dialogs +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QInputDialog, QWidget -class CtkConfigListView(BaseConfigListView, ABC): +class PySide6ConfigListView(BaseConfigListView, ABC): def __init__( self, @@ -24,11 +24,16 @@ def __init__( is_full_width: bool = False, show_toggle_button: bool = False, ): - BaseConfigListView.__init__(self, ctk_components) + BaseConfigListView.__init__(self, pyside6_components) - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) + master_lo = pyside6_components._layout(master) + master_lo.setContentsMargins( + pyside6_components.PAD, pyside6_components.PAD, + pyside6_components.PAD, pyside6_components.PAD, + ) + master_lo.setRowStretch(0, 0) + master_lo.setRowStretch(1, 1) + master_lo.setColumnStretch(0, 1) self.build( master, controller, ui_state, from_external_file, @@ -43,29 +48,51 @@ def __init__( ) def _create_top_frame(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=0, column=0, sticky="nsew") + frame = QWidget(master) + pyside6_components._layout(master).addWidget(frame, 0, 0) + pyside6_components._layout(frame).setColumnStretch(4, 1) return frame def _create_element_list_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid(row=1, column=0, sticky="nsew") + scroll, content = pyside6_components.scrollable_frame(master) + pyside6_components._layout(master).addWidget(scroll, 1, 0) if self.is_full_width: - frame.grid_columnconfigure(0, weight=1) - return frame + pyside6_components._layout(content).setColumnStretch(0, 1) + content._scroll_area = scroll + return content def _wait_for_window(self, window): - self.master.wait_window(window) + window.exec() def _remove_widget_from_layout(self, widget): - widget.grid_remove() + widget.hide() def _destroy_widget(self, widget): - with contextlib.suppress(AttributeError): - widget.destroy() + with contextlib.suppress(RuntimeError, AttributeError): + widget.hide() + widget.deleteLater() def _destroy_frame(self, frame): - frame.destroy() + with contextlib.suppress(RuntimeError, AttributeError): + scroll = getattr(frame, '_scroll_area', None) + if scroll is not None: + lo = scroll.parent().layout() if scroll.parent() else None + if lo is not None: + lo.removeWidget(scroll) + scroll.hide() + scroll.deleteLater() + else: + frame.hide() + frame.deleteLater() + + def _update_toggle_button_text(self): + if not self.show_toggle_button: + return + self._update_item_enabled_state() + if self.toggle_button is not None: + self.toggle_button.setText("Disable" if self._is_current_item_enabled else "Enable") def _show_name_dialog(self, callback): - dialogs.StringInputDialog(self.master, "name", "Name", callback) + text, ok = QInputDialog.getText(self.master, "name", "Name") + if ok and text: + callback(text) diff --git a/modules/ui/PySide6ConvertModelUIView.py b/modules/ui/PySide6ConvertModelUIView.py index 782637348..959a2ea31 100644 --- a/modules/ui/PySide6ConvertModelUIView.py +++ b/modules/ui/PySide6ConvertModelUIView.py @@ -1,34 +1,32 @@ from modules.ui.BaseConvertModelUIView import BaseConvertModelUIView from modules.ui.ConvertModelUIController import ConvertModelUIController -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QWidget -class CtkConvertModelUIView(BaseConvertModelUIView, ctk.CTkToplevel): - def __init__(self, parent, controller: ConvertModelUIController, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseConvertModelUIView.__init__(self, ctk_components) +class PySide6ConvertModelUIView(BaseConvertModelUIView, QDialog): + def __init__(self, parent, controller: ConvertModelUIController): + QDialog.__init__(self, parent) + BaseConvertModelUIView.__init__(self, pyside6_components) - ui_state = CtkUIState(self, controller.convert_model_args) + ui_state = PySide6UIState(controller.convert_model_args) - self.title("Convert models") - self.geometry("550x350") - self.resizable(True, True) + self.setWindowTitle("Convert models") + self.resize(600, 380) - self.frame = ctk.CTkFrame(self, width=600, height=300) - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) + _pad = pyside6_components.PAD + outer = QGridLayout(self) + outer.setContentsMargins(_pad, _pad, _pad, _pad) - self.build_content(self.frame, controller, ui_state) - self.frame.pack(fill="both", expand=True) + frame = QWidget(self) + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + outer.addWidget(frame, 0, 0) - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + self.build_content(frame, controller, ui_state) + lo.setRowStretch(lo.rowCount(), 1) def set_converting(self, active): - self.button.configure(state="disabled" if active else "normal") + self.button.setEnabled(not active) diff --git a/modules/ui/PySide6LoraTabView.py b/modules/ui/PySide6LoraTabView.py index 8caa1f171..842017697 100644 --- a/modules/ui/PySide6LoraTabView.py +++ b/modules/ui/PySide6LoraTabView.py @@ -1,15 +1,17 @@ - from modules.ui.BaseLoraTabView import BaseLoraTabView from modules.ui.LoraTabController import LoraTabController from modules.util.enum.ModelType import PeftType -from modules.util.ui import ctk_components +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QWidget -import customtkinter as ctk +class PySide6LoraTabView(BaseLoraTabView, QWidget): -class CtkLoraTabView(BaseLoraTabView): def __init__(self, master, controller: LoraTabController, ui_state): - BaseLoraTabView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseLoraTabView.__init__(self, pyside6_components) + self.master = master self.controller = controller self.ui_state = ui_state @@ -18,24 +20,28 @@ def __init__(self, master, controller: LoraTabController, ui_state): self.refresh_ui() def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - self.scroll_frame = ctk.CTkFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=0) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=2) + if self.scroll_frame is not None: + self.scroll_frame.hide() + self.scroll_frame.deleteLater() + + self.scroll_frame = QWidget(self) + pyside6_components._layout(self).addWidget(self.scroll_frame, 0, 0) + lo = pyside6_components._layout(self.scroll_frame) + lo.setContentsMargins(pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(2, 2) self.build(self.scroll_frame, self.controller, self.ui_state, self.setup_lora) + pyside6_components._pack_form(self.scroll_frame) def setup_lora(self, peft_type: PeftType): - if self.options_frame: - self.options_frame.destroy() - self.options_frame = ctk.CTkFrame(self.scroll_frame, fg_color="transparent") - self.options_frame.grid(row=1, column=0, columnspan=3, sticky="nsew") - master = self.options_frame - master.grid_columnconfigure(0, weight=0, uniform="a") - master.grid_columnconfigure(1, weight=1, uniform="a") - master.grid_columnconfigure(2, minsize=50, uniform="a") - master.grid_columnconfigure(3, weight=0, uniform="a") - master.grid_columnconfigure(4, weight=1, uniform="a") - self.build_lora_options(master, self.controller, self.ui_state, peft_type) + if self.options_frame is not None: + self.options_frame.hide() + self.options_frame.deleteLater() + + self.options_frame = QWidget(self.scroll_frame) + pyside6_components._layout(self.scroll_frame).addWidget(self.options_frame, 1, 0, 1, 3) + lo = pyside6_components._layout(self.options_frame) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(4, 1) + self.build_lora_options(self.options_frame, self.controller, self.ui_state, peft_type) + pyside6_components._pack_form(self.options_frame) diff --git a/modules/ui/PySide6ModelTabView.py b/modules/ui/PySide6ModelTabView.py index 5b43b7dca..0835a586f 100644 --- a/modules/ui/PySide6ModelTabView.py +++ b/modules/ui/PySide6ModelTabView.py @@ -1,48 +1,41 @@ - - from modules.ui.BaseModelTabView import BaseModelTabView from modules.ui.ModelTabController import ModelTabController -from modules.util.ui import ctk_components +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta -import customtkinter as ctk +from PySide6.QtWidgets import QWidget -class CtkModelTabView(BaseModelTabView): +class PySide6ModelTabView(BaseModelTabView, QWidget, metaclass=QtABCMeta): + def __init__(self, master, controller: ModelTabController, ui_state): - BaseModelTabView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseModelTabView.__init__(self, pyside6_components) + self.master = master self.controller = controller self.ui_state = ui_state - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - self.scroll_frame = None - self.refresh_ui() def _make_svd_frames(self, parent, row: int): - svd_label_frame = ctk.CTkFrame(parent, fg_color="transparent") - svd_label_frame.grid(row=row, column=3, sticky="nsew") - svd_entry_frame = ctk.CTkFrame(parent, fg_color="transparent") - svd_entry_frame.grid(row=row, column=4, sticky="nsew") + svd_label_frame = QWidget(parent) + pyside6_components._layout(parent).addWidget(svd_label_frame, row, 3) + svd_entry_frame = QWidget(parent) + pyside6_components._layout(parent).addWidget(svd_entry_frame, row, 4) return svd_label_frame, svd_entry_frame def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - self.scroll_frame.grid_columnconfigure(0, weight=1) + if self.scroll_frame is not None: + self.scroll_frame.hide() + self.scroll_frame.deleteLater() - base_frame = ctk.CTkFrame(master=self.scroll_frame, corner_radius=5) - base_frame.grid(row=0, column=0, padx=5, pady=5, sticky="nsew") + scroll, frame = pyside6_components.scrollable_frame(self) + pyside6_components._layout(self).addWidget(scroll, 0, 0) + self.scroll_frame = scroll - base_frame.grid_columnconfigure(0, weight=0) - base_frame.grid_columnconfigure(1, weight=10) # , minsize=500) - base_frame.grid_columnconfigure(2, minsize=50) - base_frame.grid_columnconfigure(3, weight=0) - base_frame.grid_columnconfigure(4, weight=1) + frame_lo = pyside6_components._layout(frame) + frame_lo.setColumnStretch(1, 10) + frame_lo.setColumnStretch(4, 1) - self.build_content(base_frame, self.controller, self.ui_state) + self.build_content(frame, self.controller, self.ui_state) diff --git a/modules/ui/PySide6MuonAdamWindowView.py b/modules/ui/PySide6MuonAdamWindowView.py index 3dc48ab74..0deb4945a 100644 --- a/modules/ui/PySide6MuonAdamWindowView.py +++ b/modules/ui/PySide6MuonAdamWindowView.py @@ -1,37 +1,29 @@ from modules.ui.BaseMuonAdamWindowView import BaseMuonAdamWindowView from modules.ui.MuonAdamWindowController import MuonAdamWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton -class CtkMuonAdamWindowView(BaseMuonAdamWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: MuonAdamWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseMuonAdamWindowView.__init__(self, ctk_components) +class PySide6MuonAdamWindowView(BaseMuonAdamWindowView, QDialog): + def __init__(self, parent, controller: MuonAdamWindowController, ui_state): + QDialog.__init__(self, parent) + BaseMuonAdamWindowView.__init__(self, pyside6_components) - self.title(controller.get_title()) - self.geometry("800x500") - self.resizable(True, True) + self.setWindowTitle(controller.get_title()) + self.resize(800, 500) - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(0, 1) - frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - - self.components.button(self, 1, 0, "ok", command=self.destroy) + scroll, frame = pyside6_components.scrollable_frame(self) + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + lo.setColumnMinimumWidth(2, 50) + lo.setColumnStretch(4, 1) self.build_content(frame, controller, ui_state) + outer.addWidget(scroll, 0, 0) - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) diff --git a/modules/ui/PySide6OffloadingWindowView.py b/modules/ui/PySide6OffloadingWindowView.py index b752f1602..fcc9d1f4a 100644 --- a/modules/ui/PySide6OffloadingWindowView.py +++ b/modules/ui/PySide6OffloadingWindowView.py @@ -1,31 +1,27 @@ from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView from modules.ui.OffloadingWindowController import OffloadingWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton -class CtkOffloadingWindowView(BaseOffloadingWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: OffloadingWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseOffloadingWindowView.__init__(self, ctk_components) +class PySide6OffloadingWindowView(BaseOffloadingWindowView, QDialog): + def __init__(self, parent, controller: OffloadingWindowController, ui_state): + QDialog.__init__(self, parent) + BaseOffloadingWindowView.__init__(self, pyside6_components) - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) + self.setWindowTitle("Offloading") + self.resize(800, 400) - frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + + scroll, frame = pyside6_components.scrollable_frame(self) + pyside6_components._layout(frame).setColumnStretch(0, 1) + pyside6_components._layout(frame).setColumnStretch(1, 1) self.build_content(frame, controller, ui_state) - frame.pack(fill="both", expand=1) - frame.grid(row=0, column=0, sticky='nsew') - self.components.button(self, 1, 0, "ok", self.destroy) + outer.addWidget(scroll, 0, 0) - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) diff --git a/modules/ui/PySide6OptimizerParamsWindowView.py b/modules/ui/PySide6OptimizerParamsWindowView.py index ecc1f6a38..bb6682f01 100644 --- a/modules/ui/PySide6OptimizerParamsWindowView.py +++ b/modules/ui/PySide6OptimizerParamsWindowView.py @@ -1,66 +1,54 @@ -import contextlib -from tkinter import TclError - from modules.ui.BaseOptimizerParamsWindowView import BaseOptimizerParamsWindowView -from modules.ui.CtkMuonAdamWindowView import CtkMuonAdamWindowView from modules.ui.MuonAdamWindowController import MuonAdamWindowController from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import set_window_icon +from modules.ui.PySide6MuonAdamWindowView import PySide6MuonAdamWindowView +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QWidget -class CtkOptimizerParamsWindowView(BaseOptimizerParamsWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: OptimizerParamsWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseOptimizerParamsWindowView.__init__(self, ctk_components) +class PySide6OptimizerParamsWindowView(BaseOptimizerParamsWindowView, QDialog): + def __init__(self, parent, controller: OptimizerParamsWindowController, ui_state): + QDialog.__init__(self, parent) + BaseOptimizerParamsWindowView.__init__(self, pyside6_components) self.controller = controller self.ui_state = ui_state self.optimizer_ui_state = ui_state.get_var("optimizer") self.muon_adam_button = None - self.protocol("WM_DELETE_WINDOW", self.on_window_close) + self._dynamic_frame = None - self.title("Optimizer Settings") - self.geometry("800x500") - self.resizable(True, True) + self.setWindowTitle("Optimizer Settings") + self.resize(800, 500) - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(0, 1) - self.frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - self.frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) + scroll, self._frame = pyside6_components.scrollable_frame(self) + lo = pyside6_components._layout(self._frame) + lo.setColumnStretch(1, 1) + lo.setColumnMinimumWidth(2, 50) + lo.setColumnStretch(4, 1) + outer.addWidget(scroll, 0, 0) - self.frame.grid_columnconfigure(0, weight=0) - self.frame.grid_columnconfigure(1, weight=1) - self.frame.grid_columnconfigure(2, minsize=50) - self.frame.grid_columnconfigure(3, weight=0) - self.frame.grid_columnconfigure(4, weight=1) + ok = QPushButton("ok", self) + ok.clicked.connect(self._on_close) + outer.addWidget(ok, 1, 0) - self.components.button(self, 1, 0, "ok", command=self.on_window_close) - self.build_content(self.frame, controller, ui_state, self.optimizer_ui_state, + self.build_content(self._frame, controller, ui_state, self.optimizer_ui_state, self.on_optimizer_change, self._load_defaults) self._rebuild_dynamic_ui() - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) def _rebuild_dynamic_ui(self): - with contextlib.suppress(TclError): - for widget in self.frame.winfo_children(): - grid_info = widget.grid_info() - if int(grid_info["row"]) >= 1: - widget.destroy() + if self._dynamic_frame is not None: + self._dynamic_frame.setParent(None) - if not self.winfo_exists(): - return + self._dynamic_frame = QWidget(self._frame) + pyside6_components._layout(self._frame).addWidget(self._dynamic_frame, 1, 0, 1, 5) - self.build_dynamic_content(self.frame, self.controller, self.optimizer_ui_state, + self.build_dynamic_content(self._dynamic_frame, self.controller, self.optimizer_ui_state, self.update_user_pref, self.open_muon_adam_window) self.toggle_muon_adam_button() @@ -75,17 +63,17 @@ def on_optimizer_change(self, *args): def _load_defaults(self, *args): self.controller.load_defaults(self.ui_state) - def on_window_close(self): - self.destroy() + def _on_close(self): + self.controller.on_close() + self.accept() def toggle_muon_adam_button(self): - if self.muon_adam_button and self.muon_adam_button.winfo_exists(): + if self.muon_adam_button is not None: muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() - self.muon_adam_button.configure(state="normal" if muon_with_adam else "disabled") + self.muon_adam_button.setEnabled(bool(muon_with_adam)) def open_muon_adam_window(self): adam_config, current_optimizer = self.controller.prepare_muon_adam_config() - temp_adam_ui_state = CtkUIState(self, adam_config) - window = CtkMuonAdamWindowView(self, MuonAdamWindowController(self.controller.config, current_optimizer), temp_adam_ui_state) - self.wait_window(window) + adam_ui_state = PySide6UIState(adam_config) + PySide6MuonAdamWindowView(self, MuonAdamWindowController(self.controller.config, current_optimizer), adam_ui_state).exec() self.controller.save_muon_adam_config(adam_config) diff --git a/modules/ui/PySide6ProfilingWindowView.py b/modules/ui/PySide6ProfilingWindowView.py index 15d5055a0..b8c0a41e1 100644 --- a/modules/ui/PySide6ProfilingWindowView.py +++ b/modules/ui/PySide6ProfilingWindowView.py @@ -1,49 +1,49 @@ +import contextlib from modules.ui.BaseProfilingWindowView import BaseProfilingWindowView from modules.ui.ProfilingWindowController import ProfilingWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtCore import Qt +from PySide6.QtWidgets import QGridLayout, QWidget -class CtkProfilingWindowView(BaseProfilingWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: ProfilingWindowController, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseProfilingWindowView.__init__(self, ctk_components) +class PySide6ProfilingWindowView(BaseProfilingWindowView, QWidget): + def __init__(self, parent, controller: ProfilingWindowController): + QWidget.__init__(self, parent, Qt.WindowType.Window) + BaseProfilingWindowView.__init__(self, pyside6_components) self._controller = controller - self.title("Profiling") - self.geometry("512x512") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() + self.setWindowTitle("Profiling") + self.resize(512, 512) - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=0) - self.grid_rowconfigure(2, weight=1) - self.grid_columnconfigure(0, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(2, 1) - # Bottom bar - self._bottom_bar = ctk.CTkFrame(master=self, corner_radius=0) - self._bottom_bar.grid(row=2, column=0, sticky="sew") + self._bottom_bar = QWidget(self) + QGridLayout(self._bottom_bar) + outer.addWidget(self._bottom_bar, 3, 0) self.build_content(self, self._bottom_bar, controller) - self.protocol("WM_DELETE_WINDOW", self.withdraw) - self.withdraw() - self.after(200, lambda: set_window_icon(self)) + def set_message(self, text: str): + self._message_label.setText(text) - def set_message(self, text): - self._message_label.configure(text=text) - - def set_profiling_active(self, active): + def set_profiling_active(self, active: bool): if active: - self._message_label.configure(text='Profiling active...') - self._profile_button.configure(text='End Profiling') - self._profile_button.configure(command=self._controller.end_profiler) + self._message_label.setText("Profiling active...") + self._profile_button.setText("End Profiling") + with contextlib.suppress(RuntimeError): + self._profile_button.clicked.disconnect() + self._profile_button.clicked.connect(self._controller.end_profiler) else: - self._message_label.configure(text='Inactive') - self._profile_button.configure(text='Start Profiling') - self._profile_button.configure(command=self._controller.start_profiler) + self._message_label.setText("Inactive") + self._profile_button.setText("Start Profiling") + with contextlib.suppress(RuntimeError): + self._profile_button.clicked.disconnect() + self._profile_button.clicked.connect(self._controller.start_profiler) + + def closeEvent(self, event): + event.ignore() + self.hide() diff --git a/modules/ui/PySide6SampleFrameView.py b/modules/ui/PySide6SampleFrameView.py index 167b25692..a936cccb3 100644 --- a/modules/ui/PySide6SampleFrameView.py +++ b/modules/ui/PySide6SampleFrameView.py @@ -1,43 +1,46 @@ from modules.ui.BaseSampleFrameView import BaseSampleFrameView from modules.ui.SampleFrameController import SampleFrameController -from modules.util.ui import ctk_components +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QWidget -class CtkSampleFrameView(BaseSampleFrameView, ctk.CTkFrame): +class PySide6SampleFrameView(BaseSampleFrameView, QWidget): def __init__( self, - parent, + parent: QWidget, controller: SampleFrameController, ui_state, include_prompt: bool = True, include_settings: bool = True, ): - ctk.CTkFrame.__init__(self, parent, fg_color="transparent") - BaseSampleFrameView.__init__(self, ctk_components) + QWidget.__init__(self, parent) + BaseSampleFrameView.__init__(self, pyside6_components) - if include_prompt and include_prompt: - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_columnconfigure(0, weight=1) + outer = pyside6_components._layout(self) + outer.setColumnStretch(0, 1) + if include_prompt and include_settings: + outer.setRowStretch(1, 1) top_frame = None if include_prompt: - top_frame = ctk.CTkFrame(self, fg_color="transparent") - top_frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") - - top_frame.grid_columnconfigure(0, weight=0) - top_frame.grid_columnconfigure(1, weight=1) + top_frame = QWidget(self) + top_lo = pyside6_components._layout(top_frame) + top_lo.setColumnStretch(1, 1) + outer.addWidget(top_frame, 0, 0) bottom_frame = None if include_settings: - bottom_frame = ctk.CTkFrame(self, fg_color="transparent") - bottom_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - bottom_frame.grid_columnconfigure(0, weight=0) - bottom_frame.grid_columnconfigure(1, weight=1) - bottom_frame.grid_columnconfigure(2, weight=0) - bottom_frame.grid_columnconfigure(3, weight=1) + bottom_frame = QWidget(self) + bot_lo = pyside6_components._layout(bottom_frame) + bot_lo.setColumnStretch(1, 1) + bot_lo.setColumnStretch(3, 1) + row = 1 if include_prompt else 0 + outer.addWidget(bottom_frame, row, 0) self.build_content(top_frame, bottom_frame, ui_state, controller, include_prompt, include_settings) + + if top_frame is not None: + pyside6_components._pack_form(top_frame) + if bottom_frame is not None: + pyside6_components._pack_form(bottom_frame) diff --git a/modules/ui/PySide6SampleParamsWindowView.py b/modules/ui/PySide6SampleParamsWindowView.py index 8229a19f6..054712ccc 100644 --- a/modules/ui/PySide6SampleParamsWindowView.py +++ b/modules/ui/PySide6SampleParamsWindowView.py @@ -1,32 +1,27 @@ from modules.ui.BaseSampleParamsWindowView import BaseSampleParamsWindowView -from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.PySide6SampleFrameView import PySide6SampleFrameView from modules.ui.SampleFrameController import SampleFrameController from modules.ui.SampleParamsWindowController import SampleParamsWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QWidget -class CtkSampleParamsWindowView(BaseSampleParamsWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: SampleParamsWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseSampleParamsWindowView.__init__(self, ctk_components) +class PySide6SampleParamsWindowView(BaseSampleParamsWindowView, QDialog): + def __init__(self, parent, controller: SampleParamsWindowController, ui_state): + QDialog.__init__(self, parent if isinstance(parent, QWidget) else None) + BaseSampleParamsWindowView.__init__(self, pyside6_components) - self.title("Sample") - self.geometry("800x500") - self.resizable(True, True) + self.setWindowTitle("Sample") + self.resize(800, 500) - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + outer.setColumnStretch(0, 1) - frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) - frame.grid(row=0, column=0, padx=0, pady=0, sticky="nsew") + frame = PySide6SampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) + outer.addWidget(frame, 0, 0) - self.components.button(self, 1, 0, "ok", self.destroy) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) diff --git a/modules/ui/PySide6SampleWindowView.py b/modules/ui/PySide6SampleWindowView.py index 3b67f03d5..a41777f1d 100644 --- a/modules/ui/PySide6SampleWindowView.py +++ b/modules/ui/PySide6SampleWindowView.py @@ -1,102 +1,85 @@ -import contextlib -import tkinter as tk -import traceback +import threading from modules.modelSampler.BaseModelSampler import ( ModelSamplerOutput, ) from modules.ui.BaseSampleWindowView import BaseSampleWindowView -from modules.ui.CtkSampleFrameView import CtkSampleFrameView +from modules.ui.PySide6SampleFrameView import PySide6SampleFrameView from modules.ui.SampleFrameController import SampleFrameController from modules.ui.SampleWindowController import SampleWindowController from modules.util.enum.FileType import FileType -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk -from PIL import Image +from PIL.ImageQt import ImageQt +from PySide6.QtCore import Qt, QTimer +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import QDialog, QGridLayout, QLabel, QProgressBar, QPushButton -class CtkSampleWindowView(BaseSampleWindowView, ctk.CTkToplevel): - def __init__( - self, - parent, - controller: SampleWindowController, - *args, **kwargs - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseSampleWindowView.__init__(self, ctk_components) +class PySide6SampleWindowView(BaseSampleWindowView, QDialog): + def __init__(self, parent, controller: SampleWindowController): + QDialog.__init__(self, parent) + BaseSampleWindowView.__init__(self, pyside6_components) - self.title("Sample") - self.geometry("1200x800") - self.resizable(True, True) + self.setWindowTitle("Sample") + self.resize(1200, 800) - model_type = controller.get_model_type() - self.ui_state = CtkUIState(self, controller.sample) + self.ui_state = PySide6UIState(controller.sample) if controller.use_external_model: - controller.callbacks.set_on_sample_custom(self.__update_preview) - controller.callbacks.set_on_update_sample_custom_progress(self.__update_progress) - - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_rowconfigure(3, weight=0) - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1) - - prompt_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_settings=False) - prompt_frame.grid(row=0, column=0, columnspan=2, padx=0, pady=0, sticky="nsew") - - settings_frame = CtkSampleFrameView(self, SampleFrameController(controller.sample, model_type), self.ui_state, include_prompt=False) - settings_frame.grid(row=1, column=0, padx=0, pady=0, sticky="nsew") - - # image - self.image = ctk.CTkImage( - light_image=self.__dummy_image(), - size=(512, 512) - ) + controller.callbacks.set_on_sample_custom(self._update_preview) + controller.callbacks.set_on_update_sample_custom_progress(self._update_progress) - image_label = ctk.CTkLabel(master=self, text="", image=self.image, height=512, width=512) - image_label.grid(row=1, column=1, rowspan=3, sticky="nsew") + outer = QGridLayout(self) + outer.setRowStretch(1, 1) + outer.setColumnStretch(1, 1) - self.progress = self.components.progress(self, 2, 0) - self.components.button(self, 3, 0, "sample", - lambda: controller.do_sample(self.__update_preview, self.__update_progress)) + model_type = controller.get_model_type() + frame_controller = SampleFrameController(controller.sample, model_type) + + prompt_frame = PySide6SampleFrameView(self, frame_controller, self.ui_state, include_settings=False) + outer.addWidget(prompt_frame, 0, 0, 1, 2) + + settings_frame = PySide6SampleFrameView(self, frame_controller, self.ui_state, include_prompt=False) + outer.addWidget(settings_frame, 1, 0) + + self._image_label = QLabel(self) + self._image_label.setFixedSize(512, 512) + self._image_label.setAlignment(Qt.AlignCenter) + self._image_label.setStyleSheet("background: black;") + outer.addWidget(self._image_label, 1, 1, 3, 1) + + self._progress = QProgressBar(self) + self._progress.setRange(0, 1000) + outer.addWidget(self._progress, 2, 0) + + sample_btn = QPushButton("sample", self) + # Run in a background thread so the Qt event loop stays responsive during sampling + sample_btn.clicked.connect( + lambda: threading.Thread( + target=lambda: controller.do_sample(self._update_preview, self._update_progress), + daemon=True, + ).start() + ) + outer.addWidget(sample_btn, 3, 0) - self.wait_visibility() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - def __update_preview(self, sampler_output: ModelSamplerOutput): + def schedule_on_main_thread(self, fn): + QTimer.singleShot(0, self, fn) + + def _update_preview(self, sampler_output: ModelSamplerOutput): + # Called from training thread — capture data and dispatch to main thread if sampler_output.file_type == FileType.IMAGE: image = sampler_output.data - self.image.configure( - light_image=image, - size=(image.width, image.height), - ) - - def __update_progress(self, progress: int, max_progress: int): - self.progress.set(progress / max_progress) - self.update() - - def __dummy_image(self) -> Image: - return Image.new(mode="RGB", size=(512, 512), color=(0, 0, 0)) - - def destroy(self): - try: - if hasattr(self, "_icon_image_ref"): - del self._icon_image_ref - - # Remove any pending after callbacks - for after_id in self.tk.call('after', 'info'): - with contextlib.suppress(tk.TclError, RuntimeError): - self.after_cancel(after_id) - - super().destroy() - except (tk.TclError, RuntimeError) as e: - print(f"Error destroying window: {e}") - except Exception as e: - print(f"Unexpected error destroying window: {e}") - traceback.print_exc() + self.schedule_on_main_thread(lambda: self._do_update_preview(image)) + + def _do_update_preview(self, image): + pixmap = QPixmap.fromImage(ImageQt(image.convert("RGBA"))) + self._image_label.setFixedSize(pixmap.size()) + self._image_label.setPixmap(pixmap) + + def _update_progress(self, progress: int, max_progress: int): + # Called from training thread — dispatch to main thread + value = int(progress / max_progress * 1000) + self.schedule_on_main_thread(lambda: self._progress.setValue(value)) diff --git a/modules/ui/PySide6SamplingTabView.py b/modules/ui/PySide6SamplingTabView.py index dfe1a704a..e83c4a727 100644 --- a/modules/ui/PySide6SamplingTabView.py +++ b/modules/ui/PySide6SamplingTabView.py @@ -1,16 +1,17 @@ from modules.ui.BaseSamplingTabView import BaseSampleWidgetView, BaseSamplingTabView -from modules.ui.CtkConfigListView import CtkConfigListView -from modules.ui.CtkSampleParamsWindowView import CtkSampleParamsWindowView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.ui.PySide6SampleParamsWindowView import PySide6SampleParamsWindowView from modules.ui.SamplingTabController import SamplingTabController -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta -import customtkinter as ctk +from PySide6.QtWidgets import QWidget -class CtkSamplingTabView(CtkConfigListView, BaseSamplingTabView): +class PySide6SamplingTabView(PySide6ConfigListView, BaseSamplingTabView): + def __init__(self, master, controller: SamplingTabController, ui_state): - CtkConfigListView.__init__( + PySide6ConfigListView.__init__( self, master, controller, ui_state, from_external_file=True, attr_name="sample_definition_file_name", @@ -22,29 +23,44 @@ def __init__(self, master, controller: SamplingTabController, ui_state): show_toggle_button=True, ) - def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: - return self.controller.open_element_window(self.master, self.current_config[i], ui_state, CtkSampleParamsWindowView) + def open_element_window(self, i, ui_state): + return self.controller.open_element_window(self.master, self.current_config[i], ui_state, PySide6SampleParamsWindowView) def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return CtkSampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) + return PySide6SampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) + +class PySide6SampleWidgetView(BaseSampleWidgetView, QWidget, metaclass=QtABCMeta): -class CtkSampleWidgetView(BaseSampleWidgetView, ctk.CTkFrame): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - ctk.CTkFrame.__init__(self, master=master, corner_radius=10, bg_color="transparent") - BaseSampleWidgetView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseSampleWidgetView.__init__(self, pyside6_components) - self.ui_state = CtkUIState(self, element) + from modules.util.ui.PySide6UIState import PySide6UIState + self.element = element + self.ui_state = PySide6UIState(element) - self.grid_columnconfigure(10, weight=1) + pyside6_components._layout(self).setColumnStretch(10, 1) self.build_content(self, element, self.ui_state, i, open_command, remove_command, clone_command, save_command) def _bind_save(self, save_command): - self.width_entry.bind('', lambda _: save_command()) - self.height_entry.bind('', lambda _: save_command()) - self.seed_entry.bind('', lambda _: save_command()) - self.prompt_entry.bind('', lambda _: save_command()) + self.width_entry.editingFinished.connect(save_command) + self.height_entry.editingFinished.connect(save_command) + self.seed_entry.editingFinished.connect(save_command) + self.prompt_entry.editingFinished.connect(save_command) + + def _set_enabled(self): + enabled = self.element.enabled + self.width_entry.setEnabled(enabled) + self.height_entry.setEnabled(enabled) + self.prompt_entry.setEnabled(enabled) + self.seed_entry.setEnabled(enabled) + self.button.setEnabled(enabled) def place_in_list(self): - self.grid(row=self.i, column=0, pady=5, padx=5, sticky="new") + pyside6_components._layout(self.parent()).addWidget(self, getattr(self, 'visible_index', self.i), 0) + self.show() + + def destroy(self): + self.deleteLater() diff --git a/modules/ui/PySide6SchedulerParamsWindowView.py b/modules/ui/PySide6SchedulerParamsWindowView.py index 9a7c41b96..96deaae40 100644 --- a/modules/ui/PySide6SchedulerParamsWindowView.py +++ b/modules/ui/PySide6SchedulerParamsWindowView.py @@ -1,96 +1,93 @@ from modules.ui.BaseSchedulerParamsWindowView import BaseKvParamsView, BaseSchedulerParamsWindowView -from modules.ui.CtkConfigListView import CtkConfigListView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView from modules.ui.SchedulerParamsWindowController import KvParamsController, SchedulerParamsWindowController -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QScrollArea, QWidget -class CtkKvParamsView(CtkConfigListView, BaseKvParamsView): +class PySide6KvParamsView(PySide6ConfigListView, BaseKvParamsView): def __init__(self, master, controller: KvParamsController, ui_state): - CtkConfigListView.__init__( + PySide6ConfigListView.__init__( self, master, controller, ui_state, attr_name="scheduler_params", from_external_file=False, add_button_text="add parameter", is_full_width=True, ) - BaseKvParamsView.__init__(self, ctk_components) + BaseKvParamsView.__init__(self, pyside6_components) def refresh_ui(self): self._create_element_list() def create_widget(self, master, element, i, open_command, remove_command, clone_command, save_command): - return KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) + return PySide6KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) -class KvWidget(ctk.CTkFrame): +class PySide6KvWidget(QWidget): def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): - super().__init__(master=master, bg_color="transparent") + super().__init__(master) self.element = element - self.ui_state = CtkUIState(self, element) + self.ui_state = PySide6UIState(element) self.i = i self.save_command = save_command - self.grid_columnconfigure(0, weight=0) - self.grid_columnconfigure(1, weight=1, uniform=1) - self.grid_columnconfigure(2, weight=1, uniform=1) + lo = pyside6_components._layout(self) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(2, 1) - close_button = ctk.CTkButton( - master=self, - width=20, - height=20, - text="X", - corner_radius=2, - fg_color="#C00000", - command=lambda: remove_command(self.i)) - close_button.grid(row=0, column=0) + pyside6_components.colored_icon_button(self, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) # Key - tooltip_key = "Key name for an argument in your scheduler" - self.key = ctk_components.entry(self, 0, 1, self.ui_state, "key", - tooltip=tooltip_key, wide_tooltip=True) - self.key.bind("", lambda _: save_command()) - self.key.configure(width=50) + self.key = pyside6_components.entry(self, 0, 1, self.ui_state, "key", + tooltip="Key name for an argument in your scheduler", + wide_tooltip=True, width=50) + self.key.editingFinished.connect(save_command) # Value - tooltip_val = "Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'." - self.value = ctk_components.entry(self, 0, 2, self.ui_state, "value", - tooltip=tooltip_val, wide_tooltip=True) - self.value.bind("", lambda _: save_command()) - self.value.configure(width=50) + self.value = pyside6_components.entry(self, 0, 2, self.ui_state, "value", + tooltip="Value for an argument in your scheduler. Some special values can be used, wrapped in percent signs: LR, EPOCHS, STEPS_PER_EPOCH, TOTAL_STEPS, SCHEDULER_STEPS. Note that OneTrainer calls step() after every individual learning step, not every epoch, so what Torch calls 'epoch' you should treat as 'step'.", + wide_tooltip=True, width=50) + self.value.editingFinished.connect(save_command) def place_in_list(self): - self.grid(row=self.i, column=0, padx=5, pady=5, sticky="new") + pyside6_components._layout(self.parent()).addWidget(self, getattr(self, 'visible_index', self.i), 0) + self.show() + def destroy(self): + self.deleteLater() -class CtkSchedulerParamsWindowView(BaseSchedulerParamsWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseSchedulerParamsWindowView.__init__(self, ctk_components) - self.title("Learning Rate Scheduler Settings") - self.geometry("800x400") - self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) +class PySide6SchedulerParamsWindowView(BaseSchedulerParamsWindowView, QDialog): + def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state): + QDialog.__init__(self, parent) + BaseSchedulerParamsWindowView.__init__(self, pyside6_components) - frame = ctk.CTkFrame(self) - frame.grid(row=0, column=0, sticky="nsew", padx=10, pady=10) - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) + self.setWindowTitle("Learning Rate Scheduler Settings") + self.resize(800, 500) - expand_frame = ctk.CTkFrame(frame, bg_color="transparent") - expand_frame.grid(row=1, column=0, columnspan=2, sticky="nsew") + outer = QGridLayout(self) + outer.setRowStretch(0, 1) - self.components.button(self, 1, 0, "ok", command=self.destroy) - self.build_content(frame, controller, ui_state) - CtkKvParamsView(expand_frame, KvParamsController(controller.config), ui_state) + scroll = QScrollArea(self) + scroll.setWidgetResizable(True) + inner = QWidget() + scroll.setWidget(inner) + inner_lo = pyside6_components._layout(inner) + inner_lo.setColumnStretch(1, 1) - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) + self.build_content(inner, controller, ui_state) + + expand_frame = QWidget(inner) + inner_lo.addWidget(expand_frame, inner_lo.rowCount(), 0, 1, 2) + # Must be assigned to an instance variable — PySide6ConfigListView is not a QWidget, + # so Qt won't keep it alive. Without this, the GC collects it and the button's + # clicked signal loses its connection to __add_element. + self._kv_params_view = PySide6KvParamsView(expand_frame, KvParamsController(controller.config), ui_state) + + outer.addWidget(scroll, 0, 0) + + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) diff --git a/modules/ui/PySide6TimestepDistributionWindowView.py b/modules/ui/PySide6TimestepDistributionWindowView.py index 69f2ae7d3..3cd11d9ef 100644 --- a/modules/ui/PySide6TimestepDistributionWindowView.py +++ b/modules/ui/PySide6TimestepDistributionWindowView.py @@ -1,84 +1,48 @@ - from modules.ui.BaseTimestepDistributionWindowView import BaseTimestepDistributionWindowView from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components -import customtkinter as ctk -from customtkinter import AppearanceModeTracker, ThemeManager from matplotlib import pyplot as plt -from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg - +from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton -class CtkTimestepDistributionWindowView(BaseTimestepDistributionWindowView, ctk.CTkToplevel): - def __init__( - self, - parent, - controller: TimestepDistributionWindowController, - ui_state, - *args, **kwargs, - ): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseTimestepDistributionWindowView.__init__(self, ctk_components) - self.title("Timestep Distribution") - self.geometry("900x600") - self.resizable(True, True) +class PySide6TimestepDistributionWindowView(BaseTimestepDistributionWindowView, QDialog): + def __init__(self, parent, controller: TimestepDistributionWindowController, ui_state): + QDialog.__init__(self, parent) + BaseTimestepDistributionWindowView.__init__(self, pyside6_components) - self.controller = controller - self.ax = None - self.canvas = None + self.setWindowTitle("Timestep Distribution") + self.resize(900, 600) + self._controller = controller - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) + outer = QGridLayout(self) + outer.setRowStretch(0, 1) - frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=0) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - frame.grid_rowconfigure(7, weight=1) + scroll, frame = pyside6_components.scrollable_frame(self) + lo = pyside6_components._layout(frame) + lo.setColumnStretch(3, 1) self.build_content(frame, controller, ui_state) + lo.setRowStretch(7, 1) - # matplotlib chart (CTK-only: needs winfo_rgb from the toplevel) - appearance_mode = AppearanceModeTracker.get_mode() - background_color = self.winfo_rgb(ThemeManager.theme["CTkToplevel"]["fg_color"][appearance_mode]) - text_color = self.winfo_rgb(ThemeManager.theme["CTkLabel"]["text_color"][appearance_mode]) - background_color = f"#{int(background_color[0]/256):x}{int(background_color[1]/256):x}{int(background_color[2]/256):x}" - text_color = f"#{int(text_color[0]/256):x}{int(text_color[1]/256):x}{int(text_color[2]/256):x}" - - fig, ax = plt.subplots() - self.ax = ax - self.canvas = FigureCanvasTkAgg(fig, master=frame) - self.canvas.get_tk_widget().grid(row=0, column=3, rowspan=8) - - fig.set_facecolor(background_color) - ax.set_facecolor(background_color) - ax.spines['bottom'].set_color(text_color) - ax.spines['left'].set_color(text_color) - ax.spines['top'].set_color(text_color) - ax.spines['right'].set_color(text_color) - ax.tick_params(axis='x', colors=text_color, which="both") - ax.tick_params(axis='y', colors=text_color, which="both") - ax.xaxis.label.set_color(text_color) - ax.yaxis.label.set_color(text_color) + fig, self._ax = plt.subplots() + self._canvas = FigureCanvasQTAgg(fig) + lo.addWidget(self._canvas, 0, 3, 8, 1) + self._update_preview() - self.__update_preview() + update_btn = QPushButton("Update Preview", frame) + update_btn.clicked.connect(self._update_preview) + lo.addWidget(update_btn, 8, 3) - # update button - ctk_components.button(frame, 8, 3, "Update Preview", command=self.__update_preview) + outer.addWidget(scroll, 0, 0) - frame.pack(fill="both", expand=1) - frame.grid(row=0, column=0, sticky='nsew') - ctk_components.button(self, 1, 0, "ok", self.destroy) + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) - self.wait_visibility() - self.after(200, lambda: set_window_icon(self)) - self.grab_set() - self.focus_set() - def __update_preview(self): - self.ax.cla() - self.ax.hist(self.controller.generate_preview_data(), bins=1000, range=(0, 999)) - self.canvas.draw() + def _update_preview(self): + self._ax.cla() + self._ax.hist(self._controller.generate_preview_data(), bins=1000, range=(0, 999)) + self._canvas.draw() diff --git a/modules/ui/PySide6TopBarView.py b/modules/ui/PySide6TopBarView.py index ee1bcfa75..e852723c8 100644 --- a/modules/ui/PySide6TopBarView.py +++ b/modules/ui/PySide6TopBarView.py @@ -4,13 +4,13 @@ from modules.ui.TopBarController import TopBarController from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import ctk_components, dialogs -from modules.util.ui.CtkUIState import CtkUIState +from modules.util.ui import pyside6_components -import customtkinter as ctk +from PySide6.QtWidgets import QInputDialog, QWidget -class CtkTopBarView(BaseTopBarView): +class PySide6TopBarView(BaseTopBarView, QWidget): + def __init__( self, master, @@ -20,31 +20,35 @@ def __init__( change_training_method_callback: Callable[[TrainingMethod], None], load_preset_callback: Callable[[], None], ): - BaseTopBarView.__init__(self, ctk_components) - - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=0, column=0, sticky="nsew") + QWidget.__init__(self, master) + BaseTopBarView.__init__(self, pyside6_components) + + self.frame = QWidget(self) + pyside6_components._layout(self).addWidget(self.frame, 0, 0) + pyside6_components._layout(self.frame).setContentsMargins( + pyside6_components.PAD, pyside6_components.PAD, + pyside6_components.PAD, pyside6_components.PAD, + ) - self.build(frame, master, controller, ui_state, change_model_type_callback, change_training_method_callback, load_preset_callback) + self.build(self.frame, master, controller, ui_state, + change_model_type_callback, change_training_method_callback, load_preset_callback) def _make_config_ui_state(self, master, data): - return CtkUIState(master, data) + from modules.util.ui.PySide6UIState import PySide6UIState + return PySide6UIState(data) def _get_dropdown_text(self, widget) -> str: - return widget.get() + return widget.currentText() def _setup_frame_column_weight(self): - self.frame.grid_columnconfigure(5, weight=1) + pyside6_components._layout(self.frame).setColumnStretch(5, 1) def _forget_dropdown(self): - self.configs_dropdown.grid_forget() + pyside6_components._layout(self.frame).removeWidget(self.configs_dropdown) + self.configs_dropdown.hide() + self.configs_dropdown.deleteLater() def _show_save_dialog(self, default_value: str, callback): - dialogs.StringInputDialog( - parent=self.master, - title="name", - question="Config Name", - callback=callback, - default_value=default_value, - validate_callback=lambda x: not x.startswith("#"), - ) + text, ok = QInputDialog.getText(self, "name", "Config Name", text=default_value) + if ok and not text.startswith("#"): + callback(text) diff --git a/modules/ui/PySide6TrainUIView.py b/modules/ui/PySide6TrainUIView.py index d3c1a70b5..6481a1420 100644 --- a/modules/ui/PySide6TrainUIView.py +++ b/modules/ui/PySide6TrainUIView.py @@ -3,28 +3,27 @@ from collections.abc import Callable from contextlib import suppress from pathlib import Path -from tkinter import filedialog, messagebox from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController from modules.ui.BaseTrainUIView import BaseTrainUIView from modules.ui.CloudTabController import CloudTabController from modules.ui.ConceptTabController import ConceptTabController -from modules.ui.CtkAdditionalEmbeddingsTabView import CtkAdditionalEmbeddingsTabView -from modules.ui.CtkCaptionUIView import CtkCaptionUIView -from modules.ui.CtkCloudTabView import CtkCloudTabView -from modules.ui.CtkConceptTabView import CtkConceptTabView -from modules.ui.CtkConvertModelUIView import CtkConvertModelUIView -from modules.ui.CtkLoraTabView import CtkLoraTabView -from modules.ui.CtkModelTabView import CtkModelTabView -from modules.ui.CtkProfilingWindowView import CtkProfilingWindowView -from modules.ui.CtkSampleWindowView import CtkSampleWindowView -from modules.ui.CtkSamplingTabView import CtkSamplingTabView -from modules.ui.CtkTopBarView import CtkTopBarView -from modules.ui.CtkTrainingTabView import CtkTrainingTabView -from modules.ui.CtkVideoToolUIView import CtkVideoToolUIView from modules.ui.LoraTabController import LoraTabController from modules.ui.ModelTabController import ModelTabController from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.ui.PySide6AdditionalEmbeddingsTabView import PySide6AdditionalEmbeddingsTabView +from modules.ui.PySide6CaptionUIView import PySide6CaptionUIView +from modules.ui.PySide6CloudTabView import PySide6CloudTabView +from modules.ui.PySide6ConceptTabView import PySide6ConceptTabView +from modules.ui.PySide6ConvertModelUIView import PySide6ConvertModelUIView +from modules.ui.PySide6LoraTabView import PySide6LoraTabView +from modules.ui.PySide6ModelTabView import PySide6ModelTabView +from modules.ui.PySide6ProfilingWindowView import PySide6ProfilingWindowView +from modules.ui.PySide6SampleWindowView import PySide6SampleWindowView +from modules.ui.PySide6SamplingTabView import PySide6SamplingTabView +from modules.ui.PySide6TopBarView import PySide6TopBarView +from modules.ui.PySide6TrainingTabView import PySide6TrainingTabView +from modules.ui.PySide6VideoToolUIView import PySide6VideoToolUIView from modules.ui.SamplingTabController import SamplingTabController from modules.ui.TopBarController import TopBarController from modules.ui.TrainingTabController import TrainingTabController @@ -32,12 +31,13 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ui_utils import set_window_icon +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta +from modules.util.ui.PySide6UIState import PySide6UIState -import customtkinter as ctk -from customtkinter import AppearanceModeTracker +from PySide6.QtCore import QTimer +from PySide6.QtGui import QIcon +from PySide6.QtWidgets import QFileDialog, QGridLayout, QMainWindow, QMessageBox, QTabWidget, QWidget # chunk for forcing Windows to ignore DPI scaling when moving between monitors # fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 @@ -47,121 +47,101 @@ ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE -class CtkTrainUIView(BaseTrainUIView, ctk.CTk): - set_step_progress: Callable[[int, int], None] - set_epoch_progress: Callable[[int, int], None] - - status_label: ctk.CTkLabel | None - training_button: ctk.CTkButton | None - - _TRAIN_BUTTON_STYLES = { - "idle": { - "text": "Start Training", - "state": "normal", - "fg_color": "#198754", - "hover_color": "#146c43", - "text_color": "white", - "text_color_disabled": "white", - }, - "running": { - "text": "Stop Training", - "state": "normal", - "fg_color": "#dc3545", - "hover_color": "#bb2d3b", - "text_color": "white", - }, - "stopping": { - "text": "Stopping...", - "state": "disabled", - "fg_color": "#dc3545", - "hover_color": "#dc3545", - "text_color": "white", - "text_color_disabled": "white", - }, - } - def __init__(self): - ctk.CTk.__init__(self) - BaseTrainUIView.__init__(self, ctk_components) - - self.title("OneTrainer") - self.geometry("1100x740") - self.after(100, lambda: self._set_icon()) +class PySide6TrainView(BaseTrainUIView, QMainWindow, metaclass=QtABCMeta): + def __init__(self): + QMainWindow.__init__(self) + BaseTrainUIView.__init__(self, pyside6_components) - # more efficient version of ctk.set_appearance_mode("System"), which retrieves the system theme on each main loop iteration - ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") - ctk.set_default_color_theme("blue") + self.setWindowTitle("OneTrainer") + self.setWindowIcon(QIcon("resources/icons/icon.png")) + self.resize(1100, 740) self.train_config = TrainConfig.default_values() - self.ui_state = CtkUIState(self, self.train_config) + self.ui_state = PySide6UIState(self.train_config) self.controller = TrainUIController(self.train_config) self.controller.view = self - self.grid_rowconfigure(0, weight=0) - self.grid_rowconfigure(1, weight=1) - self.grid_rowconfigure(2, weight=0) - self.grid_columnconfigure(0, weight=1) - self.status_label = None self.eta_label = None self.training_button = None self.export_button = None - self.tabview = None + self.tabview: QTabWidget | None = None + self._tab_widgets: dict[str, QWidget] = {} self.model_tab = None self.training_tab = None self.lora_tab = None self.cloud_tab = None + self.concepts_tab = None + self.sampling_tab = None self.additional_embeddings_tab = None - self.top_bar_component = self.top_bar(self) - self.content_frame(self) - self.bottom_bar(self) + central = QWidget(self) + self.setCentralWidget(central) + central_lo = QGridLayout(central) + central_lo.setContentsMargins(0, 0, 0, 0) + central_lo.setSpacing(0) + central_lo.setRowStretch(1, 1) + central_lo.setColumnStretch(0, 1) - self.controller._check_start_always_on_tensorboard() + self.top_bar_component = self._build_top_bar(central) + central_lo.addWidget(self.top_bar_component, 0, 0) + + self.tabview = QTabWidget(central) + central_lo.addWidget(self.tabview, 1, 0) - self.workspace_dir_trace_id = self.ui_state.add_var_trace("workspace_dir", self.controller._on_workspace_dir_change_trace) + bottom = self._build_bottom_bar(central) + central_lo.addWidget(bottom, 2, 0) + + self._create_tabs() + self.change_training_method(self.train_config.training_method) - # Persistent profiling window. self._profiling_controller = ProfilingWindowController() - self.profiling_window = self._profiling_controller.create_window(self, CtkProfilingWindowView) + self.profiling_window = PySide6ProfilingWindowView(self, self._profiling_controller) - self.protocol("WM_DELETE_WINDOW", self.__close) + self.controller._check_start_always_on_tensorboard() + self.workspace_dir_trace_id = self.ui_state.add_var_trace( + "workspace_dir", self.controller._on_workspace_dir_change_trace + ) - def __close(self): + def closeEvent(self, event): self.top_bar_component.save_default() self.controller._stop_always_on_tensorboard() - if hasattr(self, 'workspace_dir_trace_id'): - self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) - self.quit() + self.ui_state.remove_var_trace("workspace_dir", self.workspace_dir_trace_id) + event.accept() # --- BaseTrainUIView abstract method implementations --- def on_update_status(self, status: str): - self.status_label.configure(text=status) + # Called from training thread — defer to main thread + self.schedule_on_main_thread(lambda: self.status_label.setText(status)) def on_training_started(self): self._set_training_button_style("running") def on_training_stopped(self, error_caught: bool): - self.eta_label.configure(text="") + self.eta_label.setText("") self._set_training_button_style("idle") def on_training_stopping(self): self._set_training_button_style("stopping") def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): + # Called from training thread — defer to main thread + self.schedule_on_main_thread(lambda: self._do_update_progress(epoch_step, max_step, epoch, max_epoch, eta_str)) + + def _do_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): self.set_step_progress(epoch_step, max_step) self.set_epoch_progress(epoch, max_epoch) - if eta_str is not None: - self.eta_label.configure(text=f"ETA: {eta_str}") - else: - self.eta_label.configure(text="") + self.eta_label.setText(f"ETA: {eta_str}" if eta_str is not None else "") def schedule_on_main_thread(self, fn: Callable): - self.after(0, fn) + # The 3-argument form (msec, context, fn) is thread-safe: Qt marshals the call + # to the thread where `self` lives (the main thread), unlike the 2-arg form. + QTimer.singleShot(0, self, fn) def get_cloud_reattach(self) -> bool: return self.cloud_tab.reattach @@ -174,43 +154,37 @@ def save_default(self): def show_validation_errors(self, errors: list[str]): bullet_list = "\n".join(f"• {e}" for e in errors) - messagebox.showerror( - "Cannot Start Training", - f"Please fix the following errors before training:\n\n{bullet_list}", - ) + QMessageBox.critical(self, "Cannot Start Training", + f"Please fix the following errors before training:\n\n{bullet_list}") def open_dataset_tool(self): - self.wait_window(self.controller.open_dataset_tool(self, CtkCaptionUIView)) + self.wait_window(self.controller.open_dataset_tool(self, PySide6CaptionUIView)) def open_video_tool(self): - self.wait_window(self.controller.open_video_tool(self, CtkVideoToolUIView)) + self.wait_window(self.controller.open_video_tool(self, PySide6VideoToolUIView)) def open_convert_model_tool(self): - self.wait_window(self.controller.open_convert_model_tool(self, CtkConvertModelUIView)) + self.wait_window(self.controller.open_convert_model_tool(self, PySide6ConvertModelUIView)) def open_sampling_tool(self): - self.controller.open_sampling_tool(self, CtkSampleWindowView) + self.controller.open_sampling_tool(self, PySide6SampleWindowView) def open_manual_sample_window(self): - self.controller.open_manual_sample_window(self, CtkSampleWindowView) + self.controller.open_manual_sample_window(self, PySide6SampleWindowView) def wait_window(self, window): - ctk.CTk.wait_window(self, window) + window.exec() def show_window(self, window): - window.focus_set() + window.show() def connect_window_closed(self, window, callback): - window.bind("", lambda _: callback()) - - # --- CTK layout and frame builders --- + window.finished.connect(lambda _: callback()) - def _set_icon(self): - """Set the window icon safely after window is ready""" - set_window_icon(self) + # --- PySide6 layout builders --- - def top_bar(self, master): - return CtkTopBarView( + def _build_top_bar(self, master): + return PySide6TopBarView( master, TopBarController(self.train_config), self.ui_state, @@ -219,146 +193,143 @@ def top_bar(self, master): self.load_preset, ) - def bottom_bar(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=2, column=0, sticky="nsew") + def _build_bottom_bar(self, parent): + frame = QWidget(parent) + lo = QGridLayout(frame) + lo.setColumnStretch(0, 1) + lo.setColumnStretch(2, 2) - # status + ETA container - status_frame = ctk.CTkFrame(frame, corner_radius=0, fg_color="transparent") - status_frame.grid(row=0, column=1, sticky="w") - status_frame.grid_rowconfigure(0, weight=0) - status_frame.grid_rowconfigure(1, weight=0) - status_frame.grid_columnconfigure(0, weight=1) - - # padding - frame.grid_columnconfigure(2, weight=1) + status_frame = QWidget(frame) + status_lo = QGridLayout(status_frame) + status_lo.setContentsMargins(0, 0, 0, 0) + lo.addWidget(status_frame, 0, 1) self.build_bottom_bar_content(frame, status_frame, self.controller, self.ui_state) - self._set_training_button_style("idle") # centralized styling - + self._set_training_button_style("idle") return frame - def content_frame(self, master): - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") + def _create_scrollable_tab(self, configure_fn): + tab_page = QWidget() + tab_lo = pyside6_components._layout(tab_page) + tab_lo.setRowStretch(0, 1) + tab_lo.setColumnStretch(0, 1) + scroll, frame = pyside6_components.scrollable_frame(tab_page) + tab_lo.addWidget(scroll, 0, 0) + configure_fn(frame) + return tab_page + + def _configure_general_frame(self, frame): + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(3, 1) + self.build_general_tab_content(frame, self.controller, self.ui_state) - frame.grid_rowconfigure(0, weight=1) - frame.grid_columnconfigure(0, weight=1) + def _configure_data_frame(self, frame): + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(3, 1) + self.build_data_tab_content(frame, self.controller, self.ui_state) - self.tabview = ctk.CTkTabview(frame) - self.tabview.grid(row=0, column=0, sticky="nsew") + def _configure_backup_frame(self, frame): + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(3, 1) + self.build_backup_tab_content(frame, self.controller, self.ui_state) - self.general_tab = self.create_general_tab(self.tabview.add("general")) - self.model_tab = self.create_model_tab(self.tabview.add("model")) - self.data_tab = self.create_data_tab(self.tabview.add("data")) - self.concepts_tab = self.create_concepts_tab(self.tabview.add("concepts")) - self.training_tab = self.create_training_tab(self.tabview.add("training")) - self.sampling_tab = self.create_sampling_tab(self.tabview.add("sampling")) - self.backup_tab = self.create_backup_tab(self.tabview.add("backup")) - self.tools_tab = self.create_tools_tab(self.tabview.add("tools")) - self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) - self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) + def _configure_tools_frame(self, frame): + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + self.build_tools_tab_content(frame, self.controller, self.ui_state) - self.change_training_method(self.train_config.training_method) + def _configure_embedding_frame(self, frame): + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + self.build_embedding_tab_content(frame, self.controller, self.ui_state) - return frame + def _create_tabs(self): + general_page = self._create_scrollable_tab(self._configure_general_frame) + self.tabview.addTab(general_page, "general") + self._tab_widgets["general"] = general_page - def create_general_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, weight=0) - frame.grid_columnconfigure(3, weight=1) - self.build_general_tab_content(frame, self.controller, self.ui_state) - frame.pack(fill="both", expand=1) - return frame + self.model_tab = PySide6ModelTabView(None, ModelTabController(self.train_config), self.ui_state) + self.tabview.addTab(self.model_tab, "model") + self._tab_widgets["model"] = self.model_tab - def create_model_tab(self, master): - return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) + data_page = self._create_scrollable_tab(self._configure_data_frame) + self.tabview.addTab(data_page, "data") + self._tab_widgets["data"] = data_page - def create_data_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - self.build_data_tab_content(frame, self.controller, self.ui_state) - frame.pack(fill="both", expand=1) - return frame + concepts_page = QWidget() + self.concepts_tab = PySide6ConceptTabView(concepts_page, ConceptTabController(self.train_config), self.ui_state) + self.tabview.addTab(concepts_page, "concepts") + self._tab_widgets["concepts"] = concepts_page - def create_concepts_tab(self, master): - return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) + self.training_tab = PySide6TrainingTabView(None, TrainingTabController(self.train_config), self.ui_state) + self.tabview.addTab(self.training_tab, "training") + self._tab_widgets["training"] = self.training_tab - def create_training_tab(self, master) -> CtkTrainingTabView: - return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) + sampling_page = self.create_sampling_tab() + self.tabview.addTab(sampling_page, "sampling") + self._tab_widgets["sampling"] = sampling_page - def create_cloud_tab(self, master) -> CtkCloudTabView: - return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) + backup_page = self._create_scrollable_tab(self._configure_backup_frame) + self.tabview.addTab(backup_page, "backup") + self._tab_widgets["backup"] = backup_page - def create_sampling_tab(self, master): - master.grid_rowconfigure(0, weight=0) - master.grid_rowconfigure(1, weight=1) - master.grid_columnconfigure(0, weight=1) + tools_page = self._create_scrollable_tab(self._configure_tools_frame) + self.tabview.addTab(tools_page, "tools") + self._tab_widgets["tools"] = tools_page - top_frame = ctk.CTkFrame(master=master, corner_radius=0) - top_frame.grid(row=0, column=0, sticky="nsew") - sub_frame = ctk.CTkFrame(master=top_frame, corner_radius=0, fg_color="transparent") - sub_frame.grid(row=1, column=0, sticky="nsew", columnspan=6) + additional_embeddings_page = QWidget() + self.additional_embeddings_tab = PySide6AdditionalEmbeddingsTabView( + additional_embeddings_page, + AdditionalEmbeddingsTabController(self.train_config), + self.ui_state, + ) + self.tabview.addTab(additional_embeddings_page, "additional embeddings") + self._tab_widgets["additional embeddings"] = additional_embeddings_page + + self.cloud_tab = PySide6CloudTabView(None, CloudTabController(self.train_config, self), self.ui_state) + self.tabview.addTab(self.cloud_tab, "cloud") + self._tab_widgets["cloud"] = self.cloud_tab + + def create_sampling_tab(self): + tab_page = QWidget() + tab_lo = QGridLayout(tab_page) + tab_lo.setContentsMargins(0, 0, 0, 0) + tab_lo.setSpacing(0) + tab_lo.setRowStretch(0, 0) + tab_lo.setRowStretch(1, 1) + tab_lo.setColumnStretch(0, 1) + + top_frame = QWidget(tab_page) + tab_lo.addWidget(top_frame, 0, 0) + top_lo = pyside6_components._layout(top_frame) + top_lo.setContentsMargins(pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD) + top_lo.setColumnStretch(8, 1) + + sub_frame = QWidget(top_frame) + pyside6_components._layout(top_frame).addWidget(sub_frame, 1, 0, 1, 8) self.build_sampling_tab_header(top_frame, sub_frame, self.controller, self.ui_state) + pyside6_components._layout(sub_frame).setColumnStretch(4, 1) - frame = ctk.CTkFrame(master=master, corner_radius=0) - frame.grid(row=1, column=0, sticky="nsew") - - return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) - - def create_backup_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - self.build_backup_tab_content(frame, self.controller, self.ui_state) - frame.pack(fill="both", expand=1) - return frame - - def embedding_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - self.build_embedding_tab_content(frame, self.controller, self.ui_state) - frame.pack(fill="both", expand=1) - return frame + sampling_container = QWidget(tab_page) + tab_lo.addWidget(sampling_container, 1, 0) + self.sampling_tab = PySide6SamplingTabView( + sampling_container, SamplingTabController(self.train_config), self.ui_state + ) - def create_additional_embeddings_tab(self, master): - return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) - - def create_tools_tab(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) - frame.grid_columnconfigure(2, minsize=50) - frame.grid_columnconfigure(3, weight=0) - frame.grid_columnconfigure(4, weight=1) - self.build_tools_tab_content(frame, self.controller, self.ui_state) - frame.pack(fill="both", expand=1) - return frame + return tab_page def open_profiling_tool(self): - self.profiling_window.deiconify() + self.profiling_window.show() def change_model_type(self, model_type: ModelType): if self.model_tab: self.model_tab.refresh_ui() - if self.training_tab: self.training_tab.refresh_ui() - if self.lora_tab: self.lora_tab.refresh_ui() @@ -369,45 +340,53 @@ def change_training_method(self, training_method: TrainingMethod): if self.model_tab: self.model_tab.refresh_ui() - if training_method != TrainingMethod.LORA and "LoRA" in self.tabview._tab_dict: - self.tabview.delete("LoRA") + if training_method != TrainingMethod.LORA and 'LoRA' in self._tab_widgets: + self.tabview.removeTab(self.tabview.indexOf(self._tab_widgets['LoRA'])) + del self._tab_widgets['LoRA'] self.lora_tab = None - if training_method != TrainingMethod.EMBEDDING and "embedding" in self.tabview._tab_dict: - self.tabview.delete("embedding") - - if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) - if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: - self.embedding_tab(self.tabview.add("embedding")) + if training_method != TrainingMethod.EMBEDDING and 'embedding' in self._tab_widgets: + self.tabview.removeTab(self.tabview.indexOf(self._tab_widgets['embedding'])) + del self._tab_widgets['embedding'] + + if training_method == TrainingMethod.LORA and 'LoRA' not in self._tab_widgets: + self.lora_tab = PySide6LoraTabView(None, LoraTabController(self.train_config), self.ui_state) + self.tabview.addTab(self.lora_tab, 'LoRA') + self._tab_widgets['LoRA'] = self.lora_tab + if training_method == TrainingMethod.EMBEDDING and 'embedding' not in self._tab_widgets: + tab_page = self._create_scrollable_tab(self._configure_embedding_frame) + self.tabview.addTab(tab_page, 'embedding') + self._tab_widgets['embedding'] = tab_page def load_preset(self): - if not self.tabview: - return - if self.additional_embeddings_tab: self.additional_embeddings_tab.refresh_ui() def _set_training_button_style(self, mode: str): if not self.training_button: return - style = self._TRAIN_BUTTON_STYLES.get(mode) - if not style: - return - self.training_button.configure(**style) + styles = { + "idle": ("Start Training", True, "#198754", "white"), + "running": ("Stop Training", True, "#dc3545", "white"), + "stopping": ("Stopping...", False, "#dc3545", "white"), + } + text, enabled, bg, fg = styles.get(mode, ("Start Training", True, "#198754", "white")) + self.training_button.setText(text) + self.training_button.setEnabled(enabled) + self.training_button.setStyleSheet( + f"QPushButton {{ background-color: {bg}; color: {fg}; }}" + f"QPushButton:disabled {{ background-color: {bg}; color: {fg}; }}" + ) def export_training(self): - file_path = filedialog.asksaveasfilename(filetypes=[ - ("All Files", "*.*"), - ("json", "*.json"), - ], initialdir=".", initialfile="config.json") + file_path, _ = QFileDialog.getSaveFileName( + self, "Export Training Config", "config.json", + "JSON Files (*.json);;All Files (*.*)" + ) if file_path: self.controller.export_training(file_path) def generate_debug_package(self): - zip_path = filedialog.askdirectory( - initialdir=".", - title="Select Directory to Save Debug Package" - ) - if not zip_path: + dir_path = QFileDialog.getExistingDirectory(self, "Select Directory to Save Debug Package", ".") + if not dir_path: return - self.controller.generate_debug_package(Path(zip_path) / "OneTrainer_debug_report.zip") + self.controller.generate_debug_package(Path(dir_path) / "OneTrainer_debug_report.zip") diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py index bc29488dd..510cda361 100644 --- a/modules/ui/PySide6TrainingTabView.py +++ b/modules/ui/PySide6TrainingTabView.py @@ -1,51 +1,68 @@ - from modules.ui.BaseTrainingTabView import BaseTrainingTabView -from modules.ui.CtkOffloadingWindowView import CtkOffloadingWindowView -from modules.ui.CtkOptimizerParamsWindowView import CtkOptimizerParamsWindowView -from modules.ui.CtkSchedulerParamsWindowView import CtkSchedulerParamsWindowView -from modules.ui.CtkTimestepDistributionWindowView import CtkTimestepDistributionWindowView +from modules.ui.OffloadingWindowController import OffloadingWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.ui.PySide6OffloadingWindowView import PySide6OffloadingWindowView +from modules.ui.PySide6OptimizerParamsWindowView import PySide6OptimizerParamsWindowView +from modules.ui.PySide6SchedulerParamsWindowView import PySide6SchedulerParamsWindowView +from modules.ui.PySide6TimestepDistributionWindowView import PySide6TimestepDistributionWindowView +from modules.ui.SchedulerParamsWindowController import SchedulerParamsWindowController +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController from modules.ui.TrainingTabController import TrainingTabController -from modules.util.ui import ctk_components +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QScrollArea, QWidget -import customtkinter as ctk +class PySide6TrainingTabView(BaseTrainingTabView, QWidget, metaclass=QtABCMeta): -class CtkTrainingTabView(BaseTrainingTabView): def __init__(self, master, controller: TrainingTabController, ui_state): - BaseTrainingTabView.__init__(self, ctk_components) + QWidget.__init__(self, master) + BaseTrainingTabView.__init__(self, pyside6_components) self.master = master self.controller = controller self.ui_state = ui_state self.scroll_frame = None - - master.grid_rowconfigure(0, weight=1) - master.grid_columnconfigure(0, weight=1) - self.refresh_ui() def refresh_ui(self): - if self.scroll_frame: - self.scroll_frame.destroy() - - self.scroll_frame = ctk.CTkScrollableFrame(self.master, fg_color="transparent") - self.scroll_frame.grid(row=0, column=0, sticky="nsew") - - self.scroll_frame.grid_columnconfigure(0, weight=1) - self.scroll_frame.grid_columnconfigure(1, weight=1) - self.scroll_frame.grid_columnconfigure(2, weight=1) - - column_0 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_0.grid(row=0, column=0, sticky="nsew") - column_0.grid_columnconfigure(0, weight=1) - - column_1 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_1.grid(row=0, column=1, sticky="nsew") - column_1.grid_columnconfigure(0, weight=1) - - column_2 = ctk.CTkFrame(master=self.scroll_frame, corner_radius=0, fg_color="transparent") - column_2.grid(row=0, column=2, sticky="nsew") - column_2.grid_columnconfigure(0, weight=1) + if self.scroll_frame is not None: + self.scroll_frame.hide() + self.scroll_frame.deleteLater() + + scroll = QScrollArea(self) + scroll.setWidgetResizable(True) + pyside6_components._layout(self).addWidget(scroll, 0, 0) + + self.scroll_frame = QWidget() + scroll.setWidget(self.scroll_frame) + + lo = pyside6_components._layout(self.scroll_frame) + lo.setContentsMargins(pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD, pyside6_components.PAD) + lo.setColumnStretch(0, 1) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(2, 1) + + from PySide6.QtWidgets import QSizePolicy + + column_0 = QWidget(self.scroll_frame) + column_0.setMinimumWidth(0) + column_0.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Preferred) + pyside6_components._layout(self.scroll_frame).addWidget(column_0, 0, 0) + pyside6_components._layout(column_0).setColumnStretch(0, 1) + + column_1 = QWidget(self.scroll_frame) + column_1.setMinimumWidth(0) + column_1.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Preferred) + pyside6_components._layout(self.scroll_frame).addWidget(column_1, 0, 1) + pyside6_components._layout(column_1).setColumnStretch(0, 1) + + column_2 = QWidget(self.scroll_frame) + column_2.setMinimumWidth(0) + column_2.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Preferred) + pyside6_components._layout(self.scroll_frame).addWidget(column_2, 0, 2) + pyside6_components._layout(column_2).setColumnStretch(0, 1) callbacks = { 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), @@ -58,20 +75,23 @@ def refresh_ui(self): self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) + for col_widget in (column_0, column_1, column_2): + lo = pyside6_components._layout(col_widget) + lo.setRowStretch(lo.rowCount(), 1) + def _restore_scheduler_config(self, variable): if not hasattr(self, 'lr_scheduler_adv_comp'): return - state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" - self.lr_scheduler_adv_comp.configure(state=state) + self.lr_scheduler_adv_comp.setEnabled(self.controller.is_custom_scheduler_value(variable)) def _open_optimizer_params_window(self): - self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) + PySide6OptimizerParamsWindowView(self, OptimizerParamsWindowController(self.controller.config), self.ui_state).exec() def _open_scheduler_params_window(self): - self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) + PySide6SchedulerParamsWindowView(self, SchedulerParamsWindowController(self.controller.config), self.ui_state).exec() def _open_timestep_distribution_window(self): - self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) + PySide6TimestepDistributionWindowView(self, TimestepDistributionWindowController(self.controller.config), self.ui_state).exec() def _open_offloading_window(self): - self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) + PySide6OffloadingWindowView(self, OffloadingWindowController(self.controller.config), self.ui_state).exec() diff --git a/modules/ui/PySide6VideoToolUIView.py b/modules/ui/PySide6VideoToolUIView.py index c272c891c..56e2c2779 100644 --- a/modules/ui/PySide6VideoToolUIView.py +++ b/modules/ui/PySide6VideoToolUIView.py @@ -1,128 +1,155 @@ -from tkinter import filedialog - from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image -from modules.util.ui import ctk_components -from modules.util.ui.CtkUIState import CtkUIState - -import customtkinter as ctk - -PAD = ctk_components.PAD - - -class CtkVideoToolUIView(BaseVideoToolUIView, ctk.CTkToplevel): - def __init__(self, parent, controller: VideoToolUIController, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseVideoToolUIView.__init__(self, ctk_components) +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta +from modules.util.ui.PySide6UIState import PySide6UIState + +from PIL.ImageQt import ImageQt +from PySide6.QtCore import Qt +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import ( + QDialog, + QFileDialog, + QGridLayout, + QLabel, + QPushButton, + QScrollArea, + QTabWidget, + QTextEdit, + QWidget, +) + +_PAD = pyside6_components.PAD + + +class PySide6VideoToolUIView(BaseVideoToolUIView, QDialog, metaclass=QtABCMeta): + def __init__(self, parent, controller: VideoToolUIController): + QDialog.__init__(self, parent) + BaseVideoToolUIView.__init__(self, pyside6_components) self.controller = controller - ui_state = CtkUIState(self, controller.args) - - self.title("Video Tools") - self.geometry("600x720") - self.resizable(True, True) - self.wait_visibility() - self.focus_set() - - self.grid_rowconfigure(0, weight=1) - self.grid_rowconfigure(1, weight=0) - self.grid_columnconfigure(0, weight=1) - - tabview = ctk.CTkTabview(self) - tabview.grid(row=0, column=0, sticky="nsew") - - clip_frame = ctk.CTkScrollableFrame(tabview.add("extract clips"), fg_color="transparent") - clip_frame.grid_columnconfigure(0, weight=0, minsize=120) - clip_frame.grid_columnconfigure(1, weight=0, minsize=200) - clip_frame.grid_columnconfigure(2, weight=0) - clip_frame.grid_columnconfigure(3, weight=1) - self.build_clip_extract_tab(clip_frame, controller, ui_state) - clip_frame.pack(fill="both", expand=1) - - image_frame = ctk.CTkScrollableFrame(tabview.add("extract images"), fg_color="transparent") - image_frame.grid_columnconfigure(0, weight=0, minsize=120) - image_frame.grid_columnconfigure(1, weight=0, minsize=200) - image_frame.grid_columnconfigure(2, weight=0) - image_frame.grid_columnconfigure(3, weight=1) - self.build_image_extract_tab(image_frame, controller, ui_state) - image_frame.pack(fill="both", expand=1) - - download_frame = ctk.CTkScrollableFrame(tabview.add("download"), fg_color="transparent") - download_frame.grid_columnconfigure(0, weight=0, minsize=120) - download_frame.grid_columnconfigure(1, weight=0, minsize=200) - download_frame.grid_columnconfigure(2, weight=0) - download_frame.grid_columnconfigure(3, weight=1) - self.build_video_download_tab(download_frame, controller, ui_state) - download_frame.pack(fill="both", expand=1) - - self._build_status_bar(self) - - def _build_status_bar(self, master): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=1, column=0) - frame.grid_columnconfigure(0, weight=0, minsize=160) - frame.grid_columnconfigure(1, weight=0, minsize=300) - frame.grid_columnconfigure(2, weight=1) - - preview_path = "resources/icons/icon.png" - preview = load_image(preview_path, 'RGB') + self._status_box: QTextEdit | None = None + self._preview_label: QLabel | None = None + self._preview_caption_label: QLabel | None = None + + ui_state = PySide6UIState(controller.args) + + self.setWindowTitle("Video Tools") + self.resize(700, 750) + + outer = QGridLayout(self) + outer.setContentsMargins(0, 0, 0, 0) + outer.setSpacing(0) + outer.setRowStretch(0, 1) + outer.setRowStretch(1, 0) + + tabs = QTabWidget(self) + outer.addWidget(tabs, 0, 0) + + for name, build_fn in [ + ("extract clips", self.build_clip_extract_tab), + ("extract images", self.build_image_extract_tab), + ("download", self.build_video_download_tab), + ]: + scroll = QScrollArea() + scroll.setWidgetResizable(True) + frame = QWidget() + scroll.setWidget(frame) + lo = pyside6_components._layout(frame) + lo.setContentsMargins(_PAD, _PAD, _PAD, _PAD) + lo.setColumnMinimumWidth(0, 120) + lo.setColumnStretch(3, 1) + build_fn(frame, controller, ui_state) + lo.setRowStretch(lo.rowCount(), 1) + tabs.addTab(scroll, name) + + outer.addWidget(self._build_status_bar(), 1, 0) + + def _build_status_bar(self): + frame = QWidget(self) + lo = QGridLayout(frame) + lo.setColumnMinimumWidth(0, 160) + lo.setColumnStretch(2, 1) + + self._preview_label = QLabel(frame) + self._preview_label.setFixedSize(150, 150) + preview = load_image("resources/icons/icon.png", 'RGB') preview.thumbnail((150, 150)) - self.preview_image = ctk.CTkImage(light_image=preview, size=preview.size) - self.preview_image_label = ctk.CTkLabel( - master=frame, text="Preview image", image=self.preview_image, height=150, width=150, - compound="top") - self.preview_image_label.grid(row=0, column=0, sticky="nw", padx=5, pady=5) - - self.status_label = ctk.CTkTextbox(master=frame, width=400, height=160, wrap="word", border_width=2) - self.status_label.insert(index="1.0", text="Current status") - self.status_label.configure(state="disabled") - self.status_label.grid(row=0, column=1, sticky="ne", padx=5, pady=5) + self._preview_label.setPixmap( + QPixmap.fromImage(ImageQt(preview.convert("RGBA"))).scaled( + 150, 150, Qt.KeepAspectRatio, Qt.SmoothTransformation + ) + ) + self._preview_caption_label = QLabel("Preview image", frame) + self._preview_caption_label.setWordWrap(True) + + preview_col = QWidget(frame) + preview_lo = QGridLayout(preview_col) + preview_lo.setContentsMargins(0, 0, 0, 0) + preview_lo.addWidget(self._preview_label, 0, 0, Qt.AlignTop) + preview_lo.addWidget(self._preview_caption_label, 1, 0, Qt.AlignTop) + lo.addWidget(preview_col, 0, 0, Qt.AlignTop | Qt.AlignLeft) + + self._status_box = QTextEdit(frame) + self._status_box.setReadOnly(True) + self._status_box.setFixedHeight(160) + self._status_box.setMinimumWidth(300) + self._status_box.setPlainText("Current status") + lo.addWidget(self._status_box, 0, 1, Qt.AlignTop) + + return frame + + # --- abstract method implementations --- def _create_textbox(self, master, row, col, width, height, ui_state, var_name): var = ui_state.get_var(var_name) - textbox = ctk.CTkTextbox(master, width=width, height=height, border_width=2) - textbox.insert("1.0", var.get()) - textbox.grid(row=row, column=col, rowspan=2, sticky="w", padx=PAD, pady=PAD) - - def on_text_change(event=None): - var.set(textbox.get("1.0", "end-1c")) - - textbox.bind("", on_text_change) - return textbox + widget = QTextEdit(master) + widget.setFixedHeight(height) + widget.setMinimumWidth(width) + widget.setPlainText(var.get()) + pyside6_components._add( + pyside6_components._layout(master), widget, row, col, sticky="w", rowspan=2 + ) + widget.textChanged.connect(lambda: var.set(widget.toPlainText())) + return widget def _create_browse_dir_button(self, master, row, ui_state, var_name): def browse(): - path = filedialog.askdirectory() + path = QFileDialog.getExistingDirectory(self, "Select Directory") if path: ui_state.get_var(var_name).set(path) - self.focus_set() - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + button = QPushButton("...", master) + button.setMaximumWidth(30) + button.clicked.connect(browse) + pyside6_components._add(pyside6_components._layout(master), button, row, 1, sticky="e") return button def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): def browse(): - path = filedialog.askopenfilename(filetypes=filetypes) + filters = ";;".join(f"{label} ({pattern})" for label, pattern in filetypes) + path, _ = QFileDialog.getOpenFileName(self, "Select File", filter=filters) if path: ui_state.get_var(var_name).set(path) - self.focus_set() - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) + button = QPushButton("...", master) + button.setMaximumWidth(30) + button.clicked.connect(browse) + pyside6_components._add(pyside6_components._layout(master), button, row, 1, sticky="e") return button + # --- view interface --- + def update_status(self, status_text: str): - self.status_label.configure(state="normal") - self.status_label.insert(index="end", text=status_text + "\n") - self.status_label.configure(state="disabled") + self._status_box.append(status_text) def clear_status(self): - self.status_label.configure(state="normal") - self.status_label.delete(index1="1.0", index2="end") - self.status_label.configure(state="disabled") + self._status_box.clear() def update_preview(self, preview_image, label_text: str): - self.preview_image.configure(light_image=preview_image, size=preview_image.size) - self.preview_image_label.configure(text=label_text) + pixmap = QPixmap.fromImage(ImageQt(preview_image.convert("RGBA"))) + self._preview_label.setPixmap( + pixmap.scaled(150, 150, Qt.KeepAspectRatio, Qt.SmoothTransformation) + ) + self._preview_caption_label.setText(label_text) diff --git a/modules/util/ui/PySide6UIState.py b/modules/util/ui/PySide6UIState.py new file mode 100644 index 000000000..6bef9e919 --- /dev/null +++ b/modules/util/ui/PySide6UIState.py @@ -0,0 +1,18 @@ +from typing import Any + +from modules.util.ui.QtVar import QtVar +from modules.util.ui.UIState import BaseUIState + + +class PySide6UIState(BaseUIState): + def __init__(self, obj): + super().__init__(obj) + + def _make_str_var(self, initial_value: Any) -> QtVar: + return QtVar(initial_value) + + def _make_bool_var(self, initial_value: Any) -> QtVar: + return QtVar(initial_value) + + def _make_nested_state(self, obj: Any) -> "PySide6UIState": + return PySide6UIState(obj) diff --git a/modules/util/ui/QtVar.py b/modules/util/ui/QtVar.py new file mode 100644 index 000000000..a535c8d13 --- /dev/null +++ b/modules/util/ui/QtVar.py @@ -0,0 +1,41 @@ +from collections.abc import Callable +from typing import Any + + +class QtVar: + """Toolkit-neutral observable variable. Drop-in for tk.StringVar / tk.BooleanVar.""" + + def __init__(self, value: Any = ""): + self._value = value + self._traces: dict[int, Callable[[], None]] = {} + self._next_id = 0 + self._widget_callbacks: dict[int, Callable[[Any], None]] = {} + + def get(self) -> Any: + return self._value + + def set(self, value: Any): + self._value = value + for cb in list(self._widget_callbacks.values()): + cb(value) + for cb in list(self._traces.values()): + cb(None, None, None) + + def trace_add(self, mode: str, callback: Callable) -> int: + id_ = self._next_id + self._traces[id_] = callback + self._next_id += 1 + return id_ + + def trace_remove(self, mode: str, name: int): + self._traces.pop(name, None) + + def _bind_widget(self, push_to_widget: Callable[[Any], None]) -> int: + """Register a one-way push from var → widget. Returns an ID for _unbind_widget.""" + id_ = self._next_id + self._widget_callbacks[id_] = push_to_widget + self._next_id += 1 + return id_ + + def _unbind_widget(self, id_: int): + self._widget_callbacks.pop(id_, None) diff --git a/modules/util/ui/pyside6_abc.py b/modules/util/ui/pyside6_abc.py new file mode 100644 index 000000000..60a27dd14 --- /dev/null +++ b/modules/util/ui/pyside6_abc.py @@ -0,0 +1,7 @@ +from abc import ABCMeta + +from PySide6.QtWidgets import QWidget + + +class QtABCMeta(type(QWidget), ABCMeta): + """Combined metaclass that resolves the conflict between Qt's Shiboken metaclass and ABCMeta.""" diff --git a/modules/util/ui/pyside6_components.py b/modules/util/ui/pyside6_components.py index e462f72a1..e901886d0 100644 --- a/modules/util/ui/pyside6_components.py +++ b/modules/util/ui/pyside6_components.py @@ -1,54 +1,171 @@ import contextlib -import tkinter as tk from collections.abc import Callable from pathlib import Path -from tkinter import filedialog from typing import Any, Literal from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit from modules.util.path_util import supported_image_extensions -from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator -from modules.util.ui.CtkUIState import CtkUIState -from modules.util.ui.ToolTip import ToolTip - -import customtkinter as ctk -from customtkinter.windows.widgets.scaling import CTkScalingBaseClass -from PIL import Image +from modules.util.ui.pyside6_validation import PySide6FieldValidator, PySide6PathValidator +from modules.util.ui.UIState import BaseUIState +from modules.util.ui.validation import DEFAULT_MAX_UNDO + +from PySide6.QtCore import Qt, QTimer +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import ( + QCheckBox, + QComboBox, + QFileDialog, + QFrame, + QGridLayout, + QLabel, + QLineEdit, + QProgressBar, + QPushButton, + QScrollArea, + QSizePolicy, + QVBoxLayout, + QWidget, +) PAD = 10 -def app_title(master, row, column): - frame = ctk.CTkFrame(master) - frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") +# --------------------------------------------------------------------------- +# PySide6-only helpers +# --------------------------------------------------------------------------- - image_component = ctk.CTkImage( - Image.open("resources/icons/icon.png").resize((40, 40), Image.Resampling.LANCZOS), - size=(40, 40) - ) - image_label_component = ctk.CTkLabel(frame, image=image_component, text="") - image_label_component.grid(row=0, column=0, padx=PAD, pady=PAD) +def _layout(master: QWidget) -> QGridLayout: + lo = master.layout() + if lo is None: + lo = QGridLayout(master) + lo.setContentsMargins(0, 0, 0, 0) + lo.setSpacing(PAD) + master.setLayout(lo) + return lo + + +def _alignment(sticky: str) -> Qt.AlignmentFlag: + has_e = 'e' in sticky + has_w = 'w' in sticky + has_n = 'n' in sticky + has_s = 's' in sticky + + if has_e and has_w: + h = Qt.AlignmentFlag(0) + elif has_e: + h = Qt.AlignRight + else: + h = Qt.AlignLeft + + if has_n and has_s: + v = Qt.AlignmentFlag(0) + elif has_s: + v = Qt.AlignBottom + else: + v = Qt.AlignTop - label_component = ctk.CTkLabel(frame, text="OneTrainer", font=ctk.CTkFont(size=20, weight="bold")) - label_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD) + return h | v -def label(master, row, column, text, pad=PAD, tooltip=None, wide_tooltip=False, wraplength=0, underline=False): - component = ctk.CTkLabel(master, text=text, wraplength=wraplength) - component.grid(row=row, column=column, padx=pad, pady=pad, sticky="nw") +def _add( + layout: QGridLayout, + widget: QWidget, + row: int, + col: int, + sticky: str = "new", + padx: int = PAD, + pady: int = PAD, + rowspan: int = 1, + colspan: int = 1, +): + layout.addWidget(widget, row, col, rowspan, colspan) + align = _alignment(sticky) + if align: + layout.setAlignment(widget, align) + + +def scrollable_frame(parent: QWidget) -> tuple[QScrollArea, QWidget]: + scroll = QScrollArea(parent) + scroll.setWidgetResizable(True) + container = QWidget() + container_layout = QVBoxLayout(container) + container_layout.setContentsMargins(PAD, PAD, PAD, PAD) + container_layout.setSpacing(0) + frame = QWidget(container) + container_layout.addWidget(frame) + container_layout.addStretch(1) + scroll.setWidget(container) + return scroll, frame + + +def _pack_form(master: QWidget) -> None: + """Add a stretch row after the last content row so rows don't expand to fill available space.""" + lo = _layout(master) + lo.setRowStretch(lo.rowCount(), 1) + + +# --------------------------------------------------------------------------- +# Stateless widgets +# --------------------------------------------------------------------------- + +def app_title(master: QWidget, row: int, column: int): + frame = QFrame(master) + layout = QGridLayout(frame) + layout.setContentsMargins(5, 5, 5, 5) + _layout(master).addWidget(frame, row, column) + + pixmap = QPixmap("resources/icons/icon.png").scaled( + 40, 40, Qt.KeepAspectRatio, Qt.SmoothTransformation + ) + icon_label = QLabel(frame) + icon_label.setPixmap(pixmap) + layout.addWidget(icon_label, 0, 0) + + text_label = QLabel("OneTrainer", frame) + font = text_label.font() + font.setPointSize(14) + font.setBold(True) + text_label.setFont(font) + layout.addWidget(text_label, 0, 1) + + +def label( + master: QWidget, + row: int, + column: int, + text: str, + pad: int = PAD, + tooltip: str | None = None, + wide_tooltip: bool = False, + wraplength: int = 0, + underline: bool = False, +) -> QLabel: + component = QLabel(text, master) + if wraplength > 0: + component.setWordWrap(True) + component.setMaximumWidth(wraplength) if tooltip: - ToolTip(component, tooltip, wide=wide_tooltip) + component.setToolTip(tooltip) if underline: - component.configure(font=ctk.CTkFont(underline=True)) + font = component.font() + font.setUnderline(True) + component.setFont(font) + layout = _layout(master) + layout.addWidget(component, row, column) + layout.setAlignment(component, Qt.AlignVCenter | Qt.AlignLeft) return component +# --------------------------------------------------------------------------- +# Compound widgets +# --------------------------------------------------------------------------- + def entry( - master, - row, - column, - ui_state: CtkUIState, + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, var_name: str, command: Callable[[], None] | None = None, tooltip: str = "", @@ -56,17 +173,21 @@ def entry( width: int = 140, sticky: str = "new", max_undo: int | None = None, - validator_factory: Callable[..., FieldValidator] | None = None, + validator_factory: Callable[..., PySide6FieldValidator] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, -): +) -> QLineEdit: var = ui_state.get_var(var_name) - trace_id = None + if command: - trace_id = ui_state.add_var_trace(var_name, command) + ui_state.add_var_trace(var_name, command) + + component = QLineEdit(master) + component.setMinimumWidth(width) + _add(_layout(master), component, row, column, sticky=sticky) - component = ctk.CTkEntry(master, textvariable=var, width=width) - component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky=sticky) + if tooltip: + component.setToolTip(tooltip) if validator_factory is not None: validator = validator_factory( @@ -76,7 +197,7 @@ def entry( required=required, ) else: - validator = FieldValidator( + validator = PySide6FieldValidator( component, var, ui_state, var_name, max_undo=max_undo or DEFAULT_MAX_UNDO, extra_validate=extra_validate, @@ -85,33 +206,15 @@ def entry( validator.attach() component._validator = validator # type: ignore[attr-defined] - original_destroy = component.destroy - - def new_destroy(): - validator.detach() - - # 'temporary' fix until https://github.com/TomSchimansky/CustomTkinter/pull/2077 is merged - # unfortunately Tom has admitted to forgetting about how to maintain CTK so this likely will never be merged - if component._textvariable_callback_name: - with contextlib.suppress(tk.TclError): - component._textvariable.trace_remove("write", component._textvariable_callback_name) # type: ignore[union-attr] - component._textvariable_callback_name = "" - - if command is not None and trace_id is not None: - ui_state.remove_var_trace(var_name, trace_id) - - original_destroy() - - component.destroy = new_destroy # type: ignore[assignment] - - if tooltip: - ToolTip(component, tooltip, wide=wide_tooltip) - return component def path_entry( - master, row, column, ui_state: CtkUIState, var_name: str, + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, + var_name: str, *, mode: Literal["file", "dir"] = "file", io_type: PathIOType = PathIOType.INPUT, @@ -122,179 +225,153 @@ def path_entry( extra_validate: Callable[[str], str | None] | None = None, required: bool = False, columnspan: int = 1, -): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new", columnspan=columnspan) - - frame.grid_columnconfigure(0, weight=1) +) -> QWidget: + frame = QWidget(master) + frame_lo = QGridLayout(frame) + frame_lo.setContentsMargins(0, 0, 0, 0) + frame_lo.setSpacing(0) + frame_lo.setColumnStretch(0, 1) + _add(_layout(master), frame, row, column, sticky="new", padx=0, pady=0, colspan=columnspan) def _path_validator_factory(comp, var, state, name, **kw): - return PathValidator(comp, var, state, name, io_type=io_type, **kw) + return PySide6PathValidator(comp, var, state, name, io_type=io_type, **kw) entry_component = entry( - frame, row=0, column=0, ui_state=ui_state, var_name=var_name, + frame, 0, 0, ui_state, var_name, validator_factory=_path_validator_factory, extra_validate=extra_validate, required=required, ) - trace_ids = [] + dep_trace_ids: list[tuple] = [] if io_type in (PathIOType.OUTPUT, PathIOType.MODEL): validator = getattr(entry_component, '_validator', None) if validator is not None: for dep_var_name in ("prevent_overwrites", "output_model_format"): with contextlib.suppress(KeyError, AttributeError): dep_var = ui_state.get_var(dep_var_name) - tid = dep_var.trace_add("write", lambda *_a: validator.revalidate()) - trace_ids.append((dep_var, tid)) + tid = dep_var.trace_add("write", lambda _0, _1, _2: validator.revalidate()) + dep_trace_ids.append((dep_var, tid)) + + if dep_trace_ids: + def _cleanup_dep_traces(): + for dv, tid in dep_trace_ids: + dv.trace_remove("write", tid) + frame.destroyed.connect(_cleanup_dep_traces) use_save_dialog = io_type in (PathIOType.OUTPUT, PathIOType.MODEL) - def __open_dialog(): - # Determine currently selected filename and/or directory - current_dir, current_filename = None, None + def _open_dialog(): current_path_str = ui_state.get_var(var_name).get() or None + current_dir = "" + current_filename = "" - if current_path_str is not None: + if current_path_str: current_path = Path(current_path_str) if mode == "file": current_dir = str(current_path.parent) current_filename = str(current_path.name) elif mode == "dir": current_dir = str(current_path.parent) - current_filename = None if mode == "dir": - chosen = filedialog.askdirectory(initialdir=current_dir) + chosen = QFileDialog.getExistingDirectory(frame, "", current_dir, QFileDialog.Option.ShowDirsOnly) else: - filetypes = [ - ("All Files", "*.*"), - ] - + filters = ["All Files (*.*)"] if allow_model_files: - filetypes.extend([ - ("Diffusers", "model_index.json"), - ("Checkpoint", "*.ckpt *.pt *.bin"), - ("Safetensors", "*.safetensors"), - ]) + filters += [ + "Diffusers (model_index.json)", + "Checkpoint (*.ckpt *.pt *.bin)", + "Safetensors (*.safetensors)", + ] if allow_image_files: - filetypes.extend([ - ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), - ]) + exts = " ".join(f"*.{x}" for x in supported_image_extensions()) + filters.append(f"Image ({exts})") + filter_str = ";;".join(filters) + init_path = str(Path(current_dir) / current_filename) if current_filename else current_dir if use_save_dialog: - chosen = filedialog.asksaveasfilename(filetypes=filetypes, initialdir=current_dir, - initialfile=current_filename) + chosen, _ = QFileDialog.getSaveFileName(frame, "", init_path, filter_str) else: - chosen = filedialog.askopenfilename(filetypes=filetypes, initialdir=current_dir, - initialfile=current_filename) + chosen, _ = QFileDialog.getOpenFileName(frame, "", init_path, filter_str) if chosen: if path_modifier: chosen = path_modifier(chosen) - chosen_str = str(chosen) ui_state.get_var(var_name).set(chosen_str) - if command: command(chosen_str) - button_component = ctk.CTkButton(frame, text="...", width=40, command=__open_dialog) - button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") - - if trace_ids: - original_frame_destroy = frame.destroy - def _frame_destroy(): - for dep_var, tid in trace_ids: - with contextlib.suppress(tk.TclError, ValueError): - dep_var.trace_remove("write", tid) - original_frame_destroy() - frame.destroy = _frame_destroy # type: ignore[assignment] + btn = QPushButton("...", frame) + btn.setFixedWidth(40) + btn.clicked.connect(_open_dialog) + frame_lo.addWidget(btn, 0, 1) return frame -def time_entry(master, row, column, ui_state: CtkUIState, var_name: str, unit_var_name, supports_time_units: bool = True): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") - - frame.grid_columnconfigure(0, weight=0) - frame.grid_columnconfigure(1, weight=1) +def time_entry( + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, + var_name: str, + unit_var_name: str, + supports_time_units: bool = True, +) -> QWidget: + frame = QWidget(master) + _add(_layout(master), frame, row, column, sticky="new", padx=0, pady=0) - entry(frame, row=0, column=0, ui_state=ui_state, var_name=var_name, width=50) + entry(frame, 0, 0, ui_state, var_name, width=50) values = [str(x) for x in list(TimeUnit)] if not supports_time_units: values = [str(x) for x in list(TimeUnit) if not x.is_time_unit()] - unit_component = ctk.CTkOptionMenu( - frame, - values=values, - variable=ui_state.get_var(unit_var_name), - width=100, - ) - unit_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="new") + options(frame, 0, 1, values, ui_state, unit_var_name) return frame -def layer_filter_entry(master, row, column, ui_state: CtkUIState, preset_var_name: str, preset_label: str, preset_tooltip: str, presets, entry_var_name, entry_tooltip: str, regex_var_name, regex_tooltip: str, frame_color=None): - frame = ctk.CTkFrame(master=master, corner_radius=5, fg_color=frame_color) - frame.grid(row=row, column=column, padx=5, pady=5, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) - - layer_entry = entry( - frame, 1, 0, ui_state, entry_var_name, - tooltip=entry_tooltip - ) - layer_entry_fg_color = layer_entry.cget("fg_color") - layer_entry_text_color = layer_entry.cget("text_color") - - regex_label = label( - frame, 2, 0, "Use Regex", - tooltip=regex_tooltip, - ) - regex_switch = switch( - frame, 2, 1, ui_state, regex_var_name - ) - - # Let the user set their own layer filter - # TODO - #if self.train_config.layer_filter and self.train_config.layer_filter_preset == "custom": - # self.prior_custom = self.train_config.layer_filter - #else: - # self.prior_custom = "" - layer_entry.grid_configure(columnspan=2, sticky="ew") +def layer_filter_entry( + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, + preset_var_name: str, + preset_label: str, + preset_tooltip: str, + presets, + entry_var_name: str, + entry_tooltip: str, + regex_var_name: str, + regex_tooltip: str, + frame_color=None, +) -> QWidget: + frame = QWidget(master) + _layout(master).addWidget(frame, row, column) + + label(frame, 0, 0, preset_label, tooltip=preset_tooltip) + + layer_entry = entry(frame, 1, 0, ui_state, entry_var_name, tooltip=entry_tooltip) + _layout(frame).addWidget(layer_entry, 1, 0, 1, 2) # span 2 columns + + regex_label = label(frame, 2, 0, "Use Regex", tooltip=regex_tooltip) + regex_switch = switch(frame, 2, 1, ui_state, regex_var_name) presets_list = list(presets.keys()) + ["custom"] - - def hide_layer_entry(): - if layer_entry and layer_entry.winfo_manager(): - layer_entry.grid_remove() - - def show_layer_entry(): - if layer_entry and not layer_entry.winfo_manager(): - layer_entry.grid() - - def preset_set_layer_choice(selected: str): if not selected or selected not in presets_list: selected = presets_list[0] if selected == "custom": - # Allow editing + regex toggle - show_layer_entry() - layer_entry.configure(state="normal", fg_color=layer_entry_fg_color, text_color=layer_entry_text_color) - #layer_entry.cget('textvariable').set("") - regex_label.grid() - regex_switch.grid() + layer_entry.setVisible(True) + layer_entry.setEnabled(True) + regex_label.setVisible(True) + regex_switch.setVisible(True) else: - # Preserve custom text before overwriting - #if self.prior_selected == "custom": - # self.prior_custom = self.layer_entry.get() - - # Resolve preset definition (list[str] OR {'patterns': [...], 'regex': bool}) preset_def = presets.get(selected, []) if isinstance(preset_def, dict): patterns = preset_def.get("patterns", []) @@ -303,286 +380,348 @@ def preset_set_layer_choice(selected: str): patterns = preset_def preset_uses_regex = False - disabled_color = ("gray85", "gray17") - disabled_text_color = ("gray30", "gray70") - layer_entry.configure(state="disabled", fg_color=disabled_color, text_color=disabled_text_color) - layer_entry.cget('textvariable').set(",".join(patterns)) - + layer_entry.setEnabled(False) ui_state.get_var(entry_var_name).set(",".join(patterns)) ui_state.get_var(regex_var_name).set(preset_uses_regex) - regex_label.grid_remove() - regex_switch.grid_remove() - - if selected == "full" and not patterns: - hide_layer_entry() - else: - show_layer_entry() - -# self.prior_selected = selected - - label(frame, 0, 0, preset_label, - tooltip=preset_tooltip) + regex_label.setVisible(False) + regex_switch.setVisible(False) + layer_entry.setVisible(selected != "full" or bool(patterns)) ui_state.remove_all_var_traces(preset_var_name) layer_selector = options( frame, 0, 1, presets_list, ui_state, preset_var_name, - command=preset_set_layer_choice + command=preset_set_layer_choice, ) - def on_layer_filter_preset_change(): - if not layer_selector: - return - selected = ui_state.get_var(preset_var_name).get() - preset_set_layer_choice(selected) - - ui_state.add_var_trace( - preset_var_name, - on_layer_filter_preset_change, - ) + ui_state.add_var_trace(preset_var_name, lambda: preset_set_layer_choice( + ui_state.get_var(preset_var_name).get() + )) - preset_set_layer_choice(layer_selector.get()) + preset_set_layer_choice(layer_selector.currentText()) -def icon_button(master, row, column, text, command): - component = ctk.CTkButton(master, text=text, width=40, command=command) - component.grid(row=row, column=column, padx=PAD, pady=PAD, sticky="new") - return component + return frame -def colored_icon_button(master, row, column, text, fg_color, command, padx=0): - component = ctk.CTkButton( - master=master, width=20, height=20, text=text, - corner_radius=2, fg_color=fg_color, command=command, - ) - component.grid(row=row, column=column, padx=padx) +def icon_button(master: QWidget, row: int, column: int, text: str, command: Callable[[], None]) -> QPushButton: + component = QPushButton(text, master) + component.setFixedWidth(40) + component.clicked.connect(command) + _add(_layout(master), component, row, column, sticky="new") return component -def button(master, row, column, text, command, tooltip=None, **kwargs): - # Pop grid-specific parameters from kwargs, using PAD as the default if not provided. - padx = kwargs.pop('padx', PAD) - pady = kwargs.pop('pady', PAD) - - component = ctk.CTkButton(master, text=text, command=command, **kwargs) - component.grid(row=row, column=column, padx=padx, pady=pady, sticky="new") - if tooltip: - ToolTip(component, tooltip, x_position=25) +def colored_icon_button( + master: QWidget, + row: int, + column: int, + text: str, + fg_color, + command: Callable[[], None], + padx: int = 0, +) -> QPushButton: + color = fg_color[0] if isinstance(fg_color, (tuple, list)) else fg_color + component = QPushButton(text, master) + component.setFixedSize(20, 20) + component.setStyleSheet(f"QPushButton {{ background-color: {color}; border-radius: 2px; }}") + component.clicked.connect(command) + _add(_layout(master), component, row, column, sticky="new", padx=padx, pady=0) return component -def options(master, row, column, values, ui_state: CtkUIState, var_name: str, command: Callable[[str], None] | None = None): - component = ctk.CTkOptionMenu(master, values=values, variable=ui_state.get_var(var_name), command=command) - component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") - - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy - - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] - +def button( + master: QWidget, + row: int, + column: int, + text: str, + command: Callable[[], None], + tooltip: str | None = None, + padx: int = PAD, + pady: int = PAD, + **kwargs, +) -> QPushButton: + component = QPushButton(text, master) + component.clicked.connect(command) + if tooltip: + component.setToolTip(tooltip) + _add(_layout(master), component, row, column, sticky="new", padx=padx, pady=pady) return component -def options_adv(master, row, column, values, ui_state: CtkUIState, var_name: str, - command: Callable[[str], None] | None = None, adv_command: Callable[[], None] | None = None): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="new") +# --------------------------------------------------------------------------- +# Bound widgets +# --------------------------------------------------------------------------- - frame.grid_columnconfigure(0, weight=1) +def options( + master: QWidget, + row: int, + column: int, + values: list[str], + ui_state: BaseUIState, + var_name: str, + command: Callable[[str], None] | None = None, +) -> QComboBox: + var = ui_state.get_var(var_name) + combo = QComboBox(master) + combo.addItems(values) + combo.setCurrentText(str(var.get())) - component = ctk.CTkOptionMenu(frame, values=values, variable=ui_state.get_var(var_name), command=command) - component.grid(row=0, column=0, padx=PAD, pady=(PAD, PAD), sticky="new") + _updating = False - button_component = ctk.CTkButton(frame, text="…", width=20, command=adv_command) - button_component.grid(row=0, column=1, padx=(0, PAD), pady=PAD, sticky="nsew") + def on_combo(text: str): + nonlocal _updating + if _updating: + return + _updating = True + var.set(text) + _updating = False + if command: + command(text) + + def on_var(value): + nonlocal _updating + if _updating: + return + _updating = True + combo.setCurrentText(str(value)) + _updating = False + + combo.currentTextChanged.connect(on_combo) + cb_id = var._bind_widget(on_var) + combo.destroyed.connect(lambda: var._unbind_widget(cb_id)) + _add(_layout(master), combo, row, column) + return combo + + +def options_adv( + master: QWidget, + row: int, + column: int, + values: list[str], + ui_state: BaseUIState, + var_name: str, + command: Callable[[str], None] | None = None, + adv_command: Callable[[], None] | None = None, +) -> tuple[QWidget, dict]: + frame = QWidget(master) + frame_lo = QGridLayout(frame) + frame_lo.setContentsMargins(0, 0, 0, 0) + frame_lo.setColumnStretch(0, 1) + _add(_layout(master), frame, row, column, sticky="new", padx=0, pady=0) + + combo = options(frame, 0, 0, values, ui_state, var_name, command=command) + + adv_btn = QPushButton("…", frame) + adv_btn.setFixedWidth(20) + if adv_command: + adv_btn.clicked.connect(adv_command) + _add(frame_lo, adv_btn, 0, 1, sticky="nsew", padx=(0, PAD), pady=PAD) if command: - command(ui_state.get_var(var_name).get()) # call command once to set the initial value - - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy + command(ui_state.get_var(var_name).get()) - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] + return frame, {'component': combo, 'button_component': adv_btn} - return frame, {'component': component, 'button_component': button_component} - -def options_kv(master, row, column, values: list[tuple[str, Any]], ui_state: CtkUIState, var_name: str, - command: Callable[[Any], None] | None = None): +def options_kv( + master: QWidget, + row: int, + column: int, + values: list[tuple[str, Any]], + ui_state: BaseUIState, + var_name: str, + command: Callable[[Any], None] | None = None, +) -> QComboBox: var = ui_state.get_var(var_name) - keys = [key for key, value in values] + keys = [key for key, _ in values] + str_values = [str(v) for _, v in values] - # if the current value is not valid, select the first option - if var.get() not in [str(value) for key, value in values] and len(keys) > 0: + if var.get() not in str_values and keys: var.set(values[0][1]) - deactivate_update_var = False + _updating = False - def update_component(text): - for key, value in values: - if text == key: - nonlocal deactivate_update_var - deactivate_update_var = True - var.set(value) + def on_combo(key: str): + nonlocal _updating + if _updating: + return + _updating = True + for k, v in values: + if key == k: + var.set(v) if command: - command(value) - deactivate_update_var = False + command(v) break + _updating = False - component = ctk.CTkOptionMenu(master, values=keys, command=update_component) - component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") - - def update_var(): - if not deactivate_update_var: - for key, value in values: - if var.get() == str(value): - if component.winfo_exists(): # the component could already be destroyed - component.set(key) - if command: - command(value) - break - - var.trace_add("write", lambda _0, _1, _2: update_var()) - update_var() # call update_var once to set the initial value - - # temporary fix until https://github.com/TomSchimansky/CustomTkinter/pull/2246 is merged - def create_destroy(component): - orig_destroy = component.destroy - - def destroy(self): - orig_destroy() - CTkScalingBaseClass.destroy(self) - - return destroy - - destroy = create_destroy(component._dropdown_menu) - component._dropdown_menu.destroy = lambda: destroy(component._dropdown_menu) # type: ignore[assignment] + def on_var(value): + nonlocal _updating + if _updating: + return + _updating = True + for k, v in values: + if str(value) == str(v): + combo.setCurrentText(k) + if command: + command(v) + break + _updating = False + + combo = QComboBox(master) + combo.addItems(keys) + # set initial display from current var value + for k, v in values: + if str(var.get()) == str(v): + combo.setCurrentText(k) + break + + combo.currentTextChanged.connect(on_combo) + cb_id = var._bind_widget(on_var) + combo.destroyed.connect(lambda: var._unbind_widget(cb_id)) + _add(_layout(master), combo, row, column) + + # match CTK behavior: fire initial command with the current value + if command: + current = var.get() + for _, v in values: + if str(current) == str(v): + command(v) + break - return component + return combo def switch( - master, - row, - column, - ui_state: CtkUIState, + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, var_name: str, command: Callable[[], None] | None = None, text: str = "", width: int | None = None, -): +) -> QCheckBox: var = ui_state.get_var(var_name) - if command: - trace_id = ui_state.add_var_trace(var_name, command) + component = QCheckBox(text, master) + component.setChecked(bool(var.get())) - component = ctk.CTkSwitch(master, variable=var, text=text, command=command) - if width is not None: - component.configure(width=width) - component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="new") - - def create_destroy(component): - orig_destroy = component.destroy + if command: + ui_state.add_var_trace(var_name, command) - def destroy(self): - if command is not None: - ui_state.remove_var_trace(var_name, trace_id) + _updating = False - orig_destroy() + def on_toggle(checked: bool): + nonlocal _updating + if _updating: + return + _updating = True + var.set(checked) + _updating = False - return destroy + def on_var(value): + nonlocal _updating + if _updating: + return + _updating = True + component.setChecked(bool(value)) + _updating = False - destroy = create_destroy(component) - component.destroy = lambda: destroy(component) # type: ignore[assignment] + component.toggled.connect(on_toggle) + cb_id = var._bind_widget(on_var) + component.destroyed.connect(lambda: var._unbind_widget(cb_id)) + if width is not None: + component.setFixedWidth(width) + lo = _layout(master) + lo.addWidget(component, row, column) + lo.setAlignment(component, Qt.AlignVCenter | Qt.AlignLeft) return component -def progress(master, row, column): - component = ctk.CTkProgressBar(master) - component.grid(row=row, column=column, padx=PAD, pady=(PAD, PAD), sticky="ew") +def progress(master: QWidget, row: int, column: int) -> QProgressBar: + component = QProgressBar(master) + component.setRange(0, 1000) + component.setValue(0) + component.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + _add(_layout(master), component, row, column, sticky="ew") return component -def double_progress(master, row, column, label_1, label_2): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=column, padx=0, pady=0, sticky="nsew") - - frame.grid_rowconfigure(0, weight=1) - frame.grid_rowconfigure(1, weight=1) - frame.grid_columnconfigure(0, weight=1) - - label_1_component = ctk.CTkLabel(frame, text=label_1) - label_1_component.grid(row=0, column=0, padx=(PAD, PAD), pady=(0, 0), sticky="new") - - label_2_component = ctk.CTkLabel(frame, text=label_2) - label_2_component.grid(row=1, column=0, padx=(PAD, PAD), pady=(0, 0), sticky="sew") - - progress_1_component = ctk.CTkProgressBar(frame) - progress_1_component.grid(row=0, column=1, padx=(PAD, PAD), pady=(PAD, 0), sticky="new") - - progress_2_component = ctk.CTkProgressBar(frame) - progress_2_component.grid(row=1, column=1, padx=(PAD, PAD), pady=(0, PAD), sticky="sew") - - description_1_component = ctk.CTkLabel(frame, text="") - description_1_component.grid(row=0, column=2, padx=(PAD, PAD), pady=(0, 0), sticky="new") - - description_2_component = ctk.CTkLabel(frame, text="") - description_2_component.grid(row=1, column=2, padx=(PAD, PAD), pady=(0, 0), sticky="sew") - - def set_1(value, max_value): - progress_1_component.set(value / max_value) - description_1_component.configure(text=f"{value}/{max_value}") - - def set_2(value, max_value): - progress_2_component.set(value / max_value) - description_2_component.configure(text=f"{value}/{max_value}") +def double_progress( + master: QWidget, + row: int, + column: int, + label_1: str, + label_2: str, +) -> tuple[Callable, Callable]: + frame = QWidget(master) + lo = QGridLayout(frame) + lo.setContentsMargins(0, 0, 0, 0) + lo.setColumnStretch(1, 1) + + label_1_component = QLabel(label_1, frame) + label_2_component = QLabel(label_2, frame) + progress_1_component = QProgressBar(frame) + progress_2_component = QProgressBar(frame) + description_1_component = QLabel("", frame) + description_2_component = QLabel("", frame) + + for p in (progress_1_component, progress_2_component): + p.setRange(0, 1000) + p.setValue(0) + p.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) + + lo.addWidget(label_1_component, 0, 0) + lo.addWidget(progress_1_component, 0, 1) + lo.addWidget(description_1_component, 0, 2) + lo.addWidget(label_2_component, 1, 0) + lo.addWidget(progress_2_component, 1, 1) + lo.addWidget(description_2_component, 1, 2) + + _add(_layout(master), frame, row, column, sticky="nsew") + + def set_1(value: int | float, max_value: int | float): + progress_1_component.setValue(int(value / max_value * 1000)) + description_1_component.setText(f"{value}/{max_value}") + + def set_2(value: int | float, max_value: int | float): + progress_2_component.setValue(int(value / max_value * 1000)) + description_2_component.setText(f"{value}/{max_value}") return set_1, set_2 -def section_frame(master, row: int, col: int = 0): - frame = ctk.CTkFrame(master=master, corner_radius=5) - frame.grid(row=row, column=col, padx=PAD // 2, pady=PAD // 2, sticky="nsew") - frame.grid_columnconfigure(0, weight=1) +def section_frame(parent: QWidget, row: int, col: int = 0, colspan: int = 1) -> "QFrame": + from PySide6.QtWidgets import QFrame + frame = QFrame(parent) + frame.setFrameShape(QFrame.Shape.StyledPanel) + _layout(parent).addWidget(frame, row, col, 1, colspan) + frame_lo = _layout(frame) + frame_lo.setColumnStretch(0, 1) + frame_lo.setContentsMargins(PAD, PAD, PAD, PAD) return frame -def inline_frame(master, row: int, col: int, columnspan: int = 1): - frame = ctk.CTkFrame(master, fg_color="transparent") - frame.grid(row=row, column=col, columnspan=columnspan, sticky="ew", padx=0, pady=0) +def inline_frame(parent: QWidget, row: int, col: int, columnspan: int = 1) -> QWidget: + frame = QWidget(parent) + _layout(frame) + _layout(parent).addWidget(frame, row, col, 1, columnspan) return frame -def set_widget_enabled(widget, enabled: bool) -> None: - state = "normal" if enabled else "disabled" - if isinstance(widget, ctk.CTkFrame): - for child in widget.children.values(): - with contextlib.suppress(Exception): - child.configure(state=state) - else: - widget.configure(state=state) +# --------------------------------------------------------------------------- +# Pure helper (toolkit-neutral) +# --------------------------------------------------------------------------- + +def set_widget_enabled(widget: QWidget, enabled: bool) -> None: + widget.setEnabled(enabled) -def set_label_text(label, text: str) -> None: - label.configure(text=str(text)) +def set_label_text(label: QLabel, text: str) -> None: + label.setText(str(text)) -def call_after(widget, delay_ms: int, func) -> None: - widget.after(delay_ms, func) +def call_after(widget: QWidget, delay_ms: int, func) -> None: + QTimer.singleShot(delay_ms, widget, func) diff --git a/modules/util/ui/pyside6_validation.py b/modules/util/ui/pyside6_validation.py index 5347ba6d4..8e9e1e502 100644 --- a/modules/util/ui/pyside6_validation.py +++ b/modules/util/ui/pyside6_validation.py @@ -1,60 +1,28 @@ -from __future__ import annotations - -import contextlib -import tkinter as tk from collections.abc import Callable -from typing import TYPE_CHECKING, Any from modules.util.enum.PathIOType import PathIOType +from modules.util.ui.QtVar import QtVar +from modules.util.ui.UIState import BaseUIState from modules.util.ui.validation import ( DEBOUNCE_TYPING_MS, DEFAULT_MAX_UNDO, ERROR_BORDER_COLOR, - UNDO_DEBOUNCE_MS, BaseFieldValidator, - UndoHistory, - _active_validators, _validate_path_field, ) -if TYPE_CHECKING: - from modules.util.ui.UIState import UIState - - import customtkinter as ctk - - -class DebounceTimer: - def __init__(self, widget, delay_ms: int, callback: Callable[..., Any]): - self.widget = widget - self.delay_ms = delay_ms - self.callback = callback - self._after_id: str | None = None +from PySide6.QtCore import QTimer +from PySide6.QtWidgets import QLineEdit - def call(self, *args, **kwargs): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) +_active_qt_validators: set["PySide6FieldValidator"] = set() - def fire(): - self._after_id = None - self.callback(*args, **kwargs) - with contextlib.suppress(tk.TclError): - self._after_id = self.widget.after(self.delay_ms, fire) - - def cancel(self): - if self._after_id: - with contextlib.suppress(tk.TclError): - self.widget.after_cancel(self._after_id) - self._after_id = None - - -class FieldValidator(BaseFieldValidator): +class PySide6FieldValidator(BaseFieldValidator): def __init__( self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, + component: QLineEdit, + var: QtVar, + ui_state: BaseUIState, var_name: str, max_undo: int = DEFAULT_MAX_UNDO, extra_validate: Callable[[str], str | None] | None = None, @@ -63,209 +31,122 @@ def __init__( super().__init__(ui_state, var_name, extra_validate, required) self.component = component self.var = var - - try: - self._original_border_color = component.cget("border_color") - except Exception: - self._original_border_color = "gray50" - - self._shadow_var = tk.StringVar(master=component) - self._shadow_trace_name: str | None = None - self._real_var_trace_name: str | None = None + self._original_style = component.styleSheet() self._syncing = False self._touched = False + self._var_trace_id: int | None = None - self._debounce: DebounceTimer | None = None - self._undo_debounce: DebounceTimer | None = None - self._undo = UndoHistory(max_undo) + self._debounce = QTimer(component) + self._debounce.setSingleShot(True) + self._debounce.setInterval(DEBOUNCE_TYPING_MS) + self._debounce.timeout.connect(self._on_debounce_fire) - def attach(self) -> None: - self._shadow_var.set(self.var.get()) - self._swap_textvariable(self._shadow_var) + def _apply_error(self) -> None: + self.component.setStyleSheet(f"border: 1px solid {ERROR_BORDER_COLOR};") - self._debounce = DebounceTimer( - self.component, DEBOUNCE_TYPING_MS, self._on_debounce_fire - ) - self._undo_debounce = DebounceTimer( - self.component, UNDO_DEBOUNCE_MS, self._push_undo_snapshot - ) + def _clear_error(self) -> None: + self.component.setStyleSheet(self._original_style) - self._shadow_trace_name = self._shadow_var.trace_add("write", self._on_shadow_write) - self._real_var_trace_name = self.var.trace_add("write", self._on_real_var_write) + def attach(self) -> None: + self._syncing = True + self.component.setText(str(self.var.get())) + self._syncing = False - self.component.bind("", self._on_focus_in) - self.component.bind("", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("<>", self._on_user_input) - self.component.bind("", self._on_focus_out) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_undo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_redo) - self.component.bind("", self._on_enter) + self.component.textChanged.connect(self._on_text_changed) + self.component.editingFinished.connect(self._on_editing_finished) + self._var_trace_id = self.var.trace_add("write", self._on_real_var_write) + self.component.destroyed.connect(self._on_destroyed) self._bound = True - _active_validators.add(self) + _active_qt_validators.add(self) def detach(self) -> None: if not self._bound: return self._bound = False - _active_validators.discard(self) - + _active_qt_validators.discard(self) + self._debounce.stop() self._commit() - - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() - - if self._shadow_trace_name: - with contextlib.suppress(Exception): - self._shadow_var.trace_remove("write", self._shadow_trace_name) - self._shadow_trace_name = None - - if self._real_var_trace_name: - with contextlib.suppress(Exception): - self.var.trace_remove("write", self._real_var_trace_name) - self._real_var_trace_name = None - - self._swap_textvariable(self.var) - - def _swap_textvariable(self, new_var: tk.Variable) -> None: - comp = self.component - if comp._textvariable_callback_name: - with contextlib.suppress(Exception): - comp._textvariable.trace_remove("write", comp._textvariable_callback_name) # type: ignore[union-attr] - comp._textvariable_callback_name = "" - - comp.configure(textvariable=new_var) - - if new_var is not None: - comp._textvariable_callback_name = new_var.trace_add( - "write", comp._textvariable_callback - ) + try: + self.component.textChanged.disconnect(self._on_text_changed) + self.component.editingFinished.disconnect(self._on_editing_finished) + except RuntimeError: + pass + if self._var_trace_id is not None: + self.var.trace_remove("write", self._var_trace_id) + self._var_trace_id = None + + def _on_destroyed(self) -> None: + """Called when the Qt C++ widget is deleted; skips _commit() since widget is gone.""" + if not self._bound: + return + self._bound = False + _active_qt_validators.discard(self) + self._debounce.stop() + if self._var_trace_id is not None: + self.var.trace_remove("write", self._var_trace_id) + self._var_trace_id = None def _commit(self) -> None: - shadow_val = self._shadow_var.get() - if shadow_val != self.var.get(): + val = self.component.text() + if val != str(self.var.get()): self._syncing = True - self.var.set(shadow_val) + self.var.set(val) self._syncing = False - def _apply_error(self) -> None: - self.component.configure(border_color=ERROR_BORDER_COLOR) - - def _clear_error(self) -> None: - self.component.configure(border_color=self._original_border_color) - - def _on_shadow_write(self, *_args) -> None: - if self._syncing: - return - if not self._touched: - # external sync or initial set — commit immediately - self._commit() - if self._debounce: - self._debounce.cancel() - return - if self._debounce: - self._debounce.call() - if self._undo_debounce: - self._undo_debounce.call() - - def _on_real_var_write(self, *_args) -> None: + def _on_text_changed(self, _text: str) -> None: if self._syncing: return - # external change (preset load, file dialog, etc) — sync to shadow var - self._syncing = True - self._shadow_var.set(self.var.get()) - self._syncing = False - self._validate_and_style(self._shadow_var.get()) - - def _push_undo_snapshot(self) -> None: - self._undo.push(self._shadow_var.get()) + self._touched = True + self._debounce.start() def _on_debounce_fire(self) -> None: - val = self._shadow_var.get() + val = self.component.text() if self._validate_and_style(val): self._commit() - def _on_focus_in(self, _e=None) -> None: - self._touched = False - self._undo.push(self._shadow_var.get()) - - def _on_user_input(self, _e=None) -> None: - self._touched = True - - def _on_focus_out(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._undo_debounce: - self._undo_debounce.cancel() + def _on_editing_finished(self) -> None: + self._debounce.stop() if self._touched: - if self._validate_and_style(self._shadow_var.get()): - self._commit() - self._undo.push(self._shadow_var.get()) - - def _on_enter(self, _e=None) -> None: - if self._debounce: - self._debounce.cancel() - if self._touched: - if self._validate_and_style(self._shadow_var.get()): + val = self.component.text() + if self._validate_and_style(val): self._commit() + self._touched = False - def _set_value(self, value: str) -> None: + def _on_real_var_write(self, _0, _1, _2) -> None: + if self._syncing: + return self._syncing = True - self._shadow_var.set(value) + self.component.setText(str(self.var.get())) self._syncing = False - if self._validate_and_style(value): - self._commit() - - def _on_undo(self, _e=None) -> str: - previous = self._undo.undo(self._shadow_var.get()) - if previous is not None: - self._set_value(previous) - return "break" - - def _on_redo(self, _e=None) -> str: - next_val = self._undo.redo() - if next_val is not None: - self._set_value(next_val) - return "break" + self._validate_and_style(self.component.text()) def flush(self) -> str | None: - if self._debounce: - self._debounce.cancel() - - value = self._shadow_var.get() - error = self.validate(value) - + self._debounce.stop() + val = self.component.text() + error = self.validate(val) if error is not None: self._apply_error() else: self._clear_error() self._commit() - return error -class PathValidator(FieldValidator): - """FieldValidator with additional path-specific checks.""" - +class PySide6PathValidator(PySide6FieldValidator): def __init__( self, - component: ctk.CTkEntry, - var: tk.Variable, - ui_state: UIState, + component: QLineEdit, + var: QtVar, + ui_state: BaseUIState, var_name: str, io_type: PathIOType = PathIOType.INPUT, max_undo: int = DEFAULT_MAX_UNDO, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, ): - super().__init__(component, var, ui_state, var_name, max_undo=max_undo, extra_validate=extra_validate, required=required) + super().__init__(component, var, ui_state, var_name, max_undo=max_undo, + extra_validate=extra_validate, required=required) self.io_type = io_type def validate(self, value: str) -> str | None: @@ -277,5 +158,19 @@ def validate(self, value: str) -> str | None: return _validate_path_field(self.ui_state, self.io_type, value) def revalidate(self) -> None: - if self.component.winfo_exists(): - self._validate_and_style(self._shadow_var.get()) + self._validate_and_style(self.component.text()) + + +def flush_and_validate_all_qt() -> list[str]: + invalid: list[str] = [] + for v in list(_active_qt_validators): + v._debounce.stop() + val = v.component.text() + error = v.validate(val) + if error is not None: + v._apply_error() + invalid.append(f"{v.var_name}: {error}") + else: + v._clear_error() + v._commit() + return invalid diff --git a/requirements-global.txt b/requirements-global.txt index 91014c6ff..d5b4e029c 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -49,6 +49,7 @@ scalene==1.5.51 # ui customtkinter==5.2.2 +PySide6==6.11.0 # cloud runpod==1.7.10 diff --git a/scripts/train_ui_pyside6.py b/scripts/train_ui_pyside6.py new file mode 100644 index 000000000..e666733f2 --- /dev/null +++ b/scripts/train_ui_pyside6.py @@ -0,0 +1,43 @@ +import sys + +# Force pydantic internals into sys.modules before PySide6/shiboken installs its +# import hooks. Without this, shiboken's inspect.getsource() fires on a +# partially-initialized pydantic module, causing a circular import error. +import pydantic._internal._validators # noqa: F401 +from util.import_util import script_imports + +script_imports() + +from modules.ui.PySide6TrainUIView import PySide6TrainView + +from PySide6.QtGui import QColor, QPalette +from PySide6.QtWidgets import QApplication + + +def main(): + app = QApplication(sys.argv) + + palette = app.palette() + palette.setColor(QPalette.ColorRole.Base, QColor("white")) + palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) + app.setPalette(palette) + + app.setStyleSheet(""" + QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { + padding: 2px 2px; + } + QCheckBox::indicator { + width: 16px; + height: 16px; + } + QProgressBar { + background-color: #c8c8c8; + } + """) + window = PySide6TrainView() + window.show() + sys.exit(app.exec()) + + +if __name__ == '__main__': + main() diff --git a/scripts/video_tool_ui.py b/scripts/video_tool_ui.py index 99707506f..8d9aadac0 100644 --- a/scripts/video_tool_ui.py +++ b/scripts/video_tool_ui.py @@ -2,7 +2,7 @@ script_imports() -from modules.ui.VideoToolUI import VideoToolUI +from modules.ui.CtkVideoToolUIView import VideoToolUI def main(): From ec2ee0a9547a24e5bb8d54e4884d2c48681e602c Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 11 May 2026 20:09:48 +0200 Subject: [PATCH 13/67] refactor: replace manual entry+browse-button pairs with path_entry in VideoToolUI Removes abstract _create_browse_dir_button/_create_browse_file_button from BaseVideoToolUIView and uses the combined path_entry component instead, fixing widget alignment issues. Adds allow_video_files flag to path_entry. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseVideoToolUIView.py | 39 +++++++++---------------------- modules/ui/CtkVideoToolUIView.py | 24 ------------------- modules/util/ui/ctk_components.py | 7 +++++- 3 files changed, 17 insertions(+), 53 deletions(-) diff --git a/modules/ui/BaseVideoToolUIView.py b/modules/ui/BaseVideoToolUIView.py index b82ba5701..e60247585 100644 --- a/modules/ui/BaseVideoToolUIView.py +++ b/modules/ui/BaseVideoToolUIView.py @@ -1,8 +1,6 @@ import webbrowser from abc import ABC, abstractmethod -from modules.util.path_util import SUPPORTED_VIDEO_EXTENSIONS - class BaseVideoToolUIView(ABC): def __init__(self, components): @@ -12,9 +10,8 @@ def build_clip_extract_tab(self, frame, controller, ui_state): # single video self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.components.entry(frame, 0, 1, ui_state, "clip_single", width=190) - self._create_browse_file_button(frame, 0, ui_state, "clip_single", - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.path_entry(frame, 0, 1, ui_state, "clip_single", + mode="file", allow_model_files=False, allow_video_files=True) self.components.button(frame, 0, 2, "Extract Single", command=lambda: self._extract_clips(False, controller)) @@ -28,16 +25,14 @@ def build_clip_extract_tab(self, frame, controller, ui_state): # directory of videos self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.components.entry(frame, 2, 1, ui_state, "clip_list", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "clip_list") + self.components.path_entry(frame, 2, 1, ui_state, "clip_list", mode="dir") self.components.button(frame, 2, 2, "Extract Directory", command=lambda: self._extract_clips(True, controller)) # output directory self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted clips will be saved.") - self.components.entry(frame, 3, 1, ui_state, "clip_output", width=190) - self._create_browse_dir_button(frame, 3, ui_state, "clip_output") + self.components.path_entry(frame, 3, 1, ui_state, "clip_output", mode="dir") # output to subdirectories self.components.label(frame, 4, 0, "Output to\nSubdirectories", @@ -76,9 +71,8 @@ def build_image_extract_tab(self, frame, controller, ui_state): # single video self.components.label(frame, 0, 0, "Single Video", tooltip="Link to single video file to process.") - self.components.entry(frame, 0, 1, ui_state, "image_single", width=190) - self._create_browse_file_button(frame, 0, ui_state, "image_single", - [("Video files", " ".join(f"*{e}" for e in SUPPORTED_VIDEO_EXTENSIONS))]) + self.components.path_entry(frame, 0, 1, ui_state, "image_single", + mode="file", allow_model_files=False, allow_video_files=True) self.components.button(frame, 0, 2, "Extract Single", command=lambda: self._extract_images(False, controller)) @@ -92,16 +86,14 @@ def build_image_extract_tab(self, frame, controller, ui_state): # directory of videos self.components.label(frame, 2, 0, "Directory", tooltip="Path to directory with multiple videos to process, including in subdirectories.") - self.components.entry(frame, 2, 1, ui_state, "image_list", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "image_list") + self.components.path_entry(frame, 2, 1, ui_state, "image_list", mode="dir") self.components.button(frame, 2, 2, "Extract Directory", command=lambda: self._extract_images(True, controller)) # output directory self.components.label(frame, 3, 0, "Output", tooltip="Path to folder where extracted images will be saved.") - self.components.entry(frame, 3, 1, ui_state, "image_output", width=190) - self._create_browse_dir_button(frame, 3, ui_state, "image_output") + self.components.path_entry(frame, 3, 1, ui_state, "image_output", mode="dir") # output to subdirectories self.components.label(frame, 4, 0, "Output to\nSubdirectories", @@ -143,16 +135,15 @@ def build_video_download_tab(self, frame, controller, ui_state): # link list self.components.label(frame, 1, 0, "Link List", tooltip="Path to txt file with list of links separated by newlines.") - self.components.entry(frame, 1, 1, ui_state, "download_list", width=190) - self._create_browse_file_button(frame, 1, ui_state, "download_list", [("Text file", ".txt")]) + self.components.path_entry(frame, 1, 1, ui_state, "download_list", + mode="file", allow_model_files=False) self.components.button(frame, 1, 2, "Download List", command=lambda: self._download(True, controller)) # output directory self.components.label(frame, 2, 0, "Output", tooltip="Path to folder where downloaded videos will be saved.") - self.components.entry(frame, 2, 1, ui_state, "download_output", width=190) - self._create_browse_dir_button(frame, 2, ui_state, "download_output") + self.components.path_entry(frame, 2, 1, ui_state, "download_output", mode="dir") # additional args self.components.label(frame, 3, 0, "Additional Args", @@ -166,14 +157,6 @@ def build_video_download_tab(self, frame, controller, ui_state): def _create_textbox(self, master, row, col, width, height, ui_state, var_name): pass - @abstractmethod - def _create_browse_dir_button(self, master, row, ui_state, var_name): - pass - - @abstractmethod - def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): - pass - @abstractmethod def update_status(self, status_text: str): pass diff --git a/modules/ui/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py index c272c891c..d6190772f 100644 --- a/modules/ui/CtkVideoToolUIView.py +++ b/modules/ui/CtkVideoToolUIView.py @@ -1,5 +1,3 @@ -from tkinter import filedialog - from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image @@ -91,28 +89,6 @@ def on_text_change(event=None): textbox.bind("", on_text_change) return textbox - def _create_browse_dir_button(self, master, row, ui_state, var_name): - def browse(): - path = filedialog.askdirectory() - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - - def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): - def browse(): - path = filedialog.askopenfilename(filetypes=filetypes) - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - def update_status(self, status_text: str): self.status_label.configure(state="normal") self.status_label.insert(index="end", text=status_text + "\n") diff --git a/modules/util/ui/ctk_components.py b/modules/util/ui/ctk_components.py index e462f72a1..ec5e0ecf9 100644 --- a/modules/util/ui/ctk_components.py +++ b/modules/util/ui/ctk_components.py @@ -7,7 +7,7 @@ from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit -from modules.util.path_util import supported_image_extensions +from modules.util.path_util import supported_image_extensions, supported_video_extensions from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ToolTip import ToolTip @@ -118,6 +118,7 @@ def path_entry( path_modifier: Callable[[str], str | Path] | None = None, allow_model_files: bool = True, allow_image_files: bool = False, + allow_video_files: bool = False, command: Callable[[str], None] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, @@ -181,6 +182,10 @@ def __open_dialog(): filetypes.extend([ ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), ]) + if allow_video_files: + filetypes.extend([ + ("Video", ' '.join(f"*{e}" for e in supported_video_extensions())), + ]) if use_save_dialog: chosen = filedialog.asksaveasfilename(filetypes=filetypes, initialdir=current_dir, From e19d1496211da6e490329c417e32cfa3c0417dc3 Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 11 May 2026 20:11:41 +0200 Subject: [PATCH 14/67] refactor: mirror VideoToolUI path_entry refactor to PySide6VideoToolUIView Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6VideoToolUIView.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/modules/ui/PySide6VideoToolUIView.py b/modules/ui/PySide6VideoToolUIView.py index c272c891c..d6190772f 100644 --- a/modules/ui/PySide6VideoToolUIView.py +++ b/modules/ui/PySide6VideoToolUIView.py @@ -1,5 +1,3 @@ -from tkinter import filedialog - from modules.ui.BaseVideoToolUIView import BaseVideoToolUIView from modules.ui.VideoToolUIController import VideoToolUIController from modules.util.image_util import load_image @@ -91,28 +89,6 @@ def on_text_change(event=None): textbox.bind("", on_text_change) return textbox - def _create_browse_dir_button(self, master, row, ui_state, var_name): - def browse(): - path = filedialog.askdirectory() - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - - def _create_browse_file_button(self, master, row, ui_state, var_name, filetypes): - def browse(): - path = filedialog.askopenfilename(filetypes=filetypes) - if path: - ui_state.get_var(var_name).set(path) - self.focus_set() - - button = ctk.CTkButton(master, width=30, text="...", command=browse) - button.grid(row=row, column=1, sticky="e", padx=PAD, pady=PAD) - return button - def update_status(self, status_text: str): self.status_label.configure(state="normal") self.status_label.insert(index="end", text=status_text + "\n") From 65bdda2e07b2d9754fd4598b25fc65e3787938f6 Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 11 May 2026 20:15:45 +0200 Subject: [PATCH 15/67] copy: sync pyside6_components.py with ctk_components.py (allow_video_files) Co-Authored-By: Claude Sonnet 4.6 --- modules/util/ui/pyside6_components.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modules/util/ui/pyside6_components.py b/modules/util/ui/pyside6_components.py index e462f72a1..ec5e0ecf9 100644 --- a/modules/util/ui/pyside6_components.py +++ b/modules/util/ui/pyside6_components.py @@ -7,7 +7,7 @@ from modules.util.enum.PathIOType import PathIOType from modules.util.enum.TimeUnit import TimeUnit -from modules.util.path_util import supported_image_extensions +from modules.util.path_util import supported_image_extensions, supported_video_extensions from modules.util.ui.ctk_validation import DEFAULT_MAX_UNDO, FieldValidator, PathValidator from modules.util.ui.CtkUIState import CtkUIState from modules.util.ui.ToolTip import ToolTip @@ -118,6 +118,7 @@ def path_entry( path_modifier: Callable[[str], str | Path] | None = None, allow_model_files: bool = True, allow_image_files: bool = False, + allow_video_files: bool = False, command: Callable[[str], None] | None = None, extra_validate: Callable[[str], str | None] | None = None, required: bool = False, @@ -181,6 +182,10 @@ def __open_dialog(): filetypes.extend([ ("Image", ' '.join([f"*.{x}" for x in supported_image_extensions()])), ]) + if allow_video_files: + filetypes.extend([ + ("Video", ' '.join(f"*{e}" for e in supported_video_extensions())), + ]) if use_save_dialog: chosen = filedialog.asksaveasfilename(filetypes=filetypes, initialdir=current_dir, From c3961ba48afeacbb86c059fa1aa97ffe16d0d7f8 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 11:40:15 +0200 Subject: [PATCH 16/67] fix: declare undeclared self attrs in Base*View classes Base view concrete methods were accessing self.controller, self.ui_state, and toolkit-specific action methods that were only set by CTK subclass __init__ after the base __init__ call, with no enforcement in the base. - BaseTrainUIView: add controller/ui_state as constructor params; fix sync_cloud_secrets to use controller.train_config; add @abstractmethod for export_training, generate_debug_package, open_profiling_tool - CtkTrainUIView: reorder __init__ to create deps before base init call; replace self.train_config with self.controller.train_config - BaseCloudTabView: add controller as constructor param - CtkCloudTabView: pass controller to base __init__, drop redundant assignment - BaseCaptionUIView: add ABC + @abstractmethod for 6 action callbacks - BaseConceptTabView: remove concrete _update_filters() (accessed CTK vars); add ConceptConfig import; add concept: ConceptConfig param to BaseConceptWidgetView.__init__ - CtkConceptTabView: implement _update_filters(); pass concept to base init - BaseConceptWindowView: initialize bucket_ax/text_color/canvas to None - BaseTrainingTabView: replace callbacks dict with 6 @abstractmethod declarations; restore_optimizer_config(variable: str) matches controller - CtkTrainingTabView: implement all 6 abstract methods directly Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseCaptionUIView.py | 21 +++- modules/ui/BaseCloudTabView.py | 3 +- modules/ui/BaseConceptTabView.py | 14 ++- modules/ui/BaseConceptWindowView.py | 3 + modules/ui/BaseTrainUIView.py | 15 ++- modules/ui/BaseTrainingTabView.py | 184 +++++++++++++++------------- modules/ui/CtkCloudTabView.py | 3 +- modules/ui/CtkConceptTabView.py | 10 +- modules/ui/CtkTrainUIView.py | 32 ++--- modules/ui/CtkTrainingTabView.py | 34 +++-- 10 files changed, 184 insertions(+), 135 deletions(-) diff --git a/modules/ui/BaseCaptionUIView.py b/modules/ui/BaseCaptionUIView.py index a6561da69..c53ca84d8 100644 --- a/modules/ui/BaseCaptionUIView.py +++ b/modules/ui/BaseCaptionUIView.py @@ -1,10 +1,29 @@ import platform +from abc import ABC, abstractmethod -class BaseCaptionUIView: +class BaseCaptionUIView(ABC): def __init__(self, components): self.components = components + @abstractmethod + def open_directory(self): pass + + @abstractmethod + def open_mask_window(self): pass + + @abstractmethod + def open_caption_window(self): pass + + @abstractmethod + def open_in_explorer(self): pass + + @abstractmethod + def draw_mask_editing_mode(self, *args): pass + + @abstractmethod + def fill_mask_editing_mode(self, *args): pass + def build_top_bar(self, frame, controller, ui_state): self.components.button(frame, 0, 0, "Open", self.open_directory, tooltip="open a new directory") diff --git a/modules/ui/BaseCloudTabView.py b/modules/ui/BaseCloudTabView.py index 80be17360..3f20b5674 100644 --- a/modules/ui/BaseCloudTabView.py +++ b/modules/ui/BaseCloudTabView.py @@ -7,8 +7,9 @@ class BaseCloudTabView(ABC): - def __init__(self, components): + def __init__(self, components, controller): self.components = components + self.controller = controller @property def reattach(self): diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 077a10c24..4d2fca570 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -4,6 +4,7 @@ from modules.ui.BaseConfigListView import BaseConfigListView from modules.ui.ConceptWindowController import ConceptWindowController from modules.util import path_util +from modules.util.config.ConceptConfig import ConceptConfig from modules.util.enum.ConceptType import ConceptType from modules.util.image_util import load_image @@ -14,6 +15,11 @@ class BaseConceptTabView(BaseConfigListView): _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] + def __init__(self, search_var, filter_var, show_disabled_var): + self.search_var = search_var + self.filter_var = filter_var + self.show_disabled_var = show_disabled_var + def _element_matches_filters(self, element): if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: @@ -50,17 +56,13 @@ def _element_matches_filters(self, element): return True - def _update_filters(self): - self._create_element_list(search=self.search_var.get(), - type=self.filter_var.get(), - show_disabled=self.show_disabled_var.get()) - self._refresh_show_disabled_text() class BaseConceptWidgetView: - def __init__(self, components): + def __init__(self, components, concept: ConceptConfig): self.components = components + self.concept = concept def _get_display_name(self): if self.concept.name: diff --git a/modules/ui/BaseConceptWindowView.py b/modules/ui/BaseConceptWindowView.py index 0f851f5fb..0b94e50d1 100644 --- a/modules/ui/BaseConceptWindowView.py +++ b/modules/ui/BaseConceptWindowView.py @@ -9,6 +9,9 @@ class BaseConceptWindowView: def __init__(self, components): self.components = components + self.bucket_ax = None + self.text_color = None + self.canvas = None def build_general_tab(self, frame, controller, ui_state, text_ui_state): # name diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 7acbf0005..9a443a342 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -9,8 +9,10 @@ class BaseTrainUIView(ABC): - def __init__(self, components): + def __init__(self, components, controller, ui_state): self.components = components + self.controller = controller + self.ui_state = ui_state # --- Abstract callbacks (controller calls into view) --- @@ -51,7 +53,7 @@ def show_window(self, window): pass def connect_window_closed(self, window, callback): pass def sync_cloud_secrets(self): - self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + self.ui_state.get_var("secrets.cloud").update(self.controller.train_config.secrets.cloud) def start_training(self): self.controller.start_training() @@ -83,6 +85,15 @@ def open_sampling_tool(self): pass @abstractmethod def open_manual_sample_window(self): pass + @abstractmethod + def open_profiling_tool(self): pass + + @abstractmethod + def export_training(self): pass + + @abstractmethod + def generate_debug_package(self): pass + # --- Content builders (components calls; called by CTK view after frame creation) --- def build_bottom_bar_content(self, frame, status_frame, controller, ui_state): diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index 22a6af34e..947ca4f56 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -1,4 +1,4 @@ -from abc import ABC +from abc import ABC, abstractmethod from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode @@ -16,232 +16,250 @@ class BaseTrainingTabView(ABC): def __init__(self, components): self.components = components - def build(self, column_0, column_1, column_2, controller, ui_state, callbacks: dict): + @abstractmethod + def restore_optimizer_config(self, variable: str): pass + + @abstractmethod + def open_optimizer_params(self): pass + + @abstractmethod + def restore_scheduler(self, variable): pass + + @abstractmethod + def open_scheduler_params(self): pass + + @abstractmethod + def open_offloading(self): pass + + @abstractmethod + def open_timestep_distribution(self): pass + + def build(self, column_0, column_1, column_2, controller, ui_state): model_type = controller.config.model_type if model_type.is_stable_diffusion(): - self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_ui(column_0, column_1, column_2, controller, ui_state) if model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_3_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_stable_diffusion_xl_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_wuerstchen_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_pixart(): - self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_pixart_alpha_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_flux_1(): - self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_flux_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_flux_2(): - self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_flux_2_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_chroma(): - self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_chroma_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_qwen(): - self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_qwen_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_sana(): - self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_sana_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_hunyuan_video_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_hi_dream(): - self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_hi_dream_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_z_image(): - self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_ernie(): - self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state, callbacks) + self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state) - def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_unet_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_generalized_offset_noise=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_3_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_stable_diffusion_xl_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2) self.__create_embedding_frame(column_0, 3, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_unet_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_generalized_offset_noise=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_generalized_offset_noise=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_wuerstchen_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, supports_circular_padding=True) + self.__create_base2_frame(column_1, 0, ui_state, supports_circular_padding=True) self.__create_prior_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 0, ui_state) self.__create_loss_frame(column_2, 1, controller, ui_state) self.__create_layer_frame(column_2, 2, controller, ui_state) - def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_pixart_alpha_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state, supports_vb_loss=True) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_flux_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True, supports_sequence_length=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_flux_2_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False, supports_sequence_length=True) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_chroma_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_qwen_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_z_image_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) - self.__create_noise_frame(column_1, 2, ui_state, callbacks, supports_dynamic_timestep_shifting=True) + self.__create_noise_frame(column_1, 2, ui_state, supports_dynamic_timestep_shifting=True) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) self.__create_embedding_frame(column_0, 2, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks) + self.__create_base2_frame(column_1, 0, ui_state) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_hunyuan_video_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_embedding_frame(column_0, 4, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_base2_frame(column_1, 0, ui_state, video_training_enabled=True) self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=True) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state, callbacks): - self.__create_base_frame(column_0, 0, controller, ui_state, callbacks) + def __setup_hi_dream_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_n_frame(column_0, 1, ui_state, i=1, supports_include=True) self.__create_text_encoder_n_frame(column_0, 2, ui_state, i=2, supports_include=True) self.__create_text_encoder_n_frame(column_0, 3, ui_state, i=3, supports_include=True) self.__create_text_encoder_n_frame(column_0, 4, ui_state, i=4, supports_include=True, supports_layer_skip=False) self.__create_embedding_frame(column_0, 5, ui_state) - self.__create_base2_frame(column_1, 0, ui_state, callbacks, video_training_enabled=True) + self.__create_base2_frame(column_1, 0, ui_state, video_training_enabled=True) self.__create_transformer_frame(column_1, 1, ui_state) - self.__create_noise_frame(column_1, 2, ui_state, callbacks) + self.__create_noise_frame(column_1, 2, ui_state) self.__create_masked_frame(column_2, 1, ui_state) self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) - def __create_base_frame(self, master, row, controller, ui_state, callbacks): + def __create_base_frame(self, master, row, controller, ui_state): frame = self.components.section_frame(master, row) # optimizer self.components.label(frame, 0, 0, "Optimizer", tooltip="The type of optimizer") self.components.options_adv(frame, 0, 1, [str(x) for x in list(Optimizer)], ui_state, "optimizer.optimizer", - command=callbacks.get('restore_optimizer'), - adv_command=callbacks.get('open_optimizer_params')) + command=self.restore_optimizer_config, + adv_command=self.open_optimizer_params) # learning rate scheduler # Wackiness will ensue when reloading configs if we don't check and clear this first. @@ -252,14 +270,12 @@ def __create_base_frame(self, master, row, controller, ui_state, callbacks): tooltip="Learning rate scheduler that automatically changes the learning rate during training") _, d = self.components.options_adv(frame, 1, 1, [str(x) for x in list(LearningRateScheduler)], ui_state, "learning_rate_scheduler", - command=callbacks.get('restore_scheduler'), - adv_command=callbacks.get('open_scheduler_params')) + command=self.restore_scheduler, + adv_command=self.open_scheduler_params) self.lr_scheduler_comp = d['component'] self.lr_scheduler_adv_comp = d['button_component'] # Initial call requires the presence of self.lr_scheduler_adv_comp. - restore_scheduler = callbacks.get('restore_scheduler') - if restore_scheduler: - restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) + self.restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) # learning rate self.components.label(frame, 2, 0, "Learning Rate", @@ -308,7 +324,7 @@ def __create_base_frame(self, master, row, controller, ui_state, callbacks): tooltip="Clips the gradient norm. Leave empty to disable gradient clipping.") self.components.entry(frame, 10, 1, ui_state, "clip_grad_norm") - def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_enabled: bool = False, + def __create_base2_frame(self, master, row, ui_state, video_training_enabled: bool = False, supports_circular_padding: bool = False): frame = self.components.section_frame(master, row) row = 0 @@ -338,7 +354,7 @@ def __create_base2_frame(self, master, row, ui_state, callbacks, video_training_ tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") self.components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, "gradient_checkpointing", - adv_command=callbacks.get('open_offloading')) + adv_command=self.open_offloading) row += 1 # gradient checkpointing layer offloading @@ -588,7 +604,7 @@ def __create_transformer_frame(self, master, row, ui_state, supports_guidance_sc tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") self.components.entry(frame, 4, 1, ui_state, "transformer.guidance_scale") - def __create_noise_frame(self, master, row, ui_state, callbacks, + def __create_noise_frame(self, master, row, ui_state, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): frame = self.components.section_frame(master, row) @@ -616,7 +632,7 @@ def __create_noise_frame(self, master, row, ui_state, callbacks, wide_tooltip=True) self.components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], ui_state, "timestep_distribution", - adv_command=callbacks.get('open_timestep_distribution')) + adv_command=self.open_timestep_distribution) # min noising strength self.components.label(frame, 4, 0, "Min Noising Strength", diff --git a/modules/ui/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py index 0a5249069..ffe37451f 100644 --- a/modules/ui/CtkCloudTabView.py +++ b/modules/ui/CtkCloudTabView.py @@ -9,9 +9,8 @@ class CtkCloudTabView(BaseCloudTabView): def __init__(self, master, controller: CloudTabController, ui_state): - BaseCloudTabView.__init__(self, ctk_components) + BaseCloudTabView.__init__(self, ctk_components, controller) self.master = master - self.controller = controller self.ui_state = ui_state self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") diff --git a/modules/ui/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py index 5b3e86ac9..84d8e4508 100644 --- a/modules/ui/CtkConceptTabView.py +++ b/modules/ui/CtkConceptTabView.py @@ -86,6 +86,12 @@ def _maybe_reposition_toolbar(self, width): else: self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() + def _reset_filters(self): self.search_var.set("") self.filter_var.set("ALL") @@ -109,9 +115,7 @@ class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") - BaseConceptWidgetView.__init__(self, ctk_components) - - self.concept = concept + BaseConceptWidgetView.__init__(self, ctk_components, concept) self.ui_state = CtkUIState(self, concept) self.image_ui_state = CtkUIState(self, concept.image) self.text_ui_state = CtkUIState(self, concept.text) diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index d3c1a70b5..41af55230 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -82,7 +82,13 @@ class CtkTrainUIView(BaseTrainUIView, ctk.CTk): def __init__(self): ctk.CTk.__init__(self) - BaseTrainUIView.__init__(self, ctk_components) + + train_config = TrainConfig.default_values() + ui_state = CtkUIState(self, train_config) + controller = TrainUIController(train_config) + + BaseTrainUIView.__init__(self, ctk_components, controller, ui_state) + self.controller.view = self self.title("OneTrainer") self.geometry("1100x740") @@ -93,12 +99,6 @@ def __init__(self): ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") ctk.set_default_color_theme("blue") - self.train_config = TrainConfig.default_values() - self.ui_state = CtkUIState(self, self.train_config) - - self.controller = TrainUIController(self.train_config) - self.controller.view = self - self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_rowconfigure(2, weight=0) @@ -212,7 +212,7 @@ def _set_icon(self): def top_bar(self, master): return CtkTopBarView( master, - TopBarController(self.train_config), + TopBarController(self.controller.train_config), self.ui_state, self.change_model_type, self.change_training_method, @@ -259,7 +259,7 @@ def content_frame(self, master): self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) - self.change_training_method(self.train_config.training_method) + self.change_training_method(self.controller.train_config.training_method) return frame @@ -274,7 +274,7 @@ def create_general_tab(self, master): return frame def create_model_tab(self, master): - return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) + return CtkModelTabView(master, ModelTabController(self.controller.train_config), self.ui_state) def create_data_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -288,13 +288,13 @@ def create_data_tab(self, master): return frame def create_concepts_tab(self, master): - return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) + return CtkConceptTabView(master, ConceptTabController(self.controller.train_config), self.ui_state) def create_training_tab(self, master) -> CtkTrainingTabView: - return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) + return CtkTrainingTabView(master, TrainingTabController(self.controller.train_config), self.ui_state) def create_cloud_tab(self, master) -> CtkCloudTabView: - return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) + return CtkCloudTabView(master, CloudTabController(self.controller.train_config, parent=self), self.ui_state) def create_sampling_tab(self, master): master.grid_rowconfigure(0, weight=0) @@ -311,7 +311,7 @@ def create_sampling_tab(self, master): frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=1, column=0, sticky="nsew") - return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) + return CtkSamplingTabView(frame, SamplingTabController(self.controller.train_config), self.ui_state) def create_backup_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -336,7 +336,7 @@ def embedding_tab(self, master): return frame def create_additional_embeddings_tab(self, master): - return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.controller.train_config), self.ui_state) def create_tools_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -376,7 +376,7 @@ def change_training_method(self, training_method: TrainingMethod): self.tabview.delete("embedding") if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.controller.train_config), self.ui_state) if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: self.embedding_tab(self.tabview.add("embedding")) diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bc29488dd..bce309876 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -47,31 +47,25 @@ def refresh_ui(self): column_2.grid(row=0, column=2, sticky="nsew") column_2.grid_columnconfigure(0, weight=1) - callbacks = { - 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), - 'open_optimizer_params': self._open_optimizer_params_window, - 'restore_scheduler': self._restore_scheduler_config, - 'open_scheduler_params': self._open_scheduler_params_window, - 'open_offloading': self._open_offloading_window, - 'open_timestep_distribution': self._open_timestep_distribution_window, - } - - self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) - - def _restore_scheduler_config(self, variable): + self.build(column_0, column_1, column_2, self.controller, self.ui_state) + + def restore_optimizer_config(self, variable: str): + self.controller.restore_optimizer_config(self.ui_state) + + def open_optimizer_params(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) + + def restore_scheduler(self, variable: str): if not hasattr(self, 'lr_scheduler_adv_comp'): return state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" self.lr_scheduler_adv_comp.configure(state=state) - def _open_optimizer_params_window(self): - self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) - - def _open_scheduler_params_window(self): + def open_scheduler_params(self): self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def _open_timestep_distribution_window(self): - self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) - - def _open_offloading_window(self): + def open_offloading(self): self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) + + def open_timestep_distribution(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) From d03987122553521e733cfd858ffd5eeac5c4d867 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 11:51:00 +0200 Subject: [PATCH 17/67] refactor: move Windows DPI awareness block to CtkTrainUIView The ctypes DPI awareness call is toolkit-specific (fixes CTK transparency on Windows monitor changes). It already exists in CtkTrainUIView.py and has no place in the toolkit-agnostic controller. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/TrainUIController.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index 15d1eaff6..77c4de1a7 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -1,15 +1,12 @@ -import ctypes import datetime import json import os -import platform import subprocess import sys import threading import time import traceback import webbrowser -from contextlib import suppress from pathlib import Path import scripts.generate_debug_report @@ -28,13 +25,6 @@ import torch -# chunk for forcing Windows to ignore DPI scaling when moving between monitors -# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 -if platform.system() == "Windows": - with suppress(Exception): - # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically - ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE - class TrainUIController: def __init__(self, config: TrainConfig): From 29a62a197a5ceeba60c6b93bb61f034607afe753 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 11:52:16 +0200 Subject: [PATCH 18/67] chore: sync PySide6 copies from ctk_abstraction merge Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6CloudTabView.py | 3 +-- modules/ui/PySide6ConceptTabView.py | 10 +++++--- modules/ui/PySide6TrainUIView.py | 32 +++++++++++++------------- modules/ui/PySide6TrainingTabView.py | 34 ++++++++++++---------------- 4 files changed, 38 insertions(+), 41 deletions(-) diff --git a/modules/ui/PySide6CloudTabView.py b/modules/ui/PySide6CloudTabView.py index 0a5249069..ffe37451f 100644 --- a/modules/ui/PySide6CloudTabView.py +++ b/modules/ui/PySide6CloudTabView.py @@ -9,9 +9,8 @@ class CtkCloudTabView(BaseCloudTabView): def __init__(self, master, controller: CloudTabController, ui_state): - BaseCloudTabView.__init__(self, ctk_components) + BaseCloudTabView.__init__(self, ctk_components, controller) self.master = master - self.controller = controller self.ui_state = ui_state self.frame = ctk.CTkScrollableFrame(master, fg_color="transparent") diff --git a/modules/ui/PySide6ConceptTabView.py b/modules/ui/PySide6ConceptTabView.py index 5b3e86ac9..84d8e4508 100644 --- a/modules/ui/PySide6ConceptTabView.py +++ b/modules/ui/PySide6ConceptTabView.py @@ -86,6 +86,12 @@ def _maybe_reposition_toolbar(self, width): else: self._toolbar.grid_configure(row=0, column=4, columnspan=2, sticky="ew", padx=10) + def _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + self._refresh_show_disabled_text() + def _reset_filters(self): self.search_var.set("") self.filter_var.set("ALL") @@ -109,9 +115,7 @@ class CtkConceptWidgetView(BaseConceptWidgetView, ctk.CTkFrame): def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): ctk.CTkFrame.__init__(self, master=master, width=150, height=170, corner_radius=10, bg_color="transparent") - BaseConceptWidgetView.__init__(self, ctk_components) - - self.concept = concept + BaseConceptWidgetView.__init__(self, ctk_components, concept) self.ui_state = CtkUIState(self, concept) self.image_ui_state = CtkUIState(self, concept.image) self.text_ui_state = CtkUIState(self, concept.text) diff --git a/modules/ui/PySide6TrainUIView.py b/modules/ui/PySide6TrainUIView.py index d3c1a70b5..41af55230 100644 --- a/modules/ui/PySide6TrainUIView.py +++ b/modules/ui/PySide6TrainUIView.py @@ -82,7 +82,13 @@ class CtkTrainUIView(BaseTrainUIView, ctk.CTk): def __init__(self): ctk.CTk.__init__(self) - BaseTrainUIView.__init__(self, ctk_components) + + train_config = TrainConfig.default_values() + ui_state = CtkUIState(self, train_config) + controller = TrainUIController(train_config) + + BaseTrainUIView.__init__(self, ctk_components, controller, ui_state) + self.controller.view = self self.title("OneTrainer") self.geometry("1100x740") @@ -93,12 +99,6 @@ def __init__(self): ctk.set_appearance_mode("Light" if AppearanceModeTracker.detect_appearance_mode() == 0 else "Dark") ctk.set_default_color_theme("blue") - self.train_config = TrainConfig.default_values() - self.ui_state = CtkUIState(self, self.train_config) - - self.controller = TrainUIController(self.train_config) - self.controller.view = self - self.grid_rowconfigure(0, weight=0) self.grid_rowconfigure(1, weight=1) self.grid_rowconfigure(2, weight=0) @@ -212,7 +212,7 @@ def _set_icon(self): def top_bar(self, master): return CtkTopBarView( master, - TopBarController(self.train_config), + TopBarController(self.controller.train_config), self.ui_state, self.change_model_type, self.change_training_method, @@ -259,7 +259,7 @@ def content_frame(self, master): self.additional_embeddings_tab = self.create_additional_embeddings_tab(self.tabview.add("additional embeddings")) self.cloud_tab = self.create_cloud_tab(self.tabview.add("cloud")) - self.change_training_method(self.train_config.training_method) + self.change_training_method(self.controller.train_config.training_method) return frame @@ -274,7 +274,7 @@ def create_general_tab(self, master): return frame def create_model_tab(self, master): - return CtkModelTabView(master, ModelTabController(self.train_config), self.ui_state) + return CtkModelTabView(master, ModelTabController(self.controller.train_config), self.ui_state) def create_data_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -288,13 +288,13 @@ def create_data_tab(self, master): return frame def create_concepts_tab(self, master): - return CtkConceptTabView(master, ConceptTabController(self.train_config), self.ui_state) + return CtkConceptTabView(master, ConceptTabController(self.controller.train_config), self.ui_state) def create_training_tab(self, master) -> CtkTrainingTabView: - return CtkTrainingTabView(master, TrainingTabController(self.train_config), self.ui_state) + return CtkTrainingTabView(master, TrainingTabController(self.controller.train_config), self.ui_state) def create_cloud_tab(self, master) -> CtkCloudTabView: - return CtkCloudTabView(master, CloudTabController(self.train_config, parent=self), self.ui_state) + return CtkCloudTabView(master, CloudTabController(self.controller.train_config, parent=self), self.ui_state) def create_sampling_tab(self, master): master.grid_rowconfigure(0, weight=0) @@ -311,7 +311,7 @@ def create_sampling_tab(self, master): frame = ctk.CTkFrame(master=master, corner_radius=0) frame.grid(row=1, column=0, sticky="nsew") - return CtkSamplingTabView(frame, SamplingTabController(self.train_config), self.ui_state) + return CtkSamplingTabView(frame, SamplingTabController(self.controller.train_config), self.ui_state) def create_backup_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -336,7 +336,7 @@ def embedding_tab(self, master): return frame def create_additional_embeddings_tab(self, master): - return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.train_config), self.ui_state) + return CtkAdditionalEmbeddingsTabView(master, AdditionalEmbeddingsTabController(self.controller.train_config), self.ui_state) def create_tools_tab(self, master): frame = ctk.CTkScrollableFrame(master, fg_color="transparent") @@ -376,7 +376,7 @@ def change_training_method(self, training_method: TrainingMethod): self.tabview.delete("embedding") if training_method == TrainingMethod.LORA and "LoRA" not in self.tabview._tab_dict: - self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.train_config), self.ui_state) + self.lora_tab = CtkLoraTabView(self.tabview.add("LoRA"), LoraTabController(self.controller.train_config), self.ui_state) if training_method == TrainingMethod.EMBEDDING and "embedding" not in self.tabview._tab_dict: self.embedding_tab(self.tabview.add("embedding")) diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py index bc29488dd..bce309876 100644 --- a/modules/ui/PySide6TrainingTabView.py +++ b/modules/ui/PySide6TrainingTabView.py @@ -47,31 +47,25 @@ def refresh_ui(self): column_2.grid(row=0, column=2, sticky="nsew") column_2.grid_columnconfigure(0, weight=1) - callbacks = { - 'restore_optimizer': lambda *args: self.controller.restore_optimizer_config(self.ui_state), - 'open_optimizer_params': self._open_optimizer_params_window, - 'restore_scheduler': self._restore_scheduler_config, - 'open_scheduler_params': self._open_scheduler_params_window, - 'open_offloading': self._open_offloading_window, - 'open_timestep_distribution': self._open_timestep_distribution_window, - } - - self.build(column_0, column_1, column_2, self.controller, self.ui_state, callbacks) - - def _restore_scheduler_config(self, variable): + self.build(column_0, column_1, column_2, self.controller, self.ui_state) + + def restore_optimizer_config(self, variable: str): + self.controller.restore_optimizer_config(self.ui_state) + + def open_optimizer_params(self): + self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) + + def restore_scheduler(self, variable: str): if not hasattr(self, 'lr_scheduler_adv_comp'): return state = "normal" if self.controller.is_custom_scheduler_value(variable) else "disabled" self.lr_scheduler_adv_comp.configure(state=state) - def _open_optimizer_params_window(self): - self.master.wait_window(self.controller.open_optimizer_params_window(self.master, self.ui_state, CtkOptimizerParamsWindowView)) - - def _open_scheduler_params_window(self): + def open_scheduler_params(self): self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def _open_timestep_distribution_window(self): - self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) - - def _open_offloading_window(self): + def open_offloading(self): self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) + + def open_timestep_distribution(self): + self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) From 1bb5a567104ae53ad19bb00594c1922d95510b1e Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 12:36:28 +0200 Subject: [PATCH 19/67] refactor: remove unused __init__ from BaseConceptTabView search_var/filter_var/show_disabled_var were stored in the base but never used there after _update_filters() was removed. Subclasses manage them. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseConceptTabView.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py index 4d2fca570..6a55e515b 100644 --- a/modules/ui/BaseConceptTabView.py +++ b/modules/ui/BaseConceptTabView.py @@ -15,11 +15,6 @@ class BaseConceptTabView(BaseConfigListView): _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] - def __init__(self, search_var, filter_var, show_disabled_var): - self.search_var = search_var - self.filter_var = filter_var - self.show_disabled_var = show_disabled_var - def _element_matches_filters(self, element): if not self.filters.get("show_disabled", True): if hasattr(element, 'enabled') and not element.enabled: From a38020a3b2b18042091a1845c9753a6ae6e594d0 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 12:58:38 +0200 Subject: [PATCH 20/67] refactor: move QSizePolicy import to module level in PySide6TrainingTabView Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6TrainingTabView.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py index 2edd61a8a..0311df6a9 100644 --- a/modules/ui/PySide6TrainingTabView.py +++ b/modules/ui/PySide6TrainingTabView.py @@ -11,7 +11,7 @@ from modules.util.ui import pyside6_components from modules.util.ui.pyside6_abc import QtABCMeta -from PySide6.QtWidgets import QScrollArea, QWidget +from PySide6.QtWidgets import QScrollArea, QSizePolicy, QWidget class PySide6TrainingTabView(BaseTrainingTabView, QWidget, metaclass=QtABCMeta): @@ -44,8 +44,6 @@ def refresh_ui(self): lo.setColumnStretch(1, 1) lo.setColumnStretch(2, 1) - from PySide6.QtWidgets import QSizePolicy - column_0 = QWidget(self.scroll_frame) column_0.setMinimumWidth(0) column_0.setSizePolicy(QSizePolicy.Policy.Ignored, QSizePolicy.Policy.Preferred) From 9c05ddd85dd6b7ddafd9f9bfdc81fd1d92bb0b51 Mon Sep 17 00:00:00 2001 From: dxqb Date: Fri, 15 May 2026 21:32:03 +0200 Subject: [PATCH 21/67] fix: store str repr in options_kv var so UIState enum lookup works UIState's enum trace looks up var_type[string], so the var must hold the string repr of the value, not the value itself. Co-Authored-By: Claude Opus 4.7 --- modules/util/ui/pyside6_components.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modules/util/ui/pyside6_components.py b/modules/util/ui/pyside6_components.py index 1552cc148..8e0f936dd 100644 --- a/modules/util/ui/pyside6_components.py +++ b/modules/util/ui/pyside6_components.py @@ -543,7 +543,8 @@ def options_kv( str_values = [str(v) for _, v in values] if var.get() not in str_values and keys: - var.set(values[0][1]) + # store the str repr — UIState's enum trace looks up var_type[string] + var.set(str(values[0][1])) _updating = False @@ -554,7 +555,7 @@ def on_combo(key: str): _updating = True for k, v in values: if key == k: - var.set(v) + var.set(str(v)) if command: command(v) break From 8dd86c62ce2001739ef37a45dba9d77d8fec1a3c Mon Sep 17 00:00:00 2001 From: dxqb Date: Sat, 16 May 2026 16:18:54 +0200 Subject: [PATCH 22/67] fix: apply upstream COFT removal to CtkLoraTabView Mirrors the change from upstream commit e928ddab (Remove COFT #1447), which removed COFT from LoraTab.py. The merge didn't carry it across the rename to CtkLoraTabView.py. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/CtkLoraTabView.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index 1c73d90ce..faabf9607 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -121,19 +121,10 @@ def setup_lora(self, peft_type: PeftType): tooltip=f"The block size parameter used when creating a new {name}") components.entry(master, 1, 1, self.ui_state, "oft_block_size", required=True) - # COFT - components.label(master, 1, 3, "Constrained OFT (COFT)", - tooltip="Use the constrained variant of OFT. This constrains the learned rotation to stay very close to the identity matrix, limiting adaptation to only small changes. This improves training stability, helps prevent overfitting on small datasets, and better preserves the base models original knowledge but it may lack expressiveness for tasks requiring substantial adaptation and introduces an additional hyperparameter (COFT Epsilon) that needs tuning.") - components.switch(master, 1, 4, self.ui_state, "oft_coft") - - components.label(master, 2, 3, "COFT Epsilon", - tooltip="The control strength of COFT. Only has an effect if COFT is enabled.") - components.entry(master, 2, 4, self.ui_state, "coft_eps") - # Block Share - components.label(master, 3, 3, "Block Share", + components.label(master, 1, 3, "Block Share", tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") - components.switch(master, 3, 4, self.ui_state, "oft_block_share") + components.switch(master, 1, 4, self.ui_state, "oft_block_share") # Dropout Percentage components.label(master, 2, 0, "Dropout Probability", From 2de990a2a1bcffe1839b83aae2da9db3f3f9435a Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 17 May 2026 09:29:50 +0200 Subject: [PATCH 23/67] feat: add data prefetch and rename dataloader_threads to caching_threads Adds prefetch_next_batch option that loads the next batch on a background thread, overlapping disk reads with the current training step. Most beneficial when caching is enabled. Renames dataloader_threads to caching_threads to better reflect its purpose. The UI places Prefetch Next Batch above Clear cache before training. Co-Authored-By: Claude Sonnet 4.6 --- modules/dataLoader/ErnieBaseDataLoader.py | 2 +- modules/dataLoader/Flux2BaseDataLoader.py | 4 +- modules/dataLoader/ZImageBaseDataLoader.py | 2 +- .../dataLoader/mixin/DataLoaderMgdsMixin.py | 2 +- modules/trainer/GenericTrainer.py | 8 ++- modules/ui/TrainUI.py | 50 ++++++++------- modules/util/PrefetchIterator.py | 62 +++++++++++++++++++ modules/util/config/TrainConfig.py | 17 ++++- modules/util/create.py | 4 +- 9 files changed, 116 insertions(+), 35 deletions(-) create mode 100644 modules/util/PrefetchIterator.py diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..586c66cbb 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -32,7 +32,7 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel): image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH) - if config.dataloader_threads > 1: + if config.caching_threads > 1: apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673, unclear if Mistral is affected encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=HIDDEN_STATES_LAYER, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..8826c392f 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -43,7 +43,7 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: mistral_format_input([caption], MISTRAL_SYSTEM_MESSAGE), apply_chat_template_kwargs = {'add_generation_prompt': False}, ) - if config.dataloader_threads > 1: + if config.caching_threads > 1: apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), @@ -53,7 +53,7 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: qwen3_format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': False} ) - if config.dataloader_threads > 1: + if config.caching_threads > 1: apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=QWEN3_HIDDEN_STATES_LAYERS, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..5f1e1603d 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -38,7 +38,7 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH, apply_chat_template = lambda caption: format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': True} ) - if config.dataloader_threads > 1: + if config.caching_threads > 1: apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=-2, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/mixin/DataLoaderMgdsMixin.py b/modules/dataLoader/mixin/DataLoaderMgdsMixin.py index 8421ece11..7867ec276 100644 --- a/modules/dataLoader/mixin/DataLoaderMgdsMixin.py +++ b/modules/dataLoader/mixin/DataLoaderMgdsMixin.py @@ -44,7 +44,7 @@ def _create_mgds( settings, definition, batch_size=config.batch_size, #local batch size - state=PipelineState(config.dataloader_threads), + state=PipelineState(config.caching_threads), initial_epoch=train_progress.epoch, initial_epoch_sample=train_progress.epoch_sample, ) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index ab4926901..f38246d6f 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -29,6 +29,7 @@ from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.TimeUnit import TimeUnit from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.PrefetchIterator import PrefetchIterator from modules.util.profiling_util import TorchMemoryRecorder, TorchProfiler from modules.util.time_util import get_string_timestamp from modules.util.torch_util import torch_gc @@ -678,11 +679,12 @@ def train(self): current_epoch_length = self.data_loader.get_data_set().approximate_length() + batches = self.data_loader.get_data_loader() + if self.config.prefetch_next_batch: + batches = PrefetchIterator(batches) if multi.is_master(): - batches = step_tqdm = tqdm(self.data_loader.get_data_loader(), desc="step", total=current_epoch_length, + batches = step_tqdm = tqdm(batches, desc="step", total=current_epoch_length, initial=train_progress.epoch_step) - else: - batches = self.data_loader.get_data_loader() for batch in batches: multi.sync_commands(self.commands) if self.commands.get_stop_command(): diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index ba90d2e64..143db083e 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -300,44 +300,40 @@ def create_general_tab(self, master): components.time_entry(frame, 8, 3, self.ui_state, "validate_after", "validate_after_unit") # device - components.label(frame, 10, 0, "Dataloader Threads", - tooltip="Number of threads used for the data loader. Increase if your GPU has room during caching, decrease if it's going out of memory during caching.") - components.entry(frame, 10, 1, self.ui_state, "dataloader_threads", required=True) - - components.label(frame, 11, 0, "Train Device", + components.label(frame, 9, 0, "Train Device", tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") - components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + components.entry(frame, 9, 1, self.ui_state, "train_device", required=True) - components.label(frame, 12, 0, "Multi-GPU", + components.label(frame, 10, 0, "Multi-GPU", tooltip="Enable multi-GPU training") - components.switch(frame, 12, 1, self.ui_state, "multi_gpu") - components.label(frame, 12, 2, "Device Indexes", + components.switch(frame, 10, 1, self.ui_state, "multi_gpu") + components.label(frame, 10, 2, "Device Indexes", tooltip="Multi-GPU: A comma-separated list of device indexes. If empty, all your GPUs are used. With a list such as \"0,1,3,4\" you can omit a GPU, for example an on-board graphics GPU.") - components.entry(frame, 12, 3, self.ui_state, "device_indexes") + components.entry(frame, 10, 3, self.ui_state, "device_indexes") - components.label(frame, 13, 0, "Gradient Reduce Precision", + components.label(frame, 11, 0, "Gradient Reduce Precision", tooltip="WEIGHT_DTYPE: Reduce gradients between GPUs in your weight data type; can be imprecise, but more efficient than float32\n" "WEIGHT_DTYPE_STOCHASTIC: Sum up the gradients in your weight data type, but average them in float32 and stochastically round if your weight data type is bfloat16\n" "FLOAT_32: Reduce gradients in float32\n" "FLOAT_32_STOCHASTIC: Reduce gradients in float32; use stochastic rounding to bfloat16 if your weight data type is bfloat16", wide_tooltip=True) - components.options(frame, 13, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, + components.options(frame, 11, 1, [str(x) for x in list(GradientReducePrecision)], self.ui_state, "gradient_reduce_precision") - components.label(frame, 13, 2, "Fused Gradient Reduce", + components.label(frame, 11, 2, "Fused Gradient Reduce", tooltip="Multi-GPU: Gradient synchronisation during the backward pass. Can be more efficient, especially with Async Gradient Reduce") - components.switch(frame, 13, 3, self.ui_state, "fused_gradient_reduce") + components.switch(frame, 11, 3, self.ui_state, "fused_gradient_reduce") - components.label(frame, 14, 0, "Async Gradient Reduce", + components.label(frame, 12, 0, "Async Gradient Reduce", tooltip="Multi-GPU: Asynchroniously start the gradient reduce operations during the backward pass. Can be more efficient, but requires some VRAM.") - components.switch(frame, 14, 1, self.ui_state, "async_gradient_reduce") - components.label(frame, 14, 2, "Buffer size (MB)", + components.switch(frame, 12, 1, self.ui_state, "async_gradient_reduce") + components.label(frame, 12, 2, "Buffer size (MB)", tooltip="Multi-GPU: Maximum VRAM for \"Async Gradient Reduce\", in megabytes. A multiple of this value can be needed if combined with \"Fused Back Pass\" and/or \"Layer offload fraction\"") - components.entry(frame, 14, 3, self.ui_state, "async_gradient_reduce_buffer") + components.entry(frame, 12, 3, self.ui_state, "async_gradient_reduce_buffer") - components.label(frame, 15, 0, "Temp Device", + components.label(frame, 13, 0, "Temp Device", tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") - components.entry(frame, 15, 1, self.ui_state, "temp_device") + components.entry(frame, 13, 1, self.ui_state, "temp_device") frame.pack(fill="both", expand=1) return frame @@ -363,10 +359,20 @@ def create_data_tab(self, master): tooltip="Caching of intermediate training data that can be re-used between epochs") components.switch(frame, 1, 1, self.ui_state, "latent_caching") + # caching threads + components.label(frame, 2, 0, "Caching Threads", + tooltip="Number of threads used while building the latent and text caches. Increase if your GPU has room during caching, decrease if it's going out of memory during caching. Only affects performance while the cache is being built.") + components.entry(frame, 2, 1, self.ui_state, "caching_threads", width=100, sticky="nw", required=True) + + # prefetch next batch + components.label(frame, 3, 0, "Prefetch Next Batch", + tooltip="Load the next batch on a background thread, overlapping disk reads with the current training step. Most beneficial when caching is enabled, since the prefetch thread then only does disk reads. With caching disabled, the text encoder / VAE forward passes run concurrently with training, increasing peak VRAM.") + components.switch(frame, 3, 1, self.ui_state, "prefetch_next_batch") + # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", + components.label(frame, 4, 0, "Clear cache before training", tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + components.switch(frame, 4, 1, self.ui_state, "clear_cache_before_training") frame.pack(fill="both", expand=1) return frame diff --git a/modules/util/PrefetchIterator.py b/modules/util/PrefetchIterator.py new file mode 100644 index 000000000..dad33a3fa --- /dev/null +++ b/modules/util/PrefetchIterator.py @@ -0,0 +1,62 @@ +import queue +import threading +from collections.abc import Iterable, Iterator +from contextlib import suppress + + +class PrefetchIterator: + """Iterable wrapper that prefetches items ahead on a single background thread. + + Wrapping an iterable in PrefetchIterator lets the producer-side work + (e.g. disk reads, decoding, encoding) overlap with whatever the consumer + is doing between iterations. + """ + + def __init__(self, iterable: Iterable, queue_size: int = 1, stop_poll_interval: float = 0.1): + self._iterable = iterable + self._queue_size = queue_size + # How often the producer checks the stop signal while blocked on put. + self._stop_poll_interval = stop_poll_interval + + def __iter__(self) -> Iterator: + q: queue.Queue = queue.Queue(maxsize=self._queue_size) + stop_event = threading.Event() + + def put_or_stop(value) -> bool: + # Block on put, but periodically wake to check the stop signal so + # we can exit if the consumer has gone away. + while not stop_event.is_set(): + with suppress(queue.Full): + q.put(value, timeout=self._stop_poll_interval) + return True + return False + + def producer(): + try: + for item in self._iterable: + if not put_or_stop(item): + return + except BaseException as e: + put_or_stop(e) + return + put_or_stop(StopIteration()) + + t = threading.Thread(target=producer, daemon=True) + t.start() + + try: + while True: + item = q.get() + if isinstance(item, StopIteration): + return + if isinstance(item, BaseException): + raise item + yield item + finally: + # Signal the producer to stop and drain anything pending so it + # can wake from a blocked put and observe the stop signal. + stop_event.set() + with suppress(queue.Empty): + while True: + q.get_nowait() + t.join() diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index bbc70a030..367927f06 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -404,7 +404,8 @@ class TrainConfig(BaseConfig): ema: EMAMode ema_decay: float ema_update_step_interval: int - dataloader_threads: int + caching_threads: int + prefetch_next_batch: bool train_device: str temp_device: str train_dtype: DataType @@ -558,7 +559,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -570,6 +571,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -789,6 +791,14 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + + if "dataloader_threads" in migrated_data: + migrated_data["caching_threads"] = migrated_data.pop("dataloader_threads") + + return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -986,7 +996,8 @@ def default_values() -> 'TrainConfig': data.append(("ema", EMAMode.OFF, EMAMode, False)) data.append(("ema_decay", 0.999, float, False)) data.append(("ema_update_step_interval", 5, int, False)) - data.append(("dataloader_threads", 2, int, False)) + data.append(("caching_threads", 2, int, False)) + data.append(("prefetch_next_batch", True, bool, False)) data.append(("train_device", default_device.type, str, False)) data.append(("temp_device", "cpu", str, False)) data.append(("train_dtype", DataType.FLOAT_16, DataType, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..9bbf833aa 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -110,8 +110,8 @@ def create_data_loader( train_progress: TrainProgress | None = None, is_validation: bool = False ) -> BaseDataLoader | None: - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 and config.dataloader_threads > 1: - raise RuntimeError('layer offloading can not be activated if "dataloader_threads" > 1') + if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 and config.caching_threads > 1: + raise RuntimeError('layer offloading can not be activated if "caching_threads" > 1') if train_progress is None: train_progress = TrainProgress() From d5ac1e4d77ed1440573ff65a4d5f818c35b0665d Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 17 May 2026 09:38:14 +0200 Subject: [PATCH 24/67] refactor: split latent_caching into image_caching and text_caching Rename `config.latent_caching` to `config.image_caching` and add a new orthogonal `config.text_caching` field. Adds config migration 10 -> 11 mapping the old field to both new fields, preserving existing behavior. This enables the previously-impossible (no image cache, yes text cache) configuration. Updates modelSetup placement decisions (VAE/effnet -> image_caching, text encoder -> text_caching), dataloader mixin gating, trainer cache-clear/start-next-epoch logic, UI switches, and tracked preset JSONs. --- modules/dataLoader/ChromaBaseDataLoader.py | 6 +++--- modules/dataLoader/ErnieBaseDataLoader.py | 2 +- modules/dataLoader/Flux2BaseDataLoader.py | 2 +- modules/dataLoader/FluxBaseDataLoader.py | 5 ++++- modules/dataLoader/HiDreamBaseDataLoader.py | 10 ++++++---- .../dataLoader/HunyuanVideoBaseDataLoader.py | 5 ++++- modules/dataLoader/PixArtAlphaBaseDataLoader.py | 2 +- modules/dataLoader/QwenBaseDataLoader.py | 6 +++--- modules/dataLoader/SanaBaseDataLoader.py | 2 +- .../StableDiffusion3BaseDataLoader.py | 6 +++++- .../dataLoader/StableDiffusionBaseDataLoader.py | 2 +- .../StableDiffusionFineTuneVaeDataLoader.py | 4 ++-- .../StableDiffusionXLBaseDataLoader.py | 5 ++++- modules/dataLoader/WuerstchenBaseDataLoader.py | 2 +- modules/dataLoader/ZImageBaseDataLoader.py | 6 +++--- .../mixin/DataLoaderText2ImageMixin.py | 10 +++++----- modules/modelSetup/ChromaEmbeddingSetup.py | 5 +++-- modules/modelSetup/ChromaFineTuneSetup.py | 4 ++-- modules/modelSetup/ChromaLoRASetup.py | 4 ++-- modules/modelSetup/ErnieFineTuneSetup.py | 4 ++-- modules/modelSetup/ErnieLoRASetup.py | 4 ++-- modules/modelSetup/Flux2FineTuneSetup.py | 4 ++-- modules/modelSetup/Flux2LoRASetup.py | 4 ++-- modules/modelSetup/FluxEmbeddingSetup.py | 8 +++++--- modules/modelSetup/FluxFineTuneSetup.py | 6 +++--- modules/modelSetup/FluxLoRASetup.py | 6 +++--- modules/modelSetup/HiDreamEmbeddingSetup.py | 16 ++++++++++------ modules/modelSetup/HiDreamFineTuneSetup.py | 10 +++++----- modules/modelSetup/HiDreamLoRASetup.py | 10 +++++----- .../modelSetup/HunyuanVideoEmbeddingSetup.py | 8 +++++--- modules/modelSetup/HunyuanVideoFineTuneSetup.py | 6 +++--- modules/modelSetup/HunyuanVideoLoRASetup.py | 6 +++--- modules/modelSetup/PixArtAlphaFineTuneSetup.py | 4 ++-- modules/modelSetup/PixArtAlphaLoRASetup.py | 4 ++-- modules/modelSetup/QwenFineTuneSetup.py | 4 ++-- modules/modelSetup/QwenLoRASetup.py | 4 ++-- modules/modelSetup/SanaFineTuneSetup.py | 4 ++-- modules/modelSetup/SanaLoRASetup.py | 4 ++-- .../StableDiffusion3EmbeddingSetup.py | 13 ++++++++----- .../modelSetup/StableDiffusion3FineTuneSetup.py | 8 ++++---- modules/modelSetup/StableDiffusion3LoRASetup.py | 8 ++++---- .../modelSetup/StableDiffusionEmbeddingSetup.py | 2 +- .../modelSetup/StableDiffusionFineTuneSetup.py | 4 ++-- modules/modelSetup/StableDiffusionLoRASetup.py | 4 ++-- .../StableDiffusionXLEmbeddingSetup.py | 8 +++++--- .../StableDiffusionXLFineTuneSetup.py | 6 +++--- .../modelSetup/StableDiffusionXLLoRASetup.py | 6 +++--- modules/modelSetup/WuerstchenEmbeddingSetup.py | 2 +- modules/modelSetup/WuerstchenFineTuneSetup.py | 4 ++-- modules/modelSetup/WuerstchenLoRASetup.py | 4 ++-- modules/modelSetup/ZImageFineTuneSetup.py | 4 ++-- modules/modelSetup/ZImageLoRASetup.py | 4 ++-- modules/trainer/GenericTrainer.py | 4 ++-- modules/trainer/MultiTrainer.py | 4 ++-- modules/ui/TrainUI.py | 17 +++++++++++------ modules/util/config/TrainConfig.py | 17 ++++++++++++++--- training_presets/#sd 1.5 embedding.json | 3 ++- training_presets/#sd 2.1 embedding.json | 3 ++- training_presets/#sdxl 1.0 embedding.json | 3 ++- training_presets/#wuerstchen 2.0 embedding.json | 3 ++- 60 files changed, 192 insertions(+), 143 deletions(-) diff --git a/modules/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py index 94722bdf5..9be9076f7 100644 --- a/modules/dataLoader/ChromaBaseDataLoader.py +++ b/modules/dataLoader/ChromaBaseDataLoader.py @@ -50,7 +50,7 @@ def _preparation_modules(self, config: TrainConfig, model: ChromaModel): if not config.train_text_encoder_or_embedding(): modules.append(encode_prompt) - if config.latent_caching and not config.train_text_encoder_or_embedding(): + if config.text_caching and not config.train_text_encoder_or_embedding(): modules.append(prune_masked_tokens) return modules @@ -80,7 +80,7 @@ def _cache_modules(self, config: TrainConfig, model: ChromaModel, model_setup: B text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching = not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), ) def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup: BaseChromaSetup): @@ -110,7 +110,7 @@ def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup: train_dtype=model.train_dtype, ) - if config.latent_caching and not config.train_text_encoder_or_embedding(): + if config.text_caching and not config.train_text_encoder_or_embedding(): output_module_list = [pad_masked_tokens] + output_module_list return output_module_list diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..0ae9d4a48 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -67,7 +67,7 @@ def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: Ba text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=True, + text_caching=config.text_caching, ) def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: BaseErnieSetup): diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..1fe65271c 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -90,7 +90,7 @@ def _cache_modules(self, config: TrainConfig, model: Flux2Model, model_setup: Ba text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=True, + text_caching=config.text_caching, ) def _output_modules(self, config: TrainConfig, model: Flux2Model, model_setup: BaseFlux2Setup): diff --git a/modules/dataLoader/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py index 93cc44aa6..cd2aab061 100644 --- a/modules/dataLoader/FluxBaseDataLoader.py +++ b/modules/dataLoader/FluxBaseDataLoader.py @@ -103,7 +103,10 @@ def _cache_modules(self, config: TrainConfig, model: FluxModel, model_setup: Bas text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(), + text_caching=config.text_caching and ( + not config.train_text_encoder_or_embedding() + or not config.train_text_encoder_2_or_embedding() + ), ) def _output_modules(self, config: TrainConfig, model: FluxModel, model_setup: BaseFluxSetup): diff --git a/modules/dataLoader/HiDreamBaseDataLoader.py b/modules/dataLoader/HiDreamBaseDataLoader.py index dedea17be..5e1c94a90 100644 --- a/modules/dataLoader/HiDreamBaseDataLoader.py +++ b/modules/dataLoader/HiDreamBaseDataLoader.py @@ -132,10 +132,12 @@ def _cache_modules(self, config: TrainConfig, model: HiDreamModel, model_setup: text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding() \ - or not config.train_text_encoder_2_or_embedding() \ - or not config.train_text_encoder_3_or_embedding() \ - or not config.train_text_encoder_4_or_embedding(), + text_caching=config.text_caching and ( + not config.train_text_encoder_or_embedding() + or not config.train_text_encoder_2_or_embedding() + or not config.train_text_encoder_3_or_embedding() + or not config.train_text_encoder_4_or_embedding() + ), ) def _output_modules(self, config: TrainConfig, model: HiDreamModel, model_setup: BaseHiDreamSetup): diff --git a/modules/dataLoader/HunyuanVideoBaseDataLoader.py b/modules/dataLoader/HunyuanVideoBaseDataLoader.py index d587a5fd9..2869b72ba 100644 --- a/modules/dataLoader/HunyuanVideoBaseDataLoader.py +++ b/modules/dataLoader/HunyuanVideoBaseDataLoader.py @@ -97,7 +97,10 @@ def _cache_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_se text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(), + text_caching=config.text_caching and ( + not config.train_text_encoder_or_embedding() + or not config.train_text_encoder_2_or_embedding() + ), ) def _output_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_setup: BaseHunyuanVideoSetup): diff --git a/modules/dataLoader/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py index f2dc37857..59d101832 100644 --- a/modules/dataLoader/PixArtAlphaBaseDataLoader.py +++ b/modules/dataLoader/PixArtAlphaBaseDataLoader.py @@ -85,7 +85,7 @@ def _cache_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_set text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), ) def _output_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_setup: BasePixArtAlphaSetup): diff --git a/modules/dataLoader/QwenBaseDataLoader.py b/modules/dataLoader/QwenBaseDataLoader.py index db3ca5df3..c02c15f5a 100644 --- a/modules/dataLoader/QwenBaseDataLoader.py +++ b/modules/dataLoader/QwenBaseDataLoader.py @@ -54,7 +54,7 @@ def _preparation_modules(self, config: TrainConfig, model: QwenModel): if not config.train_text_encoder_or_embedding(): modules.append(encode_prompt) - if config.latent_caching and not config.train_text_encoder_or_embedding(): + if config.text_caching and not config.train_text_encoder_or_embedding(): modules.append(prune_masked_tokens) return modules @@ -84,7 +84,7 @@ def _cache_modules(self, config: TrainConfig, model: QwenModel, model_setup: Bas text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), ) def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: BaseQwenSetup): @@ -114,7 +114,7 @@ def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: Ba train_dtype=model.train_dtype, ) - if config.latent_caching and not config.train_text_encoder_or_embedding(): + if config.text_caching and not config.train_text_encoder_or_embedding(): output_module_list = [pad_masked_tokens] + output_module_list return output_module_list diff --git a/modules/dataLoader/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py index 38d5c31b0..36175cb16 100644 --- a/modules/dataLoader/SanaBaseDataLoader.py +++ b/modules/dataLoader/SanaBaseDataLoader.py @@ -78,7 +78,7 @@ def _cache_modules(self, config: TrainConfig, model: SanaModel, model_setup: Bas text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), ) def _output_modules(self, config: TrainConfig, model: SanaModel, model_setup: BaseSanaSetup): diff --git a/modules/dataLoader/StableDiffusion3BaseDataLoader.py b/modules/dataLoader/StableDiffusion3BaseDataLoader.py index c497a20c6..fc9c7b09d 100644 --- a/modules/dataLoader/StableDiffusion3BaseDataLoader.py +++ b/modules/dataLoader/StableDiffusion3BaseDataLoader.py @@ -117,7 +117,11 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusion3Model, mode text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding() or not config.train_text_encoder_3_or_embedding(), + text_caching=config.text_caching and ( + not config.train_text_encoder_or_embedding() + or not config.train_text_encoder_2_or_embedding() + or not config.train_text_encoder_3_or_embedding() + ), ) def _output_modules(self, config: TrainConfig, model: StableDiffusion3Model, model_setup: BaseStableDiffusion3Setup): diff --git a/modules/dataLoader/StableDiffusionBaseDataLoader.py b/modules/dataLoader/StableDiffusionBaseDataLoader.py index 781f769a6..d27398686 100644 --- a/modules/dataLoader/StableDiffusionBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionBaseDataLoader.py @@ -86,7 +86,7 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusionModel, model text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), ) def _output_modules(self, config: TrainConfig, model: StableDiffusionModel, model_setup: BaseStableDiffusionSetup): diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index ed5dd32b8..956bb70b5 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py @@ -197,7 +197,7 @@ def before_cache_fun(): modules = [] - if config.latent_caching: + if config.image_caching: modules.append(disk_cache) modules.append(variation_sorting) @@ -221,7 +221,7 @@ def __output_modules(self, config: TrainConfig): image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') - if config.latent_caching: + if config.image_caching: batch_sorting = AspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size) else: batch_sorting = InlineAspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size) diff --git a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py index ed1ad491e..340ecd848 100644 --- a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py @@ -99,7 +99,10 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusionXLModel, mod text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(), + text_caching=config.text_caching and ( + not config.train_text_encoder_or_embedding() + or not config.train_text_encoder_2_or_embedding() + ), ) def _output_modules(self, config: TrainConfig, model: StableDiffusionXLModel, model_setup: BaseStableDiffusionXLSetup): diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py index f5cf2f41a..0d13586ea 100644 --- a/modules/dataLoader/WuerstchenBaseDataLoader.py +++ b/modules/dataLoader/WuerstchenBaseDataLoader.py @@ -84,7 +84,7 @@ def before_cache_image_fun(): text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=not config.train_text_encoder_or_embedding(), + text_caching=config.text_caching and not config.train_text_encoder_or_embedding(), before_cache_image_fun=before_cache_image_fun ) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..ad7a531bd 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -51,7 +51,7 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel): modules += [tokenize_prompt, encode_prompt] - if config.latent_caching: + if config.text_caching: modules.append(prune_masked_tokens) return modules @@ -80,7 +80,7 @@ def _cache_modules(self, config: TrainConfig, model: ZImageModel, model_setup: B text_split_names=text_split_names, sort_names=sort_names, config=config, - text_caching=True, + text_caching=config.text_caching, ) def _output_modules(self, config: TrainConfig, model: ZImageModel, model_setup: BaseZImageSetup): @@ -109,7 +109,7 @@ def _output_modules(self, config: TrainConfig, model: ZImageModel, model_setup: train_dtype=model.train_dtype, ) - if config.latent_caching: + if config.text_caching: output_module_list = [pad_masked_tokens] + output_module_list return output_module_list diff --git a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py index 2654bdd19..2dde104bf 100644 --- a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py +++ b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py @@ -301,7 +301,7 @@ def prepare_vae(): ) world_size = multi.world_size() if config.multi_gpu else 1 #world_size can be 1 for validation dataloader, even if multi.world_size() returns > 1 - if config.latent_caching: + if config.image_caching: batch_sorting = AspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size * world_size) distributed_sampler = DistributedSampler(names=sort_names, world_size=world_size, rank=multi.rank()) else: @@ -358,15 +358,15 @@ def before_cache_text_fun(): modules = [] - if config.latent_caching: + if config.image_caching: modules.append(image_disk_cache) sort_names = [x for x in sort_names if x not in image_aggregate_names] sort_names = [x for x in sort_names if x not in image_split_names] - if text_caching: - modules.append(text_disk_cache) - sort_names = [x for x in sort_names if x not in text_split_names] + if text_caching: + modules.append(text_disk_cache) + sort_names = [x for x in sort_names if x not in text_split_names] if len(sort_names) > 0: variation_sorting = VariationSorting(names=sort_names, balancing_in_name='concept.balancing', balancing_strategy_in_name='concept.balancing_strategy', diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 9fa65b69f..295bc8769 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -74,9 +74,10 @@ def setup_train_device( model: ChromaModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching - model.text_encoder_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) + model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.transformer_to(self.train_device) diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 371c414ec..242822b9f 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -83,10 +83,10 @@ def setup_train_device( model: ChromaModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 1396b4a52..2003f5f6b 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -111,10 +111,10 @@ def setup_train_device( model: ChromaModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/ErnieFineTuneSetup.py b/modules/modelSetup/ErnieFineTuneSetup.py index a4a03fb02..5158648f0 100644 --- a/modules/modelSetup/ErnieFineTuneSetup.py +++ b/modules/modelSetup/ErnieFineTuneSetup.py @@ -63,8 +63,8 @@ def setup_train_device( model: ErnieModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/ErnieLoRASetup.py b/modules/modelSetup/ErnieLoRASetup.py index 108a26f92..82c043a8f 100644 --- a/modules/modelSetup/ErnieLoRASetup.py +++ b/modules/modelSetup/ErnieLoRASetup.py @@ -73,8 +73,8 @@ def setup_train_device( model: ErnieModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/Flux2FineTuneSetup.py b/modules/modelSetup/Flux2FineTuneSetup.py index 2b90618c3..8824587a5 100644 --- a/modules/modelSetup/Flux2FineTuneSetup.py +++ b/modules/modelSetup/Flux2FineTuneSetup.py @@ -63,8 +63,8 @@ def setup_train_device( model: Flux2Model, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/Flux2LoRASetup.py b/modules/modelSetup/Flux2LoRASetup.py index 3c38ebf87..1e949d902 100644 --- a/modules/modelSetup/Flux2LoRASetup.py +++ b/modules/modelSetup/Flux2LoRASetup.py @@ -75,8 +75,8 @@ def setup_train_device( model: Flux2Model, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index 0418390da..40a02fd7f 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -85,10 +85,12 @@ def setup_train_device( model: FluxModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_1_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching + text_encoder_2_on_train_device = config.train_text_encoder_2_or_embedding() or not config.text_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) + model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) + model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.transformer_to(self.train_device) diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index a2dc78057..609a590b6 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -94,14 +94,14 @@ def setup_train_device( model: FluxModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index e3de941fd..59dfa140b 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -136,14 +136,14 @@ def setup_train_device( model: FluxModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/HiDreamEmbeddingSetup.py b/modules/modelSetup/HiDreamEmbeddingSetup.py index 22c09927b..7c58f31a2 100644 --- a/modules/modelSetup/HiDreamEmbeddingSetup.py +++ b/modules/modelSetup/HiDreamEmbeddingSetup.py @@ -103,12 +103,16 @@ def setup_train_device( model: HiDreamModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.text_encoder_3_to(self.train_device if config.text_encoder_3.train_embedding else self.temp_device) - model.text_encoder_4_to(self.train_device if config.text_encoder_4.train_embedding else self.temp_device) + vae_on_train_device = not config.image_caching + text_encoder_1_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching + text_encoder_2_on_train_device = config.train_text_encoder_2_or_embedding() or not config.text_caching + text_encoder_3_on_train_device = config.train_text_encoder_3_or_embedding() or not config.text_caching + text_encoder_4_on_train_device = config.train_text_encoder_4_or_embedding() or not config.text_caching + + model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) + model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) + model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) + model.text_encoder_4_to(self.train_device if text_encoder_4_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.transformer_to(self.train_device) diff --git a/modules/modelSetup/HiDreamFineTuneSetup.py b/modules/modelSetup/HiDreamFineTuneSetup.py index 49a3adf9d..fc95e3041 100644 --- a/modules/modelSetup/HiDreamFineTuneSetup.py +++ b/modules/modelSetup/HiDreamFineTuneSetup.py @@ -119,22 +119,22 @@ def setup_train_device( model: HiDreamModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_3_on_train_device = \ config.train_text_encoder_3_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_4_on_train_device = \ config.train_text_encoder_4_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/HiDreamLoRASetup.py b/modules/modelSetup/HiDreamLoRASetup.py index 3cbcdbe8e..a2c1d949b 100644 --- a/modules/modelSetup/HiDreamLoRASetup.py +++ b/modules/modelSetup/HiDreamLoRASetup.py @@ -191,22 +191,22 @@ def setup_train_device( model: HiDreamModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_3_on_train_device = \ config.train_text_encoder_3_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_4_on_train_device = \ config.train_text_encoder_4_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py index 2dca283f8..0a3fbd681 100644 --- a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py +++ b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py @@ -83,10 +83,12 @@ def setup_train_device( model: HunyuanVideoModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_1_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching + text_encoder_2_on_train_device = config.train_text_encoder_2_or_embedding() or not config.text_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) + model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) + model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.transformer_to(self.train_device) diff --git a/modules/modelSetup/HunyuanVideoFineTuneSetup.py b/modules/modelSetup/HunyuanVideoFineTuneSetup.py index f9ab6e4b9..685e03d1d 100644 --- a/modules/modelSetup/HunyuanVideoFineTuneSetup.py +++ b/modules/modelSetup/HunyuanVideoFineTuneSetup.py @@ -93,14 +93,14 @@ def setup_train_device( model: HunyuanVideoModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/HunyuanVideoLoRASetup.py b/modules/modelSetup/HunyuanVideoLoRASetup.py index 0b83b3ad7..2ec408bcd 100644 --- a/modules/modelSetup/HunyuanVideoLoRASetup.py +++ b/modules/modelSetup/HunyuanVideoLoRASetup.py @@ -139,14 +139,14 @@ def setup_train_device( model: HunyuanVideoModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index 277f39004..c95d48291 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -87,11 +87,11 @@ def setup_train_device( model: PixArtAlphaModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 5c865b224..10290e67c 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -107,11 +107,11 @@ def setup_train_device( model: PixArtAlphaModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/QwenFineTuneSetup.py b/modules/modelSetup/QwenFineTuneSetup.py index c9bc2cae5..c4ced4d94 100644 --- a/modules/modelSetup/QwenFineTuneSetup.py +++ b/modules/modelSetup/QwenFineTuneSetup.py @@ -68,10 +68,10 @@ def setup_train_device( model: QwenModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/QwenLoRASetup.py b/modules/modelSetup/QwenLoRASetup.py index 4ad8f3eea..689e37102 100644 --- a/modules/modelSetup/QwenLoRASetup.py +++ b/modules/modelSetup/QwenLoRASetup.py @@ -97,10 +97,10 @@ def setup_train_device( model: QwenModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index 11faa747a..3835abe2e 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -81,11 +81,11 @@ def setup_train_device( model: SanaModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index f481cd445..c0472d3b5 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -107,11 +107,11 @@ def setup_train_device( model: SanaModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index 1ff229928..e8aba2a94 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -96,11 +96,14 @@ def setup_train_device( model: StableDiffusion3Model, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.text_encoder_3_to(self.train_device if config.text_encoder_3.train_embedding else self.temp_device) + vae_on_train_device = not config.image_caching + text_encoder_1_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching + text_encoder_2_on_train_device = config.train_text_encoder_2_or_embedding() or not config.text_caching + text_encoder_3_on_train_device = config.train_text_encoder_3_or_embedding() or not config.text_caching + + model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) + model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) + model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.transformer_to(self.train_device) diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 3bbe41821..4a389341e 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -105,18 +105,18 @@ def setup_train_device( model: StableDiffusion3Model, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_3_on_train_device = \ config.train_text_encoder_3_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index ae326cbbe..509a4377a 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -163,18 +163,18 @@ def setup_train_device( model: StableDiffusion3Model, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_3_on_train_device = \ config.train_text_encoder_3_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index 6e3a41c2b..1e10674a8 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -75,7 +75,7 @@ def setup_train_device( model: StableDiffusionModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching model.text_encoder_to(self.train_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index 8c42e59df..17365cf99 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -89,11 +89,11 @@ def setup_train_device( model: StableDiffusionModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index 9b9e9a613..17b390c4b 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -112,11 +112,11 @@ def setup_train_device( model: StableDiffusionModel, config: TrainConfig, ): - vae_on_train_device = self.debug_mode or not config.latent_caching + vae_on_train_device = self.debug_mode or not config.image_caching text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index 7d9ef442f..c8896462c 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -85,10 +85,12 @@ def setup_train_device( model: StableDiffusionXLModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_1_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching + text_encoder_2_on_train_device = config.train_text_encoder_2_or_embedding() or not config.text_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) + model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) + model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) model.unet_to(self.train_device) diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index ec6dbe16c..1d75f5c93 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -100,16 +100,16 @@ def setup_train_device( model: StableDiffusionXLModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.text_encoder_2.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 6b193edc5..1011cb60f 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -134,13 +134,13 @@ def setup_train_device( model: StableDiffusionXLModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching text_encoder_1_on_train_device = \ config.train_text_encoder_or_embedding()\ - or not config.latent_caching + or not config.text_caching text_encoder_2_on_train_device = \ config.train_text_encoder_2_or_embedding() \ - or not config.latent_caching + or not config.text_caching model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 37dc7f85f..0ae8a4eb2 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -75,7 +75,7 @@ def setup_train_device( model: WuerstchenModel, config: TrainConfig, ): - effnet_on_train_device = not config.latent_caching + effnet_on_train_device = not config.image_caching if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder_to(self.temp_device) diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 1291125dc..6e7b792b1 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -84,7 +84,7 @@ def setup_train_device( model: WuerstchenModel, config: TrainConfig, ): - effnet_on_train_device = not config.latent_caching + effnet_on_train_device = not config.image_caching if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder_to(self.temp_device) @@ -95,7 +95,7 @@ def setup_train_device( text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.prior_text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.prior_prior_to(self.train_device) diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index b26f57436..848d0a3dc 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -111,7 +111,7 @@ def setup_train_device( model: WuerstchenModel, config: TrainConfig, ): - effnet_on_train_device = not config.latent_caching + effnet_on_train_device = not config.image_caching if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder_to(self.temp_device) @@ -122,7 +122,7 @@ def setup_train_device( text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ - or not config.latent_caching + or not config.text_caching model.prior_text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.prior_prior_to(self.train_device) diff --git a/modules/modelSetup/ZImageFineTuneSetup.py b/modules/modelSetup/ZImageFineTuneSetup.py index a4c15c2b3..8fa63451b 100644 --- a/modules/modelSetup/ZImageFineTuneSetup.py +++ b/modules/modelSetup/ZImageFineTuneSetup.py @@ -63,8 +63,8 @@ def setup_train_device( model: ZImageModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/modelSetup/ZImageLoRASetup.py b/modules/modelSetup/ZImageLoRASetup.py index 9355a9ad0..f9ff8e4e2 100644 --- a/modules/modelSetup/ZImageLoRASetup.py +++ b/modules/modelSetup/ZImageLoRASetup.py @@ -76,8 +76,8 @@ def setup_train_device( model: ZImageModel, config: TrainConfig, ): - vae_on_train_device = not config.latent_caching - text_encoder_on_train_device = not config.latent_caching + vae_on_train_device = not config.image_caching + text_encoder_on_train_device = not config.text_caching model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) model.vae_to(self.train_device if vae_on_train_device else self.temp_device) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index ab4926901..83939ce50 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -82,7 +82,7 @@ def start(self): if multi.is_master(): self.__save_config_to_workspace() - if self.config.clear_cache_before_training and self.config.latent_caching: + if self.config.clear_cache_before_training and (self.config.image_caching or self.config.text_caching): self.__clear_cache() if self.config.train_dtype.enable_tf(): @@ -642,7 +642,7 @@ def train(self): #call start_next_epoch with only one process at first, because it might write to the cache. All subsequent processes can read in parallel: for _ in multi.master_first(): - if self.config.latent_caching: + if self.config.image_caching or self.config.text_caching: self.data_loader.get_data_set().start_next_epoch() self.model_setup.setup_train_device(self.model, self.config) else: diff --git a/modules/trainer/MultiTrainer.py b/modules/trainer/MultiTrainer.py index 1cf0cc57d..95df6808b 100644 --- a/modules/trainer/MultiTrainer.py +++ b/modules/trainer/MultiTrainer.py @@ -20,8 +20,8 @@ def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: Tra super().__init__(config, callbacks, commands) if config.samples_to_tensorboard: print("Warning: If 'Samples To Tensorboard' is enabled, only one GPU is used for sampling!") - if not config.latent_caching: - print("Warning: Latent caching is disabled, but recommended for multi-GPU training!") + if not config.image_caching: + print("Warning: Image Caching is disabled, but it is recommended for multi-GPU training!") def start(self): os.environ.setdefault('MASTER_ADDR', 'localhost') diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index ba90d2e64..c8e0c3ea8 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -358,15 +358,20 @@ def create_data_tab(self, master): tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") components.switch(frame, 0, 1, self.ui_state, "aspect_ratio_bucketing") - # latent caching - components.label(frame, 1, 0, "Latent Caching", - tooltip="Caching of intermediate training data that can be re-used between epochs") - components.switch(frame, 1, 1, self.ui_state, "latent_caching") + # image caching + components.label(frame, 1, 0, "Image Caching", + tooltip="Caches image latents (VAE outputs) so they can be re-used between epochs") + components.switch(frame, 1, 1, self.ui_state, "image_caching") + + # text caching + components.label(frame, 2, 0, "Text Caching", + tooltip="Caches text encoder outputs so they can be re-used between epochs") + components.switch(frame, 2, 1, self.ui_state, "text_caching") # clear cache before training - components.label(frame, 2, 0, "Clear cache before training", + components.label(frame, 3, 0, "Clear cache before training", tooltip="Clears the cache directory before starting to train. Only disable this if you want to continue using the same cached data. Disabling this can lead to errors, if other settings are changed during a restart") - components.switch(frame, 2, 1, self.ui_state, "clear_cache_before_training") + components.switch(frame, 3, 1, self.ui_state, "clear_cache_before_training") frame.pack(fill="both", expand=1) return frame diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index bbc70a030..72bf4c8e1 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -385,7 +385,8 @@ class TrainConfig(BaseConfig): concept_file_name: str concepts: list[ConceptConfig] aspect_ratio_bucketing: bool - latent_caching: bool + image_caching: bool + text_caching: bool clear_cache_before_training: bool # training settings @@ -558,7 +559,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -570,6 +571,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -789,6 +791,14 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + if "latent_caching" in migrated_data: + latent_caching = migrated_data.pop("latent_caching") + migrated_data["image_caching"] = latent_caching + migrated_data["text_caching"] = latent_caching + return migrated_data + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -969,7 +979,8 @@ def default_values() -> 'TrainConfig': data.append(("concept_file_name", "training_concepts/concepts.json", str, False)) data.append(("concepts", None, list[ConceptConfig], True)) data.append(("aspect_ratio_bucketing", True, bool, False)) - data.append(("latent_caching", True, bool, False)) + data.append(("image_caching", True, bool, False)) + data.append(("text_caching", True, bool, False)) data.append(("clear_cache_before_training", True, bool, False)) # training settings diff --git a/training_presets/#sd 1.5 embedding.json b/training_presets/#sd 1.5 embedding.json index 6970d96d3..441471b4c 100644 --- a/training_presets/#sd 1.5 embedding.json +++ b/training_presets/#sd 1.5 embedding.json @@ -1,7 +1,8 @@ { "backup_after": 10, "base_model_name": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "latent_caching": false, + "image_caching": false, + "text_caching": false, "learning_rate": 0.0003, "learning_rate_warmup_steps": 20, "model_type": "STABLE_DIFFUSION_15", diff --git a/training_presets/#sd 2.1 embedding.json b/training_presets/#sd 2.1 embedding.json index 369636f50..912b0a1c6 100644 --- a/training_presets/#sd 2.1 embedding.json +++ b/training_presets/#sd 2.1 embedding.json @@ -1,7 +1,8 @@ { "backup_after": 10, "base_model_name": "sd2-community/stable-diffusion-2-1", - "latent_caching": false, + "image_caching": false, + "text_caching": false, "learning_rate": 0.0003, "learning_rate_warmup_steps": 20, "model_type": "STABLE_DIFFUSION_21", diff --git a/training_presets/#sdxl 1.0 embedding.json b/training_presets/#sdxl 1.0 embedding.json index 449037c65..aa16cb91c 100644 --- a/training_presets/#sdxl 1.0 embedding.json +++ b/training_presets/#sdxl 1.0 embedding.json @@ -1,7 +1,8 @@ { "backup_after": 10, "base_model_name": "stabilityai/stable-diffusion-xl-base-1.0", - "latent_caching": false, + "image_caching": false, + "text_caching": false, "learning_rate": 0.0003, "learning_rate_warmup_steps": 20, "model_type": "STABLE_DIFFUSION_XL_10_BASE", diff --git a/training_presets/#wuerstchen 2.0 embedding.json b/training_presets/#wuerstchen 2.0 embedding.json index 5b5a75eaa..ac057bc2c 100644 --- a/training_presets/#wuerstchen 2.0 embedding.json +++ b/training_presets/#wuerstchen 2.0 embedding.json @@ -21,7 +21,8 @@ "model_name": "warp-ai/EfficientNetEncoder", "weight_dtype": "FLOAT_16" }, - "latent_caching": false, + "image_caching": false, + "text_caching": false, "learning_rate": 0.0003, "learning_rate_warmup_steps": 20, "model_type": "WUERSTCHEN_2", From 895c75a2ad7a715a88fa4de7ea37a9f7a70ce9c5 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 17 May 2026 11:18:55 +0200 Subject: [PATCH 25/67] prefetch: use dedicated CUDA stream to avoid serialising with training loop Tensor uploads to the GPU in OutputPipelineModule were enqueued on the default CUDA stream, so each H2D transfer had to wait for the current training step's GPU work to finish before it could start. Running the producer under its own stream lets uploads proceed independently, allowing the prefetch queue to stay ahead of the training loop. Co-Authored-By: Claude Sonnet 4.6 --- modules/util/PrefetchIterator.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/modules/util/PrefetchIterator.py b/modules/util/PrefetchIterator.py index dad33a3fa..aec97f3ee 100644 --- a/modules/util/PrefetchIterator.py +++ b/modules/util/PrefetchIterator.py @@ -1,7 +1,9 @@ import queue import threading from collections.abc import Iterable, Iterator -from contextlib import suppress +from contextlib import nullcontext, suppress + +import torch class PrefetchIterator: @@ -10,6 +12,9 @@ class PrefetchIterator: Wrapping an iterable in PrefetchIterator lets the producer-side work (e.g. disk reads, decoding, encoding) overlap with whatever the consumer is doing between iterations. + + The producer runs on a dedicated CUDA stream so tensor uploads to the GPU + don't have to wait for in-flight training work on the default stream. """ def __init__(self, iterable: Iterable, queue_size: int = 1, stop_poll_interval: float = 0.1): @@ -22,6 +27,8 @@ def __iter__(self) -> Iterator: q: queue.Queue = queue.Queue(maxsize=self._queue_size) stop_event = threading.Event() + stream_ctx = torch.cuda.stream(torch.cuda.Stream()) if torch.cuda.is_available() else nullcontext() + def put_or_stop(value) -> bool: # Block on put, but periodically wake to check the stop signal so # we can exit if the consumer has gone away. @@ -32,14 +39,15 @@ def put_or_stop(value) -> bool: return False def producer(): - try: - for item in self._iterable: - if not put_or_stop(item): - return - except BaseException as e: - put_or_stop(e) - return - put_or_stop(StopIteration()) + with stream_ctx: + try: + for item in self._iterable: + if not put_or_stop(item): + return + except BaseException as e: + put_or_stop(e) + return + put_or_stop(StopIteration()) t = threading.Thread(target=producer, daemon=True) t.start() From 271f51e40f5504f7e44095d1323eca2c3fbba012 Mon Sep 17 00:00:00 2001 From: dxqb Date: Sun, 24 May 2026 17:17:05 +0200 Subject: [PATCH 26/67] Upgrade transformers to 5.9 and huggingface-hub to 1.16 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Bump requirements: transformers 4.57.6 → 5.9, huggingface-hub 0.34.4 → 1.16.1 - Remove HF_HUB_DISABLE_XET workaround from startup scripts; Xet is stable in hub 1.16 - Remove _prepare_sub_modules / snapshot_download prefetching; hub 1.16 fetches lazily on demand - Delete thread_safety.py and apply_thread_safe_forward calls; workaround for transformers#42673 was fixed upstream in v5 - Replace _remove_added_embeddings_from_tokenizer (relied on internal Trie, removed in v5) with orig_tokenizer deep-copies stored at load time; model savers pass use_original_tokenizers=True to create_pipeline() so saved checkpoints use the unmodified tokenizer - Switch ErnieModelLoader to AutoTokenizer; eliminates the tokenization-logger suppress workaround - Suppress httpx INFO logs; hub 1.16 uses httpx internally and logs every HTTP request Co-Authored-By: Claude Sonnet 4.6 --- modules/dataLoader/ErnieBaseDataLoader.py | 3 -- modules/dataLoader/Flux2BaseDataLoader.py | 5 --- modules/dataLoader/ZImageBaseDataLoader.py | 3 -- modules/model/ChromaModel.py | 6 ++- modules/model/FluxModel.py | 10 +++-- modules/model/HiDreamModel.py | 10 ++--- modules/model/HunyuanVideoModel.py | 6 +-- modules/model/PixArtAlphaModel.py | 9 ++-- modules/model/SanaModel.py | 6 ++- modules/model/StableDiffusion3Model.py | 14 ++++-- modules/model/StableDiffusionModel.py | 11 +++-- modules/model/StableDiffusionXLModel.py | 10 +++-- modules/model/WuerstchenModel.py | 11 +++-- modules/modelLoader/ErnieModelLoader.py | 28 +----------- modules/modelLoader/Flux2ModelLoader.py | 13 ------ modules/modelLoader/ZImageModelLoader.py | 13 ------ .../modelLoader/chroma/ChromaModelLoader.py | 14 +----- modules/modelLoader/flux/FluxModelLoader.py | 22 +++------- .../modelLoader/hiDream/HiDreamModelLoader.py | 28 ------------ .../hunyuanVideo/HunyuanVideoModelLoader.py | 18 -------- .../modelLoader/mixin/HFModelLoaderMixin.py | 25 +++-------- .../pixartAlpha/PixArtAlphaModelLoader.py | 2 + modules/modelLoader/qwen/QwenModelLoader.py | 12 ------ modules/modelLoader/sana/SanaModelLoader.py | 2 + .../StableDiffusionModelLoader.py | 4 ++ .../StableDiffusion3ModelLoader.py | 9 +++- .../StableDiffusionXLModelLoader.py | 7 +++ .../wuerstchen/WuerstchenModelLoader.py | 2 + modules/modelSampler/HiDreamSampler.py | 2 +- modules/modelSampler/HunyuanVideoSampler.py | 2 +- modules/modelSaver/chroma/ChromaModelSaver.py | 2 +- modules/modelSaver/flux/FluxModelSaver.py | 2 +- .../modelSaver/hidream/HiDreamModelSaver.py | 2 +- .../hunyuanVideo/HunyuanVideoModelSaver.py | 2 +- .../pixartAlpha/PixArtAlphaModelSaver.py | 2 +- modules/modelSaver/sana/SanaModelSaver.py | 2 +- .../StableDiffusionModelSaver.py | 2 +- .../StableDiffusion3ModelSaver.py | 2 +- .../StableDiffusionXLModelSaver.py | 2 +- .../wuerstchen/WuerstchenModelSaver.py | 2 +- modules/modelSetup/ChromaEmbeddingSetup.py | 1 - modules/modelSetup/ChromaFineTuneSetup.py | 1 - modules/modelSetup/ChromaLoRASetup.py | 1 - modules/modelSetup/FluxEmbeddingSetup.py | 2 - modules/modelSetup/FluxFineTuneSetup.py | 2 - modules/modelSetup/FluxLoRASetup.py | 2 - .../modelSetup/PixArtAlphaEmbeddingSetup.py | 1 - .../modelSetup/PixArtAlphaFineTuneSetup.py | 1 - modules/modelSetup/PixArtAlphaLoRASetup.py | 1 - modules/modelSetup/SanaEmbeddingSetup.py | 1 - modules/modelSetup/SanaFineTuneSetup.py | 1 - modules/modelSetup/SanaLoRASetup.py | 1 - .../StableDiffusion3EmbeddingSetup.py | 3 -- .../StableDiffusion3FineTuneSetup.py | 3 -- .../modelSetup/StableDiffusion3LoRASetup.py | 3 -- .../StableDiffusionEmbeddingSetup.py | 1 - .../StableDiffusionFineTuneSetup.py | 1 - .../modelSetup/StableDiffusionLoRASetup.py | 1 - .../StableDiffusionXLEmbeddingSetup.py | 2 - .../StableDiffusionXLFineTuneSetup.py | 2 - .../modelSetup/StableDiffusionXLLoRASetup.py | 2 - .../modelSetup/WuerstchenEmbeddingSetup.py | 1 - modules/modelSetup/WuerstchenFineTuneSetup.py | 1 - modules/modelSetup/WuerstchenLoRASetup.py | 1 - .../mixin/ModelSetupEmbeddingMixin.py | 14 +----- modules/util/thread_safety.py | 43 ------------------- requirements-global.txt | 4 +- run-cmd.sh | 5 --- start-ui.bat | 8 ---- start-ui.sh | 5 --- 70 files changed, 113 insertions(+), 329 deletions(-) delete mode 100644 modules/util/thread_safety.py diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..edd634e79 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -7,7 +7,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -32,8 +31,6 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel): image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673, unclear if Mistral is affected encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=HIDDEN_STATES_LAYER, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..0aa4522d7 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -14,7 +14,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -43,8 +42,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: mistral_format_input([caption], MISTRAL_SYSTEM_MESSAGE), apply_chat_template_kwargs = {'add_generation_prompt': False}, ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), hidden_state_output_index=MISTRAL_HIDDEN_STATES_LAYERS, @@ -53,8 +50,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: qwen3_format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': False} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=QWEN3_HIDDEN_STATES_LAYERS, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..66cfea16f 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -9,7 +9,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -38,8 +37,6 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH, apply_chat_template = lambda caption: format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': True} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=-2, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state') diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index d62a8d21b..adfd82c54 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -40,6 +40,7 @@ def __init__( class ChromaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,13 +143,13 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return ChromaPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, ) def add_text_encoder_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 4b85f1272..9c58c90cf 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -50,7 +50,9 @@ def __init__( class FluxModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: T5Tokenizer | None + orig_tokenizer_2: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: T5EncoderModel | None @@ -86,7 +88,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -177,15 +181,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return FluxPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index ca888f4f0..9c6959565 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -267,19 +267,19 @@ def eval(self): self.text_encoder_4.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HiDreamImagePipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.orig_tokenizer_3 if use_original_modules else self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, text_encoder_4=self.text_encoder_4, - tokenizer_4=self.orig_tokenizer_4 if use_original_modules else self.tokenizer_4, + tokenizer_4=self.orig_tokenizer_4 if use_original_tokenizers else self.tokenizer_4, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 146569abc..ae3c93053 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -196,15 +196,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HunyuanVideoPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 248b308bb..33fd421ed 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -42,6 +42,7 @@ def __init__( class PixArtAlphaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,11 +143,12 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer match self.model_type: case ModelType.PIXART_ALPHA: return PixArtAlphaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, @@ -153,7 +156,7 @@ def create_pipeline(self) -> DiffusionPipeline: ) case ModelType.PIXART_SIGMA: return PixArtSigmaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index d00d75fa2..7ff5501ea 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -41,6 +41,7 @@ def __init__( class SanaModel(BaseModel): # base model data tokenizer: GemmaTokenizer | None + orig_tokenizer: GemmaTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: Gemma2Model | None vae: AutoencoderDC | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -143,9 +145,9 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return SanaPipeline( - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 65f29ac49..4b3a5a9c5 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -58,8 +58,11 @@ def __init__( class StableDiffusion3Model(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None tokenizer_3: T5Tokenizer | None + orig_tokenizer_3: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModelWithProjection | None text_encoder_2: CLIPTextModelWithProjection | None @@ -98,8 +101,11 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.tokenizer_3 = None + self.orig_tokenizer_3 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -208,17 +214,17 @@ def eval(self): self.text_encoder_3.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusion3Pipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index fa6e1ad59..79cf1f6d1 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -43,6 +43,7 @@ def __init__( class StableDiffusionModel(BaseModel): # base model data tokenizer: CLIPTokenizer | None + orig_tokenizer: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: CLIPTextModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -136,12 +138,13 @@ def eval(self): self.text_encoder.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer if self.model_type.has_depth_input(): return StableDiffusionDepth2ImgPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, depth_estimator=self.depth_estimator, @@ -151,7 +154,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionInpaintPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, @@ -162,7 +165,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index 5f4f37dcf..f286e44dc 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -45,7 +45,9 @@ def __init__( class StableDiffusionXLModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: CLIPTextModelWithProjection | None @@ -81,7 +83,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -166,13 +170,13 @@ def eval(self): self.text_encoder_2.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusionXLPipeline( vae=self.vae, text_encoder=self.text_encoder_1, text_encoder_2=self.text_encoder_2, - tokenizer=self.tokenizer_1, - tokenizer_2=self.tokenizer_2, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, unet=self.unet, scheduler=self.noise_scheduler, ) diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index 391c65db5..9bd9f5dd6 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -69,6 +69,7 @@ class WuerstchenModel(BaseModel): decoder_vqgan: PaellaVQModel | None effnet_encoder: WuerstchenEfficientNetEncoder | None prior_tokenizer: CLIPTokenizer | None + orig_prior_tokenizer: CLIPTokenizer | None prior_text_encoder: CLIPTextModel | None prior_noise_scheduler: DDPMWuerstchenScheduler | None prior_prior: WuerstchenPrior | StableCascadeUNet | None @@ -105,6 +106,7 @@ def __init__( self.decoder_vqgan = None self.effnet_encoder = None self.prior_tokenizer = None + self.orig_prior_tokenizer = None self.prior_text_encoder = None self.prior_noise_scheduler = None self.prior_prior = None @@ -179,7 +181,8 @@ def eval(self): self.prior_text_encoder.eval() self.prior_prior.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + prior_tokenizer = self.orig_prior_tokenizer if use_original_tokenizers else self.prior_tokenizer if self.model_type.is_wuerstchen_v2(): return WuerstchenCombinedPipeline( tokenizer=self.decoder_tokenizer, @@ -187,19 +190,19 @@ def create_pipeline(self) -> DiffusionPipeline: decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, ) elif self.model_type.is_stable_cascade(): return StableCascadeCombinedPipeline( - tokenizer=self.prior_tokenizer, + tokenizer=prior_tokenizer, text_encoder=self.prior_text_encoder, decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, diff --git a/modules/modelLoader/ErnieModelLoader.py b/modules/modelLoader/ErnieModelLoader.py index 8e8365981..c9d623463 100644 --- a/modules/modelLoader/ErnieModelLoader.py +++ b/modules/modelLoader/ErnieModelLoader.py @@ -1,4 +1,3 @@ -import logging import os import traceback @@ -22,8 +21,7 @@ FlowMatchEulerDiscreteScheduler, GGUFQuantizationConfig, ) -from transformers import Mistral3Model, MistralConfig, PreTrainedTokenizerFast -from transformers.models.auto.configuration_auto import CONFIG_MAPPING +from transformers import AutoTokenizer, Mistral3Model class ErnieModelLoader( @@ -60,23 +58,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - # transformers < 5.x doesn't register "ministral3"; patch it so Mistral3Config can parse its text_config - if "ministral3" not in CONFIG_MAPPING: - CONFIG_MAPPING.register("ministral3", MistralConfig) - - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = ErnieImageTransformer2DModel.from_single_file( transformer_model_name, @@ -98,15 +79,10 @@ def __load_diffusers( quantization, ) - # TokenizersBackend is the Rust tokenizers library backend, not a transformers class — warning is a false alarm - tokenization_logger = logging.getLogger("transformers.tokenization_utils_base") - prev_level = tokenization_logger.level - tokenization_logger.setLevel(logging.ERROR) - tokenizer = PreTrainedTokenizerFast.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - tokenization_logger.setLevel(prev_level) text_encoder = self._load_transformers_sub_module( Mistral3Model, diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 414e43232..670bacab3 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -62,19 +62,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = Flux2Transformer2DModel.from_single_file( transformer_model_name, diff --git a/modules/modelLoader/ZImageModelLoader.py b/modules/modelLoader/ZImageModelLoader.py index 0fe6e3c00..104d04e7b 100644 --- a/modules/modelLoader/ZImageModelLoader.py +++ b/modules/modelLoader/ZImageModelLoader.py @@ -60,19 +60,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/chroma/ChromaModelLoader.py b/modules/modelLoader/chroma/ChromaModelLoader.py index 9c5a5d543..7dcbef794 100644 --- a/modules/modelLoader/chroma/ChromaModelLoader.py +++ b/modules/modelLoader/chroma/ChromaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -52,18 +53,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", @@ -120,6 +109,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/flux/FluxModelLoader.py b/modules/modelLoader/flux/FluxModelLoader.py index 7747a9f3e..d4f21ea2a 100644 --- a/modules/modelLoader/flux/FluxModelLoader.py +++ b/modules/modelLoader/flux/FluxModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -58,23 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -156,7 +140,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -231,7 +217,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/hiDream/HiDreamModelLoader.py b/modules/modelLoader/hiDream/HiDreamModelLoader.py index 1912cf20c..ae6fa2d1a 100644 --- a/modules/modelLoader/hiDream/HiDreamModelLoader.py +++ b/modules/modelLoader/hiDream/HiDreamModelLoader.py @@ -67,34 +67,6 @@ def __load_diffusers( include_text_encoder_4: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if include_text_encoder_3: - transformers_sub.append("text_encoder_3") - if include_text_encoder_4: - if text_encoder_4_model_name: - self._prepare_sub_modules( - text_encoder_4_model_name, - transformers_modules=[""], - diffusers_modules=[], - ) - else: - transformers_sub.append("text_encoder_4") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py index 0e9d577cc..85c91699b 100644 --- a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py +++ b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py @@ -59,24 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = LlamaTokenizerFast.from_pretrained( base_model_name, diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index 9180014bb..d7efa6f82 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -1,7 +1,9 @@ import json + +# huggingface_hub 1.16+ uses httpx, which logs every HTTP request/response at INFO level. +import logging import os import re -import traceback from abc import ABCMeta from itertools import repeat @@ -20,6 +22,8 @@ from huggingface_hub.utils import EntryNotFoundError from safetensors.torch import load_file +logging.getLogger("httpx").setLevel(logging.WARNING) + class HFModelLoaderMixin(metaclass=ABCMeta): def __init__(self): @@ -298,22 +302,3 @@ def _convert_diffusers_sub_module_to_dtype( None, quantization, ) - - def _prepare_sub_modules(self, pretrained_model_name_or_path: str, diffusers_modules: list[str], transformers_modules: list[str]): - is_local = os.path.isdir(pretrained_model_name_or_path) - if is_local: - return - - diffusers_paths = [((folder + "/") if folder else "") + "diffusion_pytorch_model*" for folder in diffusers_modules] - transformers_paths = [((folder + "/") if folder else "") + "model*" for folder in transformers_modules] - transformers_paths.extend([((folder + "/") if folder else "") + "pytorch_model*" for folder in transformers_modules]) - try: - huggingface_hub.snapshot_download( - pretrained_model_name_or_path, - allow_patterns=diffusers_paths + transformers_paths, - ) - except huggingface_hub.errors.HFValidationError: - pass - except Exception: - traceback.print_exc() - print("Error during bulk preloading of Huggingface model repository, proceeding without preloading") diff --git a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py index f8e75907d..467c29c8a 100644 --- a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py +++ b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/qwen/QwenModelLoader.py b/modules/modelLoader/qwen/QwenModelLoader.py index 498b925af..953f15bfb 100644 --- a/modules/modelLoader/qwen/QwenModelLoader.py +++ b/modules/modelLoader/qwen/QwenModelLoader.py @@ -52,18 +52,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/sana/SanaModelLoader.py b/modules/modelLoader/sana/SanaModelLoader.py index 508579ba2..a904e3996 100644 --- a/modules/modelLoader/sana/SanaModelLoader.py +++ b/modules/modelLoader/sana/SanaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py index 4333897ae..aa610f485 100644 --- a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py +++ b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -132,6 +133,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -206,6 +208,7 @@ def __load_ckpt( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -262,6 +265,7 @@ def __load_safetensors( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py index 17e55696d..47d87da74 100644 --- a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py +++ b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -50,8 +51,6 @@ def __load_diffusers( include_text_encoder_3: bool, quantization: QuantizationConfig, ): - #no call to self._prepare_sub_modules, because SAI polluted their sd3 / sd3.5 medium repo text encoders with fp16 files - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -141,8 +140,11 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -230,8 +232,11 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py index 6d31516fe..afbab6581 100644 --- a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py +++ b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -125,7 +126,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -167,7 +170,9 @@ def __load_ckpt( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -227,7 +232,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py index 8fec87b29..188107a2c 100644 --- a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py +++ b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py @@ -1,3 +1,4 @@ +import copy import json import os.path import traceback @@ -192,6 +193,7 @@ def __load_diffusers( model.decoder_vqgan = decoder_vqgan model.effnet_encoder = effnet_encoder model.prior_tokenizer = prior_tokenizer + model.orig_prior_tokenizer = copy.deepcopy(prior_tokenizer) model.prior_text_encoder = prior_text_encoder model.prior_noise_scheduler = prior_noise_scheduler model.prior_prior = prior_prior diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index d02197a35..dbafdae81 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -31,7 +31,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index 275626341..020df6a39 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -32,7 +32,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSaver/chroma/ChromaModelSaver.py b/modules/modelSaver/chroma/ChromaModelSaver.py index f8129558d..2dee69bfc 100644 --- a/modules/modelSaver/chroma/ChromaModelSaver.py +++ b/modules/modelSaver/chroma/ChromaModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/flux/FluxModelSaver.py b/modules/modelSaver/flux/FluxModelSaver.py index dc632f2ae..b5ac966d5 100644 --- a/modules/modelSaver/flux/FluxModelSaver.py +++ b/modules/modelSaver/flux/FluxModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hidream/HiDreamModelSaver.py b/modules/modelSaver/hidream/HiDreamModelSaver.py index 4bfe56ebc..b243166ec 100644 --- a/modules/modelSaver/hidream/HiDreamModelSaver.py +++ b/modules/modelSaver/hidream/HiDreamModelSaver.py @@ -24,7 +24,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py index 866d68617..3fad1c7a0 100644 --- a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py +++ b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py index 83e8af4a9..d6ef4d580 100644 --- a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py +++ b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/sana/SanaModelSaver.py b/modules/modelSaver/sana/SanaModelSaver.py index 1dede8c8f..29301adbb 100644 --- a/modules/modelSaver/sana/SanaModelSaver.py +++ b/modules/modelSaver/sana/SanaModelSaver.py @@ -22,7 +22,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py index e0131f52f..f24e6b5f9 100644 --- a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py +++ b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py index bcfabab67..095b9e75f 100644 --- a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py +++ b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py index 09170e67b..24625fc52 100644 --- a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py +++ b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py @@ -26,7 +26,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py index 3bbe1b55d..cee5e28b4 100644 --- a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py +++ b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline().prior_pipe + pipeline = model.create_pipeline(use_original_tokenizers=True).prior_pipe original_device = pipeline.device pipeline.to("cpu") pipeline_copy = copy.deepcopy(pipeline) diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 9fa65b69f..63a6d75af 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -61,7 +61,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 371c414ec..c78e4e700 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -70,7 +70,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 1396b4a52..514eca8e8 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index 0418390da..98dd12579 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index a2dc78057..3afbdf31c 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -80,8 +80,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index e3de941fd..e6204565f 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -122,8 +122,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py index 5d94c076d..5b842e267 100644 --- a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py +++ b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index 277f39004..31b92378e 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -74,7 +74,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 5c865b224..1e18dc673 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaEmbeddingSetup.py b/modules/modelSetup/SanaEmbeddingSetup.py index 7704c3afd..c31c2aa11 100644 --- a/modules/modelSetup/SanaEmbeddingSetup.py +++ b/modules/modelSetup/SanaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index 11faa747a..1c334c44f 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -68,7 +68,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index f481cd445..45cb2f7bb 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index 1ff229928..e9209bdc4 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -81,9 +81,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 3bbe41821..3ce400a59 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -90,9 +90,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index ae326cbbe..0d802c18a 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -148,9 +148,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index 6e3a41c2b..f8cc956f2 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index 8c42e59df..9f6fff8f3 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -76,7 +76,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index 9b9e9a613..d82147611 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -99,7 +99,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index 7d9ef442f..62f78cba2 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index ec6dbe16c..56a5bfe40 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -86,8 +86,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 6b193edc5..a0db29b07 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -120,8 +120,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 37dc7f85f..d5dd8a83b 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( ): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 1291125dc..111ce016f 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -71,7 +71,6 @@ def setup_model( if config.train_any_embedding(): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index b26f57436..812bb43e6 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( model.prior_prior_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.prior_prior_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py index 33b6d34cd..24f726230 100644 --- a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py +++ b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py @@ -13,27 +13,15 @@ CLIPTextModelWithProjection, Gemma2Model, LlamaModel, + PreTrainedTokenizer, T5EncoderModel, ) -from transformers.tokenization_utils import PreTrainedTokenizer, Trie class ModelSetupEmbeddingMixin(metaclass=ABCMeta): def __init__(self): super().__init__() - def _remove_added_embeddings_from_tokenizer( - self, - tokenizer: PreTrainedTokenizer, - ): - if tokenizer: - added_tokens = list(filter(lambda item: not item[1].special, tokenizer._added_tokens_decoder.items())) - for key, added_token in added_tokens: - tokenizer._added_tokens_decoder.pop(key) - tokenizer._added_tokens_encoder.pop(added_token.content) - tokenizer.tokens_trie = Trie() - tokenizer._update_trie() - def _create_new_embedding( self, model: BaseModel, diff --git a/modules/util/thread_safety.py b/modules/util/thread_safety.py deleted file mode 100644 index 365076c2a..000000000 --- a/modules/util/thread_safety.py +++ /dev/null @@ -1,43 +0,0 @@ -import functools -import threading - -import torch - -_THREAD_SAFE_FORWARD_ATTR = "_thread_safe_forward_lock" - - -def apply_thread_safe_forward(model: torch.nn.Module) -> None: - """ - Wrap ``model.forward()`` with a per-instance ``threading.Lock`` to - serialize concurrent calls. - - This is a workaround for a thread-safety bug in the transformers library's - ``check_model_inputs`` decorator, which monkey-patches child module - ``.forward()`` methods during execution and is not safe for concurrent use - from multiple dataloader threads. - - See: https://github.com/huggingface/transformers/issues/42673 - Fix: https://github.com/huggingface/transformers/pull/43765 (v5 only) - - This patch can be removed when upgrading to transformers v5+. - - The lock is per-model-instance so different model instances do not block - each other. The function is idempotent: calling it twice on the same model - is a no-op. - - Args: - model: The ``nn.Module`` whose ``forward()`` should be made thread-safe. - """ - if hasattr(model, _THREAD_SAFE_FORWARD_ATTR): - return - - lock = threading.Lock() - original_forward = model.forward - - @functools.wraps(original_forward) - def locked_forward(*args, **kwargs): - with lock: - return original_forward(*args, **kwargs) - - model.forward = locked_forward - setattr(model, _THREAD_SAFE_FORWARD_ATTR, lock) diff --git a/requirements-global.txt b/requirements-global.txt index 91014c6ff..299f765b0 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -5,7 +5,7 @@ pillow==12.2.0 imagesize==1.4.1 #for concept statistics tqdm==4.67.1 PyYAML==6.0.2 -huggingface-hub==0.34.4 +huggingface-hub==1.16.1 scipy==1.15.3 matplotlib==3.10.3 av==16.1.0 @@ -22,7 +22,7 @@ tensorboard==2.20.0 # diffusion models -e git+https://github.com/huggingface/diffusers.git@0f1abc4#egg=diffusers gguf==0.17.1 -transformers==4.57.6 +transformers==5.9.0 sentencepiece==0.2.1 # transitive dependency of transformers for tokenizer loading omegaconf==2.3.0 # needed to load stable diffusion from single ckpt files invisible-watermark==0.2.0 # needed for the SDXL pipeline diff --git a/run-cmd.sh b/run-cmd.sh index 2b1c9d14f..8fc77cb02 100755 --- a/run-cmd.sh +++ b/run-cmd.sh @@ -2,11 +2,6 @@ set -e -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - source "${BASH_SOURCE[0]%/*}/lib.include.sh" # Fetch and validate the name of the target script. diff --git a/start-ui.bat b/start-ui.bat index 5edc8e6c4..05881410b 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -35,14 +35,6 @@ set PYTHON="%VENV_DIR%\Scripts\python.exe" -X utf8 if defined PROFILE (set PYTHON=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) echo Using Python %PYTHON% -REM Disable HF_HUB_DISABLE_XET, buggy; default disables Xet (set to 0 to enable) - https://github.com/Nerogar/OneTrainer/issues/949 -if not defined HF_HUB_DISABLE_XET ( - set "HF_HUB_DISABLE_XET=1" -) -echo HF_HUB_DISABLE_XET=%HF_HUB_DISABLE_XET% -echo. -echo NOTE: Xet disabled, to enable it set as 0 before launch - :check_python_version echo Checking Python version... %PYTHON% --version diff --git a/start-ui.sh b/start-ui.sh index 2f0eecc0d..b2960c262 100755 --- a/start-ui.sh +++ b/start-ui.sh @@ -4,11 +4,6 @@ set -e source "${BASH_SOURCE[0]%/*}/lib.include.sh" -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - prepare_runtime_environment run_python_in_active_env "scripts/train_ui.py" "$@" From 6913fc66977062c180c3f60ae313c9ae25bdcfa3 Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 25 May 2026 16:01:16 +0200 Subject: [PATCH 27/67] Upgrade torch to 2.12.0 (CUDA 13.0 / ROCm 7.2) Co-Authored-By: Claude Sonnet 4.6 --- requirements-cuda.txt | 6 +++--- requirements-default.txt | 4 ++-- requirements-rocm.txt | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/requirements-cuda.txt b/requirements-cuda.txt index e53a76215..496ea1416 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -1,7 +1,7 @@ # pytorch ---extra-index-url https://download.pytorch.org/whl/cu128 -torch==2.11.0+cu128 -torchvision==0.26.0+cu128 +--extra-index-url https://download.pytorch.org/whl/cu130 +torch==2.12.0+cu130 +torchvision==0.27.0+cu130 onnxruntime-gpu==1.23.2 nvidia-nccl-cu12==2.27.5; sys_platform == "linux" triton-windows==3.5.1.post24; sys_platform == "win32" diff --git a/requirements-default.txt b/requirements-default.txt index 7581e2ac4..0bf5ee76d 100644 --- a/requirements-default.txt +++ b/requirements-default.txt @@ -1,6 +1,6 @@ # pytorch -torch==2.11.0 -torchvision==0.26.0 +torch==2.12.0 +torchvision==0.27.0 onnxruntime==1.23.2 # optimizers diff --git a/requirements-rocm.txt b/requirements-rocm.txt index 4f92044cb..c068fc082 100644 --- a/requirements-rocm.txt +++ b/requirements-rocm.txt @@ -2,9 +2,9 @@ # please open an issue or pull request on github # pytorch ---extra-index-url https://download.pytorch.org/whl/rocm6.3 -torch==2.11.0+rocm6.3 -torchvision==0.26.0+rocm6.3 +--extra-index-url https://download.pytorch.org/whl/rocm7.2 +torch==2.12.0+rocm7.2 +torchvision==0.27.0+rocm7.2 onnxruntime==1.23.2 # optimizers From 6e4c2e59d0af9cc6bed58dc7a7cc6d19d4b62e79 Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 25 May 2026 16:04:20 +0200 Subject: [PATCH 28/67] Switch nccl pin to cu13 (matches torch 2.12+cu130 dependency) Co-Authored-By: Claude Sonnet 4.6 --- requirements-cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 496ea1416..82ac8d9d2 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -3,7 +3,7 @@ torch==2.12.0+cu130 torchvision==0.27.0+cu130 onnxruntime-gpu==1.23.2 -nvidia-nccl-cu12==2.27.5; sys_platform == "linux" +nvidia-nccl-cu13==2.29.7; sys_platform == "linux" triton-windows==3.5.1.post24; sys_platform == "win32" # optimizers From b3c64629851626d01d63b7dc8a5b42aa9318006c Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 25 May 2026 16:05:21 +0200 Subject: [PATCH 29/67] Update triton-windows to 3.7.0.post26 (matches triton 3.7.0) Co-Authored-By: Claude Sonnet 4.6 --- requirements-cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 82ac8d9d2..5297caf00 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -4,7 +4,7 @@ torch==2.12.0+cu130 torchvision==0.27.0+cu130 onnxruntime-gpu==1.23.2 nvidia-nccl-cu13==2.29.7; sys_platform == "linux" -triton-windows==3.5.1.post24; sys_platform == "win32" +triton-windows==3.7.0.post26; sys_platform == "win32" # optimizers bitsandbytes==0.49.1 # bitsandbytes for 8-bit optimizers and weight quantization From 532625e22190cef43213e754b366a469d7d2b79a Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 25 May 2026 18:11:02 +0200 Subject: [PATCH 30/67] Make gradient checkpointing and offloading per-component; centralize model composition in ModelType - Gradient checkpointing and layer offloading are now configured per component (text encoder, transformer, VAE) rather than globally - ModelType centralizes model composition and training method associations Co-Authored-By: Claude Sonnet 4.6 --- modules/model/ChromaModel.py | 6 +- modules/model/ErnieModel.py | 6 +- modules/model/Flux2Model.py | 6 +- modules/model/FluxModel.py | 6 +- modules/model/HiDreamModel.py | 9 +- modules/model/HunyuanVideoModel.py | 6 +- modules/model/PixArtAlphaModel.py | 6 +- modules/model/QwenModel.py | 6 +- modules/model/SanaModel.py | 6 +- modules/model/StableDiffusion3Model.py | 6 +- modules/model/ZImageModel.py | 6 +- modules/modelSetup/BaseChromaSetup.py | 10 +- modules/modelSetup/BaseErnieSetup.py | 5 +- modules/modelSetup/BaseFlux2Setup.py | 17 +- modules/modelSetup/BaseFluxSetup.py | 14 +- modules/modelSetup/BaseHiDreamSetup.py | 23 +- modules/modelSetup/BaseHunyuanVideoSetup.py | 14 +- modules/modelSetup/BasePixArtAlphaSetup.py | 10 +- modules/modelSetup/BaseQwenSetup.py | 10 +- modules/modelSetup/BaseSanaSetup.py | 10 +- .../modelSetup/BaseStableDiffusion3Setup.py | 18 +- .../modelSetup/BaseStableDiffusionSetup.py | 10 +- .../modelSetup/BaseStableDiffusionXLSetup.py | 10 +- modules/modelSetup/BaseWuerstchenSetup.py | 8 +- modules/modelSetup/BaseZImageSetup.py | 10 +- modules/ui/ModelTab.py | 335 +++--------------- modules/ui/OffloadingWindow.py | 75 ---- modules/ui/TopBar.py | 38 +- modules/ui/TrainUI.py | 4 + modules/ui/TrainingTab.py | 161 ++++++--- modules/util/LayerOffloadConductor.py | 11 +- modules/util/checkpointing_util.py | 136 ++++--- modules/util/config/TrainConfig.py | 76 +++- modules/util/create.py | 9 +- .../util/enum/GradientCheckpointingMethod.py | 17 - modules/util/enum/ModelType.py | 55 +++ training_presets/#chroma Finetune 16GB.json | 7 +- training_presets/#chroma Finetune 8GB.json | 7 +- training_presets/#chroma LoRA 8GB.json | 7 +- training_presets/#ernie LoRA 8GB.json | 7 +- training_presets/#flux2 Finetune 16GB.json | 7 +- training_presets/#flux2 LoRA 8GB.json | 10 +- training_presets/#hidream LoRA.json | 11 +- training_presets/#hunyuan video LoRA.json | 8 +- training_presets/#qwen Finetune 16GB.json | 5 +- training_presets/#qwen Finetune 24GB.json | 5 +- training_presets/#qwen LoRA 16GB.json | 5 +- training_presets/#qwen LoRA 24GB.json | 5 +- .../#z-image DeTurbo LoRA 8GB.json | 5 +- training_presets/#z-image Finetune 16GB.json | 5 +- training_presets/#z-image LoRA 8GB.json | 5 +- 51 files changed, 525 insertions(+), 729 deletions(-) delete mode 100644 modules/ui/OffloadingWindow.py delete mode 100644 modules/util/enum/GradientCheckpointingMethod.py diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index d62a8d21b..228ccc424 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -111,8 +111,7 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -121,8 +120,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index fb44425e8..3d15e5385 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -72,15 +72,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index 004c70bee..a501bb116 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -122,15 +122,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 4b85f1272..22baeec67 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -145,8 +145,7 @@ def text_encoder_1_to(self, device: torch.device): def text_encoder_2_to(self, device: torch.device): if self.text_encoder_2 is not None: - if self.text_encoder_2_offload_conductor is not None and \ - self.text_encoder_2_offload_conductor.layer_offload_activated(): + if self.text_encoder_2_offload_conductor is not None: self.text_encoder_2_offload_conductor.to(device) else: self.text_encoder_2.to(device=device) @@ -155,8 +154,7 @@ def text_encoder_2_to(self, device: torch.device): self.text_encoder_2_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index ca888f4f0..1b45d1488 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -220,8 +220,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -231,8 +230,7 @@ def text_encoder_3_to(self, device: torch.device): def text_encoder_4_to(self, device: torch.device): if self.text_encoder_4 is not None: - if self.text_encoder_4_offload_conductor is not None and \ - self.text_encoder_4_offload_conductor.layer_offload_activated(): + if self.text_encoder_4_offload_conductor is not None: self.text_encoder_4_offload_conductor.to(device) else: self.text_encoder_4.to(device=device) @@ -241,8 +239,7 @@ def text_encoder_4_to(self, device: torch.device): self.text_encoder_4_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 146569abc..075c406c7 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -157,8 +157,7 @@ def text_encoder_to(self, device: torch.device): def text_encoder_1_to(self, device: torch.device): if self.text_encoder_1 is not None: - if self.text_encoder_1_offload_conductor is not None and \ - self.text_encoder_1_offload_conductor.layer_offload_activated(): + if self.text_encoder_1_offload_conductor is not None: self.text_encoder_1_offload_conductor.to(device) else: self.text_encoder_1.to(device=device) @@ -174,8 +173,7 @@ def text_encoder_2_to(self, device: torch.device): self.text_encoder_2_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 248b308bb..ab0155fd3 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -112,8 +112,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -122,8 +121,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index afa6c24fe..1aadcd467 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -81,8 +81,7 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -91,8 +90,7 @@ def text_encoder_to(self, device: torch.device): #TODO share more code between m self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index d00d75fa2..39ac83310 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -114,8 +114,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -124,8 +123,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 65f29ac49..378bb78f8 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -174,8 +174,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -184,8 +183,7 @@ def text_encoder_3_to(self, device: torch.device): self.text_encoder_3_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index 7fd9e52cb..167b58f47 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -83,15 +83,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..f8cc33187 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -49,12 +49,10 @@ def setup_optimizations( model: ChromaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_chroma_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..72f8ee2d8 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -42,9 +42,8 @@ def setup_optimizations( model: ErnieModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_ernie_transformer(model.transformer, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..6e6e7320f 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -45,16 +45,13 @@ def setup_optimizations( model: Flux2Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux2_transformer(model.transformer, config) - if model.text_encoder is not None: - if model.is_dev(): - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config) - else: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + if model.is_dev(): + model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder) + else: + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..e6638ddb2 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -49,14 +49,12 @@ def setup_optimizations( model: FluxModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - model.text_encoder_2_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + model.text_encoder_2_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 17fbcc0d6..b5a856b73 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -49,19 +49,16 @@ def setup_optimizations( model: HiDreamModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hi_dream_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) - if model.text_encoder_4 is not None: - model.text_encoder_4_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hi_dream_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) + if model.text_encoder_4 is not None and config.text_encoder_4.checkpointing_or_offloading_enabled(): + model.text_encoder_4_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config, config.text_encoder_4) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index b072bf4ba..2e32165e7 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -49,14 +49,12 @@ def setup_optimizations( model: HunyuanVideoModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - model.text_encoder_1_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_1_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 3069b4884..4935ebbfd 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -51,12 +51,10 @@ def setup_optimizations( model: PixArtAlphaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, offload_enabled=True) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a8a7be8f6..82004d8ac 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -46,12 +46,10 @@ def setup_optimizations( model: QwenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_qwen_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c9ea6da0..c26acbc82 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -52,12 +52,10 @@ def setup_optimizations( config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - # model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_sana_transformer(model.transformer, config) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_gemma_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_sana_transformer(model.transformer, config, config.transformer) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_gemma_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index de5dc04e8..efe761427 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -48,16 +48,14 @@ def setup_optimizations( model: StableDiffusion3Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 8cf63ac07..a16c72cb8 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -51,11 +51,13 @@ def setup_optimizations( model: StableDiffusionModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + if config.vae.checkpointing_enabled(): + model.vae.enable_gradient_checkpointing() + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 61cb1e457..7fe35e1e1 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -48,11 +48,13 @@ def setup_optimizations( model: StableDiffusionXLModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 48cd05942..abc77bcdd 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -57,13 +57,13 @@ def setup_optimizations( model: WuerstchenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.prior.checkpointing_or_offloading_enabled(): if model.model_type.is_wuerstchen_v2(): model.prior_prior.enable_gradient_checkpointing() - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config) elif model.model_type.is_stable_cascade(): - enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config) - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config) + enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config, config.prior) + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.decoder_vqgan) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b822d7304..d64e7e88a 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -47,12 +47,10 @@ def setup_optimizations( model: ZImageModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_z_image_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index ff17ea3ba..5bc3092d3 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -45,307 +45,66 @@ def refresh_ui(self): base_frame.grid_columnconfigure(3, weight=0) base_frame.grid_columnconfigure(4, weight=1) - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ - TrainingMethod.FINE_TUNE, - TrainingMethod.FINE_TUNE_VAE, - ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_3_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_2_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + self.__setup_ui(base_frame) - def __setup_z_image_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + def __setup_ui(self, frame): + model_type = self.train_config.model_type + training_method = self.train_config.training_method + parts = model_type.model_parts() - def __setup_ernie_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + # The transformer override path exists only for these architectures; SD3, PixArt, Sana + # and HiDream have a transformer but expose no override field. + allow_override_transformer = ( + model_type.is_flux() + or model_type.is_z_image() + or model_type.is_ernie() + or model_type.is_chroma() + or model_type.is_qwen() + or model_type.is_hunyuan_video() ) - def __setup_chroma_ui(self, frame): row = 0 row = self.__create_base_dtype_components(frame, row) row = self.__create_base_components( frame, row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + has_unet="unet" in parts, + has_prior="prior" in parts, + allow_override_prior=model_type.is_stable_cascade(), + has_transformer="transformer" in parts, + allow_override_transformer=allow_override_transformer, + has_text_encoder=not model_type.has_multiple_text_encoders(), + has_text_encoder_1=model_type.has_multiple_text_encoders(), + has_text_encoder_2="text_encoder_2" in parts, + has_text_encoder_3="text_encoder_3" in parts, + has_text_encoder_4="text_encoder_4" in parts, + allow_override_text_encoder_4="text_encoder_4" in parts, + has_vae="vae" in parts, + ) + if "effnet_encoder" in parts: + row = self.__create_effnet_encoder_components(frame, row) + if "decoder" in parts: + row = self.__create_decoder_components(frame, row, "decoder_text_encoder" in parts) + + if model_type.is_sana(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE + elif model_type.is_wuerstchen(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE \ + or model_type.is_stable_cascade() + else: + allow_safetensors = True + + if model_type.is_stable_diffusion(): + allow_diffusers = training_method in [TrainingMethod.FINE_TUNE, TrainingMethod.FINE_TUNE_VAE] + else: + allow_diffusers = training_method == TrainingMethod.FINE_TUNE - def __setup_qwen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_xl_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_wuerstchen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), - has_text_encoder=True, - ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_pixart_alpha_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_sana_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hunyuan_video_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hi_dream_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_text_encoder_4=True, - allow_override_text_encoder_4=True, - has_vae=True, - ) row = self.__create_output_components( frame, row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_safetensors=allow_safetensors, + allow_diffusers=allow_diffusers, + allow_legacy_safetensors=training_method == TrainingMethod.LORA, ) def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: diff --git a/modules/ui/OffloadingWindow.py b/modules/ui/OffloadingWindow.py deleted file mode 100644 index 54035e121..000000000 --- a/modules/ui/OffloadingWindow.py +++ /dev/null @@ -1,75 +0,0 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index 820fdb71a..d2ca78f7b 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -112,37 +112,13 @@ def __create_training_method(self): if self.training_method: self.training_method.destroy() - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] + labels = { + TrainingMethod.FINE_TUNE: "Fine Tune", + TrainingMethod.LORA: "LoRA", + TrainingMethod.EMBEDDING: "Embedding", + TrainingMethod.FINE_TUNE_VAE: "Fine Tune VAE", + } + values = [(labels[m], m) for m in self.train_config.model_type.supported_training_methods()] # training method self.training_method = components.options_kv( diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index ba90d2e64..f822895d6 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -308,6 +308,10 @@ def create_general_tab(self, master): tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + components.label(frame, 11, 2, "Async Offloading", + tooltip="Overlaps CPU<->GPU transfers with computation using CUDA streams. Applies to every offloaded component") + components.switch(frame, 11, 3, self.ui_state, "async_offloading") + components.label(frame, 12, 0, "Multi-GPU", tooltip="Enable multi-GPU training") components.switch(frame, 12, 1, self.ui_state, "multi_gpu") diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index bcca11ae9..2d24d968c 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -1,4 +1,3 @@ -from modules.ui.OffloadingWindow import OffloadingWindow from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow @@ -6,13 +5,13 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.LearningRateScaler import LearningRateScaler from modules.util.enum.LearningRateScheduler import LearningRateScheduler from modules.util.enum.LossScaler import LossScaler from modules.util.enum.LossWeight import LossWeight from modules.util.enum.Optimizer import Optimizer from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.optimizer_util import change_optimizer from modules.util.ui import components from modules.util.ui.UIState import UIState @@ -99,6 +98,8 @@ def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): self.__create_unet_frame(column_1, 1) self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + if self.train_config.training_method == TrainingMethod.FINE_TUNE_VAE: + self.__create_vae_frame(column_2, 0) self.__create_masked_frame(column_2, 1) self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) @@ -375,19 +376,6 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") row += 1 - # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) - row += 1 - - # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") - row += 1 - # train dtype components.label(frame, row, 0, "Train Data Type", tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") @@ -434,38 +422,70 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, tooltip="Enables circular padding for all conv layers to better train seamless images") components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=True, supports_activation_offloading=False): + if supports_checkpointing: + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for this component. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, f"{part}.gradient_checkpointing") + row += 1 + + components.label(frame, row, 0, "Layer Offload Fraction", + tooltip="Fraction of this component's layers to offload to CPU to reduce VRAM usage. Increases training time and RAM usage. 0=disabled, 1=all layers") + components.entry(frame, row, 1, self.ui_state, f"{part}.offload_fraction") + row += 1 + + if supports_activation_offloading: + components.label(frame, row, 0, "Offload Activations", + tooltip="Offloads this component's activations to CPU during training to reduce VRAM usage") + components.switch(frame, row, 1, self.ui_state, f"{part}.activation_offloading") + row += 1 + + return row + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", + components.label(frame, row, 0, "Train Text Encoder", tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + components.switch(frame, row, 1, self.ui_state, "text_encoder.train") + row += 1 + else: + # no Train switch to act as the frame's header, so add an explicit one + components.label(frame, row, 0, "Text Encoder") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "text_encoder", supports_checkpointing=supports_training) # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", + components.label(frame, row, 0, "Caption Dropout Probability", tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + components.entry(frame, row, 1, self.ui_state, "text_encoder.dropout_probability") + row += 1 if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + components.time_entry(frame, row, 1, self.ui_state, "text_encoder.stop_training_after", "text_encoder.stop_training_after_unit", supports_time_units=False) + row += 1 # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", + components.label(frame, row, 0, "Text Encoder Learning Rate", tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + components.entry(frame, row, 1, self.ui_state, "text_encoder.learning_rate") + row += 1 if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", + components.label(frame, row, 0, "Clip Skip", tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + components.entry(frame, row, 1, self.ui_state, "text_encoder_layer_skip") + row += 1 if supports_sequence_length: # text encoder sequence length @@ -503,6 +523,8 @@ def __create_text_encoder_n_frame( components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") row += 1 + row = self.__create_offloading_widgets(frame, row, f"text_encoder{suffix}") + # train text encoder embedding components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", tooltip=f"Enables training embeddings for the text encoder {i} model") @@ -560,82 +582,119 @@ def __create_unet_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train unet - components.label(frame, 0, 0, "Train UNet", + components.label(frame, row, 0, "Train UNet", tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") + components.switch(frame, row, 1, self.ui_state, "unet.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "unet", supports_activation_offloading=True) # train unet epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", supports_time_units=False) + row += 1 # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", + components.label(frame, row, 0, "UNet Learning Rate", tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + components.entry(frame, row, 1, self.ui_state, "unet.learning_rate") + row += 1 # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + rescale_label = components.label(frame, row, 0, "Rescale Noise Scheduler + V-pred", tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + components.switch(frame, row, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + row += 1 + + def __create_vae_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + components.label(frame, row, 0, "Train VAE", + tooltip="Enables training the VAE model") + components.switch(frame, row, 1, self.ui_state, "vae.train") + row += 1 + + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for the VAE. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, "vae.gradient_checkpointing") + row += 1 def __create_prior_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train prior - components.label(frame, 0, 0, "Train Prior", + components.label(frame, row, 0, "Train Prior", tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") + components.switch(frame, row, 1, self.ui_state, "prior.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "prior", supports_activation_offloading=True) # train prior epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", supports_time_units=False) + row += 1 # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", + components.label(frame, row, 0, "Prior Learning Rate", tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + components.entry(frame, row, 1, self.ui_state, "prior.learning_rate") + row += 1 def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train transformer - components.label(frame, 0, 0, "Train Transformer", + components.label(frame, row, 0, "Train Transformer", tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") + components.switch(frame, row, 1, self.ui_state, "transformer.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "transformer", supports_activation_offloading=True) # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", supports_time_units=False) + row += 1 # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", + components.label(frame, row, 0, "Transformer Learning Rate", tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + components.entry(frame, row, 1, self.ui_state, "transformer.learning_rate") + row += 1 if supports_force_attention_mask: # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", + components.label(frame, row, 0, "Force Attention Mask", tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + components.switch(frame, row, 1, self.ui_state, "transformer.attention_mask") + row += 1 if supports_guidance_scale: # guidance scale - components.label(frame, 4, 0, "Guidance Scale", + components.label(frame, row, 0, "Guidance Scale", tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + components.entry(frame, row, 1, self.ui_state, "transformer.guidance_scale") + row += 1 def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): frame = ctk.CTkFrame(master=master, corner_radius=5) @@ -838,10 +897,6 @@ def __open_timestep_distribution_window(self): window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) self.master.wait_window(window) - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - def __restore_optimizer_config(self, *args): optimizer_config = change_optimizer(self.train_config) self.ui_state.get_var("optimizer").update(optimizer_config) diff --git a/modules/util/LayerOffloadConductor.py b/modules/util/LayerOffloadConductor.py index 0b7e6c7ca..062d1391e 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -2,7 +2,7 @@ import random from typing import Any -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.quantization_util import get_offload_tensor_bytes, offload_quantized from modules.util.torch_util import ( create_stream_context, @@ -566,6 +566,7 @@ def __init__( self, module: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): super().__init__() @@ -573,16 +574,16 @@ def __init__( self.__layers = [] self.__layer_device_map = [] - self.__layer_offload_fraction = config.layer_offload_fraction + self.__layer_offload_fraction = part.offload_fraction self.__layer_activations_included_offload_param_indices_map = [] self.__train_device = torch.device(config.train_device) self.__temp_device = torch.device(config.temp_device) - self.__offload_activations = config.gradient_checkpointing.offload() and config.enable_activation_offloading - self.__offload_layers = config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 - self.__async_transfer = self.__train_device.type == "cuda" and config.enable_async_offloading + self.__offload_activations = part.activation_offloading + self.__offload_layers = part.offload_fraction > 0 + self.__async_transfer = self.__train_device.type == "cuda" and config.async_offloading if self.__async_transfer: self.__train_stream = torch.cuda.default_stream(self.__train_device) diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index ba3f7fec2..d5b9f765c 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -3,7 +3,7 @@ from typing import Any from modules.util.compile_util import init_compile -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.LayerOffloadConductor import LayerOffloadConductor from modules.util.torch_util import add_dummy_grad_fn_, has_grad_fn @@ -74,21 +74,25 @@ def __init__(self, *args, **kwargs): class CheckpointLayer(BaseCheckpointLayer): - def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device): + def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, checkpointing: bool = True): super().__init__() assert (orig_module is None or orig_forward is None) and not (orig_module is None and orig_forward is None) self.checkpoint = orig_module self.orig_forward = orig_forward + self.checkpointing = checkpointing # dummy tensor that requires grad is needed for checkpointing to work when training a LoRA self.dummy = torch.zeros((1,), device=train_device, requires_grad=True) - def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + def __orig(self, *args, **kwargs): return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + return self.__orig(*args, **kwargs) + def forward(self, *args, **kwargs): - if torch.is_grad_enabled(): + if self.checkpointing and torch.is_grad_enabled(): return torch.utils.checkpoint.checkpoint( self.__checkpointing_forward, self.dummy, @@ -97,7 +101,7 @@ def forward(self, *args, **kwargs): use_reentrant=False ) else: - return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + return self.__orig(*args, **kwargs) class OffloadCheckpointLayer(BaseCheckpointLayer): def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, conductor: LayerOffloadConductor, layer_index: int): @@ -153,6 +157,7 @@ def create_checkpoint( train_device: torch.device, include_from_offload_param_names: list[str] = None, conductor: LayerOffloadConductor | None = None, + checkpointing: bool = True, layer_index: int = 0, compile: bool = False, ) -> Callable: @@ -164,6 +169,9 @@ def create_checkpoint( conductor.add_layer(orig_module, included_offload_param_indices) if conductor is not None and conductor.offload_activated(): + # offloading is structurally coupled to use_reentrant=True checkpointing during the back pass + # (the recompute is what fires before_layer/after_layer in the backward direction), so the offload + # layer always checkpoints when grad is enabled, regardless of the part's gradient_checkpointing flag. if compile: layer = OffloadCheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, conductor=conductor, layer_index=layer_index) #don't compile the checkpointing layer - offloading cannot be compiled: @@ -176,12 +184,12 @@ def create_checkpoint( return orig_module else: if compile: - layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device) + layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, checkpointing=checkpointing) #do compile the checkpointing layer - slightly faster layer.compile(fullgraph=True) return layer else: - layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device) + layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, checkpointing=checkpointing) orig_module.forward = layer.forward return orig_module @@ -189,6 +197,7 @@ def _create_checkpoints_for_module_list( module_list: nn.ModuleList, include_from_offload_param_names: list[str], conductor: LayerOffloadConductor, + checkpointing: bool, train_device: torch.device, layer_index: int, compile: bool, @@ -200,7 +209,7 @@ def _create_checkpoints_for_module_list( module_list[i] = create_checkpoint( layer, train_device, include_from_offload_param_names, - conductor, layer_index, compile=compile, + conductor, checkpointing, layer_index, compile=compile, ) layer_index += 1 return layer_index @@ -213,11 +222,15 @@ def _remove_checkpoint_keys(module, state_dict, prefix, local_metadata): def enable_checkpointing( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, compile: bool, lists, # if there are multiple entries in this list, they must be in the exact order they are executed - otherwise offloading fails offload_enabled: bool = True, -) -> LayerOffloadConductor: - conductor = LayerOffloadConductor(model, config) +) -> LayerOffloadConductor | None: + # a conductor exists iff this part actually offloads (and the component supports conductor offloading) + offload = offload_enabled and part.offloading_enabled() + conductor = LayerOffloadConductor(model, config, part) if offload else None + checkpointing = part.checkpointing_enabled() layer_index = 0 for type_or_list, param_names in lists: @@ -228,7 +241,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -242,7 +256,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -253,9 +268,10 @@ def enable_checkpointing( def enable_checkpointing_for_basic_transformer_blocks( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, offload_enabled: bool, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (BasicTransformerBlock , []), ], offload_enabled = offload_enabled, @@ -264,16 +280,18 @@ def enable_checkpointing_for_basic_transformer_blocks( def enable_checkpointing_for_clip_encoder_layers( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): - return enable_checkpointing(model, config, False, [ + return enable_checkpointing(model, config, part, False, [ (CLIPEncoderLayer, []), # No activation offloading for text encoders, because the output might be taken from the middle of the network ]) def enable_checkpointing_for_stable_cascade_blocks( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (SDCascadeResBlock, []), (SDCascadeAttnBlock, []), (SDCascadeTimestepBlock, []), @@ -282,8 +300,9 @@ def enable_checkpointing_for_stable_cascade_blocks( def enable_checkpointing_for_t5_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (T5Block, []), ]) @@ -291,8 +310,9 @@ def enable_checkpointing_for_t5_encoder_layers( def enable_checkpointing_for_gemma_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Gemma2DecoderLayer, []), ]) @@ -300,17 +320,19 @@ def enable_checkpointing_for_gemma_layers( def enable_checkpointing_for_llama_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (LlamaDecoderLayer, []), ]) def enable_checkpointing_for_mistral_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (MistralDecoderLayer, []), + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (MistralDecoderLayer, []), # no activation offloading: this encoder is never trained ]) @@ -318,32 +340,36 @@ def enable_checkpointing_for_mistral_encoder_layers( def enable_checkpointing_for_qwen25vl_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Qwen2_5_VLDecoderLayer, []), # TODO No activation offloading for other encoders, see above. But clip skip is not implemented for QwenVL. Then do activation offloading? ]) def enable_checkpointing_for_qwen3_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (Qwen3DecoderLayer, []), # No activation offloading, because hidden states are taken from the middle of the network by Flux2 + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (Qwen3DecoderLayer, []), # no activation offloading: this encoder is never trained ]) def enable_checkpointing_for_stable_diffusion_3_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (JointTransformerBlock, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_flux_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -351,8 +377,9 @@ def enable_checkpointing_for_flux_transformer( def enable_checkpointing_for_flux2_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -361,8 +388,9 @@ def enable_checkpointing_for_flux2_transformer( def enable_checkpointing_for_chroma_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -371,16 +399,18 @@ def enable_checkpointing_for_chroma_transformer( def enable_checkpointing_for_qwen_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_z_image_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.noise_refiner, ["x"]), (model.context_refiner, ["x"]), (model.layers, ["x"]), @@ -390,16 +420,18 @@ def enable_checkpointing_for_z_image_transformer( def enable_checkpointing_for_sana_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (SanaTransformerBlock, ["hidden_states"]), ]) def enable_checkpointing_for_hunyuan_video_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (HunyuanVideoIndividualTokenRefinerBlock, ["hidden_states" ]), (HunyuanVideoTransformerBlock, ["hidden_states", "encoder_hidden_states"]), (HunyuanVideoSingleTransformerBlock, ["hidden_states" ]), @@ -408,8 +440,9 @@ def enable_checkpointing_for_hunyuan_video_transformer( def enable_checkpointing_for_hi_dream_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (HiDreamImageTransformerBlock, ["hidden_states", "encoder_hidden_states"]), (HiDreamImageSingleTransformerBlock, ["hidden_states" ]), ]) @@ -417,7 +450,8 @@ def enable_checkpointing_for_hi_dream_transformer( def enable_checkpointing_for_ernie_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.layers, ["x"]), ]) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index bbc70a030..125b53669 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -13,7 +13,6 @@ from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.GradientReducePrecision import GradientReducePrecision from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.LearningRateScaler import LearningRateScaler @@ -267,10 +266,27 @@ class TrainModelPartConfig(BaseConfig): train_embedding: bool attention_mask: bool guidance_scale: float + gradient_checkpointing: bool + offload_fraction: float + activation_offloading: bool def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) + def offloading_enabled(self) -> bool: + # a conductor should exist iff this is True. Layer offloading applies even to frozen parts (to fit + # them in VRAM), but activation offloading only does work during a backward pass, so it only applies + # when the part is trained -- even if activation_offloading is True in the config. + return self.offload_fraction > 0 or (self.activation_offloading and self.train) + + def checkpointing_enabled(self) -> bool: + # the inner torch checkpoint() should run iff this is True + return self.gradient_checkpointing and self.train + + def checkpointing_or_offloading_enabled(self) -> bool: + # whether the checkpoint layer wrapper needs to be installed for this part at all + return self.checkpointing_enabled() or self.offloading_enabled() + @staticmethod def default_values(): data = [] @@ -287,6 +303,9 @@ def default_values(): data.append(("train_embedding", True, bool, False)) data.append(("attention_mask", False, bool, False)) data.append(("guidance_scale", 1.0, float, False)) + data.append(("gradient_checkpointing", True, bool, False)) + data.append(("offload_fraction", 0.0, float, False)) + data.append(("activation_offloading", True, bool, False)) return TrainModelPartConfig(data) @@ -374,10 +393,7 @@ class TrainConfig(BaseConfig): output_dtype: DataType output_model_format: ModelFormat output_model_destination: str - gradient_checkpointing: GradientCheckpointingMethod - enable_async_offloading: bool - enable_activation_offloading: bool - layer_offload_fraction: float + async_offloading: bool force_circular_padding: bool compile: bool @@ -558,7 +574,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -570,6 +586,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -716,12 +733,14 @@ def __migration_3(self, data: dict) -> dict: def __migration_4(self, data: dict) -> dict: migrated_data = data.copy() + # Translate the old bool form of gradient_checkpointing into the v5..v10 + # string/enum form. __migration_10 later fans this out per-component. gradient_checkpointing = migrated_data.pop("gradient_checkpointing", True) if gradient_checkpointing: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.ON + migrated_data["gradient_checkpointing"] = "ON" else: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.OFF + migrated_data["gradient_checkpointing"] = "OFF" return migrated_data @@ -789,6 +808,42 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + + # Fan the four old global offload/checkpointing settings out per-component. + # After __migration_4 gradient_checkpointing is a string "OFF"/"ON"/"CPU_OFFLOADED". + gc = migrated_data.pop("gradient_checkpointing", "ON") + act = migrated_data.pop("enable_activation_offloading", True) + frac = migrated_data.pop("layer_offload_fraction", 0.0) + migrated_data["async_offloading"] = migrated_data.pop("enable_async_offloading", True) + + def fan_out(part: str): + if part in migrated_data: + migrated_data[part]["gradient_checkpointing"] = gc != "OFF" + migrated_data[part]["activation_offloading"] = (gc == "CPU_OFFLOADED") and act + migrated_data[part]["offload_fraction"] = frac if gc == "CPU_OFFLOADED" else 0.0 + + fan_out("unet") + fan_out("prior") + fan_out("transformer") + fan_out("text_encoder") + fan_out("text_encoder_2") + fan_out("text_encoder_3") + fan_out("text_encoder_4") + fan_out("vae") + fan_out("effnet_encoder") + fan_out("decoder") + fan_out("decoder_text_encoder") + fan_out("decoder_vqgan") + + return migrated_data + + def model_part_configs(self) -> list[TrainModelPartConfig]: + # the per-part configs for the components this model_type actually has. Avoids "phantom" parts whose + # fields keep their defaults (train=True) or migrated offload values but don't exist in the model. + return [getattr(self, name) for name in self.model_type.model_parts()] + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -958,10 +1013,7 @@ def default_values() -> 'TrainConfig': data.append(("output_dtype", DataType.FLOAT_32, DataType, False)) data.append(("output_model_format", ModelFormat.SAFETENSORS, ModelFormat, False)) data.append(("output_model_destination", "models/model.safetensors", str, False)) - data.append(("gradient_checkpointing", GradientCheckpointingMethod.ON, GradientCheckpointingMethod, False)) - data.append(("enable_async_offloading", True, bool, False)) - data.append(("enable_activation_offloading", True, bool, False)) - data.append(("layer_offload_fraction", 0.0, float, False)) + data.append(("async_offloading", True, bool, False)) data.append(("force_circular_padding", False, bool, False)) data.append(("compile", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..0c0adba16 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -110,7 +110,11 @@ def create_data_loader( train_progress: TrainProgress | None = None, is_validation: bool = False ) -> BaseDataLoader | None: - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 and config.dataloader_threads > 1: + # Layer offloading uses a non-thread-safe conductor. This check is too broad: it trips whenever any model + # part does layer offloading, even though only a component that is actually cached really runs in the + # dataloader worker threads. + # TODO: narrow this to the cached components only. + if config.dataloader_threads > 1 and any(part.offload_fraction > 0 for part in config.model_part_configs()): raise RuntimeError('layer offloading can not be activated if "dataloader_threads" > 1') if train_progress is None: @@ -133,7 +137,8 @@ def create_optimizer( if optimizer_config.optimizer is None: return None - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0: + # a trained, layer-offloaded part has its params evicted during the back pass, so it needs fused_back_pass + if any(part.offload_fraction > 0 and part.train for part in config.model_part_configs()): if (not optimizer_config.optimizer.supports_fused_back_pass() or not optimizer_config.fused_back_pass) \ and config.training_method == TrainingMethod.FINE_TUNE: raise RuntimeError('layer offloading can only be used for fine tuning when using an optimizer that supports "fused_back_pass"') diff --git a/modules/util/enum/GradientCheckpointingMethod.py b/modules/util/enum/GradientCheckpointingMethod.py deleted file mode 100644 index d3f05666a..000000000 --- a/modules/util/enum/GradientCheckpointingMethod.py +++ /dev/null @@ -1,17 +0,0 @@ -from enum import Enum - - -class GradientCheckpointingMethod(Enum): - OFF = 'OFF' - ON = 'ON' - CPU_OFFLOADED = 'CPU_OFFLOADED' - - def __str__(self): - return self.value - - def enabled(self): - return self == GradientCheckpointingMethod.ON \ - or self == GradientCheckpointingMethod.CPU_OFFLOADED - - def offload(self): - return self == GradientCheckpointingMethod.CPU_OFFLOADED diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index 5b6658501..4df90ba52 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -1,5 +1,7 @@ from enum import Enum +from modules.util.enum.TrainingMethod import TrainingMethod + class ModelType(Enum): STABLE_DIFFUSION_15 = 'STABLE_DIFFUSION_15' @@ -166,6 +168,59 @@ def is_flow_matching(self) -> bool: def is_video_model(self) -> bool: return self.is_hunyuan_video() #incase we add more video models in the future + def model_parts(self) -> tuple[str, ...]: + return _MODEL_PARTS[self] + + def supported_training_methods(self) -> tuple[TrainingMethod, ...]: + if self.is_stable_diffusion(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING, TrainingMethod.FINE_TUNE_VAE) + if self.is_stable_diffusion_3() \ + or self.is_stable_diffusion_xl() \ + or self.is_wuerstchen() \ + or self.is_pixart() \ + or self.is_flux_1() \ + or self.is_sana() \ + or self.is_hunyuan_video() \ + or self.is_hi_dream() \ + or self.is_chroma(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING) + if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA) + raise ValueError(f"No supported training methods defined for model type {self}") + + +# The first text encoder is always "text_encoder" here (matching the config field), even for +# multi-encoder models that refer to it as "text_encoder_1" elsewhere in the code. +_MODEL_PARTS: dict[ModelType, tuple[str, ...]] = { + ModelType.STABLE_DIFFUSION_15: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_15_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_DEPTH: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_3: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_35: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE: ("text_encoder", "text_encoder_2", "unet", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING: ("text_encoder", "text_encoder_2", "unet", "vae"), + # Only Würstchen v2's decoder has its own text encoder; Stable Cascade's decoder does not. + ModelType.WUERSTCHEN_2: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_text_encoder", "decoder_vqgan"), + ModelType.STABLE_CASCADE_1: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_vqgan"), + ModelType.PIXART_ALPHA: ("text_encoder", "transformer", "vae"), + ModelType.PIXART_SIGMA: ("text_encoder", "transformer", "vae"), + ModelType.FLUX_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_FILL_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_2: ("text_encoder", "transformer", "vae"), + ModelType.SANA: ("text_encoder", "transformer", "vae"), + ModelType.HUNYUAN_VIDEO: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.HI_DREAM_FULL: ("text_encoder", "text_encoder_2", "text_encoder_3", "text_encoder_4", "transformer", "vae"), + ModelType.CHROMA_1: ("text_encoder", "transformer", "vae"), + ModelType.QWEN: ("text_encoder", "transformer", "vae"), + ModelType.Z_IMAGE: ("text_encoder", "transformer", "vae"), + ModelType.ERNIE: ("text_encoder", "transformer", "vae"), +} + class PeftType(Enum): LORA = 'LORA' diff --git a/training_presets/#chroma Finetune 16GB.json b/training_presets/#chroma Finetune 16GB.json index 2dacbee20..145ab5b04 100644 --- a/training_presets/#chroma Finetune 16GB.json +++ b/training_presets/#chroma Finetune 16GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.4, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.4 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma Finetune 8GB.json b/training_presets/#chroma Finetune 8GB.json index 508410995..29b36c84b 100644 --- a/training_presets/#chroma Finetune 8GB.json +++ b/training_presets/#chroma Finetune 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.85, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.85 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma LoRA 8GB.json b/training_presets/#chroma LoRA 8GB.json index 78027aac4..437ac71d7 100644 --- a/training_presets/#chroma LoRA 8GB.json +++ b/training_presets/#chroma LoRA 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 0.0003, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "LORA", "vae": { diff --git a/training_presets/#ernie LoRA 8GB.json b/training_presets/#ernie LoRA 8GB.json index bc1e4a82b..a63d34ece 100644 --- a/training_presets/#ernie LoRA 8GB.json +++ b/training_presets/#ernie LoRA 8GB.json @@ -7,7 +7,8 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, @@ -26,7 +27,5 @@ "layer_filter_preset": "blocks" }, "timestep_distribution": "LOGIT_NORMAL", - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#flux2 Finetune 16GB.json b/training_presets/#flux2 Finetune 16GB.json index 6e3addbe8..55feaf992 100644 --- a/training_presets/#flux2 Finetune 16GB.json +++ b/training_presets/#flux2 Finetune 16GB.json @@ -7,11 +7,12 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { @@ -25,8 +26,6 @@ }, "timestep_distribution": "LOGIT_NORMAL", "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "optimizer": { "optimizer": "ADAFACTOR" }, diff --git a/training_presets/#flux2 LoRA 8GB.json b/training_presets/#flux2 LoRA 8GB.json index dfdc77839..dd3992123 100644 --- a/training_presets/#flux2 LoRA 8GB.json +++ b/training_presets/#flux2 LoRA 8GB.json @@ -7,11 +7,13 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.7 }, "training_method": "LORA", "vae": { @@ -26,7 +28,5 @@ "layer_filter_preset": "blocks" }, "timestep_distribution": "LOGIT_NORMAL", - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#hidream LoRA.json b/training_presets/#hidream LoRA.json index 2eb588763..8b025a628 100644 --- a/training_presets/#hidream LoRA.json +++ b/training_presets/#hidream LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "HiDream-ai/HiDream-I1-Full", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HI_DREAM_FULL", @@ -16,7 +14,8 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, @@ -28,11 +27,13 @@ }, "text_encoder_3": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_4": { "model_name": "meta-llama/Llama-3.1-8B-Instruct", "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 } } diff --git a/training_presets/#hunyuan video LoRA.json b/training_presets/#hunyuan video LoRA.json index 754550155..96bbb80e7 100644 --- a/training_presets/#hunyuan video LoRA.json +++ b/training_presets/#hunyuan video LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "hunyuanvideo-community/HunyuanVideo", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HUNYUAN_VIDEO", @@ -16,11 +14,13 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_2": { "train": false, diff --git a/training_presets/#qwen Finetune 16GB.json b/training_presets/#qwen Finetune 16GB.json index 811d7e0b1..2224993f4 100644 --- a/training_presets/#qwen Finetune 16GB.json +++ b/training_presets/#qwen Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.75, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.75 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen Finetune 24GB.json b/training_presets/#qwen Finetune 24GB.json index 8bee3cd3f..1cf9dc09f 100644 --- a/training_presets/#qwen Finetune 24GB.json +++ b/training_presets/#qwen Finetune 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.55, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.55 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 16GB.json b/training_presets/#qwen LoRA 16GB.json index b4e0d7e88..0eda34d5d 100644 --- a/training_presets/#qwen LoRA 16GB.json +++ b/training_presets/#qwen LoRA 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 24GB.json b/training_presets/#qwen LoRA 24GB.json index 696648a42..cd7b7216e 100644 --- a/training_presets/#qwen LoRA 24GB.json +++ b/training_presets/#qwen LoRA 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image DeTurbo LoRA 8GB.json b/training_presets/#z-image DeTurbo LoRA 8GB.json index cc38e60eb..957f08b80 100644 --- a/training_presets/#z-image DeTurbo LoRA 8GB.json +++ b/training_presets/#z-image DeTurbo LoRA 8GB.json @@ -4,13 +4,12 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, "weight_dtype": "INT_W8A8", - "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors" + "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image Finetune 16GB.json b/training_presets/#z-image Finetune 16GB.json index 0d23d3992..3911c866c 100644 --- a/training_presets/#z-image Finetune 16GB.json +++ b/training_presets/#z-image Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image LoRA 8GB.json b/training_presets/#z-image LoRA 8GB.json index 78b4b05cc..8dbce7b21 100644 --- a/training_presets/#z-image LoRA 8GB.json +++ b/training_presets/#z-image LoRA 8GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, From 85e850ea5da12643f275fb1aab77ee7ed4e6bca5 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 26 May 2026 05:00:33 +0300 Subject: [PATCH 31/67] Remove cep_enabled bool, and fix vali loss --- modules/modelSetup/BaseChromaSetup.py | 2 +- modules/modelSetup/BaseFlux2Setup.py | 2 +- modules/modelSetup/BaseFluxSetup.py | 2 +- modules/modelSetup/BaseHiDreamSetup.py | 2 +- modules/modelSetup/BaseHunyuanVideoSetup.py | 2 +- modules/modelSetup/BasePixArtAlphaSetup.py | 2 +- modules/modelSetup/BaseQwenSetup.py | 2 +- modules/modelSetup/BaseSanaSetup.py | 2 +- modules/modelSetup/BaseStableDiffusion3Setup.py | 2 +- modules/modelSetup/BaseStableDiffusionSetup.py | 2 +- modules/modelSetup/BaseStableDiffusionXLSetup.py | 2 +- modules/modelSetup/BaseWuerstchenSetup.py | 2 +- modules/modelSetup/BaseZImageSetup.py | 2 +- modules/ui/TrainingTab.py | 12 +++--------- modules/util/config/TrainConfig.py | 4 +--- 15 files changed, 17 insertions(+), 25 deletions(-) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index ea4948343..0a403cf16 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -185,7 +185,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b1e4a459d..7f0f0703c 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -104,7 +104,7 @@ def predict( text_encoder_output=batch.get('text_encoder_hidden_state'), text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 2e77c112c..492c791fb 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -231,7 +231,7 @@ def predict( apply_attention_mask=config.transformer.attention_mask, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index c264ecfce..b5d3a3c77 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -332,7 +332,7 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_3_output = self._apply_conditional_embedding_perturbation( text_encoder_3_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index c79255b30..b977b384a 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -229,7 +229,7 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index cf384cd19..d00f3e963 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -178,7 +178,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index ab97025b6..7d39c31c3 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -102,7 +102,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index aeb1f783d..ad71a1ba9 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -188,7 +188,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index b0358c54a..022632e73 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -284,7 +284,7 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index b41dc4dc7..8f89593d2 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -169,7 +169,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index bbcbe9dc1..8d6e14e7f 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -220,7 +220,7 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.dropout_probability if not deterministic else None, )) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 327d6811c..3f68aad93 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -250,7 +250,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_embedding = self._apply_conditional_embedding_perturbation( text_embedding, config.cep_gamma, generator ) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b23542a8c..373f0936c 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -103,7 +103,7 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) - if config.cep_enabled: + if config.cep_gamma > 0 and not deterministic: text_encoder_output = self._apply_conditional_embedding_perturbation( text_encoder_output, config.cep_gamma, generator ) diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 5880701e3..7464db02e 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -703,15 +703,9 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") - # Conditional Embedding Perturbation (CEP) - cep_label = components.label(frame, 10, 0, "Conditional Embedding Perturbation (CEP)", - tooltip="Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images.") - cep_label.configure(wraplength=130, justify="left") - components.switch(frame, 10, 1, self.ui_state, "cep_enabled") - - components.label(frame, 11, 0, "CEP Gamma", - tooltip="Gamma controls perturbation noise magnitude, paper's default is 1. Only has an effect if CEP is enabled") - components.entry(frame, 11, 1, self.ui_state, "cep_gamma") + components.label(frame, 10, 0, "CEP Gamma", + tooltip="Conditional Embedding Perturbation. Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images. Gamma controls perturbation noise magnitude, paper's default is 1.") + components.entry(frame, 10, 1, self.ui_state, "cep_gamma") def __create_masked_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index a3c7a5234..258f8fe2d 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -441,7 +441,6 @@ class TrainConfig(BaseConfig): timestep_distribution: TimestepDistribution min_noising_strength: float max_noising_strength: float - cep_enabled: bool cep_gamma: float noising_weight: float @@ -1024,8 +1023,7 @@ def default_values() -> 'TrainConfig': data.append(("noising_bias", 0.0, float, False)) data.append(("timestep_shift", 1.0, float, False)) data.append(("dynamic_timestep_shifting", False, bool, False)) - data.append(("cep_enabled", False, bool, False)) - data.append(("cep_gamma", 1.0, float, False)) + data.append(("cep_gamma", 0.0, float, False)) # unet From 065f9227115c55b8c3dc82b7befe112bb0099de0 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 26 May 2026 05:03:38 +0300 Subject: [PATCH 32/67] add Ernie support --- modules/modelSetup/BaseErnieSetup.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..832690738 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -94,6 +94,11 @@ def predict( text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, ) + if config.cep_gamma > 0 and not deterministic: + text_encoder_output = self._apply_conditional_embedding_perturbation( + text_encoder_output, config.cep_gamma, generator + ) + # Patchify: [B, 32, H, W] -> [B, 128, H/2, W/2] latent_image = model.patchify_latents(batch['latent_image'].float()) latent_height = latent_image.shape[-2] From 7609ab5abec036485f8c144df484ae4a11e92c58 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 26 May 2026 05:04:36 +0300 Subject: [PATCH 33/67] fix paper formula --- modules/modelSetup/mixin/ModelSetupNoiseMixin.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 18961a12a..8ab487ddb 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -128,15 +128,15 @@ def _apply_conditional_embedding_perturbation( Applies Conditional Embedding Perturbation (CEP) as per Equation (8). Paper: "Slight Corruption in Pre-training Data Makes Better Diffusion Models" - delta ~ U(-sqrt(gamma/d), sqrt(gamma/d)) or N(0, sqrt(gamma/d)) + delta ~ U(-(gamma/sqrt(d), gamma/sqrt(d)) or N(0, gamma/sqrt(d)) """ def _perturb_cep(tensor: Tensor) -> Tensor: # d denotes the dimension of c_theta(y) d = tensor.shape[-1] # gamma controls perturbation magnitude (Paper uses gamma=1.0 as default baseline) - # Calculate scaling factor: sqrt(gamma / d) - scale = math.sqrt(gamma / d) + # Calculate scaling factor: gamma / sqrt(d) + scale = gamma / math.sqrt(d) # CEP-U (Uniform) scheme noise = torch.rand( From 88a0406f96de591c9565c25230ea7fa37b016d9d Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Tue, 26 May 2026 05:08:35 +0300 Subject: [PATCH 34/67] @staticmethod and no self --- modules/modelSetup/mixin/ModelSetupNoiseMixin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 8ab487ddb..54abe5bdf 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -118,8 +118,8 @@ def _create_noise( return noise + @staticmethod def _apply_conditional_embedding_perturbation( - self, embedding: Tensor | list, gamma: float, generator: Generator From 0fe37fc9d1a08161dcef7f51d09d30c5bc26bcfd Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 27 May 2026 15:57:33 +0200 Subject: [PATCH 35/67] Add a pull request template --- .github/pull_request_template.md | 36 ++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 .github/pull_request_template.md diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 000000000..f5434cec8 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,36 @@ + + +## Summary + + + +## Test plan + + + +- [ ] `pre-commit run --all-files` passes +- [ ] Launched the affected UI or script and exercised the change +- [ ] Tested with at least one real preset / config when relevant (note which: ____) + +## AI assistance + + + +- [ ] No AI involvement +- [ ] AI-assisted — I have read every line in this diff and can defend each change +- [ ] Early AI prototype — opened for discussion, **not ready for review** + +--- + +By opening this PR (and not marking it as an early prototype, aka a draft) I confirm I have read every +line in this diff and have tested it locally. \ No newline at end of file From 7ad6b6c8a9019938459b0457df336f2ae675c039 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 27 May 2026 15:59:54 +0200 Subject: [PATCH 36/67] Add CLAUDE.md --- CLAUDE.md | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..b23961163 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,3 @@ +# CLAUDE.md + +This project's contributor rules for AI coding agents live in **[AGENTS.md](AGENTS.md)**. \ No newline at end of file From 6d3356eb4ef87827df9c441025ed666e5a985492 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 27 May 2026 16:01:09 +0200 Subject: [PATCH 37/67] Add AGENTS.md --- AGENTS.md | 202 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..e77bb5f4c --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,202 @@ +# AGENTS.md + +OneTrainer trains diffusion models with full fine-tune / LoRA / embedding / +FINE_TUNE_VAE methods. Built on `diffusers` + PyTorch + +[`mgds`](https://github.com/Nerogar/MGDS). Supported families: +`modules/util/enum/ModelType.py`. + +## Run + +``` +pre-commit install # install OUTSIDE the project venv +pre-commit run --all-files # mandatory before every PR +python scripts/train.py --config-path [--secrets-path secrets.json] +python scripts/train_ui.py # GUI (CustomTkinter) +``` + +Every `scripts/*.py` (except `generate_debug_report.py`) starts with `script_imports()` +from `scripts/util/import_util.py`: it fixes `sys.path`, filters xformers/Triton noise, +and bootstraps ZLUDA. Don't try to replicate it from a bare `python -c`. + +There is **no automated test suite.** Verification is launching scripts and exercising +the change. + +## Layout + +``` +modules/ + model/ BaseModel subclasses — state (weights, optim, EMA, embeddings) + modelLoader/ BaseModelLoader — ckpt/safetensors/diffusers → model + modelSetup/ BaseModelSetup — optimizer, LR, grad checkpointing, device + dataLoader/ BaseDataLoader — wraps mgds dataset; bucketing, aug, captions + modelSampler/ BaseModelSampler — preview/inference sampling + modelSaver/ BaseModelSaver — model → disk + trainer/ BaseTrainer, GenericTrainer (main loop), CloudTrainer (RunPod), MultiTrainer + module/ + EMA / LoRA / OFT / AdditionalEmbedding training wrappers + quantized/ Fp8 / Nf4 / W8A8 / GGUFA8 / SVD linear layers + Blip / ClipSeg / WD / Rembg / HPSv2 / AestheticScore backends for caption_ui / generate_* + ui/ CustomTkinter; one *UI/*Tab/*Window per screen. + TrainUI extends ctk.CTk; others extend ctk.CTkToplevel. + util/ + create.py ★ factory — TrainConfig → pipeline + factory.py auto-discovery (import_dir + class registry) + config/ TrainConfig + sub-configs (Sample, Concept, Cloud, Secrets) + enum/ ModelType, Optimizer, LearningRateScheduler, NoiseScheduler, + TrainingMethod, DataType, ... + callbacks/ TrainCallbacks (training-loop event hooks) + commands/ TrainCommands (UI → trainer signals) + optimizer/ CAME / adam / adamw / adafactor / muon extensions + ui/ UIState, components.py, validation.py + LayerOffloadConductor.py ⚠ fragile VRAM offloading; incompatible with dataloader_threads > 1 + zluda/ Windows AMD shim; disables cuDNN / flash_sdp / mem_efficient_sdp / cudnn_sdp + +scripts/ + util/import_util.py ★ script_imports() — sys.path + ZLUDA bootstrap; called first by every script + train.py / train_ui.py / train_remote.py / sample.py + caption_ui.py / convert_model_ui.py / video_tool_ui.py + generate_captions.py / generate_masks.py / convert_model.py / calculate_loss.py + create_train_files.py / install_zluda.py / generate_debug_report.py + +training_presets/ ~50 partial-override JSONs of TrainConfig defaults, "# [ ].json" +embedding_templates/ plain-text prompt files (`` placeholder); .gitignore excludes user files +resources/ icons + sd_model_spec/*.json +docs/ ProjectStructure / Contributing / QuickStartGuide / CliTraining / + EmbeddingTraining / CaptioningAndMasking / RamOffloading / DockerImage +LAUNCH-SCRIPTS.md OT_* env vars, venv/conda selection +lib.include.sh runtime/venv/conda bootstrap; forces PYTORCH_ENABLE_MPS_FALLBACK=1 on macOS +``` + +## Architecture + +### Factory + auto-discovery (`modules/util/create.py`) + +`create_trainer(config, callbacks, commands)` is the top-level entry: returns +`CloudTrainer` if `config.cloud.enabled`, `MultiTrainer` if `config.multi_gpu`, else +`GenericTrainer` (the only branch that calls `ZLUDA.initialize_devices`). + +`GenericTrainer.start()` then calls per-layer factories in order: +`create_model_loader → create_model_setup → create_data_loader → create_model_saver → +create_model_sampler`. + +The selection mechanism for those five "per-model" layers is **auto-discovery**: + +```python +factory.import_dir("modules/modelSampler", "modules.modelSampler") +factory.import_dir("modules/modelLoader", "modules.modelLoader") +factory.import_dir("modules/modelSaver", "modules.modelSaver") +factory.import_dir("modules/modelSetup", "modules.modelSetup") +factory.import_dir("modules/dataLoader", "modules.dataLoader") +``` + +Dropping a file with the right base class in any of those five directories registers +it — you do **not** edit `create.py` to register new per-model classes. You **do** edit +`create.py` for `create_optimizer`, `create_ema`, `create_lr_scheduler`, +`create_noise_scheduler`, `create_trainer` — those are explicit `match`/`case` branches +keyed off enums. + +### UI + +CustomTkinter. One file per screen in `modules/ui/`. `TrainUI` (the root window) +extends `ctk.CTk`; `CaptionUI`, `ConvertModelUI`, `VideoToolUI` extend `ctk.CTkToplevel`. +`TrainUI` composes tabs (`ModelTab`, `TrainingTab`, `SamplingTab`, `LoraTab`, +`ConceptTab`, `AdditionalEmbeddingsTab`, `CloudTab`); sub-windows are opened by whichever +component needs them. + +Reactive state: `modules/util/ui/UIState.py` is `UIState(master, obj)` — it introspects +typed attributes on any config-shaped object and two-way-binds them to tkinter +vars. Shared widgets: `components.py`. Validation: `validation.py` + +`validation_helpers.py`. + +### Config + +`TrainConfig` (`modules/util/config/TrainConfig.py`) is **not** a `@dataclass`; it's a +`BaseConfig` subclass using class-level annotations. Serialization is via `to_dict()` / +`from_dict()` on `BaseConfig` — JSON conversion happens at the call site +(`scripts/train.py` does `json.load` → `from_dict`). `TrainConfig.py` also declares +several sub-configs inline (`TrainOptimizerConfig`, `TrainModelPartConfig`, +`TrainEmbeddingConfig`, `QuantizationConfig`); standalone sub-configs: +`SampleConfig.py`, `ConceptConfig.py`, `CloudConfig.py`, `SecretsConfig.py`. + +Validation lives in `modules/util/ui/validation.py` — the **UI** layer. CLI scripts +bypass it. Validate at load time too if a field can be invalid from JSON. + +Presets in `training_presets/` are partial overrides applied on top of +`TrainConfig.default_values()` (only fields that differ from defaults are present). + +## Recipes + +### Add a new model family + +Pick the closest existing family as a template (`Flux2`, `Qwen`, `Sana`, `HiDream`, +`Chroma`, `HunyuanVideo`, `PixArtAlpha`, `StableDiffusion3`, etc.) — file-naming +conventions vary slightly per family, so copy the analog rather than guess. + +1. Drop subclasses (auto-registered via `factory.import_dir`): + - `modules/model/Model.py` ← `BaseModel` + - `modules/modelLoader/{FineTune,LoRA,Embedding}ModelLoader.py` ← `BaseModelLoader` + (some families instead use a single `ModelLoader.py`) + - `modules/modelSetup/BaseSetup.py` + per-method `{FineTune,LoRA,Embedding}Setup.py` + ← `BaseModelSetup` (only the per-method classes auto-register) + - `modules/dataLoader/BaseDataLoader.py` ← `BaseDataLoader` + - `modules/modelSampler/Sampler.py` ← `BaseModelSampler` (one file per family) + - `modules/modelSaver/{FineTune,LoRA,Embedding}ModelSaver.py` ← `BaseModelSaver` +2. Add `ModelType.` to `modules/util/enum/ModelType.py`; mirror any sibling `is_()` predicate. +3. Sweep `modules/ui/ModelTab.py` and `modules/ui/TrainingTab.py` for `if/elif model_type == ...` chains. Partial registration = silent feature absence. +4. Add a starter `training_presets/# LoRA.json` (clone an analog's preset). +5. Optional: `resources/sd_model_spec/.json`, `resources/icons/.png`. + +Cross-family shared logic lives in `modules/modelSetup/mixin/` and `modules/modelLoader/mixin/` — extend a mixin rather than duplicating. + +### Add a new optimizer + +1. `modules/util/enum/Optimizer.py` — add enum entry; update `is_adaptive` / `is_schedule_free` / `supports_fused_back_pass` predicates if applicable. +2. `modules/util/create.py::create_optimizer` — add `case Optimizer.:`. +3. `modules/util/optimizer_util.py::OPTIMIZER_DEFAULT_PARAMETERS` — register default hyperparams. +4. If torch internals need patching: add a module under `modules/util/optimizer/` modelled on the existing `*_extensions.py`. +5. UI exposure: `modules/ui/OptimizerParamsWindow.py`. +6. Pin any new package in `requirements-global.txt` (or appropriate platform file). + +### Add a new LR scheduler + +1. `modules/util/enum/LearningRateScheduler.py` — add enum entry. +2. `modules/util/lr_scheduler_util.py` — add `lr_lambda_(...)`. +3. `modules/util/create.py::create_lr_scheduler` — add `case` branch wiring the lambda. +4. UI exposure: `modules/ui/SchedulerParamsWindow.py` if it needs parameters. + +### Add a new noise scheduler + +1. `modules/util/enum/NoiseScheduler.py` — add enum entry. +2. `modules/util/create.py::create_noise_scheduler` — add `case` branch returning a `diffusers` scheduler. Mirror an existing case for the signature. +3. UI exposure: `modules/ui/SamplingTab.py`. + +### Add a `TrainConfig` field + +1. Declare a class-level annotation (`my_field: int`) on the appropriate `BaseConfig` subclass; set its default in that class's `default_values()`. +2. UI: bind via `UIState`. Add validation in `modules/util/ui/validation.py` if it can be invalid. +3. Presets keep loading (partial overrides). Migrate only if a default they relied on changes. +4. CLI bypasses UI validation — if a malformed JSON value would crash deeper in, validate at load. + +### Hook into the training loop + +- **Callbacks** (training → caller): `modules/util/callbacks/TrainCallbacks.py`. Instantiated in `scripts/train.py` / `TrainUI`, passed to `create_trainer(...)`. +- **Commands** (caller → training): `modules/util/commands/TrainCommands.py`. Available: `stop`, `sample_default`, `sample_custom(SampleConfig)`, `backup`, `save` (each paired with `get_and_reset_*` polled by the trainer). Multi-GPU sync: `modules/util/multi_gpu_util.py::sync_commands` + `TrainCommands.merge`. + +## Footguns + +- `modules/util/LayerOffloadConductor.py` — **incompatible with `dataloader_threads > 1`** when `gradient_checkpointing.offload()` and `layer_offload_fraction > 0` (hard-raised in `create.py`). +- ZLUDA (`modules/zluda/ZLUDA.py`) — `ZLUDA.initialize_devices(config)` runs only in the `GenericTrainer` branch of `create_trainer`. On Win-AMD it disables `cuDNN`, `flash_sdp`, `mem_efficient_sdp`, `cudnn_sdp`. Silent CPU fallback if its self-test fails. +- macOS — `lib.include.sh` exports `PYTORCH_ENABLE_MPS_FALLBACK=1`. Slow workflow on Mac? Suspect this first. +- `OT_*` env vars (`LAUNCH-SCRIPTS.md`) change install / venv selection / low-mem mode / platform requirements. +- CustomTkinter trace-removal workaround in `modules/util/ui/components.py` (references upstream CTK PR #2077, unlikely to merge). Don't simplify without re-verifying. +- `requirements-global.txt` has `git+https` commit pins for `diffusers`, `mgds`, `muon`. Don't bump silently. +- Adding a `ModelType` without sweeping every `if/elif model_type == ...` causes silent feature absence (no crash). +- CLI scripts bypass UI validation — a field validated only in `modules/util/ui/validation.py` can arrive malformed from JSON. + +## Before opening a PR + +- `pre-commit run --all-files` clean +- Launched the affected UI / script and exercised the change +- Touched the training path? Ran a short training job with a real preset +- No silent new top-level dependencies +- Can defend every line in the diff \ No newline at end of file From 4d00a21f4776531231419303b55234101796e35a Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Wed, 27 May 2026 16:13:05 +0200 Subject: [PATCH 38/67] Make doc references visible in PR template --- .github/pull_request_template.md | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index f5434cec8..cb945f3b2 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,8 +1,6 @@ - +Please read [AGENTS.md](../AGENTS.md) (project rules for any AI-assisted work) before opening this PR, +and [docs/Contributing.md](../docs/Contributing.md) / [docs/ProjectStructure.md](../docs/ProjectStructure.md) for the wider project guide. ## Summary From 3a86c173172517b1f6d1da89b8659442496f78f1 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Wed, 27 May 2026 17:11:25 +0200 Subject: [PATCH 39/67] Fix missing end-of-file newlines --- .github/pull_request_template.md | 2 +- AGENTS.md | 2 +- CLAUDE.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index cb945f3b2..6e7bc5f6b 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -31,4 +31,4 @@ your co-author. Pick exactly one option below. --- By opening this PR (and not marking it as an early prototype, aka a draft) I confirm I have read every -line in this diff and have tested it locally. \ No newline at end of file +line in this diff and have tested it locally. diff --git a/AGENTS.md b/AGENTS.md index e77bb5f4c..35f4a0660 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -199,4 +199,4 @@ Cross-family shared logic lives in `modules/modelSetup/mixin/` and `modules/mode - Launched the affected UI / script and exercised the change - Touched the training path? Ran a short training job with a real preset - No silent new top-level dependencies -- Can defend every line in the diff \ No newline at end of file +- Can defend every line in the diff diff --git a/CLAUDE.md b/CLAUDE.md index b23961163..49afc658c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,3 +1,3 @@ # CLAUDE.md -This project's contributor rules for AI coding agents live in **[AGENTS.md](AGENTS.md)**. \ No newline at end of file +This project's contributor rules for AI coding agents live in **[AGENTS.md](AGENTS.md)**. From d62455677410d897eb574731680fc65792f51131 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 27 May 2026 22:47:37 +0200 Subject: [PATCH 40/67] fix: gate autocast on train dtype instead of dtype-set uniformity create_autocast_context and disable_fp16_autocast_context disabled autocast whenever all configured dtypes were uniform. But autocast also keeps precision-sensitive ops (norms, softmax, reductions) in fp32 and casts matmul/conv operands to the compute dtype, so all-bf16 configs silently ran those ops in bf16 (Nerogar/OneTrainer#1156) and DoRA's internally-built fp32 weight crashed against bf16 activations (Nerogar/OneTrainer#1479). Gate autocast on train_dtype instead: enable autocast at the training dtype so a weight stored at a different dtype is cast on the fly rather than mismatching. Only CUDA (incl. ROCm, which also reports device type "cuda") accepts float32 as an autocast dtype and upcasts lower-precision weights, so float32/tfloat32 enables autocast on CUDA but explicitly disables it (full precision at the weight dtype, with a warning) on cpu/mps/xpu. The "mixed precision untested" warning now covers any non-CUDA device. Remove the dead allow_mixed_precision helper and the now-unused weight_dtypes arguments. disable_bf16_on_fp16_autocast_context (Wuerstchen / Stable Cascade effnet only) is left unchanged. Co-Authored-By: Claude Opus 4.7 --- modules/modelSetup/BaseChromaSetup.py | 15 +--- modules/modelSetup/BaseErnieSetup.py | 13 +-- modules/modelSetup/BaseFlux2Setup.py | 13 +-- modules/modelSetup/BaseFluxSetup.py | 16 +--- modules/modelSetup/BaseHiDreamSetup.py | 23 +----- modules/modelSetup/BaseHunyuanVideoSetup.py | 16 +--- modules/modelSetup/BasePixArtAlphaSetup.py | 15 +--- modules/modelSetup/BaseQwenSetup.py | 13 +-- modules/modelSetup/BaseSanaSetup.py | 18 +---- .../modelSetup/BaseStableDiffusion3Setup.py | 17 +--- .../modelSetup/BaseStableDiffusionSetup.py | 10 +-- .../modelSetup/BaseStableDiffusionXLSetup.py | 14 +--- modules/modelSetup/BaseWuerstchenSetup.py | 17 +--- modules/modelSetup/BaseZImageSetup.py | 13 +-- modules/util/dtype_util.py | 79 +++++++++++-------- 15 files changed, 72 insertions(+), 220 deletions(-) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..77bb5c952 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -17,7 +17,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -56,24 +55,14 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..95e4baeb7 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -12,7 +12,6 @@ from modules.util.checkpointing_util import enable_checkpointing_for_ernie_transformer from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -46,22 +45,14 @@ def setup_optimizations( model.transformer_offload_conductor = \ enable_checkpointing_for_ernie_transformer(model.transformer, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..2fc0d02dd 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -17,7 +17,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -56,22 +55,14 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..da4f0faed 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -18,7 +18,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -58,25 +57,14 @@ def setup_optimizations( model.text_encoder_2_offload_conductor = \ enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_2_autocast_context, model.text_encoder_2_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 17fbcc0d6..26c5f39af 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -18,7 +18,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -63,27 +62,14 @@ def setup_optimizations( model.text_encoder_4_offload_conductor = \ enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().text_encoder_3, - config.weight_dtypes().text_encoder_4, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder_3, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) @@ -92,11 +78,6 @@ def setup_optimizations( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().transformer, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index b072bf4ba..b5ab548ef 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -18,7 +18,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -58,25 +57,14 @@ def setup_optimizations( if model.text_encoder_2 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.transformer_autocast_context, model.transformer_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().transformer, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 3069b4884..b33be767f 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -17,7 +17,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -58,23 +57,13 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a8a7be8f6..19aa0f684 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -15,7 +15,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -53,22 +52,14 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c9ea6da0..4e18f81d4 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -17,7 +17,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -59,23 +58,13 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_gemma_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_autocast_context, model.text_encoder_train_dtype = disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) @@ -83,9 +72,6 @@ def setup_optimizations( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().vae, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index de5dc04e8..d6c8b7eed 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -18,7 +18,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -59,26 +58,14 @@ def setup_optimizations( model.text_encoder_3_offload_conductor = \ enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().text_encoder_3, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \ disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder_3, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 8cf63ac07..40ef9acf3 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.conv_util import apply_circular_padding_to_conv2d from modules.util.dtype_util import create_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -63,13 +62,8 @@ def setup_optimizations( if model.unet_lora is not None: apply_circular_padding_to_conv2d(model.unet_lora) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().unet, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) quantize_layers(model.text_encoder, self.train_device, model.train_dtype, config) quantize_layers(model.vae, self.train_device, model.train_dtype, config) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 61cb1e457..51f46ee77 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.conv_util import apply_circular_padding_to_conv2d from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -60,22 +59,13 @@ def setup_optimizations( if model.unet_lora is not None: apply_circular_padding_to_conv2d(model.unet_lora) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().unet, - config.weight_dtypes().text_encoder, - config.weight_dtypes().text_encoder_2, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) model.vae_autocast_context, model.vae_train_dtype = disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().vae, - ], config.enable_autocast_cache, ) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 48cd05942..a9c97616e 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -22,7 +22,6 @@ disable_bf16_on_fp16_autocast_context, disable_fp16_autocast_context, ) -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -72,26 +71,14 @@ def setup_optimizations( if model.prior_prior_lora is not None: apply_circular_padding_to_conv2d(model.prior_prior_lora) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().decoder_text_encoder, - config.weight_dtypes().decoder, - config.weight_dtypes().decoder_vqgan, - config.weight_dtypes().effnet_encoder, - config.weight_dtypes().text_encoder, - config.weight_dtypes().prior, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - config.weight_dtypes().embedding if config.train_any_embedding() else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) if model.model_type.is_stable_cascade(): model.prior_autocast_context, model.prior_train_dtype = disable_fp16_autocast_context( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().prior, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache, ) else: diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b822d7304..f40945dac 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -16,7 +16,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -54,12 +53,8 @@ def setup_optimizations( model.text_encoder_offload_conductor = \ enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) - model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ - config.weight_dtypes().transformer, - config.weight_dtypes().text_encoder, - config.weight_dtypes().vae, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache) + model.autocast_context, model.train_dtype = create_autocast_context( + self.train_device, config.train_dtype, config.enable_autocast_cache) #TODO necessary if we don't train it? model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ @@ -67,10 +62,6 @@ def setup_optimizations( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().text_encoder, - config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, - ], config.enable_autocast_cache, ) diff --git a/modules/util/dtype_util.py b/modules/util/dtype_util.py index d0df10c59..1ad80a4ef 100644 --- a/modules/util/dtype_util.py +++ b/modules/util/dtype_util.py @@ -1,20 +1,11 @@ from contextlib import nullcontext -from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.DataType import DataType import torch from torch.nn import Parameter -def allow_mixed_precision(train_config: TrainConfig): - all_dtypes = list(train_config.weight_dtypes().all_dtypes() + [train_config.train_dtype]) - all_dtypes = list(filter(lambda dtype: dtype != DataType.NONE, all_dtypes)) - all_dtypes = set(all_dtypes) - - return len(all_dtypes) != 1 - - def enable_grad_scaling(train_dtype: DataType, parameters: list[Parameter]): trainable_parameter_dtype = list({parameter.dtype for parameter in parameters}) return train_dtype == DataType.FLOAT_16 and all(dtype == torch.float32 for dtype in trainable_parameter_dtype) @@ -28,46 +19,61 @@ def create_grad_scaler(): def create_autocast_context( device: torch.device, train_dtype: DataType | None, - weight_dtypes: list[DataType | None], enable_autocast_cache: bool, ) -> tuple[torch.autocast | nullcontext, DataType]: - if torch.backends.mps.is_available(): - if any(train_dtype != dt for dt in weight_dtypes if dt is not None): - print("Warning: Mixed precision training is untested on macOS. Consider setting all dtypes to be the same.") - else: - return nullcontext(), train_dtype - - weight_dtypes = list(weight_dtypes) - weight_dtypes = list(filter(lambda dtype: dtype != DataType.NONE and dtype is not None, weight_dtypes)) - weight_dtypes = list(set(weight_dtypes)) - - if len(weight_dtypes) == 1 and train_dtype == weight_dtypes[0]: - return torch.autocast(device_type=device.type, enabled=False), train_dtype - else: - return torch.autocast(device_type=device.type, dtype=train_dtype.torch_dtype(), + torch_train_dtype = train_dtype.torch_dtype() + + if torch_train_dtype in (torch.float16, torch.bfloat16): + # fp16/bf16 autocast is supported on every backend. autocast casts the operands + # of matmul/conv-type ops to train_dtype (precision-sensitive ops like norms stay + # in fp32), so a weight stored at a different dtype is cast on the fly rather than + # mismatching in the matmul. + if device.type != "cuda": + # CUDA (incl. ROCm) is the tested backend. fp16/bf16 autocast works on + # other backends too (mps, xpu, cpu, ...) but is untested here; bf16 on + # MPS additionally needs macOS >= 14. + print(f"Warning: Mixed precision training is untested on device type '{device.type}'.") + return torch.autocast(device_type=device.type, dtype=torch_train_dtype, cache_enabled=enable_autocast_cache), train_dtype + elif device.type == "cuda": + # float32/tfloat32 on CUDA (and ROCm, which also reports device type "cuda"): + # CUDA accepts float32 as an autocast dtype and upcasts lower-precision weights + # on the fly (this is undocumented but works). + return torch.autocast(device_type=device.type, dtype=torch_train_dtype, + cache_enabled=enable_autocast_cache), train_dtype + else: + # float32/tfloat32 on a non-CUDA backend (cpu, mps, xpu, ...): those backends + # reject fp32 autocast, so disable autocast and let the model run at its weight + # dtype. Disable explicitly (not nullcontext) so any enclosing autocast is + # suppressed too. + print("Warning: float32 training does not upcast lower-precision weights on this device " + "(only CUDA can autocast to float32); the model runs at its weight dtype. " + "Set the weight data types to float32 for full precision.") + return torch.autocast(device_type=device.type, enabled=False), train_dtype def disable_fp16_autocast_context( device: torch.device, train_dtype: DataType | None, fallback_train_dtype: DataType | None, - weight_dtypes: list[DataType | None], enable_autocast_cache: bool, ) -> tuple[torch.autocast | nullcontext, DataType]: - weight_dtypes = list(filter(lambda dtype: dtype != DataType.NONE and dtype is not None, weight_dtypes)) - weight_dtypes = list(set(weight_dtypes)) - if train_dtype != DataType.FLOAT_16: - # train dtype is not fp16 -> nothing to disable + # the main autocast context isn't fp16 -> nothing to override, defer to it return nullcontext(), train_dtype - if len(weight_dtypes) == 1 and fallback_train_dtype == weight_dtypes[0]: - # fallback_train_dtype is the same as all weights -> disable autocast - return torch.autocast(device_type=device.type, enabled=False), weight_dtypes[0] - - return torch.autocast(device_type=device.type, dtype=fallback_train_dtype.torch_dtype(), - cache_enabled=enable_autocast_cache), fallback_train_dtype + # fp16 training but this component is unstable in fp16 -> override the outer fp16 + # autocast and run it at the fallback precision. A bf16 fallback works on every + # backend; a float32 fallback can only be applied via autocast on CUDA, so on + # other devices disable autocast and let the component run at its weight dtype. + fallback_torch_dtype = fallback_train_dtype.torch_dtype() + if fallback_torch_dtype in (torch.float16, torch.bfloat16) or device.type == "cuda": + return torch.autocast(device_type=device.type, dtype=fallback_torch_dtype, + cache_enabled=enable_autocast_cache), fallback_train_dtype + else: + print("Warning: the float32 fallback for fp16-unstable layers is not applied on device type " + f"'{device.type}' (only CUDA can autocast to float32); these layers run at their weight dtype.") + return torch.autocast(device_type=device.type, enabled=False), fallback_train_dtype def disable_bf16_on_fp16_autocast_context( @@ -76,6 +82,9 @@ def disable_bf16_on_fp16_autocast_context( weight_dtypes: list[DataType | None], enable_autocast_cache: bool, ) -> tuple[torch.autocast | nullcontext, DataType]: + # Only used for the Wuerstchen / Stable Cascade effnet encoder. The rationale for + # this special case is unknown, so its original behavior is deliberately kept + # unchanged rather than migrated to the create_autocast_context approach above. weight_dtypes = list(filter(lambda dtype: dtype != DataType.NONE and dtype is not None, weight_dtypes)) weight_dtypes = list(set(weight_dtypes)) From 3b8214c8454389b5d739a2eaf8e7d60d94c40424 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Thu, 28 May 2026 12:27:41 +0200 Subject: [PATCH 41/67] Address PR review feedback on AGENTS.md and PR template --- .github/pull_request_template.md | 11 ++++---- AGENTS.md | 46 +++----------------------------- docs/recipes/AddOptimizer.md | 8 ++++++ 3 files changed, 16 insertions(+), 49 deletions(-) create mode 100644 docs/recipes/AddOptimizer.md diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 6e7bc5f6b..2fec144e3 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,6 +1,8 @@ + ## Summary @@ -28,7 +30,4 @@ your co-author. Pick exactly one option below. - [ ] AI-assisted — I have read every line in this diff and can defend each change - [ ] Early AI prototype — opened for discussion, **not ready for review** ---- - -By opening this PR (and not marking it as an early prototype, aka a draft) I confirm I have read every -line in this diff and have tested it locally. + diff --git a/AGENTS.md b/AGENTS.md index 35f4a0660..380ca6c56 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -126,49 +126,9 @@ Presets in `training_presets/` are partial overrides applied on top of ## Recipes -### Add a new model family - -Pick the closest existing family as a template (`Flux2`, `Qwen`, `Sana`, `HiDream`, -`Chroma`, `HunyuanVideo`, `PixArtAlpha`, `StableDiffusion3`, etc.) — file-naming -conventions vary slightly per family, so copy the analog rather than guess. - -1. Drop subclasses (auto-registered via `factory.import_dir`): - - `modules/model/Model.py` ← `BaseModel` - - `modules/modelLoader/{FineTune,LoRA,Embedding}ModelLoader.py` ← `BaseModelLoader` - (some families instead use a single `ModelLoader.py`) - - `modules/modelSetup/BaseSetup.py` + per-method `{FineTune,LoRA,Embedding}Setup.py` - ← `BaseModelSetup` (only the per-method classes auto-register) - - `modules/dataLoader/BaseDataLoader.py` ← `BaseDataLoader` - - `modules/modelSampler/Sampler.py` ← `BaseModelSampler` (one file per family) - - `modules/modelSaver/{FineTune,LoRA,Embedding}ModelSaver.py` ← `BaseModelSaver` -2. Add `ModelType.` to `modules/util/enum/ModelType.py`; mirror any sibling `is_()` predicate. -3. Sweep `modules/ui/ModelTab.py` and `modules/ui/TrainingTab.py` for `if/elif model_type == ...` chains. Partial registration = silent feature absence. -4. Add a starter `training_presets/# LoRA.json` (clone an analog's preset). -5. Optional: `resources/sd_model_spec/.json`, `resources/icons/.png`. - -Cross-family shared logic lives in `modules/modelSetup/mixin/` and `modules/modelLoader/mixin/` — extend a mixin rather than duplicating. - -### Add a new optimizer - -1. `modules/util/enum/Optimizer.py` — add enum entry; update `is_adaptive` / `is_schedule_free` / `supports_fused_back_pass` predicates if applicable. -2. `modules/util/create.py::create_optimizer` — add `case Optimizer.:`. -3. `modules/util/optimizer_util.py::OPTIMIZER_DEFAULT_PARAMETERS` — register default hyperparams. -4. If torch internals need patching: add a module under `modules/util/optimizer/` modelled on the existing `*_extensions.py`. -5. UI exposure: `modules/ui/OptimizerParamsWindow.py`. -6. Pin any new package in `requirements-global.txt` (or appropriate platform file). - -### Add a new LR scheduler - -1. `modules/util/enum/LearningRateScheduler.py` — add enum entry. -2. `modules/util/lr_scheduler_util.py` — add `lr_lambda_(...)`. -3. `modules/util/create.py::create_lr_scheduler` — add `case` branch wiring the lambda. -4. UI exposure: `modules/ui/SchedulerParamsWindow.py` if it needs parameters. - -### Add a new noise scheduler - -1. `modules/util/enum/NoiseScheduler.py` — add enum entry. -2. `modules/util/create.py::create_noise_scheduler` — add `case` branch returning a `diffusers` scheduler. Mirror an existing case for the signature. -3. UI exposure: `modules/ui/SamplingTab.py`. +Detailed implementation guides live in [`docs/recipes/`](docs/recipes/): + +- **New optimizer** — see [`docs/recipes/AddOptimizer.md`](docs/recipes/AddOptimizer.md). ### Add a `TrainConfig` field diff --git a/docs/recipes/AddOptimizer.md b/docs/recipes/AddOptimizer.md new file mode 100644 index 000000000..1033f19f4 --- /dev/null +++ b/docs/recipes/AddOptimizer.md @@ -0,0 +1,8 @@ +# Adding a New Optimizer + +1. `modules/util/enum/Optimizer.py` — add enum entry; update `is_adaptive` / `is_schedule_free` / `supports_fused_back_pass` predicates if applicable. +2. `modules/util/create.py::create_optimizer` — add `case Optimizer.:`. +3. `modules/util/optimizer_util.py::OPTIMIZER_DEFAULT_PARAMETERS` — register default hyperparams. +4. If torch internals need patching: add a module under `modules/util/optimizer/` modelled on the existing `*_extensions.py`. +5. UI exposure: `modules/ui/OptimizerParamsWindow.py`. +6. Pin any new package in `requirements-global.txt` (or appropriate platform file). From bcc7d5dae77ee19a4396d537e6dc3921a16a6a0a Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 28 May 2026 23:50:33 +0200 Subject: [PATCH 42/67] raise instead of silently degrading float32 fallback on non-CUDA The fp16-unstable-layer fallback can only be applied via autocast to float32 on CUDA. On other backends, silently running the layer at its weight dtype could drop it to fp16 - the exact instability the fallback exists to prevent - so refuse explicitly instead. Co-Authored-By: Claude Opus 4.8 --- modules/util/dtype_util.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/modules/util/dtype_util.py b/modules/util/dtype_util.py index 1ad80a4ef..b6d6b2643 100644 --- a/modules/util/dtype_util.py +++ b/modules/util/dtype_util.py @@ -64,16 +64,16 @@ def disable_fp16_autocast_context( # fp16 training but this component is unstable in fp16 -> override the outer fp16 # autocast and run it at the fallback precision. A bf16 fallback works on every - # backend; a float32 fallback can only be applied via autocast on CUDA, so on - # other devices disable autocast and let the component run at its weight dtype. + # backend; a float32 fallback can only be applied via autocast on CUDA. fallback_torch_dtype = fallback_train_dtype.torch_dtype() if fallback_torch_dtype in (torch.float16, torch.bfloat16) or device.type == "cuda": return torch.autocast(device_type=device.type, dtype=fallback_torch_dtype, cache_enabled=enable_autocast_cache), fallback_train_dtype else: - print("Warning: the float32 fallback for fp16-unstable layers is not applied on device type " - f"'{device.type}' (only CUDA can autocast to float32); these layers run at their weight dtype.") - return torch.autocast(device_type=device.type, enabled=False), fallback_train_dtype + raise RuntimeError( + f"A float32 fallback for fp16-unstable layers cannot be applied on device type " + f"'{device.type}' (only CUDA can autocast to float32). Use a bfloat16 fallback dtype." + ) def disable_bf16_on_fp16_autocast_context( From 96ccec01825601c49772bd8bde53b18b44255f2c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 29 May 2026 06:52:20 +0300 Subject: [PATCH 43/67] Resolve delta docstring --- modules/modelSetup/mixin/ModelSetupNoiseMixin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 54abe5bdf..13fb7b00d 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -128,7 +128,7 @@ def _apply_conditional_embedding_perturbation( Applies Conditional Embedding Perturbation (CEP) as per Equation (8). Paper: "Slight Corruption in Pre-training Data Makes Better Diffusion Models" - delta ~ U(-(gamma/sqrt(d), gamma/sqrt(d)) or N(0, gamma/sqrt(d)) + delta ~ U(-(gamma/sqrt(d), gamma/sqrt(d)) """ def _perturb_cep(tensor: Tensor) -> Tensor: # d denotes the dimension of c_theta(y) From ca91c00b3d6db562e46e8820e7b93ff4d40cc9b4 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 29 May 2026 06:54:27 +0300 Subject: [PATCH 44/67] row += 1 for cep gamma --- modules/ui/TrainingTab.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 7464db02e..3df890c84 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -697,15 +697,20 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b tooltip="Shift the timestep distribution. Use the preview to see more details.") components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True) + row = 9 + if supports_dynamic_timestep_shifting: # dynamic timestep shifting - components.label(frame, 9, 0, "Dynamic Timestep Shifting", + components.label(frame, row, 0, "Dynamic Timestep Shifting", tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) - components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting") + components.switch(frame, row, 1, self.ui_state, "dynamic_timestep_shifting") + row += 1 - components.label(frame, 10, 0, "CEP Gamma", + components.label(frame, row, 0, "CEP Gamma", tooltip="Conditional Embedding Perturbation. Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images. Gamma controls perturbation noise magnitude, paper's default is 1.") - components.entry(frame, 10, 1, self.ui_state, "cep_gamma") + components.entry(frame, row, 1, self.ui_state, "cep_gamma") + row += 1 + def __create_masked_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) From 2045afd21a6a056a710f4e83ae4322f24b6feced Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Fri, 29 May 2026 06:55:07 +0300 Subject: [PATCH 45/67] required=True for cep gamma --- modules/ui/TrainingTab.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 3df890c84..1bc85eaa7 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -708,7 +708,7 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b components.label(frame, row, 0, "CEP Gamma", tooltip="Conditional Embedding Perturbation. Inject a slight noise into the TEs outputs to enhance the quality, diversity, and fidelity of the generated images. Gamma controls perturbation noise magnitude, paper's default is 1.") - components.entry(frame, row, 1, self.ui_state, "cep_gamma") + components.entry(frame, row, 1, self.ui_state, "cep_gamma", required=True) row += 1 From c60245812c52e8be686dbf779e62ac5db870f901 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 19:41:48 +0200 Subject: [PATCH 46/67] Sync Ctk* copies with upstream changes merged into Base* files Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/CtkLoraTabView.py | 70 ++++++++++++++++++++++++++++++++ modules/ui/CtkTrainUIView.py | 22 ++++++++-- modules/ui/CtkTrainingTabView.py | 26 +++++++----- modules/ui/TrainUIController.py | 22 ++++++++-- 4 files changed, 124 insertions(+), 16 deletions(-) diff --git a/modules/ui/CtkLoraTabView.py b/modules/ui/CtkLoraTabView.py index faabf9607..6e9975476 100644 --- a/modules/ui/CtkLoraTabView.py +++ b/modules/ui/CtkLoraTabView.py @@ -39,6 +39,7 @@ def refresh_ui(self): ("LoRA", PeftType.LORA), ("LoHa", PeftType.LOHA), ("OFT v2", PeftType.OFT_2), + ("LoKr", PeftType.LOKR), ], self.ui_state, "peft_type", command=self.setup_lora) def setup_lora(self, peft_type: PeftType): @@ -46,6 +47,8 @@ def setup_lora(self, peft_type: PeftType): name = "LoHa" elif peft_type == PeftType.OFT_2: name = "OFT v2" + elif peft_type == PeftType.LOKR: + name = "LoKr" else: name = "LoRA" @@ -126,6 +129,11 @@ def setup_lora(self, peft_type: PeftType): tooltip="Share the OFT parameters between blocks. A single rotation matrix is shared across all blocks within a layer, drastically cutting the number of trainable parameters and yielding very compact adapter files, potentially improving generalization but at the cost of significant expressiveness, which can lead to underfitting on more complex or diverse tasks.") components.switch(master, 1, 4, self.ui_state, "oft_block_share") + # Scaled OFT (SOFT) + components.label(master, 2, 3, "Scaled OFT (SOFT)", + tooltip="Applies a scaling factor to the learned weights. This ensures that the effective learning rate remains consistent across different block sizes. Without this, different block sizes require significantly different learning rates.") + components.switch(master, 2, 4, self.ui_state, "oft_scaled") + # Dropout Percentage components.label(master, 2, 0, "Dropout Probability", tooltip="Dropout probability. This percentage of the rotated adapter nodes that will be randomly restored to the base model initial statue. Helps with overfitting. 0 disables, 1 maximum.") @@ -143,3 +151,65 @@ def setup_lora(self, peft_type: PeftType): components.label(master, 4, 0, "Bundle Embeddings", tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") components.switch(master, 4, 1, self.ui_state, "bundle_additional_embeddings") + + # LoKr + elif peft_type == PeftType.LOKR: + # LoKr Main Settings + components.label(master, 1, 0, f"{name} dimension", + tooltip="The dimension parameter used for the secondary decomposition. Analogous to rank in LoRA.") + components.entry(master, 1, 1, self.ui_state, "lokr_dim") + + components.label(master, 2, 0, "Decomposition Factor", + tooltip="Factor for Kronecker product decomposition. -1 for auto, which is recommended. Changing this drastically affects parameter count.") + components.entry(master, 2, 1, self.ui_state, "lokr_decompose_factor") + + # alpha + components.label(master, 3, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + components.entry(master, 3, 1, self.ui_state, "lora_alpha") + + # Dropout Percentage + components.label(master, 4, 0, "Dropout Probability", + tooltip="Dropout probability. This percentage of model nodes will be randomly ignored at each training step. Helps with overfitting. 0 disables, 1 maximum.") + components.entry(master, 4, 1, self.ui_state, "dropout_probability") + + # LoKr weight dtype + components.label(master, 5, 0, f"{name} Weight Data Type", + tooltip=f"The {name} weight data type used for training. This can reduce memory consumption, but reduces precision") + components.options_kv(master, 5, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], self.ui_state, "lora_weight_dtype") + + # LoKr Vectorization trick + components.label(master, 6, 0, "Kronecker-Vec Trick", + tooltip="Uses an accelerated path that bypasses the materialization of the full Kronecker product. This delivers a massive speedup to the LoKr without sacrificing precision. Highly recommended.") + components.switch(master, 6, 1, self.ui_state, "lokr_vec_trick") + + #LoKr Decomposition Settings + components.label(master, 1, 3, "Decompose Both Matrices", + tooltip="Perform rank decomposition on both Kronecker product matrices (W1 and W2). Only effective for very small dimensions.") + components.switch(master, 1, 4, self.ui_state, "lokr_decompose_both") + + components.label(master, 2, 3, "Use Tucker Decomposition (Conv)", + tooltip="Use Tucker decomposition for convolutional layers. Can be more efficient for some architectures.") + components.switch(master, 2, 4, self.ui_state, "lokr_use_tucker") + + components.label(master, 3, 3, "Force Full Matrix (W2)", + tooltip="Forces the second Kronecker matrix (W2) to be a full matrix, ignoring the dimension setting. For expert use.") + components.switch(master, 3, 4, self.ui_state, "lokr_full_matrix") + + # LoKr DoRA Settings + components.label(master, 4, 3, "Decompose Weights (DoRA)", + tooltip="Apply weight decomposition (DoRA) on top of the LoKr update.") + components.switch(master, 4, 4, self.ui_state, "lokr_weight_decompose") + + components.label(master, 5, 3, "Apply DoRA on Output Axis", + tooltip="Apply the DoRA weight decomposition on the output axis instead of the input axis.") + components.switch(master, 5, 4, self.ui_state, "lokr_dora_on_output") + + + # Additional embeddings + components.label(master, 6, 3, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + components.switch(master, 6, 4, self.ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index ba90d2e64..91c327694 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -132,6 +132,9 @@ def __init__(self): self.training_callbacks = None self.training_commands = None + self.start_time = None + self.start_total_steps = None + self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir self._check_start_always_on_tensorboard() @@ -600,14 +603,21 @@ def open_tensorboard(self): webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + assert self.start_time is not None and self.start_total_steps is not None + spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step + + # calculate steps done in THIS SESSION only + current_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + steps_done_this_session = current_total_steps - self.start_total_steps + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - if train_progress.global_step <= 30: + if steps_done_this_session <= 30: return "Estimating ..." + total_eta = spent_total / steps_done_this_session * remaining_steps + td = datetime.timedelta(seconds=total_eta) days = td.days hours, remainder = divmod(td.seconds, 3600) @@ -632,6 +642,10 @@ def delete_eta_label(self): self.eta_label.configure(text="") def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + # capture session start on first progress update + if self.start_total_steps is None: + self.start_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + self.set_step_progress(train_progress.epoch_step, max_step) self.set_epoch_progress(train_progress.epoch, max_epoch) self.set_eta_label(train_progress, max_step, max_epoch) @@ -715,6 +729,8 @@ def __training_thread_function(self): if self.train_config.cloud.enabled: self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + # Reset session tracking - actual values captured on first progress callback + self.start_total_steps = None self.start_time = time.monotonic() trainer.train() except Exception: diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bcca11ae9..c839abf73 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -438,34 +438,40 @@ def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supp frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", + components.label(frame, row, 0, "Train Text Encoder", tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + components.switch(frame, row, 1, self.ui_state, "text_encoder.train") + row += 1 # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", + components.label(frame, row, 0, "Caption Dropout Probability", tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + components.entry(frame, row, 1, self.ui_state, "text_encoder.dropout_probability") + row += 1 if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + components.time_entry(frame, row, 1, self.ui_state, "text_encoder.stop_training_after", "text_encoder.stop_training_after_unit", supports_time_units=False) + row += 1 # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", + components.label(frame, row, 0, "Text Encoder Learning Rate", tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + components.entry(frame, row, 1, self.ui_state, "text_encoder.learning_rate") + row += 1 if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", + components.label(frame, row, 0, "Clip Skip", tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + components.entry(frame, row, 1, self.ui_state, "text_encoder_layer_skip") + row += 1 if supports_sequence_length: # text encoder sequence length diff --git a/modules/ui/TrainUIController.py b/modules/ui/TrainUIController.py index ba90d2e64..91c327694 100644 --- a/modules/ui/TrainUIController.py +++ b/modules/ui/TrainUIController.py @@ -132,6 +132,9 @@ def __init__(self): self.training_callbacks = None self.training_commands = None + self.start_time = None + self.start_total_steps = None + self.always_on_tensorboard_subprocess = None self.current_workspace_dir = self.train_config.workspace_dir self._check_start_always_on_tensorboard() @@ -600,14 +603,21 @@ def open_tensorboard(self): webbrowser.open("http://localhost:" + str(self.train_config.tensorboard_port), new=0, autoraise=False) def _calculate_eta_string(self, train_progress: TrainProgress, max_step: int, max_epoch: int) -> str | None: + assert self.start_time is not None and self.start_total_steps is not None + spent_total = time.monotonic() - self.start_time - steps_done = train_progress.epoch * max_step + train_progress.epoch_step + + # calculate steps done in THIS SESSION only + current_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + steps_done_this_session = current_total_steps - self.start_total_steps + remaining_steps = (max_epoch - train_progress.epoch - 1) * max_step + (max_step - train_progress.epoch_step) - total_eta = spent_total / steps_done * remaining_steps - if train_progress.global_step <= 30: + if steps_done_this_session <= 30: return "Estimating ..." + total_eta = spent_total / steps_done_this_session * remaining_steps + td = datetime.timedelta(seconds=total_eta) days = td.days hours, remainder = divmod(td.seconds, 3600) @@ -632,6 +642,10 @@ def delete_eta_label(self): self.eta_label.configure(text="") def on_update_train_progress(self, train_progress: TrainProgress, max_step: int, max_epoch: int): + # capture session start on first progress update + if self.start_total_steps is None: + self.start_total_steps = train_progress.epoch * max_step + train_progress.epoch_step + self.set_step_progress(train_progress.epoch_step, max_step) self.set_epoch_progress(train_progress.epoch, max_epoch) self.set_eta_label(train_progress, max_step, max_epoch) @@ -715,6 +729,8 @@ def __training_thread_function(self): if self.train_config.cloud.enabled: self.ui_state.get_var("secrets.cloud").update(self.train_config.secrets.cloud) + # Reset session tracking - actual values captured on first progress callback + self.start_total_steps = None self.start_time = time.monotonic() trainer.train() except Exception: From 02f1032c2fdc52b190a9085a9e2d276c31bde214 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:09:06 +0200 Subject: [PATCH 47/67] Make Qt UI the default; rename train_ui scripts to _ctk/_qt Co-Authored-By: Claude Sonnet 4.6 --- scripts/{train_ui.py => train_ui_ctk.py} | 0 scripts/{train_ui_pyside6.py => train_ui_qt.py} | 0 start-ui.bat | 6 +++--- start-ui.sh | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) rename scripts/{train_ui.py => train_ui_ctk.py} (100%) rename scripts/{train_ui_pyside6.py => train_ui_qt.py} (100%) diff --git a/scripts/train_ui.py b/scripts/train_ui_ctk.py similarity index 100% rename from scripts/train_ui.py rename to scripts/train_ui_ctk.py diff --git a/scripts/train_ui_pyside6.py b/scripts/train_ui_qt.py similarity index 100% rename from scripts/train_ui_pyside6.py rename to scripts/train_ui_qt.py diff --git a/start-ui.bat b/start-ui.bat index 5edc8e6c4..f5b9d5aec 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -4,8 +4,8 @@ REM Avoid footgun by explictly navigating to the directory containing the batch cd /d "%~dp0" REM Verify that OneTrainer is our current working directory -if not exist "scripts\train_ui.py" ( - echo Error: train_ui.py does not exist, you have done something very wrong. Reclone the repository. +if not exist "scripts\train_ui_qt.py" ( + echo Error: train_ui_qt.py does not exist, you have done something very wrong. Reclone the repository. goto :end ) @@ -60,7 +60,7 @@ if errorlevel 1 ( :launch echo Starting UI... -%PYTHON% scripts\train_ui.py +%PYTHON% scripts\train_ui_qt.py if errorlevel 1 ( echo Error: UI script exited with code %ERRORLEVEL% ) diff --git a/start-ui.sh b/start-ui.sh index 2f0eecc0d..dd56ce97a 100755 --- a/start-ui.sh +++ b/start-ui.sh @@ -11,4 +11,4 @@ fi prepare_runtime_environment -run_python_in_active_env "scripts/train_ui.py" "$@" +run_python_in_active_env "scripts/train_ui_qt.py" "$@" From a3a8d2d2afce225433edd8683b0088d57af00f37 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:17:04 +0200 Subject: [PATCH 48/67] Fix over-wide widgets: extend _pack_form to absorb extra horizontal space _pack_form now adds a stretch column after the last content column, the same way it already adds a stretch row. This stops buttons and entries from expanding to fill the full frame width when the column they sit in carries the only stretch factor. Tools and embedding frames had setColumnStretch on their content column; drop that and let _pack_form place the absorber column after the build call instead. General/data/backup frames keep their intentional entry- column stretches (form fields should expand) and just gain the row+column gutter from _pack_form. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6TrainUIView.py | 9 +++++---- modules/util/ui/pyside6_components.py | 4 +++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/modules/ui/PySide6TrainUIView.py b/modules/ui/PySide6TrainUIView.py index 2c1fba90c..ddb1b803f 100644 --- a/modules/ui/PySide6TrainUIView.py +++ b/modules/ui/PySide6TrainUIView.py @@ -223,28 +223,29 @@ def _configure_general_frame(self, frame): lo.setColumnStretch(1, 1) lo.setColumnStretch(3, 1) self.build_general_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) def _configure_data_frame(self, frame): lo = pyside6_components._layout(frame) lo.setColumnStretch(1, 1) lo.setColumnStretch(3, 1) self.build_data_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) def _configure_backup_frame(self, frame): lo = pyside6_components._layout(frame) lo.setColumnStretch(1, 1) lo.setColumnStretch(3, 1) self.build_backup_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) def _configure_tools_frame(self, frame): - lo = pyside6_components._layout(frame) - lo.setColumnStretch(1, 1) self.build_tools_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) def _configure_embedding_frame(self, frame): - lo = pyside6_components._layout(frame) - lo.setColumnStretch(1, 1) self.build_embedding_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) def _create_tabs(self): general_page = self._create_scrollable_tab(self._configure_general_frame) diff --git a/modules/util/ui/pyside6_components.py b/modules/util/ui/pyside6_components.py index 8e0f936dd..c95add49a 100644 --- a/modules/util/ui/pyside6_components.py +++ b/modules/util/ui/pyside6_components.py @@ -100,9 +100,11 @@ def scrollable_frame(parent: QWidget) -> tuple[QScrollArea, QWidget]: def _pack_form(master: QWidget) -> None: - """Add a stretch row after the last content row so rows don't expand to fill available space.""" + # Add a stretch row and column after the last content cell so extra space + # goes to the empty gutter rather than stretching content widgets. lo = _layout(master) lo.setRowStretch(lo.rowCount(), 1) + lo.setColumnStretch(lo.columnCount(), 1) # --------------------------------------------------------------------------- From 51636c50d8e1bd76d030b498d4ac7f4e462a0856 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:37:19 +0200 Subject: [PATCH 49/67] Hint ctk fallback in unported Qt6 views Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6CaptionUIView.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui/PySide6CaptionUIView.py b/modules/ui/PySide6CaptionUIView.py index 69e2b5286..620e4faab 100644 --- a/modules/ui/PySide6CaptionUIView.py +++ b/modules/ui/PySide6CaptionUIView.py @@ -6,7 +6,7 @@ def __init__(self, parent, controller): super().__init__(parent) self.setWindowTitle("Dataset Tool") lo = QVBoxLayout(self) - lo.addWidget(QLabel("The dataset tool has not been ported to Qt6 yet.")) + lo.addWidget(QLabel("The dataset tool has not been ported to Qt6 yet.\nYou can still use it by launching the CustomTkinter UI: scripts/train_ui_ctk.py")) ok = QPushButton("OK") ok.clicked.connect(self.accept) lo.addWidget(ok) From 01cca79010fc7fe862172f463304b4c688eb43a3 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:42:53 +0200 Subject: [PATCH 50/67] Move Rolling Backup Count under the Rolling Backup toggle Was sitting on the same row at columns 3-4; now on its own row directly below the toggle for clearer visual grouping. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseTrainUIView.py | 22 +++++++++++----------- modules/ui/CtkTrainUIView.py | 22 +++++++++++----------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 91c327694..b9fa0c04a 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -449,32 +449,32 @@ def create_backup_tab(self, master): components.switch(frame, 1, 1, self.ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + components.entry(frame, 2, 1, self.ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + components.switch(frame, 3, 1, self.ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + components.time_entry(frame, 4, 1, self.ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + components.button(frame, 4, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + components.entry(frame, 5, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + components.entry(frame, 6, 1, self.ui_state, "save_filename_prefix") frame.pack(fill="both", expand=1) return frame diff --git a/modules/ui/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py index 91c327694..b9fa0c04a 100644 --- a/modules/ui/CtkTrainUIView.py +++ b/modules/ui/CtkTrainUIView.py @@ -449,32 +449,32 @@ def create_backup_tab(self, master): components.switch(frame, 1, 1, self.ui_state, "rolling_backup") # rolling backup count - components.label(frame, 1, 3, "Rolling Backup Count", + components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 1, 4, self.ui_state, "rolling_backup_count") + components.entry(frame, 2, 1, self.ui_state, "rolling_backup_count") # backup before save - components.label(frame, 2, 0, "Backup Before Save", + components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - components.switch(frame, 2, 1, self.ui_state, "backup_before_save") + components.switch(frame, 3, 1, self.ui_state, "backup_before_save") # save after - components.label(frame, 3, 0, "Save Every", + components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 3, 1, self.ui_state, "save_every", "save_every_unit") + components.time_entry(frame, 4, 1, self.ui_state, "save_every", "save_every_unit") # save now - components.button(frame, 3, 3, "save now", self.save_now) + components.button(frame, 4, 3, "save now", self.save_now) # skip save - components.label(frame, 4, 0, "Skip First", + components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 4, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") + components.entry(frame, 5, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - components.label(frame, 5, 0, "Save Filename Prefix", + components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - components.entry(frame, 5, 1, self.ui_state, "save_filename_prefix") + components.entry(frame, 6, 1, self.ui_state, "save_filename_prefix") frame.pack(fill="both", expand=1) return frame From 647afc27b158c320e0ed9a7d7af04abbb07a36c7 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 29 May 2026 20:45:06 +0200 Subject: [PATCH 51/67] Move Rolling Backup Count under the Rolling Backup toggle Was sitting on the same row at columns 3-4; now on its own row directly below the toggle for clearer visual grouping. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/BaseTrainUIView.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py index 9a443a342..ff3acd2b7 100644 --- a/modules/ui/BaseTrainUIView.py +++ b/modules/ui/BaseTrainUIView.py @@ -277,32 +277,32 @@ def build_backup_tab_content(self, frame, controller, ui_state): self.components.switch(frame, 1, 1, ui_state, "rolling_backup") # rolling backup count - self.components.label(frame, 1, 3, "Rolling Backup Count", + self.components.label(frame, 2, 0, "Rolling Backup Count", tooltip="Defines the number of backups to keep if rolling backups are enabled") - self.components.entry(frame, 1, 4, ui_state, "rolling_backup_count") + self.components.entry(frame, 2, 1, ui_state, "rolling_backup_count") # backup before save - self.components.label(frame, 2, 0, "Backup Before Save", + self.components.label(frame, 3, 0, "Backup Before Save", tooltip="Create a full backup before saving the final model") - self.components.switch(frame, 2, 1, ui_state, "backup_before_save") + self.components.switch(frame, 3, 1, ui_state, "backup_before_save") # save after - self.components.label(frame, 3, 0, "Save Every", + self.components.label(frame, 4, 0, "Save Every", tooltip="The interval used when automatically saving the model during training") - self.components.time_entry(frame, 3, 1, ui_state, "save_every", "save_every_unit") + self.components.time_entry(frame, 4, 1, ui_state, "save_every", "save_every_unit") # save now - self.components.button(frame, 3, 3, "save now", self.save_now) + self.components.button(frame, 4, 3, "save now", self.save_now) # skip save - self.components.label(frame, 4, 0, "Skip First", + self.components.label(frame, 5, 0, "Skip First", tooltip="Start saving automatically after this interval has elapsed") - self.components.entry(frame, 4, 1, ui_state, "save_skip_first", width=50, sticky="nw") + self.components.entry(frame, 5, 1, ui_state, "save_skip_first", width=50, sticky="nw") # save filename prefix - self.components.label(frame, 5, 0, "Save Filename Prefix", + self.components.label(frame, 6, 0, "Save Filename Prefix", tooltip="The prefix for filenames used when saving the model during training") - self.components.entry(frame, 5, 1, ui_state, "save_filename_prefix") + self.components.entry(frame, 6, 1, ui_state, "save_filename_prefix") def build_embedding_tab_content(self, frame, controller, ui_state): # embedding model name From 4c9c01261391e9bc4cdd2c66003a21dda6b8e70b Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 29 May 2026 22:17:00 +0200 Subject: [PATCH 52/67] Remove AI assistance checkboxes to fix task counter pollution --- .github/pull_request_template.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 2fec144e3..95d818a84 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -23,11 +23,11 @@ Describe what you actually did, not what you "would do". -- [ ] No AI involvement -- [ ] AI-assisted — I have read every line in this diff and can defend each change -- [ ] Early AI prototype — opened for discussion, **not ready for review** +- No AI involvement +- AI-assisted — I have read every line in this diff and can defend each change +- Early AI prototype — opened for discussion, **not ready for review** From ea59c93a9180c5137f0a405fcf3eda232dc92047 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 30 May 2026 10:51:27 +0200 Subject: [PATCH 53/67] anima: Anima model support (LoRA + Fine-Tune) Co-Authored-By: dxqb <183307934+dxqb@users.noreply.github.com> --- modules/dataLoader/AnimaBaseDataLoader.py | 161 +++++++++++ modules/model/AnimaModel.py | 254 ++++++++++++++++++ modules/modelLoader/AnimaModelLoader.py | 199 ++++++++++++++ modules/modelSampler/AnimaSampler.py | 173 ++++++++++++ modules/modelSaver/AnimaFineTuneModelSaver.py | 11 + modules/modelSaver/AnimaLoRAModelSaver.py | 11 + modules/modelSaver/anima/AnimaLoRASaver.py | 37 +++ modules/modelSaver/anima/AnimaModelSaver.py | 90 +++++++ modules/modelSetup/AnimaFineTuneSetup.py | 93 +++++++ modules/modelSetup/AnimaLoRASetup.py | 105 ++++++++ modules/modelSetup/BaseAnimaSetup.py | 185 +++++++++++++ modules/ui/ConvertModelUI.py | 1 + modules/ui/ModelTab.py | 21 ++ modules/ui/TopBar.py | 2 + modules/ui/TrainingTab.py | 14 + modules/util/config/SampleConfig.py | 8 + modules/util/enum/ModelType.py | 6 + modules/util/optimizer/muon_util.py | 2 +- requirements-global.txt | 4 +- resources/sd_model_spec/anima-lora.json | 6 + resources/sd_model_spec/anima.json | 6 + training_presets/#anima Finetune.json | 46 ++++ training_presets/#anima LoRA.json | 31 +++ 23 files changed, 1463 insertions(+), 3 deletions(-) create mode 100644 modules/dataLoader/AnimaBaseDataLoader.py create mode 100644 modules/model/AnimaModel.py create mode 100644 modules/modelLoader/AnimaModelLoader.py create mode 100644 modules/modelSampler/AnimaSampler.py create mode 100644 modules/modelSaver/AnimaFineTuneModelSaver.py create mode 100644 modules/modelSaver/AnimaLoRAModelSaver.py create mode 100644 modules/modelSaver/anima/AnimaLoRASaver.py create mode 100644 modules/modelSaver/anima/AnimaModelSaver.py create mode 100644 modules/modelSetup/AnimaFineTuneSetup.py create mode 100644 modules/modelSetup/AnimaLoRASetup.py create mode 100644 modules/modelSetup/BaseAnimaSetup.py create mode 100644 resources/sd_model_spec/anima-lora.json create mode 100644 resources/sd_model_spec/anima.json create mode 100644 training_presets/#anima Finetune.json create mode 100644 training_presets/#anima LoRA.json diff --git a/modules/dataLoader/AnimaBaseDataLoader.py b/modules/dataLoader/AnimaBaseDataLoader.py new file mode 100644 index 000000000..9c1447a67 --- /dev/null +++ b/modules/dataLoader/AnimaBaseDataLoader.py @@ -0,0 +1,161 @@ +import os + +from modules.dataLoader.BaseDataLoader import BaseDataLoader +from modules.dataLoader.mixin.DataLoaderText2ImageMixin import DataLoaderText2ImageMixin +from modules.model.AnimaModel import PROMPT_MAX_LENGTH, AnimaModel +from modules.model.BaseModel import BaseModel +from modules.modelSetup.BaseAnimaSetup import BaseAnimaSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.TrainProgress import TrainProgress + +from mgds.pipelineModules.DecodeTokens import DecodeTokens +from mgds.pipelineModules.DecodeVAE import DecodeVAE +from mgds.pipelineModules.EncodeAnimaText import EncodeAnimaText +from mgds.pipelineModules.EncodeVAE import EncodeVAE +from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels +from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution +from mgds.pipelineModules.SaveImage import SaveImage +from mgds.pipelineModules.SaveText import SaveText +from mgds.pipelineModules.ScaleImage import ScaleImage +from mgds.pipelineModules.Tokenize import Tokenize + + +class AnimaBaseDataLoader( + BaseDataLoader, + DataLoaderText2ImageMixin, +): + def _preparation_modules(self, config: TrainConfig, model: AnimaModel): + rescale_image = RescaleImageChannels(image_in_name='image', image_out_name='image', in_range_min=0, in_range_max=1, out_range_min=-1, out_range_max=1) + encode_image = EncodeVAE(in_name='image', out_name='latent_image_distribution', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') + downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) + # Anima has no chat template — tokenize raw prompt with both tokenizers + tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH) + tokenize_t5 = Tokenize(in_name='prompt', tokens_out_name='t5_tokens', mask_out_name='t5_tokens_mask', tokenizer=model.t5_tokenizer, max_token_length=PROMPT_MAX_LENGTH) + # EncodeAnimaText runs Qwen3 encoder + AnimaTextConditioner; output is fixed (512, 1024) + encode_prompt = EncodeAnimaText( + tokens_name='tokens', tokens_attention_mask_name='tokens_mask', + t5_tokens_name='t5_tokens', t5_tokens_attention_mask_name='t5_tokens_mask', + hidden_state_out_name='text_encoder_hidden_state', + text_encoder=model.text_encoder, text_conditioner=model.text_conditioner, + autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), + ) + + modules = [rescale_image, encode_image, image_sample] + if config.masked_training or config.model_type.has_mask_input(): + modules.append(downscale_mask) + + modules += [tokenize_prompt, tokenize_t5] + + if not config.train_text_encoder_or_embedding(): + modules.append(encode_prompt) + + return modules + + def _cache_modules(self, config: TrainConfig, model: AnimaModel, model_setup: BaseAnimaSetup): + image_split_names = ['latent_image', 'original_resolution', 'crop_offset'] + + if config.masked_training or config.model_type.has_mask_input(): + image_split_names.append('latent_mask') + + image_aggregate_names = ['crop_resolution', 'image_path'] + + text_split_names = [] + + sort_names = image_aggregate_names + image_split_names + [ + 'prompt', 'tokens', 'tokens_mask', 't5_tokens', 't5_tokens_mask', 'text_encoder_hidden_state', + 'concept' + ] + + if not config.train_text_encoder_or_embedding(): + text_split_names += ['tokens', 'tokens_mask', 't5_tokens', 't5_tokens_mask', 'text_encoder_hidden_state'] + + return self._cache_modules_from_names( + model, model_setup, + image_split_names=image_split_names, + image_aggregate_names=image_aggregate_names, + text_split_names=text_split_names, + sort_names=sort_names, + config=config, + text_caching=not config.train_text_encoder_or_embedding(), + ) + + def _output_modules(self, config: TrainConfig, model: AnimaModel, model_setup: BaseAnimaSetup): + output_names = [ + 'image_path', 'latent_image', + 'prompt', + 'tokens', + 'tokens_mask', + 't5_tokens', + 't5_tokens_mask', + 'original_resolution', 'crop_resolution', 'crop_offset', + ] + + if config.masked_training or config.model_type.has_mask_input(): + output_names.append('latent_mask') + + if not config.train_text_encoder_or_embedding(): + output_names.append('text_encoder_hidden_state') + + return self._output_modules_from_out_names( + model, model_setup, + output_names=output_names, + config=config, + use_conditioning_image=False, + vae=model.vae, + autocast_context=[model.autocast_context], + train_dtype=model.train_dtype, + ) + + def _debug_modules(self, config: TrainConfig, model: AnimaModel): #TODO clean up + debug_dir = os.path.join(config.debug_dir, "dataloader") + + def before_save_fun(): + model.vae_to(self.train_device) + + decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8) + decode_prompt = DecodeTokens(in_name='tokens', out_name='decoded_prompt', tokenizer=model.tokenizer) + + #FIXME https://github.com/Nerogar/OneTrainer/issues/1015 + #save_image = SaveImage(image_in_name='decoded_image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1, before_save_fun=before_save_fun) + + # SaveImage(image_in_name='latent_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun) + save_mask = SaveImage(image_in_name='decoded_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun) + save_prompt = SaveText(text_in_name='decoded_prompt', original_path_in_name='image_path', path=debug_dir, before_save_fun=before_save_fun) + + # These modules don't really work, since they are inserted after a sorting operation that does not include this data + # SaveImage(image_in_name='mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1), + # SaveImage(image_in_name='image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1), + + modules = [decode_image] + + #FIXME https://github.com/Nerogar/OneTrainer/issues/1015 + #modules.append(save_image) + + if config.masked_training or config.model_type.has_mask_input(): + modules += [upscale_mask, save_mask] + + modules += [decode_prompt, save_prompt] + + return modules + + def _create_dataset( + self, + config: TrainConfig, + model: BaseModel, + model_setup: BaseModelSetup, + train_progress: TrainProgress, + is_validation: bool = False, + ): + return DataLoaderText2ImageMixin._create_dataset(self, + config, model, model_setup, train_progress, is_validation, + aspect_bucketing_quantization=64, + allow_video_files=False, #don't allow video files, but... + vae_frame_dim=True, #...Anima has a video-capable VAE. convert images to video dimensions + ) + +factory.register(BaseDataLoader, AnimaBaseDataLoader, ModelType.ANIMA) diff --git a/modules/model/AnimaModel.py b/modules/model/AnimaModel.py new file mode 100644 index 000000000..dee1afcb1 --- /dev/null +++ b/modules/model/AnimaModel.py @@ -0,0 +1,254 @@ +import math +from contextlib import nullcontext +from random import Random + +from modules.model.BaseModel import BaseModel +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.enum.DataType import DataType +from modules.util.enum.ModelType import ModelType +from modules.util.LayerOffloadConductor import LayerOffloadConductor + +import torch +from torch import Tensor + +from diffusers import ( + AnimaAutoBlocks, + AnimaTextConditioner, + AutoencoderKLQwenImage, + CosmosTransformer3DModel, + FlowMatchEulerDiscreteScheduler, +) +from transformers import Qwen2Tokenizer, Qwen3Model, T5TokenizerFast + +PROMPT_MAX_LENGTH = 512 + + +# Maps the diffusers CosmosTransformer3DModel state dict back to the original Anima checkpoint keys. +# This is the exact inverse of the forward conversion in diffusers' scripts/convert_anima_to_diffusers.py +# (which delegates the transformer to convert_cosmos_to_diffusers.convert_transformer with +# TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0). +# The original keys carry a "net." prefix; the conversion is a flat 1:1 rename, no tensor fusion. +def diffusers_to_original(): + return [ + ("patch_embed.proj", "net.x_embedder.proj.1"), + ("time_embed.t_embedder", "net.t_embedder.1"), + ("time_embed.norm", "net.t_embedding_norm"), + ("norm_out.linear_1", "net.final_layer.adaln_modulation.1"), + ("norm_out.linear_2", "net.final_layer.adaln_modulation.2"), + ("proj_out", "net.final_layer.linear"), + ("transformer_blocks.{i}", "net.blocks.{i}", [ + ("norm1.linear_1", "adaln_modulation_self_attn.1"), + ("norm1.linear_2", "adaln_modulation_self_attn.2"), + ("attn1.norm_q", "self_attn.q_norm"), + ("attn1.norm_k", "self_attn.k_norm"), + ("attn1.to_q", "self_attn.q_proj"), + ("attn1.to_k", "self_attn.k_proj"), + ("attn1.to_v", "self_attn.v_proj"), + ("attn1.to_out.0", "self_attn.output_proj"), + ("norm2.linear_1", "adaln_modulation_cross_attn.1"), + ("norm2.linear_2", "adaln_modulation_cross_attn.2"), + ("attn2.norm_q", "cross_attn.q_norm"), + ("attn2.norm_k", "cross_attn.k_norm"), + ("attn2.to_q", "cross_attn.q_proj"), + ("attn2.to_k", "cross_attn.k_proj"), + ("attn2.to_v", "cross_attn.v_proj"), + ("attn2.to_out.0", "cross_attn.output_proj"), + ("norm3.linear_1", "adaln_modulation_mlp.1"), + ("norm3.linear_2", "adaln_modulation_mlp.2"), + ("ff.net.0.proj", "mlp.layer1"), + ("ff.net.2", "mlp.layer2"), + ]), + ] + +diffusers_checkpoint_to_original = diffusers_to_original() + + +class AnimaModel(BaseModel): + # base model data + tokenizer: Qwen2Tokenizer | None + t5_tokenizer: T5TokenizerFast | None + noise_scheduler: FlowMatchEulerDiscreteScheduler | None + text_encoder: Qwen3Model | None + text_conditioner: AnimaTextConditioner | None + vae: AutoencoderKLQwenImage | None + transformer: CosmosTransformer3DModel | None + + # autocast context + text_encoder_autocast_context: torch.autocast | nullcontext + + text_encoder_train_dtype: DataType + + text_encoder_offload_conductor: LayerOffloadConductor | None + transformer_offload_conductor: LayerOffloadConductor | None + + # persistent lora training data + text_encoder_lora: LoRAModuleWrapper | None + transformer_lora: LoRAModuleWrapper | None + lora_state_dict: dict | None + + def __init__( + self, + model_type: ModelType, + ): + super().__init__( + model_type=model_type, + ) + + self.tokenizer = None + self.t5_tokenizer = None + self.noise_scheduler = None + self.text_encoder = None + self.text_conditioner = None + self.vae = None + self.transformer = None + + self.text_encoder_autocast_context = nullcontext() + + self.text_encoder_train_dtype = DataType.FLOAT_32 + + self.text_encoder_offload_conductor = None + self.transformer_offload_conductor = None + + self.text_encoder_lora = None + self.transformer_lora = None + self.lora_state_dict = None + + def adapters(self) -> list[LoRAModuleWrapper]: + return [a for a in [ + self.text_encoder_lora, + self.transformer_lora, + ] if a is not None] + + def vae_to(self, device: torch.device): + self.vae.to(device=device) + + def text_encoder_to(self, device: torch.device): #TODO share more code between models + if self.text_encoder is not None: + if self.text_encoder_offload_conductor is not None and \ + self.text_encoder_offload_conductor.layer_offload_activated(): + self.text_encoder_offload_conductor.to(device) + else: + self.text_encoder.to(device=device) + self.text_conditioner.to(device=device) + + if self.text_encoder_lora is not None: + self.text_encoder_lora.to(device) + + def transformer_to(self, device: torch.device): + if self.transformer_offload_conductor is not None and \ + self.transformer_offload_conductor.layer_offload_activated(): + self.transformer_offload_conductor.to(device) + else: + self.transformer.to(device=device) + + if self.transformer_lora is not None: + self.transformer_lora.to(device) + + def to(self, device: torch.device): + self.vae_to(device) + self.text_encoder_to(device) + self.transformer_to(device) + + def eval(self): + self.vae.eval() + if self.text_encoder is not None: + self.text_encoder.eval() + self.text_conditioner.eval() + self.transformer.eval() + + def create_pipeline(self): + pipe = AnimaAutoBlocks().init_pipeline() + pipe.update_components( + text_encoder=self.text_encoder, + tokenizer=self.tokenizer, + t5_tokenizer=self.t5_tokenizer, + text_conditioner=self.text_conditioner, + transformer=self.transformer, + vae=self.vae, + scheduler=self.noise_scheduler, + ) + return pipe + + def encode_text( + self, + train_device: torch.device, + batch_size: int = 1, + rand: Random | None = None, + text: str | list[str] = None, + tokens: Tensor = None, + tokens_mask: Tensor = None, + text_encoder_layer_skip: int = 0, + text_encoder_dropout_probability: float | None = None, + text_encoder_output: Tensor = None, + ) -> Tensor: + # Two-stage encoding: Qwen3 text encoder → AnimaTextConditioner (with T5 token ids as queries). + # text_encoder_output, when provided from cache, is already the conditioner output. + if tokens is None and text is not None: + if isinstance(text, str): + text = [text] + + tokenizer_output = self.tokenizer( + text, + max_length=PROMPT_MAX_LENGTH, + padding='max_length', + truncation=True, + return_tensors="pt", + ) + tokens = tokenizer_output.input_ids.to(self.text_encoder.device) + tokens_mask = tokenizer_output.attention_mask.to(self.text_encoder.device) + + t5_output = self.t5_tokenizer( + text, + max_length=PROMPT_MAX_LENGTH, + padding='max_length', + truncation=True, + return_tensors="pt", + ) + t5_ids = t5_output.input_ids.to(self.text_encoder.device) + t5_mask = t5_output.attention_mask.to(self.text_encoder.device) + + if text_encoder_output is None and self.text_encoder is not None: + with self.text_encoder_autocast_context: + qwen_hidden = self.text_encoder( + tokens, + attention_mask=tokens_mask.float(), + output_hidden_states=False, + ).last_hidden_state + # zero out padding positions (mirrors diffusers AnimaTextEncoderStep) + qwen_hidden = qwen_hidden * tokens_mask.to(qwen_hidden).unsqueeze(-1) + text_encoder_output = self.text_conditioner( + source_hidden_states=qwen_hidden.to(dtype=self.text_conditioner.dtype), + target_input_ids=t5_ids, + target_attention_mask=t5_mask, + source_attention_mask=tokens_mask, + ) + + if text_encoder_dropout_probability is not None and text_encoder_dropout_probability > 0.0: + raise NotImplementedError # https://github.com/Nerogar/OneTrainer/issues/957 + + # conditioner output is always (B, 512, 1024) and fully dense (zeros for padding positions); + # the Cosmos transformer takes encoder_hidden_states with no separate text attention mask. + return text_encoder_output + + def scale_latents(self, latents: Tensor) -> Tensor: + latents_mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=latents.dtype).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=latents.dtype).view(1, self.vae.config.z_dim, 1, 1, 1) + return (latents - latents_mean) * latents_std + + def unscale_latents(self, latents: Tensor) -> Tensor: + latents_mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=latents.dtype).view(1, self.vae.config.z_dim, 1, 1, 1) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=latents.dtype).view(1, self.vae.config.z_dim, 1, 1, 1) + return latents / latents_std + latents_mean + + def calculate_timestep_shift(self, latent_width: int, latent_height: int): + base_seq_len = self.noise_scheduler.config.base_image_seq_len + max_seq_len = self.noise_scheduler.config.max_image_seq_len + base_shift = self.noise_scheduler.config.base_shift + max_shift = self.noise_scheduler.config.max_shift + patch_size = 2 + + image_seq_len = (latent_width // patch_size) * (latent_height // patch_size) + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return math.exp(mu) diff --git a/modules/modelLoader/AnimaModelLoader.py b/modules/modelLoader/AnimaModelLoader.py new file mode 100644 index 000000000..066fe7578 --- /dev/null +++ b/modules/modelLoader/AnimaModelLoader.py @@ -0,0 +1,199 @@ +import os +import traceback + +from modules.model.AnimaModel import AnimaModel +from modules.model.BaseModel import BaseModel +from modules.modelLoader.GenericFineTuneModelLoader import make_fine_tune_model_loader +from modules.modelLoader.GenericLoRAModelLoader import make_lora_model_loader +from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin +from modules.modelLoader.mixin.LoRALoaderMixin import LoRALoaderMixin +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelType import ModelType +from modules.util.ModelNames import ModelNames +from modules.util.ModelWeightDtypes import ModelWeightDtypes + +import torch + +from diffusers import ( + AnimaTextConditioner, + AutoencoderKLQwenImage, + CosmosTransformer3DModel, + FlowMatchEulerDiscreteScheduler, + GGUFQuantizationConfig, +) +from transformers import Qwen2Tokenizer, Qwen3Model, T5TokenizerFast + + +class AnimaModelLoader( + HFModelLoaderMixin, +): + def __init__(self): + super().__init__() + + def __load_internal( + self, + model: AnimaModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + ): + if os.path.isfile(os.path.join(base_model_name, "meta.json")): + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, quantization, + ) + else: + raise Exception("not an internal model") + + def __load_diffusers( + self, + model: AnimaModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + ): + tokenizer = Qwen2Tokenizer.from_pretrained( + base_model_name, + subfolder="tokenizer", + ) + + t5_tokenizer = T5TokenizerFast.from_pretrained( + base_model_name, + subfolder="t5_tokenizer", + ) + + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + base_model_name, + subfolder="scheduler", + ) + + text_encoder = self._load_transformers_sub_module( + Qwen3Model, + weight_dtypes.text_encoder, + weight_dtypes.fallback_train_dtype, + base_model_name, + "text_encoder", + ) + + # conditioner is always bfloat16 — small adapter, no user dtype control + text_conditioner = AnimaTextConditioner.from_pretrained( + base_model_name, + subfolder="text_conditioner", + torch_dtype=torch.bfloat16, + ) + + if vae_model_name: #TODO simplify + vae = self._load_diffusers_sub_module( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + vae_model_name, + ) + else: + vae = self._load_diffusers_sub_module( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + "vae", + ) + + if transformer_model_name: + transformer = CosmosTransformer3DModel.from_single_file( + transformer_model_name, + config=base_model_name, + subfolder="transformer", + #avoid loading the transformer in float32: + torch_dtype=torch.bfloat16 if weight_dtypes.transformer.torch_dtype() is None else weight_dtypes.transformer.torch_dtype(), + quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer.is_gguf() else None, + ) + transformer = self._convert_diffusers_sub_module_to_dtype( + transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, + ) + else: + transformer = self._load_diffusers_sub_module( + CosmosTransformer3DModel, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, + ) + + model.model_type = model_type + model.tokenizer = tokenizer + model.t5_tokenizer = t5_tokenizer + model.noise_scheduler = noise_scheduler + model.text_encoder = text_encoder + model.text_conditioner = text_conditioner + model.vae = vae + model.transformer = transformer + + def load( #TODO share code between models + self, + model: AnimaModel, + model_type: ModelType, + model_names: ModelNames, + weight_dtypes: ModelWeightDtypes, + quantization: QuantizationConfig, + ): + stacktraces = [] + + try: + self.__load_internal( + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + try: + self.__load_diffusers( + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + for stacktrace in stacktraces: + print(stacktrace) + raise Exception("could not load model: " + model_names.base_model) + + +class AnimaLoRALoader( + LoRALoaderMixin, +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: BaseModel) -> list[LoraConversionKeySet] | None: + return None + + def load( + self, + model: AnimaModel, + model_names: ModelNames, + ): + return self._load(model, model_names) + + +AnimaLoRAModelLoader = make_lora_model_loader( + model_spec_map={ModelType.ANIMA: "resources/sd_model_spec/anima-lora.json"}, + model_class=AnimaModel, + model_loader_class=AnimaModelLoader, + embedding_loader_class=None, + lora_loader_class=AnimaLoRALoader, +) + +AnimaFineTuneModelLoader = make_fine_tune_model_loader( + model_spec_map={ModelType.ANIMA: "resources/sd_model_spec/anima.json"}, + model_class=AnimaModel, + model_loader_class=AnimaModelLoader, + embedding_loader_class=None, +) diff --git a/modules/modelSampler/AnimaSampler.py b/modules/modelSampler/AnimaSampler.py new file mode 100644 index 000000000..b2589f97c --- /dev/null +++ b/modules/modelSampler/AnimaSampler.py @@ -0,0 +1,173 @@ +import copy +import inspect +from collections.abc import Callable + +from modules.model.AnimaModel import AnimaModel +from modules.modelSampler.BaseModelSampler import BaseModelSampler, ModelSamplerOutput +from modules.util import factory +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.AudioFormat import AudioFormat +from modules.util.enum.FileType import FileType +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.enum.VideoFormat import VideoFormat +from modules.util.torch_util import torch_gc + +import torch + +from diffusers import VaeImageProcessor + +import numpy as np +from tqdm import tqdm + + +class AnimaSampler(BaseModelSampler): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + model: AnimaModel, + model_type: ModelType, + ): + super().__init__(train_device, temp_device) + + self.model = model + self.model_type = model_type + self.image_processor = VaeImageProcessor(vae_scale_factor=8) + + @torch.no_grad() + def __sample_base( + self, + prompt: str, + negative_prompt: str, + height: int, + width: int, + seed: int, + random_seed: bool, + diffusion_steps: int, + cfg_scale: float, + noise_scheduler: NoiseScheduler, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ) -> ModelSamplerOutput: + with self.model.autocast_context: + generator = torch.Generator(device=self.train_device) + if random_seed: + generator.seed() + else: + generator.manual_seed(seed) + + noise_scheduler = copy.deepcopy(self.model.noise_scheduler) + + transformer = self.model.transformer + vae = self.model.vae + vae_scale_factor = 8 + num_latent_channels = 16 + + # prepare prompt + self.model.text_encoder_to(self.train_device) + + batch_size = 2 if cfg_scale > 1.0 else 1 + combined_prompt_embedding = self.model.encode_text( + text=[prompt, negative_prompt] if cfg_scale > 1.0 else prompt, + batch_size=batch_size, + train_device=self.train_device, + ) + + self.model.text_encoder_to(self.temp_device) + torch_gc() + + # prepare latent image + latent_image = torch.randn( + size=(1, num_latent_channels, 1, height // vae_scale_factor, width // vae_scale_factor), + generator=generator, + device=self.train_device, + dtype=torch.float32, + ) + + sigmas = np.linspace(1.0, 1.0 / diffusion_steps, diffusion_steps) + noise_scheduler.set_timesteps(sigmas=sigmas, device=self.train_device) + timesteps = noise_scheduler.timesteps + + padding_mask = latent_image.new_zeros( + 1, 1, height, width, dtype=transformer.dtype, + ) + + # denoising loop + extra_step_kwargs = {} + if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): + extra_step_kwargs["generator"] = generator + + self.model.transformer_to(self.train_device) + for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): + latent_model_input = torch.cat([latent_image] * batch_size) + expanded_timestep = timestep.expand(batch_size) / noise_scheduler.config.num_train_timesteps + noise_pred = transformer( + hidden_states=latent_model_input.to(dtype=transformer.dtype), + timestep=expanded_timestep, + encoder_hidden_states=combined_prompt_embedding.to(dtype=transformer.dtype), + padding_mask=padding_mask, + return_dict=False, + )[0] + + if cfg_scale > 1.0: + noise_pred_positive, noise_pred_negative = noise_pred.chunk(2) + noise_pred = noise_pred_negative + cfg_scale * (noise_pred_positive - noise_pred_negative) + + latent_image = noise_scheduler.step( + noise_pred, timestep, latent_image, return_dict=False, **extra_step_kwargs + )[0] + + on_update_progress(i + 1, len(timesteps)) + + self.model.transformer_to(self.temp_device) + torch_gc() + + # decode + self.model.vae_to(self.train_device) + + latents = self.model.unscale_latents(latent_image) + image = vae.decode(latents, return_dict=False)[0][:, :, 0] + + do_denormalize = [True] * image.shape[0] + image = self.image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) + + self.model.vae_to(self.temp_device) + torch_gc() + + return ModelSamplerOutput( + file_type=FileType.IMAGE, + data=image[0], + ) + + def sample( + self, + sample_config: SampleConfig, + destination: str, + image_format: ImageFormat | None = None, + video_format: VideoFormat | None = None, + audio_format: AudioFormat | None = None, + on_sample: Callable[[ModelSamplerOutput], None] = lambda _: None, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ): + sampler_output = self.__sample_base( + prompt=sample_config.prompt, + negative_prompt=sample_config.negative_prompt, + height=self.quantize_resolution(sample_config.height, 64), + width=self.quantize_resolution(sample_config.width, 64), + seed=sample_config.seed, + random_seed=sample_config.random_seed, + diffusion_steps=sample_config.diffusion_steps, + cfg_scale=sample_config.cfg_scale, + noise_scheduler=sample_config.noise_scheduler, + on_update_progress=on_update_progress, + ) + + self.save_sampler_output( + sampler_output, destination, + image_format, video_format, audio_format, + ) + + on_sample(sampler_output) + +factory.register(BaseModelSampler, AnimaSampler, ModelType.ANIMA) diff --git a/modules/modelSaver/AnimaFineTuneModelSaver.py b/modules/modelSaver/AnimaFineTuneModelSaver.py new file mode 100644 index 000000000..2cc51a711 --- /dev/null +++ b/modules/modelSaver/AnimaFineTuneModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.AnimaModel import AnimaModel +from modules.modelSaver.anima.AnimaModelSaver import AnimaModelSaver +from modules.modelSaver.GenericFineTuneModelSaver import make_fine_tune_model_saver +from modules.util.enum.ModelType import ModelType + +AnimaFineTuneModelSaver = make_fine_tune_model_saver( + ModelType.ANIMA, + model_class=AnimaModel, + model_saver_class=AnimaModelSaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/AnimaLoRAModelSaver.py b/modules/modelSaver/AnimaLoRAModelSaver.py new file mode 100644 index 000000000..b948ce3a5 --- /dev/null +++ b/modules/modelSaver/AnimaLoRAModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.AnimaModel import AnimaModel +from modules.modelSaver.anima.AnimaLoRASaver import AnimaLoRASaver +from modules.modelSaver.GenericLoRAModelSaver import make_lora_model_saver +from modules.util.enum.ModelType import ModelType + +AnimaLoRAModelSaver = make_lora_model_saver( + ModelType.ANIMA, + model_class=AnimaModel, + lora_saver_class=AnimaLoRASaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/anima/AnimaLoRASaver.py b/modules/modelSaver/anima/AnimaLoRASaver.py new file mode 100644 index 000000000..dc85035dd --- /dev/null +++ b/modules/modelSaver/anima/AnimaLoRASaver.py @@ -0,0 +1,37 @@ +from modules.model.AnimaModel import AnimaModel +from modules.modelSaver.mixin.LoRASaverMixin import LoRASaverMixin +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelFormat import ModelFormat + +import torch +from torch import Tensor + + +class AnimaLoRASaver( + LoRASaverMixin, +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: AnimaModel) -> list[LoraConversionKeySet] | None: + return None + + def _get_state_dict( + self, + model: AnimaModel, + ) -> dict[str, Tensor]: + state_dict = {} + if model.transformer_lora is not None: + state_dict |= model.transformer_lora.state_dict() + if model.lora_state_dict is not None: + state_dict |= model.lora_state_dict + return state_dict + + def save( + self, + model: AnimaModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + self._save(model, output_model_format, output_model_destination, dtype) diff --git a/modules/modelSaver/anima/AnimaModelSaver.py b/modules/modelSaver/anima/AnimaModelSaver.py new file mode 100644 index 000000000..6105c46a6 --- /dev/null +++ b/modules/modelSaver/anima/AnimaModelSaver.py @@ -0,0 +1,90 @@ +import copy +import os.path +from pathlib import Path + +from modules.model.AnimaModel import AnimaModel, diffusers_checkpoint_to_original +from modules.modelSaver.mixin.DtypeModelSaverMixin import DtypeModelSaverMixin +from modules.util.convert_util import convert +from modules.util.enum.ModelFormat import ModelFormat + +import torch + +from safetensors.torch import save_file + + +class AnimaModelSaver( + DtypeModelSaverMixin, +): + def __init__(self): + super().__init__() + + def __save_diffusers( + self, + model: AnimaModel, + destination: str, + dtype: torch.dtype | None, + ): + # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. + pipeline = model.create_pipeline() + pipeline.to("cpu") + if dtype is not None: + #TODO is this code necessary for all models? in that case, share code + # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. + # the tokenizer tries to reload from the file system otherwise + tokenizer = pipeline.tokenizer + tokenizer.__deepcopy__ = lambda memo: tokenizer + + save_pipeline = copy.deepcopy(pipeline) + save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) + + delattr(tokenizer, '__deepcopy__') + else: + save_pipeline = pipeline + + os.makedirs(Path(destination).absolute(), exist_ok=True) + save_pipeline.save_pretrained(destination) + + if dtype is not None: + del save_pipeline + + def __save_safetensors( + self, + model: AnimaModel, + destination: str, + dtype: torch.dtype | None, + ): + # convert the diffusers transformer keys back to the original Anima format (net.*) + state_dict = convert(model.transformer.state_dict(), diffusers_checkpoint_to_original) + # the original checkpoint bundles the text conditioner under net.llm_adapter.*; its keys are + # identical to the diffusers module, so only a prefix is needed. + for key, value in model.text_conditioner.state_dict().items(): + state_dict["net.llm_adapter." + key] = value + + save_state_dict = self._convert_state_dict_dtype(state_dict, dtype) + self._convert_state_dict_to_contiguous(save_state_dict) + + os.makedirs(Path(destination).parent.absolute(), exist_ok=True) + + save_file(save_state_dict, destination, self._create_safetensors_header(model, save_state_dict)) + + def __save_internal( + self, + model: AnimaModel, + destination: str, + ): + self.__save_diffusers(model, destination, None) + + def save( + self, + model: AnimaModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + match output_model_format: + case ModelFormat.DIFFUSERS: + self.__save_diffusers(model, output_model_destination, dtype) + case ModelFormat.SAFETENSORS: + self.__save_safetensors(model, output_model_destination, dtype) + case ModelFormat.INTERNAL: + self.__save_internal(model, output_model_destination) diff --git a/modules/modelSetup/AnimaFineTuneSetup.py b/modules/modelSetup/AnimaFineTuneSetup.py new file mode 100644 index 000000000..02407a9b6 --- /dev/null +++ b/modules/modelSetup/AnimaFineTuneSetup.py @@ -0,0 +1,93 @@ +from modules.model.AnimaModel import AnimaModel +from modules.modelSetup.BaseAnimaSetup import BaseAnimaSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModuleFilter import ModuleFilter +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +class AnimaFineTuneSetup( + BaseAnimaSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: AnimaModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer, config.transformer, freeze=ModuleFilter.create(config), debug=config.debug_mode) + + if config.train_any_embedding() or config.train_any_output_embedding(): + raise NotImplementedError("Embeddings not implemented for Anima") + + return parameter_group_collection + + def __setup_requires_grad( + self, + model: AnimaModel, + config: TrainConfig, + ): + self._setup_model_part_requires_grad("transformer", model.transformer, config.transformer, model.train_progress) + + model.text_encoder.requires_grad_(False) + model.vae.requires_grad_(False) + + + def setup_model( + self, + model: AnimaModel, + config: TrainConfig, + ): + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: AnimaModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + + model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device) + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + model.text_encoder.eval() + + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: AnimaModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) + +factory.register(BaseModelSetup, AnimaFineTuneSetup, ModelType.ANIMA, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/AnimaLoRASetup.py b/modules/modelSetup/AnimaLoRASetup.py new file mode 100644 index 000000000..4eb7e2d4f --- /dev/null +++ b/modules/modelSetup/AnimaLoRASetup.py @@ -0,0 +1,105 @@ +from modules.model.AnimaModel import AnimaModel +from modules.modelSetup.BaseAnimaSetup import BaseAnimaSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +class AnimaLoRASetup( + BaseAnimaSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: AnimaModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer_lora, config.transformer) + + if config.train_any_embedding() or config.train_any_output_embedding(): + raise NotImplementedError("Embeddings not implemented for Anima") + + return parameter_group_collection + + def __setup_requires_grad( + self, + model: AnimaModel, + config: TrainConfig, + ): + model.text_encoder.requires_grad_(False) + model.transformer.requires_grad_(False) + model.vae.requires_grad_(False) + + self._setup_model_part_requires_grad("transformer", model.transformer_lora, config.transformer, model.train_progress) + + def setup_model( + self, + model: AnimaModel, + config: TrainConfig, + ): + model.transformer_lora = LoRAModuleWrapper( + model.transformer, "transformer", config, config.layer_filter.split(",") + ) + + if model.lora_state_dict: + model.transformer_lora.load_state_dict(model.lora_state_dict) + model.lora_state_dict = None + + model.transformer_lora.set_dropout(config.dropout_probability) + model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) + model.transformer_lora.hook_to_module() + + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: AnimaModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + + model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device) + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + model.text_encoder.eval() + + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: AnimaModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) + +factory.register(BaseModelSetup, AnimaLoRASetup, ModelType.ANIMA, TrainingMethod.LORA) diff --git a/modules/modelSetup/BaseAnimaSetup.py b/modules/modelSetup/BaseAnimaSetup.py new file mode 100644 index 000000000..379ae9d95 --- /dev/null +++ b/modules/modelSetup/BaseAnimaSetup.py @@ -0,0 +1,185 @@ +from abc import ABCMeta +from random import Random + +import modules.util.multi_gpu_util as multi +from modules.model.AnimaModel import AnimaModel +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.modelSetup.mixin.ModelSetupDebugMixin import ModelSetupDebugMixin +from modules.modelSetup.mixin.ModelSetupDiffusionLossMixin import ModelSetupDiffusionLossMixin +from modules.modelSetup.mixin.ModelSetupFlowMatchingMixin import ModelSetupFlowMatchingMixin +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin +from modules.modelSetup.mixin.ModelSetupText2ImageMixin import ModelSetupText2ImageMixin +from modules.util.checkpointing_util import ( + enable_checkpointing_for_qwen3_encoder_layers, + enable_checkpointing_for_qwen_transformer, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.quantization_util import quantize_layers +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress + +import torch +from torch import Tensor + + +#TODO share more code with other models +class BaseAnimaSetup( + BaseModelSetup, + ModelSetupDiffusionLossMixin, + ModelSetupDebugMixin, + ModelSetupNoiseMixin, + ModelSetupFlowMatchingMixin, + ModelSetupText2ImageMixin, + metaclass=ABCMeta +): + # CosmosTransformerBlock has attn1 (self-attn), attn2 (cross-attn), ff (feedforward) + LAYER_PRESETS = { + "attn-mlp": ["attn1", "attn2", "ff"], + "attn-only": ["attn1", "attn2"], + "blocks": ["transformer_block"], + "full": [], + } + + def setup_optimizations( + self, + model: AnimaModel, + config: TrainConfig, + ): + if config.gradient_checkpointing.enabled(): + model.transformer_offload_conductor = \ + enable_checkpointing_for_qwen_transformer(model.transformer, config) + if model.text_encoder is not None: + model.text_encoder_offload_conductor = \ + enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + + model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ + config.weight_dtypes().transformer, + config.weight_dtypes().text_encoder, + config.weight_dtypes().vae, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], config.enable_autocast_cache) + + model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ + disable_fp16_autocast_context( + self.train_device, + config.train_dtype, + config.fallback_train_dtype, + [ + config.weight_dtypes().text_encoder, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], + config.enable_autocast_cache, + ) + + quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) + quantize_layers(model.vae, self.train_device, model.train_dtype, config) + quantize_layers(model.transformer, self.train_device, model.train_dtype, config) + + def predict( + self, + model: AnimaModel, + batch: dict, + config: TrainConfig, + train_progress: TrainProgress, + *, + deterministic: bool = False, + ) -> dict: + with model.autocast_context: + batch_seed = 0 if deterministic else train_progress.global_step * multi.world_size() + multi.rank() + generator = torch.Generator(device=config.train_device) + generator.manual_seed(batch_seed) + rand = Random(batch_seed) + + # Anima encode_text returns a plain Tensor (no mask); conditioner output is (B,512,1024). + text_encoder_output = model.encode_text( + train_device=self.train_device, + batch_size=batch['latent_image'].shape[0], + rand=rand, + tokens=batch.get("tokens"), + tokens_mask=batch.get("tokens_mask"), + text_encoder_output=batch['text_encoder_hidden_state'] \ + if 'text_encoder_hidden_state' in batch and not config.train_text_encoder_or_embedding() else None, + text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, + ) + + latent_image = batch['latent_image'] + scaled_latent_image = model.scale_latents(latent_image) + latent_noise = self._create_noise(scaled_latent_image, config, generator) + + shift = model.calculate_timestep_shift(scaled_latent_image.shape[-2], scaled_latent_image.shape[-1]) + timestep = self._get_timestep_discrete( + model.noise_scheduler.config['num_train_timesteps'], + deterministic, + generator, + scaled_latent_image.shape[0], + config, + shift=shift if config.dynamic_timestep_shifting else config.timestep_shift, + ) + + scaled_noisy_latent_image, sigma = self._add_noise_discrete( + scaled_latent_image, + latent_noise, + timestep, + model.noise_scheduler.timesteps, + ) + + # Anima latents are 5D (B,16,1,H/8,W/8) — no pack/unpack needed. + # CosmosTransformer3DModel requires padding_mask in pixel space (1,1,H,W). + latent_h, latent_w = scaled_noisy_latent_image.shape[-2], scaled_noisy_latent_image.shape[-1] + padding_mask = scaled_noisy_latent_image.new_zeros( + 1, 1, latent_h * 8, latent_w * 8, + ).to(dtype=model.train_dtype.torch_dtype()) + + predicted_flow = model.transformer( + hidden_states=scaled_noisy_latent_image.to(dtype=model.train_dtype.torch_dtype()), + timestep=timestep / 1000, + encoder_hidden_states=text_encoder_output.to(dtype=model.train_dtype.torch_dtype()), + padding_mask=padding_mask, + return_dict=False, + )[0] + + flow = latent_noise - scaled_latent_image + model_output_data = { + 'loss_type': 'target', + 'timestep': timestep, + 'predicted': predicted_flow, + 'target': flow, + } + + if config.debug_mode: + with torch.no_grad(): + predicted_scaled_latent_image = scaled_noisy_latent_image - predicted_flow * sigma + self._save_tokens("7-prompt", batch['tokens'], model.tokenizer, config, train_progress) + self._save_latent("1-noise", latent_noise, config, train_progress) + self._save_latent("2-noisy_image", scaled_noisy_latent_image, config, train_progress) + self._save_latent("3-predicted_flow", predicted_flow, config, train_progress) + self._save_latent("4-flow", flow, config, train_progress) + self._save_latent("5-predicted_image", predicted_scaled_latent_image, config, train_progress) + self._save_latent("6-image", scaled_latent_image, config, train_progress) + + return model_output_data + + def calculate_loss( + self, + model: AnimaModel, + batch: dict, + data: dict, + config: TrainConfig, + ) -> Tensor: + return self._flow_matching_losses( + batch=batch, + data=data, + config=config, + train_device=self.train_device, + sigmas=model.noise_scheduler.sigmas, + ).mean() + + def prepare_text_caching(self, model: AnimaModel, config: TrainConfig): + model.to(self.temp_device) + + model.text_encoder_to(self.train_device) + + model.eval() + torch_gc() diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py index 6cb1b507a..56e39b125 100644 --- a/modules/ui/ConvertModelUI.py +++ b/modules/ui/ConvertModelUI.py @@ -70,6 +70,7 @@ def main_frame(self, master): ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here + ("Anima", ModelType.ANIMA), ("ZImage", ModelType.Z_IMAGE), ], self.ui_state, "model_type") diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index ff17ea3ba..d8f94efcf 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -65,6 +65,8 @@ def refresh_ui(self): self.__setup_chroma_ui(base_frame) elif self.train_config.model_type.is_qwen(): self.__setup_qwen_ui(base_frame) + elif self.train_config.model_type.is_anima(): + self.__setup_anima_ui(base_frame) elif self.train_config.model_type.is_sana(): self.__setup_sana_ui(base_frame) elif self.train_config.model_type.is_hunyuan_video(): @@ -230,6 +232,25 @@ def __setup_qwen_ui(self, frame): allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, ) + def __setup_anima_ui(self, frame): + row = 0 + row = self.__create_base_dtype_components(frame, row) + row = self.__create_base_components( + frame, + row, + has_transformer=True, + allow_override_transformer=True, + has_text_encoder_1=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + allow_safetensors=True, + allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + ) + def __setup_stable_diffusion_xl_ui(self, frame): row = 0 row = self.__create_base_dtype_components(frame, row) diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index 820fdb71a..167f8fde2 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -100,6 +100,7 @@ def __init__( ("HiDream Full", ModelType.HI_DREAM_FULL), ("Chroma1", ModelType.CHROMA_1), ("QwenImage", ModelType.QWEN), + ("Anima", ModelType.ANIMA), ("Z-Image", ModelType.Z_IMAGE), ("Ernie Image", ModelType.ERNIE), ], @@ -136,6 +137,7 @@ def __create_training_method(self): ("Embedding", TrainingMethod.EMBEDDING), ] elif self.train_config.model_type.is_qwen() \ + or self.train_config.model_type.is_anima() \ or self.train_config.model_type.is_z_image() \ or self.train_config.model_type.is_flux_2() \ or self.train_config.model_type.is_ernie(): diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index c839abf73..490961995 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -78,6 +78,8 @@ def refresh_ui(self): self.__setup_chroma_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_qwen(): self.__setup_qwen_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_anima(): + self.__setup_anima_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_sana(): self.__setup_sana_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_hunyuan_video(): @@ -209,6 +211,18 @@ def __setup_qwen_ui(self, column_0, column_1, column_2): self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) + def __setup_anima_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + def __setup_z_image_ui(self, column_0, column_1, column_2): self.__create_base_frame(column_0, 0) self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False) diff --git a/modules/util/config/SampleConfig.py b/modules/util/config/SampleConfig.py index 38e99e182..50cd956f3 100644 --- a/modules/util/config/SampleConfig.py +++ b/modules/util/config/SampleConfig.py @@ -86,6 +86,14 @@ def _get_model_defaults(model_type) -> dict: "diffusion_steps": 25, "cfg_scale": 3.5, }) + elif model_type.is_anima(): + defaults.update({ + "width": 1024, + "height": 1024, + "diffusion_steps": 25, + "cfg_scale": 4.0, + "negative_prompt": "worst quality, low quality, score_1, score_2, score_3, artist name", + }) elif model_type.is_z_image(): defaults.update({ "width": 1024, diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index a3ad940ec..b9bea4311 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -37,6 +37,8 @@ class ModelType(Enum): QWEN = 'QWEN' + ANIMA = 'ANIMA' + Z_IMAGE = 'Z_IMAGE' ERNIE = 'ERNIE' @@ -97,6 +99,9 @@ def is_chroma(self): def is_qwen(self): return self == ModelType.QWEN + def is_anima(self): + return self == ModelType.ANIMA + def is_sana(self): return self == ModelType.SANA @@ -157,6 +162,7 @@ def is_flow_matching(self) -> bool: or self.is_flux() \ or self.is_chroma() \ or self.is_qwen() \ + or self.is_anima() \ or self.is_sana() \ or self.is_hunyuan_video() \ or self.is_hi_dream() \ diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index b2630c16b..74df5802f 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -35,7 +35,7 @@ def build_muon_adam_key_fn( 'block', # UNet 'text_model.encoder.layers', # TEs (CLIPs) ] - case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA: + case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA | ModelType.ANIMA: default_patterns = [ 'transformer_blocks', 'encoder.block', # TE (T5) diff --git a/requirements-global.txt b/requirements-global.txt index 299f765b0..7a55f2128 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -20,7 +20,7 @@ safetensors==0.8.0rc0 tensorboard==2.20.0 # diffusion models --e git+https://github.com/huggingface/diffusers.git@0f1abc4#egg=diffusers +-e git+https://github.com/huggingface/diffusers.git@b003a47#egg=diffusers gguf==0.17.1 transformers==5.9.0 sentencepiece==0.2.1 # transitive dependency of transformers for tokenizer loading @@ -32,7 +32,7 @@ pooch==1.8.2 open-clip-torch==2.32.0 # data loader --e git+https://github.com/Nerogar/mgds.git@9320a69#egg=mgds +-e git+https://github.com/dxqb/mgds.git@fa6ae65#egg=mgds # optimizers dadaptation==3.2 # dadaptation optimizers diff --git a/resources/sd_model_spec/anima-lora.json b/resources/sd_model_spec/anima-lora.json new file mode 100644 index 000000000..4aa9ccd27 --- /dev/null +++ b/resources/sd_model_spec/anima-lora.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Anima/lora", + "modelspec.implementation": "https://github.com/huggingface/diffusers", + "modelspec.title": "Anima LoRA" +} diff --git a/resources/sd_model_spec/anima.json b/resources/sd_model_spec/anima.json new file mode 100644 index 000000000..c4ce6b736 --- /dev/null +++ b/resources/sd_model_spec/anima.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Anima", + "modelspec.implementation": "https://github.com/huggingface/diffusers", + "modelspec.title": "Anima" +} diff --git a/training_presets/#anima Finetune.json b/training_presets/#anima Finetune.json new file mode 100644 index 000000000..07043de6b --- /dev/null +++ b/training_presets/#anima Finetune.json @@ -0,0 +1,46 @@ +{ + "base_model_name": "circlestone-labs/Anima-Base-v1.0-Diffusers", + "batch_size": 2, + "learning_rate": 1e-6, + "model_type": "ANIMA", + "resolution": "512", + "compile": true, + "dataloader_threads": 1, + "transformer": { + "train": true, + "weight_dtype": "BFLOAT_16" + }, + "text_encoder": { + "train": false, + "weight_dtype": "FLOAT_8" + }, + "training_method": "FINE_TUNE", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "weight_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "timestep_distribution": "LOGIT_NORMAL", + "optimizer": { + "optimizer": "ADAFACTOR" + }, + "optimizer_defaults": { + "ADAFACTOR": { + "optimizer": "ADAFACTOR", + "fused_back_pass": true, + "beta1": null, + "clip_threshold": 1.0, + "decay_rate": -0.8, + "eps": 1e-30, + "eps2": 0.001, + "relative_step": false, + "scale_parameter": false, + "stochastic_rounding": true, + "warmup_init": false, + "weight_decay": 0.0 + } + }, + "layer_filter": "transformer_block", + "layer_filter_preset": "blocks" +} diff --git a/training_presets/#anima LoRA.json b/training_presets/#anima LoRA.json new file mode 100644 index 000000000..9b936a7ad --- /dev/null +++ b/training_presets/#anima LoRA.json @@ -0,0 +1,31 @@ +{ + "base_model_name": "circlestone-labs/Anima-Base-v1.0-Diffusers", + "batch_size": 2, + "learning_rate": 3e-05, + "model_type": "ANIMA", + "resolution": "512", + "compile": true, + "dataloader_threads": 1, + "transformer": { + "train": true, + "weight_dtype": "INT_W8A8" + }, + "text_encoder": { + "train": false, + "weight_dtype": "FLOAT_8" + }, + "training_method": "LORA", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "weight_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "timestep_distribution": "LOGIT_NORMAL", + "layer_filter": "attn1,attn2,ff", + "layer_filter_preset": "attn-mlp", + "quantization": { + "layer_filter": "attn1,attn2,ff", + "layer_filter_preset": "attn-mlp" + } +} From 25366c8a08c8d392a4c5bed0c82b9bc82d00314f Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Sun, 31 May 2026 21:58:52 +0200 Subject: [PATCH 54/67] Remove PR template, moved to separate PR --- .github/pull_request_template.md | 33 -------------------------------- 1 file changed, 33 deletions(-) delete mode 100644 .github/pull_request_template.md diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md deleted file mode 100644 index 95d818a84..000000000 --- a/.github/pull_request_template.md +++ /dev/null @@ -1,33 +0,0 @@ - - -## Summary - - - -## Test plan - - - -- [ ] `pre-commit run --all-files` passes -- [ ] Launched the affected UI or script and exercised the change -- [ ] Tested with at least one real preset / config when relevant (note which: ____) - -## AI assistance - - - -- No AI involvement -- AI-assisted — I have read every line in this diff and can defend each change -- Early AI prototype — opened for discussion, **not ready for review** - - From e54623e5cd68d2b8b7f090d2da35d49b60eb4626 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 3 Jun 2026 08:22:52 +0200 Subject: [PATCH 55/67] Remove Windows DPI-awareness hack from PySide6 view SetProcessDpiAwareness was intended for CustomTkinter only; Qt handles DPI natively and the call is not needed here. Suggested by PR #1488. Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/PySide6TrainUIView.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/modules/ui/PySide6TrainUIView.py b/modules/ui/PySide6TrainUIView.py index ddb1b803f..d189d05b9 100644 --- a/modules/ui/PySide6TrainUIView.py +++ b/modules/ui/PySide6TrainUIView.py @@ -1,7 +1,4 @@ -import ctypes -import platform from collections.abc import Callable -from contextlib import suppress from pathlib import Path from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController @@ -39,15 +36,6 @@ from PySide6.QtGui import QIcon from PySide6.QtWidgets import QFileDialog, QGridLayout, QMainWindow, QMessageBox, QTabWidget, QWidget -# chunk for forcing Windows to ignore DPI scaling when moving between monitors -# fixes the long standing transparency bug https://github.com/Nerogar/OneTrainer/issues/90 -if platform.system() == "Windows": - with suppress(Exception): - # https://learn.microsoft.com/en-us/windows/win32/hidpi/setting-the-default-dpi-awareness-for-a-process#setting-default-awareness-programmatically - ctypes.windll.shcore.SetProcessDpiAwareness(1) # PROCESS_SYSTEM_DPI_AWARE - - - class PySide6TrainView(BaseTrainUIView, QMainWindow, metaclass=QtABCMeta): def __init__(self): From 711783315cc93199c766a9abe3206593dbddd219 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:28:51 +0200 Subject: [PATCH 56/67] fix: call init_compile() in training thread to apply dynamo config torch._dynamo.config overrides are thread-local. The existing call in checkpointing_util runs in the main thread and is invisible to the training thread spawned by the UI. This caused compiled optimizers (e.g. AdamW_adv with compiled_optimizer=True) to hit the default recompile_limit of 8 and abort with FailOnRecompileLimitHit when training models with more than 8 distinct parameter shapes. Fix: call init_compile() from GenericTrainer.__init__, which runs in whichever thread/process owns training (UI thread, CLI main thread, or torch.multiprocessing.spawn subprocess for multi-GPU). Co-Authored-By: Claude Sonnet 4.6 --- modules/trainer/GenericTrainer.py | 3 +++ modules/util/compile_util.py | 1 + 2 files changed, 4 insertions(+) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index ab4926901..f7c97c132 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -20,6 +20,7 @@ from modules.util.bf16_stochastic_rounding import set_seed as bf16_stochastic_rounding_set_seed from modules.util.callbacks.TrainCallbacks import TrainCallbacks from modules.util.commands.TrainCommands import TrainCommands +from modules.util.compile_util import init_compile from modules.util.config.SampleConfig import SampleConfig from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_grad_scaler, enable_grad_scaling @@ -66,6 +67,8 @@ class GenericTrainer(BaseTrainer): def __init__(self, config: TrainConfig, callbacks: TrainCallbacks, commands: TrainCommands): super().__init__(config, callbacks, commands) + # torch._dynamo.config overrides are thread-local, so init_compile() must be called in the training thread/process. + init_compile() if multi.is_master(): tensorboard_log_dir = os.path.join(config.workspace_dir, "tensorboard") diff --git a/modules/util/compile_util.py b/modules/util/compile_util.py index aeb5980e4..3fd310cd2 100644 --- a/modules/util/compile_util.py +++ b/modules/util/compile_util.py @@ -57,5 +57,6 @@ def Mod_patched_eval(cls, p, q): def init_compile(): + # cache_size_limit and recompile_limit are aliases for the same dynamo config value. torch._dynamo.config.cache_size_limit = 8192 torch.utils._sympy.functions.Mod.eval = Mod_patched_eval From 5a5411dcc02a19add0e5e2bb69737d4093bbed9e Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:53:16 +0200 Subject: [PATCH 57/67] Use decorator form for factory.register() factory.register() now detects whether its second argument is a type (direct call form for loops) or an enum key (decorator form), allowing implementations to declare themselves with @factory.register() directly on the class definition instead of imperative calls at the bottom of the file. Co-Authored-By: Claude Sonnet 4.6 --- modules/dataLoader/ChromaBaseDataLoader.py | 3 +-- modules/dataLoader/ErnieBaseDataLoader.py | 4 +--- modules/dataLoader/Flux2BaseDataLoader.py | 4 +--- modules/dataLoader/FluxBaseDataLoader.py | 5 ++--- modules/dataLoader/HiDreamBaseDataLoader.py | 3 +-- .../dataLoader/HunyuanVideoBaseDataLoader.py | 3 +-- modules/dataLoader/PixArtAlphaBaseDataLoader.py | 5 ++--- modules/dataLoader/QwenBaseDataLoader.py | 3 +-- modules/dataLoader/SanaBaseDataLoader.py | 3 +-- .../StableDiffusion3BaseDataLoader.py | 3 +-- .../dataLoader/StableDiffusionBaseDataLoader.py | 17 ++++++++--------- .../StableDiffusionFineTuneVaeDataLoader.py | 17 ++++++++--------- .../StableDiffusionXLBaseDataLoader.py | 4 ++-- modules/dataLoader/WuerstchenBaseDataLoader.py | 5 ++--- modules/dataLoader/ZImageBaseDataLoader.py | 3 +-- modules/modelSampler/ChromaSampler.py | 3 +-- modules/modelSampler/ErnieSampler.py | 4 +--- modules/modelSampler/Flux2Sampler.py | 3 +-- modules/modelSampler/FluxSampler.py | 5 ++--- modules/modelSampler/HiDreamSampler.py | 3 +-- modules/modelSampler/HunyuanVideoSampler.py | 3 +-- modules/modelSampler/PixArtAlphaSampler.py | 5 ++--- modules/modelSampler/QwenSampler.py | 3 +-- modules/modelSampler/SanaSampler.py | 3 +-- modules/modelSampler/StableDiffusion3Sampler.py | 5 ++--- modules/modelSampler/StableDiffusionSampler.py | 17 ++++++++--------- .../modelSampler/StableDiffusionVaeSampler.py | 17 ++++++++--------- .../modelSampler/StableDiffusionXLSampler.py | 5 ++--- modules/modelSampler/WuerstchenSampler.py | 5 ++--- modules/modelSampler/ZImageSampler.py | 3 +-- modules/modelSetup/ChromaEmbeddingSetup.py | 3 +-- modules/modelSetup/ChromaFineTuneSetup.py | 3 +-- modules/modelSetup/ChromaLoRASetup.py | 3 +-- modules/modelSetup/ErnieFineTuneSetup.py | 4 +--- modules/modelSetup/ErnieLoRASetup.py | 4 +--- modules/modelSetup/Flux2FineTuneSetup.py | 3 +-- modules/modelSetup/Flux2LoRASetup.py | 3 +-- modules/modelSetup/FluxEmbeddingSetup.py | 5 ++--- modules/modelSetup/FluxFineTuneSetup.py | 5 ++--- modules/modelSetup/FluxLoRASetup.py | 5 ++--- modules/modelSetup/HiDreamEmbeddingSetup.py | 3 +-- modules/modelSetup/HiDreamFineTuneSetup.py | 3 +-- modules/modelSetup/HiDreamLoRASetup.py | 3 +-- .../modelSetup/HunyuanVideoEmbeddingSetup.py | 3 +-- modules/modelSetup/HunyuanVideoFineTuneSetup.py | 3 +-- modules/modelSetup/HunyuanVideoLoRASetup.py | 3 +-- modules/modelSetup/PixArtAlphaEmbeddingSetup.py | 5 ++--- modules/modelSetup/PixArtAlphaFineTuneSetup.py | 5 ++--- modules/modelSetup/PixArtAlphaLoRASetup.py | 5 ++--- modules/modelSetup/QwenFineTuneSetup.py | 3 +-- modules/modelSetup/QwenLoRASetup.py | 3 +-- modules/modelSetup/SanaEmbeddingSetup.py | 3 +-- modules/modelSetup/SanaFineTuneSetup.py | 3 +-- modules/modelSetup/SanaLoRASetup.py | 3 +-- .../StableDiffusion3EmbeddingSetup.py | 5 ++--- .../modelSetup/StableDiffusion3FineTuneSetup.py | 5 ++--- modules/modelSetup/StableDiffusion3LoRASetup.py | 5 ++--- .../modelSetup/StableDiffusionEmbeddingSetup.py | 17 ++++++++--------- .../modelSetup/StableDiffusionFineTuneSetup.py | 17 ++++++++--------- .../StableDiffusionFineTuneVaeSetup.py | 17 ++++++++--------- modules/modelSetup/StableDiffusionLoRASetup.py | 17 ++++++++--------- .../StableDiffusionXLEmbeddingSetup.py | 5 ++--- .../StableDiffusionXLFineTuneSetup.py | 5 ++--- .../modelSetup/StableDiffusionXLLoRASetup.py | 5 ++--- modules/modelSetup/WuerstchenEmbeddingSetup.py | 5 ++--- modules/modelSetup/WuerstchenFineTuneSetup.py | 5 ++--- modules/modelSetup/WuerstchenLoRASetup.py | 5 ++--- modules/modelSetup/ZImageFineTuneSetup.py | 3 +-- modules/modelSetup/ZImageLoRASetup.py | 3 +-- modules/util/factory.py | 14 +++++++++++++- 70 files changed, 162 insertions(+), 223 deletions(-) diff --git a/modules/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py index 94722bdf5..6b89457e6 100644 --- a/modules/dataLoader/ChromaBaseDataLoader.py +++ b/modules/dataLoader/ChromaBaseDataLoader.py @@ -26,6 +26,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.CHROMA_1) class ChromaBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -154,5 +155,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - -factory.register(BaseDataLoader, ChromaBaseDataLoader, ModelType.CHROMA_1) diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..4859e24ef 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -22,6 +22,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.ERNIE) class ErnieBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -133,6 +134,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - - -factory.register(BaseDataLoader, ErnieBaseDataLoader, ModelType.ERNIE) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..942f8bb6b 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -30,6 +30,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.FLUX_2) class Flux2BaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -161,6 +162,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - - -factory.register(BaseDataLoader, Flux2BaseDataLoader, ModelType.FLUX_2) diff --git a/modules/dataLoader/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py index 8dcb98599..d23dee3a8 100644 --- a/modules/dataLoader/FluxBaseDataLoader.py +++ b/modules/dataLoader/FluxBaseDataLoader.py @@ -26,6 +26,8 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.FLUX_DEV_1) +@factory.register(BaseDataLoader, ModelType.FLUX_FILL_DEV_1) class FluxBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -181,6 +183,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - -factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_DEV_1) -factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_FILL_DEV_1) diff --git a/modules/dataLoader/HiDreamBaseDataLoader.py b/modules/dataLoader/HiDreamBaseDataLoader.py index dedea17be..f080943bf 100644 --- a/modules/dataLoader/HiDreamBaseDataLoader.py +++ b/modules/dataLoader/HiDreamBaseDataLoader.py @@ -27,6 +27,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.HI_DREAM_FULL) class HiDreamBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -219,5 +220,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - -factory.register(BaseDataLoader, HiDreamBaseDataLoader, ModelType.HI_DREAM_FULL) diff --git a/modules/dataLoader/HunyuanVideoBaseDataLoader.py b/modules/dataLoader/HunyuanVideoBaseDataLoader.py index d587a5fd9..38f7e6a8a 100644 --- a/modules/dataLoader/HunyuanVideoBaseDataLoader.py +++ b/modules/dataLoader/HunyuanVideoBaseDataLoader.py @@ -29,6 +29,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.HUNYUAN_VIDEO) class HunyuanVideoBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -183,5 +184,3 @@ def _create_dataset( allow_video_files=True, vae_frame_dim=True, ) - -factory.register(BaseDataLoader, HunyuanVideoBaseDataLoader, ModelType.HUNYUAN_VIDEO) diff --git a/modules/dataLoader/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py index f2dc37857..bb110dc2e 100644 --- a/modules/dataLoader/PixArtAlphaBaseDataLoader.py +++ b/modules/dataLoader/PixArtAlphaBaseDataLoader.py @@ -24,6 +24,8 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.PIXART_ALPHA) +@factory.register(BaseDataLoader, ModelType.PIXART_SIGMA) class PixArtAlphaBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -159,6 +161,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=16, ) - -factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_ALPHA) -factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_SIGMA) diff --git a/modules/dataLoader/QwenBaseDataLoader.py b/modules/dataLoader/QwenBaseDataLoader.py index db3ca5df3..9a4a962a9 100644 --- a/modules/dataLoader/QwenBaseDataLoader.py +++ b/modules/dataLoader/QwenBaseDataLoader.py @@ -30,6 +30,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.QWEN) class QwenBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -166,5 +167,3 @@ def _create_dataset( allow_video_files=False, #don't allow video files, but... vae_frame_dim=True, #...Qwen has a video-capable VAE. convert images to video dimensions ) - -factory.register(BaseDataLoader, QwenBaseDataLoader, ModelType.QWEN) diff --git a/modules/dataLoader/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py index 38d5c31b0..a44ff8130 100644 --- a/modules/dataLoader/SanaBaseDataLoader.py +++ b/modules/dataLoader/SanaBaseDataLoader.py @@ -23,6 +23,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.SANA) class SanaBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -152,5 +153,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=32, ) - -factory.register(BaseDataLoader, SanaBaseDataLoader, ModelType.SANA) diff --git a/modules/dataLoader/StableDiffusion3BaseDataLoader.py b/modules/dataLoader/StableDiffusion3BaseDataLoader.py index c497a20c6..f8a4a2211 100644 --- a/modules/dataLoader/StableDiffusion3BaseDataLoader.py +++ b/modules/dataLoader/StableDiffusion3BaseDataLoader.py @@ -25,6 +25,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_35) class StableDiffusion3BaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -198,5 +199,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - -factory.register(BaseDataLoader, StableDiffusion3BaseDataLoader, ModelType.STABLE_DIFFUSION_35) diff --git a/modules/dataLoader/StableDiffusionBaseDataLoader.py b/modules/dataLoader/StableDiffusionBaseDataLoader.py index 781f769a6..63ad57cac 100644 --- a/modules/dataLoader/StableDiffusionBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionBaseDataLoader.py @@ -24,6 +24,14 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE) class StableDiffusionBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -164,12 +172,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=8, ) - -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_15) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_21) -factory.register(BaseDataLoader, StableDiffusionBaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE) diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index ed5dd32b8..ad0c890b0 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py @@ -39,6 +39,14 @@ import torch +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) class StableDiffusionFineTuneVaeDataLoader(BaseDataLoader): def _setup_cache_device( self, @@ -293,12 +301,3 @@ def _create_dataset( train_progress, is_validation, ) - -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseDataLoader, StableDiffusionFineTuneVaeDataLoader, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) diff --git a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py index ed1ad491e..12739adf8 100644 --- a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py @@ -24,6 +24,8 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE) +@factory.register(BaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING) class StableDiffusionXLBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -176,5 +178,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) -factory.register(BaseDataLoader, StableDiffusionXLBaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE) -factory.register(BaseDataLoader, StableDiffusionXLBaseDataLoader, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING) diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py index f5cf2f41a..4689b09a4 100644 --- a/modules/dataLoader/WuerstchenBaseDataLoader.py +++ b/modules/dataLoader/WuerstchenBaseDataLoader.py @@ -23,6 +23,8 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.WUERSTCHEN_2) +@factory.register(BaseDataLoader, ModelType.STABLE_CASCADE_1) class WuerstchenBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -156,6 +158,3 @@ def _create_dataset( aspect_bucketing_quantization=128, supports_inpainting=False, ) - -factory.register(BaseDataLoader, WuerstchenBaseDataLoader, ModelType.WUERSTCHEN_2) -factory.register(BaseDataLoader, WuerstchenBaseDataLoader, ModelType.STABLE_CASCADE_1) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..98801a817 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -26,6 +26,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.Z_IMAGE) class ZImageBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -153,5 +154,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - -factory.register(BaseDataLoader, ZImageBaseDataLoader, ModelType.Z_IMAGE) diff --git a/modules/modelSampler/ChromaSampler.py b/modules/modelSampler/ChromaSampler.py index 4b6033623..bc89e5cab 100644 --- a/modules/modelSampler/ChromaSampler.py +++ b/modules/modelSampler/ChromaSampler.py @@ -19,6 +19,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.CHROMA_1) class ChromaSampler(BaseModelSampler): def __init__( self, @@ -189,5 +190,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, ChromaSampler, ModelType.CHROMA_1) diff --git a/modules/modelSampler/ErnieSampler.py b/modules/modelSampler/ErnieSampler.py index b1246517c..2122e5343 100644 --- a/modules/modelSampler/ErnieSampler.py +++ b/modules/modelSampler/ErnieSampler.py @@ -20,6 +20,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.ERNIE) class ErnieSampler(BaseModelSampler): def __init__( self, @@ -162,6 +163,3 @@ def sample( ) on_sample(sampler_output) - - -factory.register(BaseModelSampler, ErnieSampler, ModelType.ERNIE) diff --git a/modules/modelSampler/Flux2Sampler.py b/modules/modelSampler/Flux2Sampler.py index 38986bb6f..0a4cca9b9 100644 --- a/modules/modelSampler/Flux2Sampler.py +++ b/modules/modelSampler/Flux2Sampler.py @@ -22,6 +22,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.FLUX_2) class Flux2Sampler(BaseModelSampler): def __init__( self, @@ -190,5 +191,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, Flux2Sampler, ModelType.FLUX_2) diff --git a/modules/modelSampler/FluxSampler.py b/modules/modelSampler/FluxSampler.py index 93f8837ba..fbe532053 100644 --- a/modules/modelSampler/FluxSampler.py +++ b/modules/modelSampler/FluxSampler.py @@ -23,6 +23,8 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.FLUX_DEV_1) +@factory.register(BaseModelSampler, ModelType.FLUX_FILL_DEV_1) class FluxSampler(BaseModelSampler): def __init__( self, @@ -450,6 +452,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, FluxSampler, ModelType.FLUX_DEV_1) -factory.register(BaseModelSampler, FluxSampler, ModelType.FLUX_FILL_DEV_1) diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index d02197a35..2bd85f57a 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -19,6 +19,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.HI_DREAM_FULL) class HiDreamSampler(BaseModelSampler): def __init__( self, @@ -192,5 +193,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, HiDreamSampler, ModelType.HI_DREAM_FULL) diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index 275626341..4fa76d347 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -20,6 +20,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.HUNYUAN_VIDEO) class HunyuanVideoSampler(BaseModelSampler): def __init__( self, @@ -209,5 +210,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, HunyuanVideoSampler, ModelType.HUNYUAN_VIDEO) diff --git a/modules/modelSampler/PixArtAlphaSampler.py b/modules/modelSampler/PixArtAlphaSampler.py index 7bec3f59a..9adddf3f2 100644 --- a/modules/modelSampler/PixArtAlphaSampler.py +++ b/modules/modelSampler/PixArtAlphaSampler.py @@ -18,6 +18,8 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.PIXART_ALPHA) +@factory.register(BaseModelSampler, ModelType.PIXART_SIGMA) class PixArtAlphaSampler(BaseModelSampler): def __init__( self, @@ -191,6 +193,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, PixArtAlphaSampler, ModelType.PIXART_ALPHA) -factory.register(BaseModelSampler, PixArtAlphaSampler, ModelType.PIXART_SIGMA) diff --git a/modules/modelSampler/QwenSampler.py b/modules/modelSampler/QwenSampler.py index 46405a36b..4ca604102 100644 --- a/modules/modelSampler/QwenSampler.py +++ b/modules/modelSampler/QwenSampler.py @@ -20,6 +20,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.QWEN) class QwenSampler(BaseModelSampler): def __init__( self, @@ -188,5 +189,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, QwenSampler, ModelType.QWEN) diff --git a/modules/modelSampler/SanaSampler.py b/modules/modelSampler/SanaSampler.py index edf7510cc..6251ab87e 100644 --- a/modules/modelSampler/SanaSampler.py +++ b/modules/modelSampler/SanaSampler.py @@ -19,6 +19,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.SANA) class SanaSampler(BaseModelSampler): def __init__( self, @@ -177,5 +178,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, SanaSampler, ModelType.SANA) diff --git a/modules/modelSampler/StableDiffusion3Sampler.py b/modules/modelSampler/StableDiffusion3Sampler.py index 4cbbe4cc1..fe159e914 100644 --- a/modules/modelSampler/StableDiffusion3Sampler.py +++ b/modules/modelSampler/StableDiffusion3Sampler.py @@ -19,6 +19,8 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_3) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_35) class StableDiffusion3Sampler(BaseModelSampler): def __init__( self, @@ -191,6 +193,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, StableDiffusion3Sampler, ModelType.STABLE_DIFFUSION_3) -factory.register(BaseModelSampler, StableDiffusion3Sampler, ModelType.STABLE_DIFFUSION_35) diff --git a/modules/modelSampler/StableDiffusionSampler.py b/modules/modelSampler/StableDiffusionSampler.py index e16abeda2..92bfcd759 100644 --- a/modules/modelSampler/StableDiffusionSampler.py +++ b/modules/modelSampler/StableDiffusionSampler.py @@ -21,6 +21,14 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15_INPAINTING) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_BASE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_INPAINTING) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_DEPTH) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_21) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_21_BASE) class StableDiffusionSampler(BaseModelSampler): def __init__( self, @@ -429,12 +437,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_15) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_15_INPAINTING) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_20) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_20_BASE) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_20_INPAINTING) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_20_DEPTH) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_21) -factory.register(BaseModelSampler, StableDiffusionSampler, ModelType.STABLE_DIFFUSION_21_BASE) diff --git a/modules/modelSampler/StableDiffusionVaeSampler.py b/modules/modelSampler/StableDiffusionVaeSampler.py index afb39bcff..254c73fdd 100644 --- a/modules/modelSampler/StableDiffusionVaeSampler.py +++ b/modules/modelSampler/StableDiffusionVaeSampler.py @@ -17,6 +17,14 @@ from PIL import Image +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) class StableDiffusionVaeSampler(BaseModelSampler): def __init__( self, @@ -78,12 +86,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSampler, StableDiffusionVaeSampler, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) diff --git a/modules/modelSampler/StableDiffusionXLSampler.py b/modules/modelSampler/StableDiffusionXLSampler.py index e157d3b0a..1f066268c 100644 --- a/modules/modelSampler/StableDiffusionXLSampler.py +++ b/modules/modelSampler/StableDiffusionXLSampler.py @@ -21,6 +21,8 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE) +@factory.register(BaseModelSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING) class StableDiffusionXLSampler(BaseModelSampler): def __init__( self, @@ -501,6 +503,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, StableDiffusionXLSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE) -factory.register(BaseModelSampler, StableDiffusionXLSampler, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING) diff --git a/modules/modelSampler/WuerstchenSampler.py b/modules/modelSampler/WuerstchenSampler.py index 061a1e22b..d679757c2 100644 --- a/modules/modelSampler/WuerstchenSampler.py +++ b/modules/modelSampler/WuerstchenSampler.py @@ -19,6 +19,8 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.WUERSTCHEN_2) +@factory.register(BaseModelSampler, ModelType.STABLE_CASCADE_1) class WuerstchenSampler(BaseModelSampler): def __init__( self, @@ -365,6 +367,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, WuerstchenSampler, ModelType.WUERSTCHEN_2) -factory.register(BaseModelSampler, WuerstchenSampler, ModelType.STABLE_CASCADE_1) diff --git a/modules/modelSampler/ZImageSampler.py b/modules/modelSampler/ZImageSampler.py index 4ac837649..0cfcb46d6 100644 --- a/modules/modelSampler/ZImageSampler.py +++ b/modules/modelSampler/ZImageSampler.py @@ -19,6 +19,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.Z_IMAGE) class ZImageSampler(BaseModelSampler): def __init__( self, @@ -163,5 +164,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, ZImageSampler, ModelType.Z_IMAGE) diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 9fa65b69f..bc0a25105 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -12,6 +12,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.EMBEDDING) class ChromaEmbeddingSetup( BaseChromaSetup, ): @@ -96,5 +97,3 @@ def after_optimizer_step( if model.embedding_wrapper is not None: model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, ChromaEmbeddingSetup, ModelType.CHROMA_1, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 371c414ec..8335cffce 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.FINE_TUNE) class ChromaFineTuneSetup( BaseChromaSetup, ): @@ -116,5 +117,3 @@ def after_optimizer_step( if model.embedding_wrapper is not None: model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, ChromaFineTuneSetup, ModelType.CHROMA_1, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 1396b4a52..d03050922 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -14,6 +14,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.LORA) class ChromaLoRASetup( BaseChromaSetup, ): @@ -144,5 +145,3 @@ def after_optimizer_step( if model.embedding_wrapper is not None: model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, ChromaLoRASetup, ModelType.CHROMA_1, TrainingMethod.LORA) diff --git a/modules/modelSetup/ErnieFineTuneSetup.py b/modules/modelSetup/ErnieFineTuneSetup.py index a4a03fb02..feadd4fca 100644 --- a/modules/modelSetup/ErnieFineTuneSetup.py +++ b/modules/modelSetup/ErnieFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.ERNIE, TrainingMethod.FINE_TUNE) class ErnieFineTuneSetup( BaseErnieSetup, ): @@ -85,6 +86,3 @@ def after_optimizer_step( train_progress: TrainProgress, ): self.__setup_requires_grad(model, config) - - -factory.register(BaseModelSetup, ErnieFineTuneSetup, ModelType.ERNIE, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/ErnieLoRASetup.py b/modules/modelSetup/ErnieLoRASetup.py index 108a26f92..1a70ef269 100644 --- a/modules/modelSetup/ErnieLoRASetup.py +++ b/modules/modelSetup/ErnieLoRASetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.ERNIE, TrainingMethod.LORA) class ErnieLoRASetup( BaseErnieSetup, ): @@ -95,6 +96,3 @@ def after_optimizer_step( train_progress: TrainProgress, ): self.__setup_requires_grad(model, config) - - -factory.register(BaseModelSetup, ErnieLoRASetup, ModelType.ERNIE, TrainingMethod.LORA) diff --git a/modules/modelSetup/Flux2FineTuneSetup.py b/modules/modelSetup/Flux2FineTuneSetup.py index 2b90618c3..7ca128a09 100644 --- a/modules/modelSetup/Flux2FineTuneSetup.py +++ b/modules/modelSetup/Flux2FineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.FLUX_2, TrainingMethod.FINE_TUNE) class Flux2FineTuneSetup( BaseFlux2Setup, ): @@ -85,5 +86,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, Flux2FineTuneSetup, ModelType.FLUX_2, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/Flux2LoRASetup.py b/modules/modelSetup/Flux2LoRASetup.py index 3c38ebf87..7e31e2bae 100644 --- a/modules/modelSetup/Flux2LoRASetup.py +++ b/modules/modelSetup/Flux2LoRASetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.FLUX_2, TrainingMethod.LORA) class Flux2LoRASetup( BaseFlux2Setup, ): @@ -97,5 +98,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, Flux2LoRASetup, ModelType.FLUX_2, TrainingMethod.LORA) diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index 0418390da..f5fd4c844 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -12,6 +12,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.EMBEDDING) class FluxEmbeddingSetup( BaseFluxSetup, ): @@ -112,6 +114,3 @@ def after_optimizer_step( if model.embedding_wrapper_2 is not None: model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, FluxEmbeddingSetup, ModelType.FLUX_DEV_1, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, FluxEmbeddingSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index a2dc78057..1533de129 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -13,6 +13,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.FINE_TUNE) class FluxFineTuneSetup( BaseFluxSetup, ): @@ -140,6 +142,3 @@ def after_optimizer_step( if model.embedding_wrapper_2 is not None: model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, FluxFineTuneSetup, ModelType.FLUX_DEV_1, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, FluxFineTuneSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index e3de941fd..ad35555d1 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -14,6 +14,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.LORA) class FluxLoRASetup( BaseFluxSetup, ): @@ -182,6 +184,3 @@ def after_optimizer_step( if model.embedding_wrapper_2 is not None: model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, FluxLoRASetup, ModelType.FLUX_DEV_1, TrainingMethod.LORA) -factory.register(BaseModelSetup, FluxLoRASetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.LORA) diff --git a/modules/modelSetup/HiDreamEmbeddingSetup.py b/modules/modelSetup/HiDreamEmbeddingSetup.py index 22c09927b..659f3834e 100644 --- a/modules/modelSetup/HiDreamEmbeddingSetup.py +++ b/modules/modelSetup/HiDreamEmbeddingSetup.py @@ -12,6 +12,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.EMBEDDING) class HiDreamEmbeddingSetup( BaseHiDreamSetup, ): @@ -141,5 +142,3 @@ def after_optimizer_step( if model.embedding_wrapper_4 is not None: model.embedding_wrapper_4.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HiDreamEmbeddingSetup, ModelType.HI_DREAM_FULL, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/HiDreamFineTuneSetup.py b/modules/modelSetup/HiDreamFineTuneSetup.py index 49a3adf9d..2361e69ff 100644 --- a/modules/modelSetup/HiDreamFineTuneSetup.py +++ b/modules/modelSetup/HiDreamFineTuneSetup.py @@ -15,6 +15,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.FINE_TUNE) class HiDreamFineTuneSetup( BaseHiDreamSetup, ): @@ -192,5 +193,3 @@ def after_optimizer_step( if model.embedding_wrapper_4 is not None: model.embedding_wrapper_4.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HiDreamFineTuneSetup, ModelType.HI_DREAM_FULL, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/HiDreamLoRASetup.py b/modules/modelSetup/HiDreamLoRASetup.py index 3cbcdbe8e..1d97e6acf 100644 --- a/modules/modelSetup/HiDreamLoRASetup.py +++ b/modules/modelSetup/HiDreamLoRASetup.py @@ -16,6 +16,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.LORA) class HiDreamLoRASetup( BaseHiDreamSetup, ): @@ -264,5 +265,3 @@ def after_optimizer_step( if model.embedding_wrapper_4 is not None: model.embedding_wrapper_4.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HiDreamLoRASetup, ModelType.HI_DREAM_FULL, TrainingMethod.LORA) diff --git a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py index 2dca283f8..91022c9f6 100644 --- a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py +++ b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py @@ -12,6 +12,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.EMBEDDING) class HunyuanVideoEmbeddingSetup( BaseHunyuanVideoSetup, ): @@ -110,5 +111,3 @@ def after_optimizer_step( if model.embedding_wrapper_2 is not None: model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HunyuanVideoEmbeddingSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/HunyuanVideoFineTuneSetup.py b/modules/modelSetup/HunyuanVideoFineTuneSetup.py index f9ab6e4b9..9bb26e5d0 100644 --- a/modules/modelSetup/HunyuanVideoFineTuneSetup.py +++ b/modules/modelSetup/HunyuanVideoFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.FINE_TUNE) class HunyuanVideoFineTuneSetup( BaseHunyuanVideoSetup, ): @@ -139,5 +140,3 @@ def after_optimizer_step( if model.embedding_wrapper_2 is not None: model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HunyuanVideoFineTuneSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/HunyuanVideoLoRASetup.py b/modules/modelSetup/HunyuanVideoLoRASetup.py index 0b83b3ad7..e70a5adf4 100644 --- a/modules/modelSetup/HunyuanVideoLoRASetup.py +++ b/modules/modelSetup/HunyuanVideoLoRASetup.py @@ -16,6 +16,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.LORA) class HunyuanVideoLoRASetup( BaseHunyuanVideoSetup, ): @@ -186,5 +187,3 @@ def after_optimizer_step( model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, HunyuanVideoLoRASetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.LORA) diff --git a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py index 5d94c076d..3687a76b6 100644 --- a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py +++ b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py @@ -12,6 +12,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.EMBEDDING) class PixArtAlphaEmbeddingSetup( BasePixArtAlphaSetup, ): @@ -91,6 +93,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, PixArtAlphaEmbeddingSetup, ModelType.PIXART_ALPHA, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, PixArtAlphaEmbeddingSetup, ModelType.PIXART_SIGMA, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index 277f39004..e2bcc6b3d 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -13,6 +13,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.FINE_TUNE) class PixArtAlphaFineTuneSetup( BasePixArtAlphaSetup, ): @@ -119,6 +121,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, PixArtAlphaFineTuneSetup, ModelType.PIXART_ALPHA, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, PixArtAlphaFineTuneSetup, ModelType.PIXART_SIGMA, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 5c865b224..f35fa1c61 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -14,6 +14,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.LORA) class PixArtAlphaLoRASetup( BasePixArtAlphaSetup, ): @@ -139,6 +141,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, PixArtAlphaLoRASetup, ModelType.PIXART_ALPHA, TrainingMethod.LORA) -factory.register(BaseModelSetup, PixArtAlphaLoRASetup, ModelType.PIXART_SIGMA, TrainingMethod.LORA) diff --git a/modules/modelSetup/QwenFineTuneSetup.py b/modules/modelSetup/QwenFineTuneSetup.py index c9bc2cae5..174790e63 100644 --- a/modules/modelSetup/QwenFineTuneSetup.py +++ b/modules/modelSetup/QwenFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.QWEN, TrainingMethod.FINE_TUNE) class QwenFineTuneSetup( BaseQwenSetup, ): @@ -97,5 +98,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, QwenFineTuneSetup, ModelType.QWEN, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/QwenLoRASetup.py b/modules/modelSetup/QwenLoRASetup.py index 4ad8f3eea..ac11c3c19 100644 --- a/modules/modelSetup/QwenLoRASetup.py +++ b/modules/modelSetup/QwenLoRASetup.py @@ -14,6 +14,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.QWEN, TrainingMethod.LORA) class QwenLoRASetup( BaseQwenSetup, ): @@ -126,5 +127,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, QwenLoRASetup, ModelType.QWEN, TrainingMethod.LORA) diff --git a/modules/modelSetup/SanaEmbeddingSetup.py b/modules/modelSetup/SanaEmbeddingSetup.py index 7704c3afd..cfecad2bd 100644 --- a/modules/modelSetup/SanaEmbeddingSetup.py +++ b/modules/modelSetup/SanaEmbeddingSetup.py @@ -12,6 +12,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.EMBEDDING) class SanaEmbeddingSetup( BaseSanaSetup, ): @@ -91,5 +92,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, SanaEmbeddingSetup, ModelType.SANA, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index 11faa747a..eb9217bb9 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.FINE_TUNE) class SanaFineTuneSetup( BaseSanaSetup, ): @@ -113,5 +114,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, SanaFineTuneSetup, ModelType.SANA, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index f481cd445..ad571be82 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -14,6 +14,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.LORA) class SanaLoRASetup( BaseSanaSetup, ): @@ -139,5 +140,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, SanaLoRASetup, ModelType.SANA, TrainingMethod.LORA) diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index 1ff229928..e2319cc00 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -12,6 +12,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.EMBEDDING) class StableDiffusion3EmbeddingSetup( BaseStableDiffusion3Setup, ): @@ -130,6 +132,3 @@ def after_optimizer_step( if model.embedding_wrapper_3 is not None: model.embedding_wrapper_3.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusion3EmbeddingSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusion3EmbeddingSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 3bbe41821..febc77683 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -13,6 +13,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.FINE_TUNE) class StableDiffusion3FineTuneSetup( BaseStableDiffusion3Setup, ): @@ -166,6 +168,3 @@ def after_optimizer_step( if model.embedding_wrapper_3 is not None: model.embedding_wrapper_3.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusion3FineTuneSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusion3FineTuneSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index ae326cbbe..6eea9935b 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -14,6 +14,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.LORA) class StableDiffusion3LoRASetup( BaseStableDiffusion3Setup, ): @@ -224,6 +226,3 @@ def after_optimizer_step( if model.embedding_wrapper_3 is not None: model.embedding_wrapper_3.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusion3LoRASetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusion3LoRASetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.LORA) diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index 6e3a41c2b..deb00c5ff 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -12,6 +12,14 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.EMBEDDING) class StableDiffusionEmbeddingSetup( BaseStableDiffusionSetup, ): @@ -96,12 +104,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionEmbeddingSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index 8c42e59df..386add772 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -13,6 +13,14 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE) class StableDiffusionFineTuneSetup( BaseStableDiffusionSetup, ): @@ -122,12 +130,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionFineTuneSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py b/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py index 295ecddb2..168020b53 100644 --- a/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py @@ -12,6 +12,14 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) class StableDiffusionFineTuneVaeSetup( BaseStableDiffusionSetup, ): @@ -111,12 +119,3 @@ def after_optimizer_step( train_progress: TrainProgress ): pass - -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.FINE_TUNE_VAE) -factory.register(BaseModelSetup, StableDiffusionFineTuneVaeSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.FINE_TUNE_VAE) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index 9b9e9a613..7ba430fb4 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -14,6 +14,14 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.LORA) class StableDiffusionLoRASetup( BaseStableDiffusionSetup, ): @@ -145,12 +153,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_text_encoder_embeddings()) model.embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_20, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_20_BASE, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_20_INPAINTING, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_20_DEPTH, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_21, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionLoRASetup, ModelType.STABLE_DIFFUSION_21_BASE, TrainingMethod.LORA) diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index 7d9ef442f..12eca361c 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -12,6 +12,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.EMBEDDING) class StableDiffusionXLEmbeddingSetup( BaseStableDiffusionXLSetup, ): @@ -109,6 +111,3 @@ def after_optimizer_step( model.embedding_wrapper_1.normalize_embeddings() model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionXLEmbeddingSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, StableDiffusionXLEmbeddingSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index ec6dbe16c..7eafe84d6 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -13,6 +13,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.FINE_TUNE) class StableDiffusionXLFineTuneSetup( BaseStableDiffusionXLSetup, ): @@ -145,6 +147,3 @@ def after_optimizer_step( model.embedding_wrapper_1.normalize_embeddings() model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionXLFineTuneSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, StableDiffusionXLFineTuneSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 6b193edc5..feabf5e78 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -14,6 +14,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.LORA) class StableDiffusionXLLoRASetup( BaseStableDiffusionXLSetup, ): @@ -176,6 +178,3 @@ def after_optimizer_step( model.embedding_wrapper_1.normalize_embeddings() model.embedding_wrapper_2.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, StableDiffusionXLLoRASetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.LORA) -factory.register(BaseModelSetup, StableDiffusionXLLoRASetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.LORA) diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 37dc7f85f..453d07af5 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -12,6 +12,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.EMBEDDING) +@factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.EMBEDDING) class WuerstchenEmbeddingSetup( BaseWuerstchenSetup, ): @@ -105,6 +107,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_prior_text_encoder_embeddings()) model.prior_embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, WuerstchenEmbeddingSetup, ModelType.WUERSTCHEN_2, TrainingMethod.EMBEDDING) -factory.register(BaseModelSetup, WuerstchenEmbeddingSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.EMBEDDING) diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 1291125dc..cbda58510 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -13,6 +13,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.FINE_TUNE) +@factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.FINE_TUNE) class WuerstchenFineTuneSetup( BaseWuerstchenSetup, ): @@ -126,6 +128,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_prior_text_encoder_embeddings()) model.prior_embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, WuerstchenFineTuneSetup, ModelType.WUERSTCHEN_2, TrainingMethod.FINE_TUNE) -factory.register(BaseModelSetup, WuerstchenFineTuneSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index b26f57436..b5f0a3afa 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -14,6 +14,8 @@ import torch +@factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.LORA) +@factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.LORA) class WuerstchenLoRASetup( BaseWuerstchenSetup, ): @@ -153,6 +155,3 @@ def after_optimizer_step( self._normalize_output_embeddings(model.all_prior_text_encoder_embeddings()) model.prior_embedding_wrapper.normalize_embeddings() self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, WuerstchenLoRASetup, ModelType.WUERSTCHEN_2, TrainingMethod.LORA) -factory.register(BaseModelSetup, WuerstchenLoRASetup, ModelType.STABLE_CASCADE_1, TrainingMethod.LORA) diff --git a/modules/modelSetup/ZImageFineTuneSetup.py b/modules/modelSetup/ZImageFineTuneSetup.py index a4c15c2b3..6f2642de7 100644 --- a/modules/modelSetup/ZImageFineTuneSetup.py +++ b/modules/modelSetup/ZImageFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.Z_IMAGE, TrainingMethod.FINE_TUNE) class ZImageFineTuneSetup( BaseZImageSetup, ): @@ -85,5 +86,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, ZImageFineTuneSetup, ModelType.Z_IMAGE, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/ZImageLoRASetup.py b/modules/modelSetup/ZImageLoRASetup.py index 9355a9ad0..85cd789d6 100644 --- a/modules/modelSetup/ZImageLoRASetup.py +++ b/modules/modelSetup/ZImageLoRASetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.Z_IMAGE, TrainingMethod.LORA) class ZImageLoRASetup( BaseZImageSetup, ): @@ -98,5 +99,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, ZImageLoRASetup, ModelType.Z_IMAGE, TrainingMethod.LORA) diff --git a/modules/util/factory.py b/modules/util/factory.py index 0480286e8..08c5c3461 100644 --- a/modules/util/factory.py +++ b/modules/util/factory.py @@ -12,7 +12,7 @@ def get(base_cls, *args, **kwargs): return entry[2] return None -def register(base_cls, cls, *args, **kwargs): +def _do_register(base_cls, cls, *args, **kwargs): if get(base_cls, *args, **kwargs) is not None: raise RuntimeError(f"{cls} already registered as an implementation of {base_cls} with the same criteria {args} {kwargs}") @@ -20,6 +20,18 @@ def register(base_cls, cls, *args, **kwargs): __registry[base_cls] = [] __registry[base_cls].append((args, kwargs, cls)) +def register(base_cls, cls_or_first_key, *args, **kwargs): + if isinstance(cls_or_first_key, type): + # direct call: register(Base, Cls, key1, key2) + _do_register(base_cls, cls_or_first_key, *args, **kwargs) + return cls_or_first_key + else: + # decorator: @register(Base, key1, key2) + def decorator(cls): + _do_register(base_cls, cls, cls_or_first_key, *args, **kwargs) + return cls + return decorator + def import_dir(path: str, parent: str): for _finder, name, _ispkg in pkgutil.walk_packages([path], parent+"."): importlib.import_module(name) From c81e97012fda009feb0549214a27441023449573 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:44:47 +0000 Subject: [PATCH 58/67] wire up Ernie text encoder layer offloading Co-Authored-By: Claude Sonnet 4.6 --- modules/modelSetup/BaseErnieSetup.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..3d630b8cf 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -9,7 +9,10 @@ from modules.modelSetup.mixin.ModelSetupEmbeddingMixin import ModelSetupEmbeddingMixin from modules.modelSetup.mixin.ModelSetupFlowMatchingMixin import ModelSetupFlowMatchingMixin from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin -from modules.util.checkpointing_util import enable_checkpointing_for_ernie_transformer +from modules.util.checkpointing_util import ( + enable_checkpointing_for_ernie_transformer, + enable_checkpointing_for_mistral_encoder_layers, +) from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.enum.TrainingMethod import TrainingMethod @@ -45,6 +48,8 @@ def setup_optimizations( if config.gradient_checkpointing.enabled(): model.transformer_offload_conductor = \ enable_checkpointing_for_ernie_transformer(model.transformer, config) + model.text_encoder_offload_conductor = \ + enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, From 440051ce228c3b37b252c503e261dd9e374c6aa5 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 22:51:56 +0200 Subject: [PATCH 59/67] Force light color scheme on PySide6 UI startup Co-Authored-By: Claude Sonnet 4.6 --- scripts/train_ui_qt.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/train_ui_qt.py b/scripts/train_ui_qt.py index e666733f2..9a7c772b5 100644 --- a/scripts/train_ui_qt.py +++ b/scripts/train_ui_qt.py @@ -10,12 +10,14 @@ from modules.ui.PySide6TrainUIView import PySide6TrainView +from PySide6.QtCore import Qt from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication def main(): app = QApplication(sys.argv) + app.styleHints().setColorScheme(Qt.ColorScheme.Light) palette = app.palette() palette.setColor(QPalette.ColorRole.Base, QColor("white")) From a0bd3e44ae7d6d0ff3ffa155102f9e7cf71b3da1 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:02:10 +0200 Subject: [PATCH 60/67] Sync CtkConceptWindowView and ConceptWindowController with Base copy Co-Authored-By: Claude Sonnet 4.6 --- modules/ui/ConceptWindowController.py | 3 ++- modules/ui/CtkConceptWindowView.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/modules/ui/ConceptWindowController.py b/modules/ui/ConceptWindowController.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/ConceptWindowController.py +++ b/modules/ui/ConceptWindowController.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() diff --git a/modules/ui/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/CtkConceptWindowView.py +++ b/modules/ui/CtkConceptWindowView.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() From fc363ff817692a57c785c06dfc0cdeaf92bc3836 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:49:22 +0200 Subject: [PATCH 61/67] Remove OffloadingWindow: superseded by per-component offloading from #1476 --- modules/ui/BaseOffloadingWindowView.py | 24 ------------------ modules/ui/CtkOffloadingWindowView.py | 31 ----------------------- modules/ui/CtkTrainingTabView.py | 4 --- modules/ui/OffloadingWindowController.py | 6 ----- modules/ui/PySide6OffloadingWindowView.py | 27 -------------------- modules/ui/PySide6TrainingTabView.py | 5 ---- modules/ui/TrainingTabController.py | 4 --- 7 files changed, 101 deletions(-) delete mode 100644 modules/ui/BaseOffloadingWindowView.py delete mode 100644 modules/ui/CtkOffloadingWindowView.py delete mode 100644 modules/ui/OffloadingWindowController.py delete mode 100644 modules/ui/PySide6OffloadingWindowView.py diff --git a/modules/ui/BaseOffloadingWindowView.py b/modules/ui/BaseOffloadingWindowView.py deleted file mode 100644 index 8ec9c0cb1..000000000 --- a/modules/ui/BaseOffloadingWindowView.py +++ /dev/null @@ -1,24 +0,0 @@ -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod - - -class BaseOffloadingWindowView: - def __init__(self, components): - self.components = components - - def build_content(self, frame, controller, ui_state): - self.components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - self.components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], ui_state, - "gradient_checkpointing") - - self.components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - self.components.switch(frame, 1, 1, ui_state, "enable_async_offloading") - - self.components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - self.components.switch(frame, 2, 1, ui_state, "enable_activation_offloading") - - self.components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - self.components.entry(frame, 3, 1, ui_state, "layer_offload_fraction") diff --git a/modules/ui/CtkOffloadingWindowView.py b/modules/ui/CtkOffloadingWindowView.py deleted file mode 100644 index b752f1602..000000000 --- a/modules/ui/CtkOffloadingWindowView.py +++ /dev/null @@ -1,31 +0,0 @@ -from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView -from modules.ui.OffloadingWindowController import OffloadingWindowController -from modules.util.ui import ctk_components -from modules.util.ui.ui_utils import set_window_icon - -import customtkinter as ctk - - -class CtkOffloadingWindowView(BaseOffloadingWindowView, ctk.CTkToplevel): - def __init__(self, parent, controller: OffloadingWindowController, ui_state, *args, **kwargs): - ctk.CTkToplevel.__init__(self, parent, *args, **kwargs) - BaseOffloadingWindowView.__init__(self, ctk_components) - - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = ctk.CTkScrollableFrame(self, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - self.build_content(frame, controller, ui_state) - frame.pack(fill="both", expand=1) - frame.grid(row=0, column=0, sticky='nsew') - self.components.button(self, 1, 0, "ok", self.destroy) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) diff --git a/modules/ui/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py index bce309876..a126b2fa4 100644 --- a/modules/ui/CtkTrainingTabView.py +++ b/modules/ui/CtkTrainingTabView.py @@ -1,6 +1,5 @@ from modules.ui.BaseTrainingTabView import BaseTrainingTabView -from modules.ui.CtkOffloadingWindowView import CtkOffloadingWindowView from modules.ui.CtkOptimizerParamsWindowView import CtkOptimizerParamsWindowView from modules.ui.CtkSchedulerParamsWindowView import CtkSchedulerParamsWindowView from modules.ui.CtkTimestepDistributionWindowView import CtkTimestepDistributionWindowView @@ -64,8 +63,5 @@ def restore_scheduler(self, variable: str): def open_scheduler_params(self): self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) - def open_offloading(self): - self.master.wait_window(self.controller.open_offloading_window(self.master, self.ui_state, CtkOffloadingWindowView)) - def open_timestep_distribution(self): self.master.wait_window(self.controller.open_timestep_distribution_window(self.master, self.ui_state, CtkTimestepDistributionWindowView)) diff --git a/modules/ui/OffloadingWindowController.py b/modules/ui/OffloadingWindowController.py deleted file mode 100644 index cc9bad142..000000000 --- a/modules/ui/OffloadingWindowController.py +++ /dev/null @@ -1,6 +0,0 @@ -from modules.util.config.TrainConfig import TrainConfig - - -class OffloadingWindowController: - def __init__(self, config: TrainConfig): - self.config = config diff --git a/modules/ui/PySide6OffloadingWindowView.py b/modules/ui/PySide6OffloadingWindowView.py deleted file mode 100644 index fcc9d1f4a..000000000 --- a/modules/ui/PySide6OffloadingWindowView.py +++ /dev/null @@ -1,27 +0,0 @@ -from modules.ui.BaseOffloadingWindowView import BaseOffloadingWindowView -from modules.ui.OffloadingWindowController import OffloadingWindowController -from modules.util.ui import pyside6_components - -from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton - - -class PySide6OffloadingWindowView(BaseOffloadingWindowView, QDialog): - def __init__(self, parent, controller: OffloadingWindowController, ui_state): - QDialog.__init__(self, parent) - BaseOffloadingWindowView.__init__(self, pyside6_components) - - self.setWindowTitle("Offloading") - self.resize(800, 400) - - outer = QGridLayout(self) - outer.setRowStretch(0, 1) - - scroll, frame = pyside6_components.scrollable_frame(self) - pyside6_components._layout(frame).setColumnStretch(0, 1) - pyside6_components._layout(frame).setColumnStretch(1, 1) - self.build_content(frame, controller, ui_state) - outer.addWidget(scroll, 0, 0) - - ok = QPushButton("ok", self) - ok.clicked.connect(self.accept) - outer.addWidget(ok, 1, 0) diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py index 0311df6a9..d115c678d 100644 --- a/modules/ui/PySide6TrainingTabView.py +++ b/modules/ui/PySide6TrainingTabView.py @@ -1,7 +1,5 @@ from modules.ui.BaseTrainingTabView import BaseTrainingTabView -from modules.ui.OffloadingWindowController import OffloadingWindowController from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController -from modules.ui.PySide6OffloadingWindowView import PySide6OffloadingWindowView from modules.ui.PySide6OptimizerParamsWindowView import PySide6OptimizerParamsWindowView from modules.ui.PySide6SchedulerParamsWindowView import PySide6SchedulerParamsWindowView from modules.ui.PySide6TimestepDistributionWindowView import PySide6TimestepDistributionWindowView @@ -82,8 +80,5 @@ def open_optimizer_params(self): def open_scheduler_params(self): PySide6SchedulerParamsWindowView(self, SchedulerParamsWindowController(self.controller.config), self.ui_state).exec() - def open_offloading(self): - PySide6OffloadingWindowView(self, OffloadingWindowController(self.controller.config), self.ui_state).exec() - def open_timestep_distribution(self): PySide6TimestepDistributionWindowView(self, TimestepDistributionWindowController(self.controller.config), self.ui_state).exec() diff --git a/modules/ui/TrainingTabController.py b/modules/ui/TrainingTabController.py index 858ca0701..013c4d377 100644 --- a/modules/ui/TrainingTabController.py +++ b/modules/ui/TrainingTabController.py @@ -1,5 +1,4 @@ -from modules.ui.OffloadingWindowController import OffloadingWindowController from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController from modules.ui.SchedulerParamsWindowController import SchedulerParamsWindowController from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController @@ -34,6 +33,3 @@ def open_scheduler_params_window(self, parent, ui_state, view_cls): def open_timestep_distribution_window(self, parent, ui_state, view_cls): return view_cls(parent, TimestepDistributionWindowController(self.config), ui_state) - - def open_offloading_window(self, parent, ui_state, view_cls): - return view_cls(parent, OffloadingWindowController(self.config), ui_state) From 5de0bc263309c444943123c121a915fbf9029552 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 22:26:27 +0200 Subject: [PATCH 62/67] Add Windows-only dark mode support with Qt6 tab rendering fixes Qt's dark mode support on Windows is broken (tabs render with no contrast). apply_theme() applies explicit palette and QSS fixes on Windows only, using palette() references so it adapts to the system light/dark preference. Linux/macOS are left to Qt's native handling. Replaces the previous force-light-mode workaround in train_ui_qt.py. --- modules/util/ui/theme.py | 53 ++++++++++++++++++++++++++++++++++++++++ scripts/train_ui_qt.py | 23 ++--------------- 2 files changed, 55 insertions(+), 21 deletions(-) create mode 100644 modules/util/ui/theme.py diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py new file mode 100644 index 000000000..adc7c38b0 --- /dev/null +++ b/modules/util/ui/theme.py @@ -0,0 +1,53 @@ +import platform + +from PySide6.QtGui import QColor, QPalette +from PySide6.QtWidgets import QApplication + +# Single stylesheet using palette() references so it works in both light and dark mode +# without hardcoding any colors. +_STYLESHEET = """ + QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { + padding: 2px 2px; + } + QCheckBox::indicator { + width: 16px; + height: 16px; + } + QProgressBar { + background-color: palette(mid); + } + QTabWidget::pane { + border: 1px solid palette(mid); + } + QTabBar::tab { + background-color: palette(window); + color: palette(window-text); + border: 1px solid palette(mid); + border-bottom: none; + padding: 4px 8px; + } + QTabBar::tab:selected { + background-color: palette(highlight); + color: palette(highlighted-text); + } + QTabBar::tab:!selected:hover { + background-color: palette(button); + } +""" + + +def apply_theme(app: QApplication) -> None: + # On Linux/macOS, Qt already maps the system light/dark preference onto its + # palette correctly, so we leave it untouched. Qt's dark mode support on + # Windows is incomplete (tabs render with no contrast), so we apply explicit + # fixes there only. + if platform.system() != "Windows": + return + is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 + if not is_dark: + # Qt's light palette leaves Base as grey which looks disabled — restore to white. + palette = app.palette() + palette.setColor(QPalette.ColorRole.Base, QColor("white")) + palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) + app.setPalette(palette) + app.setStyleSheet(_STYLESHEET) diff --git a/scripts/train_ui_qt.py b/scripts/train_ui_qt.py index 9a7c772b5..1d49d3fd1 100644 --- a/scripts/train_ui_qt.py +++ b/scripts/train_ui_qt.py @@ -9,33 +9,14 @@ script_imports() from modules.ui.PySide6TrainUIView import PySide6TrainView +from modules.util.ui.theme import apply_theme -from PySide6.QtCore import Qt -from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication def main(): app = QApplication(sys.argv) - app.styleHints().setColorScheme(Qt.ColorScheme.Light) - - palette = app.palette() - palette.setColor(QPalette.ColorRole.Base, QColor("white")) - palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) - app.setPalette(palette) - - app.setStyleSheet(""" - QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { - padding: 2px 2px; - } - QCheckBox::indicator { - width: 16px; - height: 16px; - } - QProgressBar { - background-color: #c8c8c8; - } - """) + apply_theme(app) window = PySide6TrainView() window.show() sys.exit(app.exec()) From 81dfba6a6cde62a5f36c5a978281edbd3990f793 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 22:28:10 +0200 Subject: [PATCH 63/67] Trim theme comment --- modules/util/ui/theme.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py index adc7c38b0..7e7a733c9 100644 --- a/modules/util/ui/theme.py +++ b/modules/util/ui/theme.py @@ -37,10 +37,8 @@ def apply_theme(app: QApplication) -> None: - # On Linux/macOS, Qt already maps the system light/dark preference onto its - # palette correctly, so we leave it untouched. Qt's dark mode support on - # Windows is incomplete (tabs render with no contrast), so we apply explicit - # fixes there only. + # Qt's dark mode is broken on Windows (tabs render with no contrast); fix it there + # only — Linux/macOS map the system theme onto the palette correctly already. if platform.system() != "Windows": return is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 From 3a35ced5113cda9679e3515a67fd5c35240bef8e Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 22:30:22 +0200 Subject: [PATCH 64/67] Keep original light-mode stylesheet, add dark mode separately --- modules/util/ui/theme.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py index 7e7a733c9..9feee24d1 100644 --- a/modules/util/ui/theme.py +++ b/modules/util/ui/theme.py @@ -3,9 +3,23 @@ from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication -# Single stylesheet using palette() references so it works in both light and dark mode -# without hardcoding any colors. -_STYLESHEET = """ +# Light mode keeps the original stylesheet unchanged. +_LIGHT_STYLESHEET = """ + QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { + padding: 2px 2px; + } + QCheckBox::indicator { + width: 16px; + height: 16px; + } + QProgressBar { + background-color: #c8c8c8; + } +""" + +# Dark mode adds tab styling (Qt renders tabs with no contrast on Windows) and uses +# palette() references so colors follow the system theme. +_DARK_STYLESHEET = """ QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { padding: 2px 2px; } @@ -42,10 +56,11 @@ def apply_theme(app: QApplication) -> None: if platform.system() != "Windows": return is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 - if not is_dark: - # Qt's light palette leaves Base as grey which looks disabled — restore to white. + if is_dark: + app.setStyleSheet(_DARK_STYLESHEET) + else: palette = app.palette() palette.setColor(QPalette.ColorRole.Base, QColor("white")) palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) app.setPalette(palette) - app.setStyleSheet(_STYLESHEET) + app.setStyleSheet(_LIGHT_STYLESHEET) From feba96b5d4e192283108d4d7b50a05c835758464 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 22:49:15 +0200 Subject: [PATCH 65/67] Apply Windows dark mode fixes as overrides on top of base stylesheet --- modules/util/ui/theme.py | 37 ++++++++++++------------------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py index 9feee24d1..4c6bc7910 100644 --- a/modules/util/ui/theme.py +++ b/modules/util/ui/theme.py @@ -3,8 +3,7 @@ from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication -# Light mode keeps the original stylesheet unchanged. -_LIGHT_STYLESHEET = """ +_BASE_STYLESHEET = """ QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { padding: 2px 2px; } @@ -17,16 +16,7 @@ } """ -# Dark mode adds tab styling (Qt renders tabs with no contrast on Windows) and uses -# palette() references so colors follow the system theme. -_DARK_STYLESHEET = """ - QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { - padding: 2px 2px; - } - QCheckBox::indicator { - width: 16px; - height: 16px; - } +_WINDOWS_OVERRIDES = """ QProgressBar { background-color: palette(mid); } @@ -51,16 +41,13 @@ def apply_theme(app: QApplication) -> None: - # Qt's dark mode is broken on Windows (tabs render with no contrast); fix it there - # only — Linux/macOS map the system theme onto the palette correctly already. - if platform.system() != "Windows": - return - is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 - if is_dark: - app.setStyleSheet(_DARK_STYLESHEET) - else: - palette = app.palette() - palette.setColor(QPalette.ColorRole.Base, QColor("white")) - palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) - app.setPalette(palette) - app.setStyleSheet(_LIGHT_STYLESHEET) + stylesheet = _BASE_STYLESHEET + if platform.system() == "Windows": + is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 + if not is_dark: + palette = app.palette() + palette.setColor(QPalette.ColorRole.Base, QColor("white")) + palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) + app.setPalette(palette) + stylesheet += _WINDOWS_OVERRIDES + app.setStyleSheet(stylesheet) From d876a93d1c684f525abb23883861ba5c88908be3 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 22:56:48 +0200 Subject: [PATCH 66/67] Preserve forced light mode on non-Windows platforms --- modules/util/ui/theme.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py index 4c6bc7910..28c64f5df 100644 --- a/modules/util/ui/theme.py +++ b/modules/util/ui/theme.py @@ -1,5 +1,6 @@ import platform +from PySide6.QtCore import Qt from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication @@ -40,14 +41,21 @@ """ +def _apply_light_base(app: QApplication) -> None: + palette = app.palette() + palette.setColor(QPalette.ColorRole.Base, QColor("white")) + palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) + app.setPalette(palette) + + def apply_theme(app: QApplication) -> None: - stylesheet = _BASE_STYLESHEET - if platform.system() == "Windows": - is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 - if not is_dark: - palette = app.palette() - palette.setColor(QPalette.ColorRole.Base, QColor("white")) - palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) - app.setPalette(palette) - stylesheet += _WINDOWS_OVERRIDES - app.setStyleSheet(stylesheet) + if platform.system() != "Windows": + app.styleHints().setColorScheme(Qt.ColorScheme.Light) + _apply_light_base(app) + app.setStyleSheet(_BASE_STYLESHEET) + return + + is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 + if not is_dark: + _apply_light_base(app) + app.setStyleSheet(_BASE_STYLESHEET + _WINDOWS_OVERRIDES) From eb7be99b9dcf3742e4adca7d4eec3d80135bbfc7 Mon Sep 17 00:00:00 2001 From: Gabriel El Faouzi Date: Fri, 5 Jun 2026 23:13:28 +0200 Subject: [PATCH 67/67] Simplify theme: drop QSS overrides, fix Windows dark by not forcing light palette --- modules/util/ui/theme.py | 50 ++++++++-------------------------------- 1 file changed, 10 insertions(+), 40 deletions(-) diff --git a/modules/util/ui/theme.py b/modules/util/ui/theme.py index 28c64f5df..c9dc7b1b7 100644 --- a/modules/util/ui/theme.py +++ b/modules/util/ui/theme.py @@ -4,6 +4,8 @@ from PySide6.QtGui import QColor, QPalette from PySide6.QtWidgets import QApplication +IS_WINDOWS = platform.system() == "Windows" + _BASE_STYLESHEET = """ QLineEdit, QSpinBox, QDoubleSpinBox, QTextEdit, QPlainTextEdit { padding: 2px 2px; @@ -17,45 +19,13 @@ } """ -_WINDOWS_OVERRIDES = """ - QProgressBar { - background-color: palette(mid); - } - QTabWidget::pane { - border: 1px solid palette(mid); - } - QTabBar::tab { - background-color: palette(window); - color: palette(window-text); - border: 1px solid palette(mid); - border-bottom: none; - padding: 4px 8px; - } - QTabBar::tab:selected { - background-color: palette(highlight); - color: palette(highlighted-text); - } - QTabBar::tab:!selected:hover { - background-color: palette(button); - } -""" - - -def _apply_light_base(app: QApplication) -> None: - palette = app.palette() - palette.setColor(QPalette.ColorRole.Base, QColor("white")) - palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) - app.setPalette(palette) - - def apply_theme(app: QApplication) -> None: - if platform.system() != "Windows": + is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 + palette = app.palette() + if not IS_WINDOWS or not is_dark: app.styleHints().setColorScheme(Qt.ColorScheme.Light) - _apply_light_base(app) - app.setStyleSheet(_BASE_STYLESHEET) - return - - is_dark = app.palette().color(QPalette.ColorRole.Window).lightness() < 128 - if not is_dark: - _apply_light_base(app) - app.setStyleSheet(_BASE_STYLESHEET + _WINDOWS_OVERRIDES) + palette = app.palette() + palette.setColor(QPalette.ColorRole.Base, QColor("white")) + palette.setColor(QPalette.ColorGroup.Disabled, QPalette.ColorRole.Base, QColor("#e0e0e0")) + app.setPalette(palette) + app.setStyleSheet(_BASE_STYLESHEET)