diff --git a/.gitignore b/.gitignore index 0b3e5887d..df075f9ea 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ # development +PLAN.md # claude's planning file, kept local-only .idea *.bak *.pyc @@ -38,3 +39,5 @@ pixi.toml train.bat debug_report.log config_diff.txt +CLAUDE.md +PLAN.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..380ca6c56 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,162 @@ +# 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 + +Detailed implementation guides live in [`docs/recipes/`](docs/recipes/): + +- **New optimizer** — see [`docs/recipes/AddOptimizer.md`](docs/recipes/AddOptimizer.md). + +### 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 diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..49afc658c --- /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)**. 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). 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/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py index 94722bdf5..c8c5bf07b 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, @@ -50,7 +51,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 +81,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 +111,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 @@ -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 edd634e79..a6cfb8d4e 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -21,6 +21,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.ERNIE) class ErnieBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -64,7 +65,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): @@ -130,6 +131,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 0aa4522d7..56a278015 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -29,6 +29,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.FLUX_2) class Flux2BaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -85,7 +86,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): @@ -156,6 +157,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..91854791a 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, @@ -103,7 +105,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): @@ -181,6 +186,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..b57f07be4 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, @@ -132,10 +133,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): @@ -219,5 +222,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..e35decbdf 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, @@ -97,7 +98,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): @@ -183,5 +187,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..22febd256 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, @@ -85,7 +87,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): @@ -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..2ccf655b6 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, @@ -54,7 +55,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 +85,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 +115,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 @@ -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..159c44775 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, @@ -78,7 +79,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): @@ -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..76759ca84 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, @@ -117,7 +118,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): @@ -198,5 +203,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..74c8dc137 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, @@ -86,7 +94,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): @@ -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..e98ccbd02 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, @@ -197,7 +205,7 @@ def before_cache_fun(): modules = [] - if config.latent_caching: + if config.image_caching: modules.append(disk_cache) modules.append(variation_sorting) @@ -221,7 +229,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) @@ -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..21c58f230 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, @@ -99,7 +101,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): @@ -176,5 +181,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..0af019f97 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, @@ -84,7 +86,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 ) @@ -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 66cfea16f..8b0cf8ef9 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -25,6 +25,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.Z_IMAGE) class ZImageBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -48,7 +49,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 @@ -77,7 +78,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): @@ -106,7 +107,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 @@ -150,5 +151,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/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/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/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/model/ChromaModel.py b/modules/model/ChromaModel.py index 59967c8dc..3fc96d8f9 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -113,8 +113,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) @@ -123,8 +122,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 b981865c4..b790335ac 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -149,8 +149,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) @@ -159,8 +158,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 049b0e6d3..9e9684851 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 2107e55d7..bfc07ca95 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 466cc61f9..42c6621f8 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.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/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 9e8008219..c1d563cb2 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -116,8 +116,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) @@ -126,8 +125,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 8f6cf5818..a2bae894d 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -180,8 +180,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) @@ -190,8 +189,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/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/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 dbafdae81..ee7569ccf 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 020df6a39..c12056c92 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/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..06c21323c --- /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.image_caching + + model.text_encoder_to(self.temp_device if config.text_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..0f15b5082 --- /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.image_caching + + model.text_encoder_to(self.temp_device if config.text_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..99bf2cf1f --- /dev/null +++ b/modules/modelSetup/BaseAnimaSetup.py @@ -0,0 +1,176 @@ +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.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.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_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.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.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/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..c3e02fdd0 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 @@ -49,31 +48,19 @@ 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) - - 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) + 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.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, ) @@ -185,6 +172,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 + ) + latent_image = batch['latent_image'] scaled_latent_image = (latent_image - vae_shift_factor) * vae_scaling_factor diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..064b4eb4c 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -9,10 +9,12 @@ 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 from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -42,26 +44,19 @@ 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) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_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, - 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, ) @@ -94,6 +89,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] diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..329e5a506 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 @@ -45,33 +44,22 @@ 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) - - 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) + 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.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, ) @@ -104,6 +92,11 @@ 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_gamma > 0 and not deterministic: + 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] diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..bd3a869d9 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 @@ -49,34 +48,21 @@ 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) - - 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) + 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.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, ) @@ -231,6 +217,14 @@ def predict( apply_attention_mask=config.transformer.attention_mask, ) + if config.cep_gamma > 0 and not deterministic: + 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 17fbcc0d6..d9c3f8e15 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 @@ -49,41 +48,25 @@ 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) - - 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) + 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.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 +75,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, ) @@ -332,6 +310,17 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) + 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 + ) + 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 b072bf4ba..cb2de2def 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 @@ -49,34 +48,21 @@ 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) - - 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) + 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.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, ) @@ -229,6 +215,14 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.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 + ) + 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 3069b4884..2195598c7 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 @@ -51,30 +50,18 @@ 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) - - 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) + 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.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, ) @@ -178,6 +165,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 + ) + 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 a8a7be8f6..d7ff44342 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 @@ -46,29 +45,19 @@ 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) - - 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) + 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.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, ) @@ -102,6 +91,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 + ) + 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 0c9ea6da0..f372df83f 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 @@ -52,30 +51,18 @@ 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) - - 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) + 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.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 +70,6 @@ def setup_optimizations( self.train_device, config.train_dtype, config.fallback_train_dtype, - [ - config.weight_dtypes().vae, - ], config.enable_autocast_cache, ) @@ -188,6 +172,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 + ) + 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 de5dc04e8..ff4404d8a 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 @@ -48,37 +47,23 @@ 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) - - 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) + 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.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, ) @@ -284,6 +269,14 @@ def predict( apply_attention_mask=config.transformer.attention_mask, )) + if config.cep_gamma > 0 and not deterministic: + 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 8cf63ac07..c5a8f0475 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 @@ -51,11 +50,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) @@ -63,13 +64,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) @@ -169,6 +165,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 + ) + 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 61cb1e457..f8d55f3bd 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 @@ -48,11 +47,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) @@ -60,22 +61,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, ) @@ -220,6 +212,14 @@ def predict( text_encoder_2_dropout_probability=config.text_encoder_2.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 + ) + 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 48cd05942..13f243cca 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 @@ -57,13 +56,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) @@ -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: @@ -250,6 +237,14 @@ 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_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 b822d7304..bd824f749 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 @@ -47,19 +46,13 @@ 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) - - 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) + 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.enable_autocast_cache) #TODO necessary if we don't train it? model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ @@ -67,10 +60,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, ) @@ -102,6 +91,12 @@ 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_gamma > 0 and not deterministic: + 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/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 63a6d75af..1b1234824 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, ): @@ -73,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) @@ -95,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 c78e4e700..22c046b61 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, ): @@ -82,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) @@ -115,5 +116,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 514eca8e8..09f3e1166 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, ): @@ -110,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) @@ -143,5 +144,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..b8d73da85 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, ): @@ -63,8 +64,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) @@ -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..abf628776 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, ): @@ -73,8 +74,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) @@ -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..36703e79c 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, ): @@ -63,8 +64,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) @@ -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..b07afe346 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, ): @@ -75,8 +76,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) @@ -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 98dd12579..e29c91424 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, ): @@ -83,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) @@ -110,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 3afbdf31c..40b015227 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, ): @@ -92,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) @@ -138,6 +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, 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 e6204565f..78799fc5f 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, ): @@ -134,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) @@ -180,6 +182,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..04d57b389 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, ): @@ -103,12 +104,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) @@ -141,5 +146,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..8e95bf2dc 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, ): @@ -119,22 +120,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) @@ -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..8c02db5ad 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, ): @@ -191,22 +192,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) @@ -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..f1e1bf528 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, ): @@ -83,10 +84,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) @@ -110,5 +113,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..eed5807c3 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, ): @@ -93,14 +94,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) @@ -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..178d9c228 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, ): @@ -139,14 +140,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) @@ -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 5b842e267..232bf2e01 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, ): @@ -90,6 +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, 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 31b92378e..8f0cdd91b 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, ): @@ -86,11 +88,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) @@ -118,6 +120,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 1e18dc673..e9e41a4da 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, ): @@ -106,11 +108,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) @@ -138,6 +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, 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..1effeb428 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, ): @@ -68,10 +69,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) @@ -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..a6583c25f 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, ): @@ -97,10 +98,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) @@ -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 c31c2aa11..fad567b75 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, ): @@ -90,5 +91,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 1c334c44f..dbe82d1d4 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, ): @@ -80,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) @@ -112,5 +113,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 45cb2f7bb..9041bf622 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, ): @@ -106,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) @@ -138,5 +139,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 e9209bdc4..32b5644c9 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, ): @@ -93,11 +95,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) @@ -127,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 3ce400a59..7bdc4e913 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, ): @@ -102,18 +104,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) @@ -163,6 +165,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 0d802c18a..46dd90dd3 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, ): @@ -160,18 +162,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) @@ -221,6 +223,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 f8cc956f2..e67ba039a 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, ): @@ -74,7 +82,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) @@ -95,12 +103,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 9f6fff8f3..17e20834c 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, ): @@ -88,11 +96,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) @@ -121,12 +129,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 d82147611..b51808bc6 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, ): @@ -111,11 +119,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) @@ -144,12 +152,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 62f78cba2..93c5a4ea0 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, ): @@ -83,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) @@ -107,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 56a5bfe40..6849ae4cc 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, ): @@ -98,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) @@ -143,6 +145,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 a0db29b07..3f7584183 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, ): @@ -132,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) @@ -174,6 +176,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 d5dd8a83b..0afe3be87 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, ): @@ -74,7 +76,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) @@ -104,6 +106,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 111ce016f..6f455be0e 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, ): @@ -83,7 +85,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) @@ -94,7 +96,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) @@ -125,6 +127,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 812bb43e6..a9745f6dd 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, ): @@ -110,7 +112,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) @@ -121,7 +123,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) @@ -152,6 +154,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..654abe048 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, ): @@ -63,8 +64,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) @@ -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..822196b08 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, ): @@ -76,8 +77,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) @@ -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/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 680e13df9..f2dac071d 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py @@ -119,6 +119,41 @@ def _create_noise( return noise + @staticmethod + def _apply_conditional_embedding_perturbation( + embedding: Tensor | list, + gamma: float, + generator: Generator + ) -> Tensor | list: + """ + 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)) + """ + 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: gamma / sqrt(d) + scale = gamma / math.sqrt(d) + + # 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 + + if isinstance(embedding, list): + return [_perturb_cep(emb) for emb in embedding] + else: + return _perturb_cep(embedding) + def _get_timestep_discrete( self, num_train_timesteps: int, diff --git a/modules/module/quantized/LinearA8.py b/modules/module/quantized/LinearA8.py new file mode 100644 index 000000000..17a76ee38 --- /dev/null +++ b/modules/module/quantized/LinearA8.py @@ -0,0 +1,122 @@ +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 LinearIntA8Function(torch.autograd.Function): + @staticmethod + 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, compute_dtype) + + @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, None + +class LinearFpA8Function(torch.autograd.Function): + @staticmethod + 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, compute_dtype) + + @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, None + +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 = LinearIntA8Function.apply(x, self.weight, self.bias, self.compute_dtype) + else: + 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/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 605994e57..24173b367 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -30,6 +30,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 @@ -85,7 +86,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 +643,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: @@ -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/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/AdditionalEmbeddingsTab.py b/modules/ui/AdditionalEmbeddingsTab.py deleted file mode 100644 index 6a5e3fbe7..000000000 --- a/modules/ui/AdditionalEmbeddingsTab.py +++ /dev/null @@ -1,136 +0,0 @@ - -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/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 new file mode 100644 index 000000000..744f0abdc --- /dev/null +++ b/modules/ui/BaseAdditionalEmbeddingsTabView.py @@ -0,0 +1,75 @@ + +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util import path_util + +import customtkinter as ctk + + +class BaseAdditionalEmbeddingsTabView(BaseConfigListView): + + def refresh_ui(self): + if self.element_list is not None: + self._destroy_frame(self.element_list) + self.element_list = None + self.widgets_initialized = False + self._create_element_list() + + def open_element_window(self, i, ui_state) -> ctk.CTkToplevel: + pass + + +class BaseEmbeddingWidgetView: + + def __init__(self, components): + self.components = components + + 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 + + # close button + self.components.colored_icon_button(top_frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) + + # clone button + self.components.colored_icon_button(top_frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i, controller.randomize_uuid), padx=5) + + # embedding model names + 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=path_util.json_path_modifier + ) + + # 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 + 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 + 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 + 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 + 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 + 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 diff --git a/modules/ui/BaseCaptionUIView.py b/modules/ui/BaseCaptionUIView.py new file mode 100644 index 000000000..c53ca84d8 --- /dev/null +++ b/modules/ui/BaseCaptionUIView.py @@ -0,0 +1,51 @@ +import platform +from abc import ABC, abstractmethod + + +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") + self.components.button(frame, 0, 1, "Generate Masks", self.open_mask_window, + tooltip="open a dialog to automatically generate masks") + self.components.button(frame, 0, 2, "Generate Captions", self.open_caption_window, + tooltip="open a dialog to automatically generate captions") + + if platform.system() == "Windows": + self.components.button(frame, 0, 3, "Open in Explorer", self.open_in_explorer, + tooltip="open the current image in Explorer") + + self.components.switch(frame, 0, 4, ui_state, "include_subdirectories", + text="include subdirectories") + + frame.grid_columnconfigure(5, weight=1) + + self.components.button(frame, 0, 6, "Help", controller.print_help, + tooltip=controller.help_text) + + 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") + self.components.button(right_frame, 0, 1, "Fill", self.fill_mask_editing_mode, + tooltip="draw a mask using a fill tool") diff --git a/modules/ui/BaseCloudTabView.py b/modules/ui/BaseCloudTabView.py new file mode 100644 index 000000000..3f20b5674 --- /dev/null +++ b/modules/ui/BaseCloudTabView.py @@ -0,0 +1,185 @@ + +from abc import ABC, abstractmethod + +from modules.util.enum.CloudAction import CloudAction +from modules.util.enum.CloudFileSync import CloudFileSync +from modules.util.enum.CloudType import CloudType + + +class BaseCloudTabView(ABC): + def __init__(self, components, controller): + self.components = components + self.controller = controller + + @property + def reattach(self): + return self.controller.reattach + + @abstractmethod + def _make_reattach_frame(self, frame): pass + + @abstractmethod + def _make_create_frame(self, frame): pass + + @abstractmethod + def _on_set_gpu_types(self): pass + + def build_content(self, frame, controller, ui_state): + self.components.label(frame, 0, 0, "Enabled", + tooltip="Enable cloud training") + self.components.switch(frame, 0, 1, ui_state, "cloud.enabled") + + 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.") + self.components.options_kv(frame, 1, 1, [ + ("RUNPOD", CloudType.RUNPOD), + ("LINUX", CloudType.LINUX), + ], ui_state, "cloud.type") + + 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.") + self.components.options_kv(frame, 2, 1, [ + ("NATIVE_SCP", CloudFileSync.NATIVE_SCP), + ("FABRIC_SFTP", CloudFileSync.FABRIC_SFTP), + ], ui_state, "cloud.file_sync") + + 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. ") + self.components.entry(frame, 3, 1, ui_state, "secrets.cloud.api_key") + + 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.") + self.components.entry(frame, 4, 1, ui_state, "secrets.cloud.host") + + 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.") + self.components.entry(frame, 5, 1, ui_state, "secrets.cloud.port") + + 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.') + self.components.entry(frame, 6, 1, ui_state, "secrets.cloud.user") + + 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.") + self.components.path_entry(frame, 7, 1, ui_state, "secrets.cloud.key_file", mode="file") + + 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.") + self.components.entry(frame, 8, 1, ui_state, "secrets.cloud.password") + + 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.") + self.components.entry(frame, 9, 1, ui_state, "secrets.cloud.id") + + 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") + self.components.switch(frame, 10, 1, ui_state, "cloud.tensorboard_tunnel") + + self.components.label(frame, 1, 2, "Remote Directory", + tooltip="The directory on the cloud where files will be uploaded and downloaded.") + 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.") + 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.") + 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.") + 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.") + 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.") + self.components.switch(frame, 6, 3, ui_state, "cloud.update_onetrainer") + + 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.") + 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 = 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.") + 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.") + 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.") + 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.") + 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.") + 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.") + self.components.switch(frame, 16, 3, ui_state, "cloud.delete_workspace") + + 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 = 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.") + 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.") + self.components.options_kv(frame, 3, 5, [ + ("", ""), + ("Community", "COMMUNITY"), + ("Secure", "SECURE"), + ], ui_state, "cloud.sub_type") + + 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'] + + 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") + self.components.entry(frame, 5, 5, ui_state, "cloud.volume_size") + + self.components.label(frame, 6, 4, "Min download", + tooltip="Set the minimum download speed of the cloud in Mbps.") + self.components.entry(frame, 6, 5, ui_state, "cloud.min_download") + + 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.") + self.components.options_kv(frame, 8, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], ui_state, "cloud.on_finish") + + 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.") + self.components.options_kv(frame, 9, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], ui_state, "cloud.on_error") + + 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.") + self.components.options_kv(frame, 10, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], ui_state, "cloud.on_detached_finish") + + 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.") + self.components.options_kv(frame, 11, 5, [ + ("None", CloudAction.NONE), + ("Stop", CloudAction.STOP), + ("Delete", CloudAction.DELETE), + ], ui_state, "cloud.on_detached_error") diff --git a/modules/ui/BaseConceptTabView.py b/modules/ui/BaseConceptTabView.py new file mode 100644 index 000000000..6a55e515b --- /dev/null +++ b/modules/ui/BaseConceptTabView.py @@ -0,0 +1,104 @@ +import os +import pathlib + +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 + +from PIL import Image + + +class BaseConceptTabView(BaseConfigListView): + + _FILTER_TYPES = ["ALL", "STANDARD", "VALIDATION", "PRIOR_PREDICTION"] + + def _element_matches_filters(self, element): + if not self.filters.get("show_disabled", True): + if hasattr(element, 'enabled') and not element.enabled: + return False + + 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 = 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 + + + +class BaseConceptWidgetView: + + def __init__(self, components, concept: ConceptConfig): + self.components = components + self.concept = 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 _get_preview_image(self): + preview_path = "resources/icons/icon.png" + glob_pattern = "**/*.*" if getattr(self.concept, 'include_subdirectories', False) else "*.*" + + 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): + 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 _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 new file mode 100644 index 000000000..0b94e50d1 --- /dev/null +++ b/modules/ui/BaseConceptWindowView.py @@ -0,0 +1,446 @@ +import fractions +import math + +from modules.util import path_util +from modules.util.enum.BalancingStrategy import BalancingStrategy +from modules.util.enum.ConceptType import ConceptType + + +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 + self.components.label(frame, 0, 0, "Name", + tooltip="Name of the concept") + self.components.entry(frame, 0, 1, ui_state, "name") + + # enabled + self.components.label(frame, 1, 0, "Enabled", + tooltip="Enable or disable this concept") + self.components.switch(frame, 1, 1, ui_state, "enabled") + + # 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) + self.components.options(frame, 2, 1, [str(x) for x in list(ConceptType)], ui_state, "type") + + # path + self.components.label(frame, 3, 0, "Path", + tooltip="Path where the training data is located") + 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 + 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 = self.components.path_entry(frame, 4, 2, text_ui_state, "prompt_path", mode="file") + + def set_prompt_path_entry_enabled(option: str): + self.components.set_widget_enabled(prompt_path_entry, option == 'concept') + + self.components.options_kv(frame, 4, 1, [ + ("From text file per sample", 'sample'), + ("From single text file", 'concept'), + ("From image file name", 'filename'), + ], text_ui_state, "prompt_source", command=set_prompt_path_entry_enabled) + set_prompt_path_entry_enabled(controller.concept.text.prompt_source) + + # include subdirectories + self.components.label(frame, 5, 0, "Include Subdirectories", + tooltip="Includes images from subdirectories into the dataset") + self.components.switch(frame, 5, 1, ui_state, "include_subdirectories") + + # image variations + self.components.label(frame, 6, 0, "Image Variations", + tooltip="The number of different image versions to cache if latent caching is enabled.") + self.components.entry(frame, 6, 1, ui_state, "image_variations") + + # text variations + self.components.label(frame, 7, 0, "Text Variations", + tooltip="The number of different text versions to cache if latent caching is enabled.") + self.components.entry(frame, 7, 1, ui_state, "text_variations") + + # 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.") + 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 + self.components.label(frame, 9, 0, "Loss Weight", + tooltip="The loss multiplyer for this concept.") + self.components.entry(frame, 9, 1, ui_state, "loss_weight") + + def build_image_augmentation_tab(self, frame, controller, image_ui_state): + # header + self.components.label(frame, 0, 1, "Random", + tooltip="Enable this augmentation with random values") + self.components.label(frame, 0, 2, "Fixed", + tooltip="Enable this augmentation with fixed values") + + # crop jitter + self.components.label(frame, 1, 0, "Crop Jitter", + tooltip="Enables random cropping of samples") + self.components.switch(frame, 1, 1, image_ui_state, "enable_crop_jitter") + + # random flip + self.components.label(frame, 2, 0, "Random Flip", + tooltip="Randomly flip the sample during training") + 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 + self.components.label(frame, 3, 0, "Random Rotation", + tooltip="Randomly rotates the sample during training") + 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 + self.components.label(frame, 4, 0, "Random Brightness", + tooltip="Randomly adjusts the brightness of the sample during training") + 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 + self.components.label(frame, 5, 0, "Random Contrast", + tooltip="Randomly adjusts the contrast of the sample during training") + 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 + self.components.label(frame, 6, 0, "Random Saturation", + tooltip="Randomly adjusts the saturation of the sample during training") + 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 + self.components.label(frame, 7, 0, "Random Hue", + tooltip="Randomly adjusts the hue of the sample during training") + 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 + self.components.label(frame, 8, 0, "Circular Mask Generation", + tooltip="Automatically create circular masks for masked training") + self.components.switch(frame, 8, 1, image_ui_state, "enable_random_circular_mask_shrink") + + # 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") + self.components.switch(frame, 9, 1, image_ui_state, "enable_random_mask_rotate_crop") + + # circular mask generation + 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") + 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 + self.components.label(frame, 0, 0, "Tag Shuffling", + tooltip="Enables tag shuffling") + self.components.switch(frame, 0, 1, text_ui_state, "enable_tag_shuffling") + + # keep tag count + self.components.label(frame, 1, 0, "Tag Delimiter", + tooltip="The delimiter between tags") + self.components.entry(frame, 1, 1, text_ui_state, "tag_delimiter") + + # 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") + self.components.entry(frame, 2, 1, text_ui_state, "keep_tags_count") + + # tag dropout + self.components.label(frame, 3, 0, "Tag Dropout", + tooltip="Enables random dropout for tags in the captions.") + 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.") + self.components.options_kv(frame, 4, 1, [ + ("Full", 'FULL'), + ("Random", 'RANDOM'), + ("Random Weighted", 'RANDOM WEIGHTED'), + ], text_ui_state, "tag_dropout_mode", None) + self.components.label(frame, 4, 2, "Probability", + tooltip="Probability to drop tags, from 0 to 1.") + self.components.entry(frame, 4, 3, text_ui_state, "tag_dropout_probability") + + 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.") + self.components.options_kv(frame, 5, 1, [ + ("None", 'NONE'), + ("Blacklist", 'BLACKLIST'), + ("Whitelist", 'WHITELIST'), + ], 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.") + self.components.switch(frame, 6, 1, text_ui_state, "tag_dropout_special_tags_regex") + + #capitalization randomization + self.components.label(frame, 7, 0, "Randomize Capitalization", + tooltip="Enables randomization of capitalization for tags in the caption.") + 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.") + self.components.switch(frame, 7, 3, text_ui_state, "caps_randomize_lowercase") + + 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.") + 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.") + self.components.entry(frame, 8, 3, text_ui_state, "caps_randomize_probability") + + def build_concept_stats_tab(self, frame, controller): + self.concept_stats_tab = frame + + #file size + 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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.", 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 = 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 = 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 = 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 = self.components.label(frame, 0, 3, text="-", tooltip="Time taken to process concept directory") + + def _update_concept_stats(self, controller): + #file size + 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.components.set_label_text(self.dir_count_preview, controller.concept.concept_stats["directory_count"]) + + #image 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.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.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 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.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 = 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.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 = 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.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.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 = 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.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 = 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 + 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.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())] + 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 _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/ConfigList.py b/modules/ui/BaseConfigListView.py similarity index 78% rename from modules/ui/ConfigList.py rename to modules/ui/BaseConfigListView.py index 75d69252a..2165e854b 100644 --- a/modules/ui/ConfigList.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 new file mode 100644 index 000000000..178c5759b --- /dev/null +++ b/modules/ui/BaseConvertModelUIView.py @@ -0,0 +1,84 @@ +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 + + +class BaseConvertModelUIView: + def __init__(self, components): + self.components = components + + def build_content(self, frame, controller, ui_state): + # model type + self.components.label(frame, 0, 0, "Model Type", + tooltip="Type of the model") + 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), + ("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 + ("Anima", ModelType.ANIMA), + ("ZImage", ModelType.Z_IMAGE), + ], ui_state, "model_type") + + # training method + self.components.label(frame, 1, 0, "Model Type", + tooltip="The type of model to convert") + self.components.options_kv(frame, 1, 1, [ + ("Base Model", TrainingMethod.FINE_TUNE), + ("LoRA", TrainingMethod.LORA), + ("Embedding", TrainingMethod.EMBEDDING), + ], ui_state, "training_method") + + # input name + self.components.label(frame, 2, 0, "Input name", + tooltip="Filename, directory or hugging face repository of the base model") + self.components.path_entry( + frame, 2, 1, ui_state, "input_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # output data type + self.components.label(frame, 3, 0, "Output Data Type", + tooltip="Precision to use when saving the output model") + self.components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("float16", DataType.FLOAT_16), + ("bfloat16", DataType.BFLOAT_16), + ], ui_state, "output_dtype") + + # output format + self.components.label(frame, 4, 0, "Output Format", + tooltip="Format to use when saving the output model") + self.components.options_kv(frame, 4, 1, [ + ("Safetensors", ModelFormat.SAFETENSORS), + ("Diffusers", ModelFormat.DIFFUSERS), + ], ui_state, "output_model_format") + + # output model destination + self.components.label(frame, 5, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + self.components.path_entry( + frame, 5, 1, ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + self.button = self.components.button(frame, 6, 1, "Convert", controller.convert_model) diff --git a/modules/ui/BaseGenerateCaptionsWindowView.py b/modules/ui/BaseGenerateCaptionsWindowView.py new file mode 100644 index 000000000..68e24638e --- /dev/null +++ b/modules/ui/BaseGenerateCaptionsWindowView.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + + +class BaseGenerateCaptionsWindowView(ABC): + @abstractmethod + def set_progress(self, value, max_value): + pass diff --git a/modules/ui/BaseGenerateMasksWindowView.py b/modules/ui/BaseGenerateMasksWindowView.py new file mode 100644 index 000000000..f9e82231e --- /dev/null +++ b/modules/ui/BaseGenerateMasksWindowView.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + + +class BaseGenerateMasksWindowView(ABC): + @abstractmethod + def set_progress(self, value, max_value): + pass diff --git a/modules/ui/BaseLoraTabView.py b/modules/ui/BaseLoraTabView.py new file mode 100644 index 000000000..542247b82 --- /dev/null +++ b/modules/ui/BaseLoraTabView.py @@ -0,0 +1,164 @@ + +from modules.util import path_util +from modules.util.enum.ModelType import PeftType +from modules.util.ui.validation_helpers import check_range + + +class BaseLoraTabView: + def __init__(self, components): + self.components = components + + 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) + + 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: + name = "OFT v2" + elif peft_type == PeftType.LOKR: + name = "LoKr" + else: + name = "LoRA" + + # lora model name + 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, + ) + + # LoRA decomposition + if peft_type == PeftType.LORA: + 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") + + 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 + 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 + 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 + 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 + 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. + 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 + 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) + + # Block Share + self.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.") + self.components.switch(master, 1, 4, ui_state, "oft_block_share") + + # Scaled OFT (SOFT) + self.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.") + self.components.switch(master, 2, 4, ui_state, "oft_scaled") + + # Dropout Percentage + 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 + 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. + 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") + + # LoKr + elif peft_type == PeftType.LOKR: + # LoKr Main Settings + self.components.label(master, 1, 0, f"{name} dimension", + tooltip="The dimension parameter used for the secondary decomposition. Analogous to rank in LoRA.") + self.components.entry(master, 1, 1, ui_state, "lokr_dim") + + self.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.") + self.components.entry(master, 2, 1, ui_state, "lokr_decompose_factor") + + # alpha + self.components.label(master, 3, 0, f"{name} alpha", + tooltip=f"The alpha parameter used when creating a new {name}") + self.components.entry(master, 3, 1, ui_state, "lora_alpha") + + # Dropout Percentage + self.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.") + self.components.entry(master, 4, 1, ui_state, "dropout_probability") + + # LoKr weight dtype + self.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") + self.components.options_kv(master, 5, 1, controller.get_lora_weight_dtypes(), ui_state, "lora_weight_dtype") + + # LoKr Vectorization trick + self.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.") + self.components.switch(master, 6, 1, ui_state, "lokr_vec_trick") + + # LoKr Decomposition Settings + self.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.") + self.components.switch(master, 1, 4, ui_state, "lokr_decompose_both") + + self.components.label(master, 2, 3, "Use Tucker Decomposition (Conv)", + tooltip="Use Tucker decomposition for convolutional layers. Can be more efficient for some architectures.") + self.components.switch(master, 2, 4, ui_state, "lokr_use_tucker") + + self.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.") + self.components.switch(master, 3, 4, ui_state, "lokr_full_matrix") + + # LoKr DoRA Settings + self.components.label(master, 4, 3, "Decompose Weights (DoRA)", + tooltip="Apply weight decomposition (DoRA) on top of the LoKr update.") + self.components.switch(master, 4, 4, ui_state, "lokr_weight_decompose") + + self.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.") + self.components.switch(master, 5, 4, ui_state, "lokr_dora_on_output") + + # Additional embeddings + self.components.label(master, 6, 3, "Bundle Embeddings", + tooltip=f"Bundles any additional embeddings into the {name} output file, rather than as separate files") + self.components.switch(master, 6, 4, ui_state, "bundle_additional_embeddings") diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py new file mode 100644 index 000000000..21a55614d --- /dev/null +++ b/modules/ui/BaseModelTabView.py @@ -0,0 +1,728 @@ + +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 + + +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_anima(): + self.__setup_anima_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, ui_state) + row = self.__create_base_components( + frame, + row, + controller, + ui_state, + has_unet=True, + has_text_encoder=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + allow_diffusers=controller.train_config.training_method in [ + TrainingMethod.FINE_TUNE, + TrainingMethod.FINE_TUNE_VAE, + ], + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_stable_diffusion_3_ui(self, frame, controller, ui_state): + row = 0 + 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, + has_text_encoder_3=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + allow_diffusers=controller.train_config.training_method == TrainingMethod.FINE_TUNE, + allow_legacy_safetensors=controller.train_config.training_method == TrainingMethod.LORA, + ) + + def __setup_anima_ui(self, frame, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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=controller.train_config.model_type.is_stable_cascade(), + has_text_encoder=True, + ) + 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, + 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, controller, ui_state): + row = 0 + 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, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + ) + row = self.__create_output_components( + frame, + row, + 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, controller, ui_state): + row = 0 + 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, + has_text_encoder_2=True, + has_vae=True, + ) + row = self.__create_output_components( + frame, + row, + ui_state, + allow_safetensors=True, + 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, controller, ui_state): + row = 0 + 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, + 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, + ui_state, + allow_safetensors=True, + 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]]: + 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), + ("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 int", DataType.GGUF_A8_INT), + ("GGUF A8 float", DataType.GGUF_A8_FLOAT), + ] + + return options + + def __create_base_dtype_components(self, frame, row: int, ui_state) -> int: + # huggingface 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) + self.components.entry(frame, row, 1, ui_state, "secrets.huggingface_token") + + row += 1 + + # base model + self.components.label(frame, row, 0, "Base Model", + tooltip="Filename, directory or Hugging Face repository of the base model") + self.components.path_entry( + frame, row, 1, ui_state, "base_model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # compile + 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.") + self.components.switch(frame, row, 4, ui_state, "compile") + + row += 1 + + return row + + def __create_base_components( + self, + frame, + row: int, + controller, + ui_state, + 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 + self.components.label(frame, row, 3, "UNet Data Type", + tooltip="The unet weight data type") + 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 + self.components.label(frame, row, 0, "Prior Model", + tooltip="Filename, directory or Hugging Face repository of the prior model") + self.components.path_entry( + frame, row, 1, ui_state, "prior.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # prior weight dtype + self.components.label(frame, row, 3, "Prior Data Type", + tooltip="The prior weight data type") + 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 + 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") + self.components.path_entry( + frame, row, 1, ui_state, "transformer.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # transformer weight dtype + self.components.label(frame, row, 3, "Transformer Data Type", + tooltip="The transformer weight data type") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_gguf=True, include_a8=True), + ui_state, "transformer.weight_dtype") + + row += 1 + + presets = controller.get_presets() + + 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", + 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, 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.") + 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") + self.components.entry(svd_entry_frame, 1, 0, ui_state, "quantization.svd_rank") + row += 1 + + if has_text_encoder: + # text encoder weight dtype + self.components.label(frame, row, 3, "Text Encoder Data Type", + tooltip="The text encoder weight data type") + 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 + self.components.label(frame, row, 3, "Text Encoder 1 Data Type", + tooltip="The text encoder 1 weight data type") + 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 + self.components.label(frame, row, 3, "Text Encoder 2 Data Type", + tooltip="The text encoder 2 weight data type") + 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 + self.components.label(frame, row, 3, "Text Encoder 3 Data Type", + tooltip="The text encoder 3 weight data type") + 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 + self.components.label(frame, row, 0, "Text Encoder 4 Override", + tooltip="Filename, directory or Hugging Face repository of the text encoder 4 model") + 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 + self.components.label(frame, row, 3, "Text Encoder 4 Data Type", + tooltip="The text encoder 4 weight data type") + 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 + 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.") + self.components.path_entry( + frame, row, 1, ui_state, "vae.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # vae weight dtype + self.components.label(frame, row, 3, "VAE Data Type", + tooltip="The vae weight data type") + 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, ui_state) -> int: + # effnet encoder model + self.components.label(frame, row, 0, "Effnet Encoder Model", + tooltip="Filename, directory or Hugging Face repository of the effnet encoder model") + 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 + self.components.label(frame, row, 3, "Effnet Encoder Data Type", + tooltip="The effnet encoder weight data type") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "effnet_encoder.weight_dtype") + + row += 1 + + return row + + def __create_decoder_components( + self, + frame, + row: int, + ui_state, + has_text_encoder: bool, + ) -> int: + # decoder model + self.components.label(frame, row, 0, "Decoder Model", + tooltip="Filename, directory or Hugging Face repository of the decoder model") + self.components.path_entry( + frame, row, 1, ui_state, "decoder.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # decoder weight dtype + self.components.label(frame, row, 3, "Decoder Data Type", + tooltip="The decoder weight data type") + 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 + self.components.label(frame, row, 3, "Decoder Text Encoder Data Type", + tooltip="The decoder text encoder weight data type") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_text_encoder.weight_dtype") + + row += 1 + + # decoder vqgan weight dtype + self.components.label(frame, row, 3, "Decoder VQGAN Data Type", + tooltip="The decoder vqgan weight data type") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(), + ui_state, "decoder_vqgan.weight_dtype") + + row += 1 + + return row + + 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 + self.components.label(frame, row, 0, "Model Output Destination", + tooltip="Filename or directory where the output model is saved") + self.components.path_entry( + frame, row, 1, ui_state, "output_model_destination", + mode="file", + io_type=PathIOType.MODEL, + ) + + # output data type + self.components.label(frame, row, 3, "Output Data Type", + tooltip="Precision to use when saving the output model") + 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), + ], 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)) + + self.components.label(frame, row, 0, "Output Format", + tooltip="Format to use when saving the output model") + self.components.options_kv(frame, row, 1, formats, ui_state, "output_model_format") + + # 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.") + self.components.options_kv(frame, row, 4, [ + ("None", ConfigPart.NONE), + ("Settings", ConfigPart.SETTINGS), + ("All", ConfigPart.ALL), + ], ui_state, "include_train_config") + + row += 1 + + return row diff --git a/modules/ui/MuonAdamWindow.py b/modules/ui/BaseMuonAdamWindowView.py similarity index 62% rename from modules/ui/MuonAdamWindow.py rename to modules/ui/BaseMuonAdamWindowView.py index 5879ab432..90dc04eee 100644 --- a/modules/ui/MuonAdamWindow.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/OptimizerParamsWindow.py b/modules/ui/BaseOptimizerParamsWindowView.py similarity index 76% rename from modules/ui/OptimizerParamsWindow.py rename to modules/ui/BaseOptimizerParamsWindowView.py index 1d8cbe2b5..b5a5c8911 100644 --- a/modules/ui/OptimizerParamsWindow.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 not current_state: - 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 new file mode 100644 index 000000000..8e0de3b64 --- /dev/null +++ b/modules/ui/BaseProfilingWindowView.py @@ -0,0 +1,21 @@ +from abc import abstractmethod + + +class BaseProfilingWindowView: + def __init__(self, components): + self.components = components + + 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") + + @abstractmethod + def set_message(self, text): + pass + + @abstractmethod + def set_profiling_active(self, active): + pass diff --git a/modules/ui/BaseSampleFrameView.py b/modules/ui/BaseSampleFrameView.py new file mode 100644 index 000000000..c3e5860c9 --- /dev/null +++ b/modules/ui/BaseSampleFrameView.py @@ -0,0 +1,95 @@ +from modules.util.enum.NoiseScheduler import NoiseScheduler + + +class BaseSampleFrameView: + def __init__(self, components): + self.components = components + + 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 + self.components.label(top_frame, 0, 0, "prompt:") + self.components.entry(top_frame, 0, 1, ui_state, "prompt") + + # 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 + self.components.label(bottom_frame, 0, 0, "width:") + self.components.entry(bottom_frame, 0, 1, ui_state, "width") + + # height + self.components.label(bottom_frame, 0, 2, "height:") + self.components.entry(bottom_frame, 0, 3, ui_state, "height") + + if is_video_model: + # 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 + 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 + self.components.label(bottom_frame, 2, 0, "seed:") + self.components.entry(bottom_frame, 2, 1, ui_state, "seed") + + # random seed + self.components.label(bottom_frame, 2, 2, "random seed:") + self.components.switch(bottom_frame, 2, 3, ui_state, "random_seed") + + # 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: + 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), + # ("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) + ], ui_state, "noise_scheduler") + + # 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: + 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 + 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 + 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 new file mode 100644 index 000000000..bbe410a9a --- /dev/null +++ b/modules/ui/BaseSampleParamsWindowView.py @@ -0,0 +1,6 @@ + + + +class BaseSampleParamsWindowView: + def __init__(self, components): + self.components = components diff --git a/modules/ui/BaseSampleWindowView.py b/modules/ui/BaseSampleWindowView.py new file mode 100644 index 000000000..1a9c6c0c6 --- /dev/null +++ b/modules/ui/BaseSampleWindowView.py @@ -0,0 +1,8 @@ + + + + + +class BaseSampleWindowView: + def __init__(self, components): + pass diff --git a/modules/ui/BaseSamplingTabView.py b/modules/ui/BaseSamplingTabView.py new file mode 100644 index 000000000..69ec082a1 --- /dev/null +++ b/modules/ui/BaseSamplingTabView.py @@ -0,0 +1,67 @@ +from abc import ABC, abstractmethod + +from modules.ui.BaseConfigListView import BaseConfigListView + + +class BaseSamplingTabView(BaseConfigListView): + pass + + +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.i = i + self.save_command = save_command + + # close button + self.components.colored_icon_button(frame, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) + + # clone button + self.components.colored_icon_button(frame, 0, 1, "+", "#00C000", lambda: clone_command(self.i), padx=5) + + # enabled + self.enabled_switch = self.components.switch(frame, 0, 2, ui_state, "enabled", self._switch_enabled, width=40) + + # width + self.components.label(frame, 0, 3, "width:") + self.width_entry = self.components.entry(frame, 0, 4, ui_state, "width", width=50) + + # height + self.components.label(frame, 0, 5, "height:") + self.height_entry = self.components.entry(frame, 0, 6, ui_state, "height", width=50) + + # seed + self.components.label(frame, 0, 7, "seed:") + self.seed_entry = self.components.entry(frame, 0, 8, ui_state, "seed", width=80) + + # prompt + self.components.label(frame, 0, 9, "prompt:") + self.prompt_entry = self.components.entry(frame, 0, 10, ui_state, "prompt") + + # button + self.button = self.components.icon_button(frame, 0, 11, "...", lambda: open_command(self.i, ui_state)) + + self._bind_save(save_command) + self._set_enabled() + + @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() + + 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") diff --git a/modules/ui/BaseSchedulerParamsWindowView.py b/modules/ui/BaseSchedulerParamsWindowView.py new file mode 100644 index 000000000..1106d5227 --- /dev/null +++ b/modules/ui/BaseSchedulerParamsWindowView.py @@ -0,0 +1,21 @@ + +from modules.ui.BaseConfigListView import BaseConfigListView + + +class BaseSchedulerParamsWindowView: + def __init__(self, components): + self.components = components + + 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") + + +class BaseKvParamsView(BaseConfigListView): + def __init__(self, components): + self.components = components + + def open_element_window(self, i, ui_state): + pass diff --git a/modules/ui/BaseTimestepDistributionWindowView.py b/modules/ui/BaseTimestepDistributionWindowView.py new file mode 100644 index 000000000..19354bf43 --- /dev/null +++ b/modules/ui/BaseTimestepDistributionWindowView.py @@ -0,0 +1,46 @@ + + + + + +class BaseTimestepDistributionWindowView: + def __init__(self, components): + self.components = components + + def build_content(self, frame, controller, ui_state): + # timestep distribution + self.components.label(frame, 0, 0, "Timestep Distribution", + tooltip="Selects the function to sample timesteps during training", + wide_tooltip=True) + self.components.options(frame, 0, 1, controller.get_distribution_options(), ui_state, + "timestep_distribution") + + # 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") + self.components.entry(frame, 1, 1, ui_state, "min_noising_strength") + + # 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") + self.components.entry(frame, 2, 1, ui_state, "max_noising_strength") + + # 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.") + self.components.entry(frame, 3, 1, ui_state, "noising_weight") + + # 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.") + self.components.entry(frame, 4, 1, ui_state, "noising_bias") + + # timestep shift + self.components.label(frame, 5, 0, "Timestep Shift", + tooltip="Shift the timestep distribution. Use the preview to see more details.") + self.components.entry(frame, 5, 1, ui_state, "timestep_shift") + + # 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) + self.components.switch(frame, 6, 1, ui_state, "dynamic_timestep_shifting") diff --git a/modules/ui/BaseTopBarView.py b/modules/ui/BaseTopBarView.py new file mode 100644 index 000000000..f50c1ce14 --- /dev/null +++ b/modules/ui/BaseTopBarView.py @@ -0,0 +1,190 @@ +import json +import os +import traceback +from abc import abstractmethod +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 + + +class BaseTopBarView: + def __init__(self, components): + self.components = components + + @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, + 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.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 = self._make_config_ui_state(master, self.config_ui_data) + + self.configs = controller.load_available_config_names(self.dir) + + self.current_config = [] + + self.training_method = None + + # title + self.components.app_title(self.frame, 0, 0) + + # dropdown + self.configs_dropdown = None + self.__create_configs_dropdown() + + # remove button + # TODO + # self.components.icon_button(self.frame, 0, 2, "-", self.__remove_config) + + # Wiki button + self.components.button(self.frame, 0, 4, "Wiki", self.open_wiki, width=50) + + # save button + 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._setup_frame_column_weight() + + # model type + self.components.options_kv( + master=self.frame, + row=0, + column=6, + values=controller.get_model_types(), + ui_state=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 = 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, + 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._forget_dropdown() + + self.configs_dropdown = self.components.options_kv( + self.frame, 0, 1, self.configs, self.config_ui_state, "config_name", self.__load_current_config + ) + + def __save_config(self): + default_value = self._get_dropdown_text(self.configs_dropdown) + while default_value.startswith('#'): + default_value = default_value[1:] + + self._show_save_dialog(default_value, self.__save_new_config) + + def __save_new_config(self, name): + path = self.controller.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 __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.controller.train_config.from_dict(loaded_config.to_dict()) + self.ui_state.update(loaded_config) + + optimizer_config = change_optimizer(self.controller.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 open_wiki(self): + self.controller.open_wiki() + + def save_default(self): + self.controller.save_default() diff --git a/modules/ui/BaseTrainUIView.py b/modules/ui/BaseTrainUIView.py new file mode 100644 index 000000000..5ac9e1b42 --- /dev/null +++ b/modules/ui/BaseTrainUIView.py @@ -0,0 +1,382 @@ +from abc import ABC, abstractmethod +from collections.abc import Callable + +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.PathIOType import PathIOType + + +class BaseTrainUIView(ABC): + def __init__(self, components, controller, ui_state): + self.components = components + self.controller = controller + self.ui_state = ui_state + + # --- Abstract callbacks (controller calls into view) --- + + @abstractmethod + def on_update_status(self, status: str): pass + + @abstractmethod + def on_training_started(self): pass + + @abstractmethod + def on_training_stopped(self, error_caught: bool): pass + + @abstractmethod + def on_training_stopping(self): pass + + @abstractmethod + def on_update_progress(self, epoch_step: int, max_step: int, epoch: int, max_epoch: int, eta_str: str | None): pass + + @abstractmethod + def schedule_on_main_thread(self, fn: Callable): pass + + @abstractmethod + def get_cloud_reattach(self) -> bool: pass + + @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.controller.train_config.secrets.cloud) + + def start_training(self): + self.controller.start_training() + + def open_tensorboard(self): + self.controller.open_tensorboard() + + def sample_now(self): + self.controller.sample_now() + + def backup_now(self): + self.controller.backup_now() + + def save_now(self): + self.controller.save_now() + + @abstractmethod + def open_dataset_tool(self): pass + + @abstractmethod + def open_video_tool(self): pass + + @abstractmethod + def open_convert_model_tool(self): pass + + @abstractmethod + 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): + 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 + self.components.label(frame, 0, 0, "Workspace Directory", + tooltip="The directory where all files of this training run are saved") + self.components.path_entry(frame, 0, 1, ui_state, "workspace_dir", mode="dir", command=controller._on_workspace_dir_change) + + # cache dir + self.components.label(frame, 0, 2, "Cache Directory", + tooltip="The directory where cached data is saved") + self.components.path_entry(frame, 0, 3, ui_state, "cache_dir", mode="dir") + + # continue from previous backup + self.components.label(frame, 2, 0, "Continue from last backup", + tooltip="Automatically continues training from the last backup saved in /backup") + self.components.switch(frame, 2, 1, ui_state, "continue_last_backup") + + # only cache + self.components.label(frame, 2, 2, "Only Cache", + tooltip="Only populate the cache, without any training") + self.components.switch(frame, 2, 3, ui_state, "only_cache") + + # TODO: In Phase 4 rework the general tab. + # 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") + self.components.switch(frame, 3, 1, ui_state, "prevent_overwrites") + + # debug + self.components.label(frame, 4, 0, "Debug mode", + tooltip="Save debug information during the training into the debug directory") + self.components.switch(frame, 4, 1, ui_state, "debug_mode") + + self.components.label(frame, 4, 2, "Debug Directory", + tooltip="The directory where debug data is saved") + self.components.path_entry(frame, 4, 3, ui_state, "debug_dir", mode="dir", io_type=PathIOType.OUTPUT) + + # tensorboard + self.components.label(frame, 6, 0, "Tensorboard", + tooltip="Starts the Tensorboard Web UI during training") + self.components.switch(frame, 6, 1, ui_state, "tensorboard") + + self.components.label(frame, 6, 2, "Always-On Tensorboard", + tooltip="Keep Tensorboard accessible even when not training. Useful for monitoring completed training sessions.") + self.components.switch(frame, 6, 3, ui_state, "tensorboard_always_on", command=controller._on_always_on_tensorboard_toggle) + + self.components.label(frame, 7, 0, "Expose Tensorboard", + tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)") + 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") + self.components.entry(frame, 7, 3, ui_state, "tensorboard_port") + + # validation + self.components.label(frame, 8, 0, "Validation", + tooltip="Enable validation steps and add new graph in tensorboard") + self.components.switch(frame, 8, 1, ui_state, "validation") + + self.components.label(frame, 8, 2, "Validate after", + tooltip="The interval used when validate training") + self.components.time_entry(frame, 8, 3, ui_state, "validate_after", "validate_after_unit") + + # device + self.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.") + self.components.entry(frame, 9, 1, ui_state, "train_device", required=True) + + self.components.label(frame, 9, 2, "Async Offloading", + tooltip="Overlaps CPU<->GPU transfers with computation using CUDA streams. Applies to every offloaded component") + self.components.switch(frame, 9, 3, ui_state, "async_offloading") + + self.components.label(frame, 10, 0, "Multi-GPU", + tooltip="Enable multi-GPU training") + self.components.switch(frame, 10, 1, ui_state, "multi_gpu") + self.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.") + self.components.entry(frame, 10, 3, ui_state, "device_indexes") + + self.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) + self.components.options(frame, 11, 1, [str(x) for x in list(GradientReducePrecision)], ui_state, + "gradient_reduce_precision") + + self.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") + self.components.switch(frame, 11, 3, ui_state, "fused_gradient_reduce") + + self.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.") + self.components.switch(frame, 12, 1, ui_state, "async_gradient_reduce") + self.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\"") + self.components.entry(frame, 12, 3, ui_state, "async_gradient_reduce_buffer") + + self.components.label(frame, 13, 0, "Temp Device", + tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\"") + self.components.entry(frame, 13, 1, ui_state, "temp_device") + + def build_data_tab_content(self, frame, controller, ui_state): + # aspect ratio bucketing + self.components.label(frame, 0, 0, "Aspect Ratio Bucketing", + tooltip="Aspect ratio bucketing enables training on images with different aspect ratios") + self.components.switch(frame, 0, 1, ui_state, "aspect_ratio_bucketing") + + # image caching + self.components.label(frame, 1, 0, "Image Caching", + tooltip="Caches image latents (VAE outputs) so they can be re-used between epochs") + self.components.switch(frame, 1, 1, ui_state, "image_caching") + + # text caching + self.components.label(frame, 2, 0, "Text Caching", + tooltip="Caches text encoder outputs so they can be re-used between epochs") + self.components.switch(frame, 2, 1, ui_state, "text_caching") + + # caching threads + self.components.label(frame, 3, 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.") + self.components.entry(frame, 3, 1, ui_state, "caching_threads", width=100, sticky="nw", required=True) + + # prefetch next batch + self.components.label(frame, 4, 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.") + self.components.switch(frame, 4, 1, ui_state, "prefetch_next_batch") + + # clear cache before training + self.components.label(frame, 5, 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") + self.components.switch(frame, 5, 1, ui_state, "clear_cache_before_training") + + 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") + self.components.time_entry(top_frame, 0, 1, ui_state, "sample_after", "sample_after_unit") + + self.components.label(top_frame, 0, 2, "Skip First", + tooltip="Start sampling automatically after this interval has elapsed.") + self.components.entry(top_frame, 0, 3, ui_state, "sample_skip_first", width=50, sticky="nw") + + self.components.label(top_frame, 0, 4, "Format", + tooltip="File Format used when saving samples") + self.components.options_kv(top_frame, 0, 5, [ + ("PNG", ImageFormat.PNG), + ("JPG", ImageFormat.JPG), + ], ui_state, "sample_image_format") + + self.components.button(top_frame, 0, 6, "sample now", self.sample_now) + + self.components.button(top_frame, 0, 7, "manual sample", self.open_manual_sample_window) + + self.components.label(sub_frame, 0, 0, "Non-EMA Sampling", + tooltip="Whether to include non-ema sampling when using ema.") + self.components.switch(sub_frame, 0, 1, ui_state, "non_ema_sampling") + + self.components.label(sub_frame, 0, 2, "Samples to Tensorboard", + tooltip="Whether to include sample images in the Tensorboard output.") + self.components.switch(sub_frame, 0, 3, ui_state, "samples_to_tensorboard") + + def build_backup_tab_content(self, frame, controller, ui_state): + # backup after + self.components.label(frame, 0, 0, "Backup After", + tooltip="The interval used when automatically creating model backups during training") + self.components.time_entry(frame, 0, 1, ui_state, "backup_after", "backup_after_unit") + + # backup now + self.components.button(frame, 0, 3, "backup now", self.backup_now) + + # rolling backup + self.components.label(frame, 1, 0, "Rolling Backup", + tooltip="If rolling backups are enabled, older backups are deleted automatically") + self.components.switch(frame, 1, 1, ui_state, "rolling_backup") + + # 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, 2, 1, ui_state, "rolling_backup_count") + + # 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, 3, 1, ui_state, "backup_before_save") + + # save after + self.components.label(frame, 4, 0, "Save Every", + tooltip="The interval used when automatically saving the model during training") + self.components.time_entry(frame, 4, 1, ui_state, "save_every", "save_every_unit") + + # save now + self.components.button(frame, 4, 3, "save now", self.save_now) + + # skip save + self.components.label(frame, 5, 0, "Skip First", + tooltip="Start saving automatically after this interval has elapsed") + self.components.entry(frame, 5, 1, ui_state, "save_skip_first", width=50, sticky="nw") + + # 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, 6, 1, ui_state, "save_filename_prefix") + + def build_embedding_tab_content(self, frame, controller, ui_state): + # embedding model name + self.components.label(frame, 0, 0, "Base embedding", + tooltip="The base embedding to train on. Leave empty to create a new embedding") + self.components.path_entry( + frame, 0, 1, ui_state, "embedding.model_name", + mode="file", path_modifier=path_util.json_path_modifier + ) + + # 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.") + self.components.entry(frame, 1, 1, ui_state, "embedding.token_count") + + # initial embedding text + self.components.label(frame, 2, 0, "Initial embedding text", + tooltip="The initial embedding text used when creating a new embedding") + self.components.entry(frame, 2, 1, ui_state, "embedding.initial_embedding_text") + + # embedding weight dtype + 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") + self.components.options_kv(frame, 3, 1, [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ], ui_state, "embedding_weight_dtype") + + # placeholder + self.components.label(frame, 4, 0, "Placeholder", + tooltip="The placeholder used when using the embedding in a prompt") + self.components.entry(frame, 4, 1, ui_state, "embedding.placeholder") + + # 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.") + self.components.switch(frame, 5, 1, ui_state, "embedding.is_output_embedding") + + def build_tools_tab_content(self, frame, controller, ui_state): + # dataset + self.components.label(frame, 0, 0, "Dataset Tools", + tooltip="Open the captioning tool") + self.components.button(frame, 0, 1, "Open", self.open_dataset_tool) + + # video tools + self.components.label(frame, 1, 0, "Video Tools", + tooltip="Open the video tools") + self.components.button(frame, 1, 1, "Open", self.open_video_tool) + + # convert model + self.components.label(frame, 2, 0, "Convert Model Tools", + tooltip="Open the model conversion tool") + self.components.button(frame, 2, 1, "Open", self.open_convert_model_tool) + + # sample + self.components.label(frame, 3, 0, "Sampling Tool", + tooltip="Open the model sampling tool") + self.components.button(frame, 3, 1, "Open", self.open_sampling_tool) + + self.components.label(frame, 4, 0, "Profiling Tool", + tooltip="Open the profiling tools.") + self.components.button(frame, 4, 1, "Open", self.open_profiling_tool) diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py new file mode 100644 index 000000000..11f6f857b --- /dev/null +++ b/modules/ui/BaseTrainingTabView.py @@ -0,0 +1,850 @@ +from abc import ABC, abstractmethod + +from modules.util.enum.DataType import DataType +from modules.util.enum.EMAMode import EMAMode +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.ui.validation_helpers import check_range, validate_resolution + + +class BaseTrainingTabView(ABC): + def __init__(self, components): + self.components = components + + @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_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) + if model_type.is_stable_diffusion_3(): + 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) + elif model_type.is_wuerstchen(): + 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) + elif model_type.is_flux_1(): + 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) + elif model_type.is_chroma(): + 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) + elif model_type.is_anima(): + self.__setup_anima_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) + elif model_type.is_hunyuan_video(): + 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) + elif model_type.is_z_image(): + 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) + + 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, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + self.__create_noise_frame(column_1, 2, ui_state, supports_generalized_offset_noise=True) + + if controller.config.training_method == TrainingMethod.FINE_TUNE_VAE: + self.__create_vae_frame(column_2, 0, 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_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) + self.__create_transformer_frame(column_1, 1, ui_state) + 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): + 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, supports_circular_padding=True) + self.__create_unet_frame(column_1, 1, ui_state) + 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): + 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, supports_circular_padding=True) + self.__create_prior_frame(column_1, 1, ui_state) + 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): + 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) + self.__create_transformer_frame(column_1, 1, ui_state) + 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): + 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) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=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): + 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) + 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, 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): + 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) + 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) + + 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): + 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) + 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, 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_anima_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) + 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, 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): + 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) + 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, 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): + 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) + 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, 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): + 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) + self.__create_transformer_frame(column_1, 1, ui_state) + 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): + 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, 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) + + 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): + 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, video_training_enabled=True) + self.__create_transformer_frame(column_1, 1, ui_state) + 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): + 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=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. + if hasattr(self, "lr_scheduler_comp"): + delattr(self, "lr_scheduler_comp") + delattr(self, "lr_scheduler_adv_comp") + 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=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. + self.restore_scheduler(ui_state.get_var("learning_rate_scheduler").get()) + + # learning rate + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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 + 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, video_training_enabled: bool = False, + supports_circular_padding: bool = False): + frame = self.components.section_frame(master, row) + row = 0 + + # 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 + 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 + 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 + + # train dtype + 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), + ], ui_state, "train_dtype") + row += 1 + + # fallback train dtype + 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), + ], ui_state, "fallback_train_dtype") + row += 1 + + # 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 + 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: + 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: + 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_offloading_widgets(self, frame, row, ui_state, part, supports_checkpointing=True, + supports_activation_offloading=False): + if supports_checkpointing: + self.components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for this component. Reduces VRAM usage at the cost of training speed") + self.components.switch(frame, row, 1, ui_state, f"{part}.gradient_checkpointing") + row += 1 + + self.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") + self.components.entry(frame, row, 1, ui_state, f"{part}.offload_fraction") + row += 1 + + if supports_activation_offloading: + self.components.label(frame, row, 0, "Offload Activations", + tooltip="Offloads this component's activations to CPU during training to reduce VRAM usage") + self.components.switch(frame, row, 1, ui_state, f"{part}.activation_offloading") + row += 1 + + return row + + 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) + row = 0 + + if supports_training: + self.components.label(frame, row, 0, "Train Text Encoder", + tooltip="Enables training the text encoder model") + self.components.switch(frame, row, 1, ui_state, "text_encoder.train") + row += 1 + else: + # no Train switch to act as the frame's header, so add an explicit one + self.components.label(frame, row, 0, "Text Encoder") + row += 1 + + row = self.__create_offloading_widgets(frame, row, ui_state, "text_encoder", supports_checkpointing=supports_training) + + # dropout + self.components.label(frame, row, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + self.components.entry(frame, row, 1, ui_state, "text_encoder.dropout_probability") + row += 1 + + if supports_training: + # train text encoder epochs + self.components.label(frame, row, 0, "Stop Training After", + tooltip="When to stop training the text encoder") + self.components.time_entry(frame, row, 1, ui_state, "text_encoder.stop_training_after", + "text_encoder.stop_training_after_unit", supports_time_units=False) + row += 1 + + # text encoder learning rate + self.components.label(frame, row, 0, "Text Encoder Learning Rate", + tooltip="The learning rate of the text encoder. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, "text_encoder.learning_rate") + row += 1 + + if supports_clip_skip: + # text encoder layer skip (clip skip) + self.components.label(frame, row, 0, "Clip Skip", + tooltip="The number of additional clip layers to skip. 0 = the model default") + self.components.entry(frame, row, 1, ui_state, "text_encoder_layer_skip") + row += 1 + + if supports_sequence_length: + # 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 = self.components.section_frame(master, row) + row = 0 + + suffix = f"_{i}" if i > 1 else "" + + if supports_include: + # include text encoder + 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 + 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 + + row = self.__create_offloading_widgets(frame, row, ui_state, f"text_encoder{suffix}") + + # train text encoder 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 + 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 + 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 + 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) + 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 + 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, ui_state): + frame = self.components.section_frame(master, row) + + # 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 + 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, ui_state): + frame = self.components.section_frame(master, row) + row = 0 + + # train unet + self.components.label(frame, row, 0, "Train UNet", + tooltip="Enables training the UNet model") + self.components.switch(frame, row, 1, ui_state, "unet.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, ui_state, "unet", supports_activation_offloading=True) + + # train unet epochs + self.components.label(frame, row, 0, "Stop Training After", + tooltip="When to stop training the UNet") + self.components.time_entry(frame, row, 1, ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + supports_time_units=False) + row += 1 + + # unet learning rate + self.components.label(frame, row, 0, "UNet Learning Rate", + tooltip="The learning rate of the UNet. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, "unet.learning_rate") + row += 1 + + # rescale noise scheduler to zero terminal SNR + self.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", + wraplength=130) + self.components.switch(frame, row, 1, ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + row += 1 + + def __create_vae_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) + row = 0 + + self.components.label(frame, row, 0, "Train VAE", + tooltip="Enables training the VAE model") + self.components.switch(frame, row, 1, ui_state, "vae.train") + row += 1 + + self.components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for the VAE. Reduces VRAM usage at the cost of training speed") + self.components.switch(frame, row, 1, ui_state, "vae.gradient_checkpointing") + row += 1 + + def __create_prior_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) + row = 0 + + # train prior + self.components.label(frame, row, 0, "Train Prior", + tooltip="Enables training the Prior model") + self.components.switch(frame, row, 1, ui_state, "prior.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, ui_state, "prior", supports_activation_offloading=True) + + # train prior epochs + self.components.label(frame, row, 0, "Stop Training After", + tooltip="When to stop training the Prior") + self.components.time_entry(frame, row, 1, ui_state, "prior.stop_training_after", + "prior.stop_training_after_unit", supports_time_units=False) + row += 1 + + # prior learning rate + self.components.label(frame, row, 0, "Prior Learning Rate", + tooltip="The learning rate of the Prior. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, "prior.learning_rate") + row += 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) + row = 0 + + # train transformer + self.components.label(frame, row, 0, "Train Transformer", + tooltip="Enables training the Transformer model") + self.components.switch(frame, row, 1, ui_state, "transformer.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, ui_state, "transformer", supports_activation_offloading=True) + + # train transformer epochs + self.components.label(frame, row, 0, "Stop Training After", + tooltip="When to stop training the Transformer") + self.components.time_entry(frame, row, 1, ui_state, "transformer.stop_training_after", + "transformer.stop_training_after_unit", supports_time_units=False) + row += 1 + + # transformer learning rate + self.components.label(frame, row, 0, "Transformer Learning Rate", + tooltip="The learning rate of the Transformer. Overrides the base learning rate") + self.components.entry(frame, row, 1, ui_state, "transformer.learning_rate") + row += 1 + + if supports_force_attention_mask: + self.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.") + self.components.switch(frame, row, 1, ui_state, "transformer.attention_mask") + row += 1 + + if supports_guidance_scale: + # guidance scale + self.components.label(frame, row, 0, "Guidance Scale", + tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") + self.components.entry(frame, row, 1, ui_state, "transformer.guidance_scale") + row += 1 + + 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) + + # 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 + 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 + 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 + 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=self.open_timestep_distribution) + + # min noising strength + 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 + 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 + 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 + 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 + 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) + + row = 9 + + if supports_dynamic_timestep_shifting: + # dynamic timestep shifting + self.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) + self.components.switch(frame, row, 1, ui_state, "dynamic_timestep_shifting") + row += 1 + + self.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.") + self.components.entry(frame, row, 1, ui_state, "cep_gamma", required=True) + row += 1 + + def __create_masked_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) + + # 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 + 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 + 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 + 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 + 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 + 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, controller, ui_state, + supports_vb_loss: bool = False): + frame = self.components.section_frame(master, row) + + # MSE Strength + 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 + 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 + 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 + 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 + 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 + 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 + 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 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 + 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, 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 new file mode 100644 index 000000000..e60247585 --- /dev/null +++ b/modules/ui/BaseVideoToolUIView.py @@ -0,0 +1,179 @@ +import webbrowser +from abc import ABC, abstractmethod + + +class BaseVideoToolUIView(ABC): + def __init__(self, components): + self.components = components + + 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.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)) + + # 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.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 + self.components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + 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.path_entry(frame, 3, 1, ui_state, "clip_output", mode="dir") + + # output to subdirectories + 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.components.switch(frame, 4, 1, ui_state, "output_subdir_clip") + + # 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.components.switch(frame, 5, 1, ui_state, "split_cuts") + + # maximum length + 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.components.entry(frame, 6, 1, ui_state, "clip_length", width=220) + + # 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.components.entry(frame, 7, 1, ui_state, "clip_fps", width=220) + + # Remove borders + self.components.label(frame, 8, 0, "Remove Borders", + tooltip="Remove black borders from output clip") + self.components.switch(frame, 8, 1, ui_state, "clip_bordercrop") + + # 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.components.entry(frame, 9, 1, ui_state, "clip_crop", width=220) + + 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.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)) + + # 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.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 + self.components.label(frame, 2, 0, "Directory", + tooltip="Path to directory with multiple videos to process, including in subdirectories.") + 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.path_entry(frame, 3, 1, ui_state, "image_output", mode="dir") + + # output to subdirectories + 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.components.switch(frame, 4, 1, ui_state, "output_subdir_img") + + # image capture rate + 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.components.entry(frame, 5, 1, ui_state, "capture_rate", width=220) + + # 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.components.entry(frame, 6, 1, ui_state, "blur_threshold", width=220) + + # Remove borders + self.components.label(frame, 7, 0, "Remove Borders", + tooltip="Remove black borders from output image") + self.components.switch(frame, 7, 1, ui_state, "image_bordercrop") + + # 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.components.entry(frame, 8, 1, ui_state, "image_crop", width=220) + + def build_video_download_tab(self, frame, controller, ui_state): + # 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.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 + self.components.label(frame, 1, 0, "Link List", + tooltip="Path to txt file with list of links separated by newlines.") + 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.path_entry(frame, 2, 1, ui_state, "download_output", mode="dir") + + # 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._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)) + + @abstractmethod + def _create_textbox(self, master, row, col, width, height, ui_state, var_name): + 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/CaptionUI.py b/modules/ui/CaptionUIController.py similarity index 56% rename from modules/ui/CaptionUI.py rename to modules/ui/CaptionUIController.py index e6cc0551e..0495322a3 100644 --- a/modules/ui/CaptionUI.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/CloudTab.py b/modules/ui/CloudTab.py deleted file mode 100644 index 99057e428..000000000 --- a/modules/ui/CloudTab.py +++ /dev/null @@ -1,221 +0,0 @@ - -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/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/ConceptTab.py b/modules/ui/ConceptTab.py deleted file mode 100644 index 0b6505694..000000000 --- a/modules/ui/ConceptTab.py +++ /dev/null @@ -1,286 +0,0 @@ -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/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/ConceptWindow.py b/modules/ui/ConceptWindow.py deleted file mode 100644 index 824ff43e4..000000000 --- a/modules/ui/ConceptWindow.py +++ /dev/null @@ -1,935 +0,0 @@ -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: - 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() - - 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/ConceptWindowController.py b/modules/ui/ConceptWindowController.py new file mode 100644 index 000000000..12577891c --- /dev/null +++ b/modules/ui/ConceptWindowController.py @@ -0,0 +1,298 @@ +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.image_util import load_image + +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 huggingface_hub +from PIL import Image + + +class ConceptWindowController: + def __init__(self, train_config: TrainConfig, concept: ConceptConfig): + self.train_config = train_config + self.concept = concept + self.cancel_scan_flag = threading.Event() + self.scan_thread = None + + @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: + 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() + + 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, preview_augmentations: bool) -> str: + empty_msg = "[Empty prompt]" + try: + with open(file_path, "r") as f: + 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() + 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, 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 "*.*" + + 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 == 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), preview_augmentations) if file_path else "[Empty prompt]" + + modules = [] + if preview_augmentations: + 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 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() + 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}") + view.components.call_after(view.concept_stats_tab, 0, view._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() + view.components.call_after(view.concept_stats_tab, 0, lambda: view._update_concept_stats(self)) + + self.cancel_scan_flag.clear() + 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, 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 auto_update_concept_stats(self, view): + try: + 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(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(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 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 diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py deleted file mode 100644 index 6cb1b507a..000000000 --- a/modules/ui/ConvertModelUI.py +++ /dev/null @@ -1,170 +0,0 @@ -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/ConvertModelUIController.py b/modules/ui/ConvertModelUIController.py new file mode 100644 index 000000000..c3cdebd31 --- /dev/null +++ b/modules/ui/ConvertModelUIController.py @@ -0,0 +1,70 @@ +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.TrainingMethod import TrainingMethod +from modules.util.ModelNames import EmbeddingName, ModelNames +from modules.util.torch_util import torch_gc + + +class ConvertModelUIController: + def __init__(self): + self.convert_model_args = ConvertModelArgs.default_values() + self.view = None + + def create_window(self, parent, view_cls): + self.view = view_cls(parent, self) + return self.view + + def convert_model(self): + try: + 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 + ) + 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.view.set_converting(False) diff --git a/modules/ui/CtkAdditionalEmbeddingsTabView.py b/modules/ui/CtkAdditionalEmbeddingsTabView.py new file mode 100644 index 000000000..fc24c61d1 --- /dev/null +++ b/modules/ui/CtkAdditionalEmbeddingsTabView.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/CtkCaptionUIView.py b/modules/ui/CtkCaptionUIView.py new file mode 100644 index 000000000..281036912 --- /dev/null +++ b/modules/ui/CtkCaptionUIView.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/CtkCloudTabView.py b/modules/ui/CtkCloudTabView.py new file mode 100644 index 000000000..ffe37451f --- /dev/null +++ b/modules/ui/CtkCloudTabView.py @@ -0,0 +1,44 @@ + + +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, controller) + self.master = master + 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/CtkConceptTabView.py b/modules/ui/CtkConceptTabView.py new file mode 100644 index 000000000..84d8e4508 --- /dev/null +++ b/modules/ui/CtkConceptTabView.py @@ -0,0 +1,180 @@ +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 _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 _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, 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/CtkConceptWindowView.py b/modules/ui/CtkConceptWindowView.py new file mode 100644 index 000000000..60c0f57fe --- /dev/null +++ b/modules/ui/CtkConceptWindowView.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/CtkConfigListView.py b/modules/ui/CtkConfigListView.py new file mode 100644 index 000000000..72995bfcc --- /dev/null +++ b/modules/ui/CtkConfigListView.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/CtkConvertModelUIView.py b/modules/ui/CtkConvertModelUIView.py new file mode 100644 index 000000000..782637348 --- /dev/null +++ b/modules/ui/CtkConvertModelUIView.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/GenerateCaptionsWindow.py b/modules/ui/CtkGenerateCaptionsWindowView.py similarity index 81% rename from modules/ui/GenerateCaptionsWindow.py rename to modules/ui/CtkGenerateCaptionsWindowView.py index 1690879f1..09d82f74b 100644 --- a/modules/ui/GenerateCaptionsWindow.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/GenerateMasksWindow.py b/modules/ui/CtkGenerateMasksWindowView.py similarity index 84% rename from modules/ui/GenerateMasksWindow.py rename to modules/ui/CtkGenerateMasksWindowView.py index daff0d3d5..631179fac 100644 --- a/modules/ui/GenerateMasksWindow.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 new file mode 100644 index 000000000..8caa1f171 --- /dev/null +++ b/modules/ui/CtkLoraTabView.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/CtkModelTabView.py b/modules/ui/CtkModelTabView.py new file mode 100644 index 000000000..5b43b7dca --- /dev/null +++ b/modules/ui/CtkModelTabView.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/CtkMuonAdamWindowView.py b/modules/ui/CtkMuonAdamWindowView.py new file mode 100644 index 000000000..3dc48ab74 --- /dev/null +++ b/modules/ui/CtkMuonAdamWindowView.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/CtkOptimizerParamsWindowView.py b/modules/ui/CtkOptimizerParamsWindowView.py new file mode 100644 index 000000000..ecc1f6a38 --- /dev/null +++ b/modules/ui/CtkOptimizerParamsWindowView.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/CtkProfilingWindowView.py b/modules/ui/CtkProfilingWindowView.py new file mode 100644 index 000000000..15d5055a0 --- /dev/null +++ b/modules/ui/CtkProfilingWindowView.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/CtkSampleFrameView.py b/modules/ui/CtkSampleFrameView.py new file mode 100644 index 000000000..167b25692 --- /dev/null +++ b/modules/ui/CtkSampleFrameView.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/CtkSampleParamsWindowView.py b/modules/ui/CtkSampleParamsWindowView.py new file mode 100644 index 000000000..8229a19f6 --- /dev/null +++ b/modules/ui/CtkSampleParamsWindowView.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/CtkSampleWindowView.py b/modules/ui/CtkSampleWindowView.py new file mode 100644 index 000000000..3b67f03d5 --- /dev/null +++ b/modules/ui/CtkSampleWindowView.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/CtkSamplingTabView.py b/modules/ui/CtkSamplingTabView.py new file mode 100644 index 000000000..dfe1a704a --- /dev/null +++ b/modules/ui/CtkSamplingTabView.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/SchedulerParamsWindow.py b/modules/ui/CtkSchedulerParamsWindowView.py similarity index 53% rename from modules/ui/SchedulerParamsWindow.py rename to modules/ui/CtkSchedulerParamsWindowView.py index f96ed4876..9a7c41b96 100644 --- a/modules/ui/SchedulerParamsWindow.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 new file mode 100644 index 000000000..69f2ae7d3 --- /dev/null +++ b/modules/ui/CtkTimestepDistributionWindowView.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/CtkTopBarView.py b/modules/ui/CtkTopBarView.py new file mode 100644 index 000000000..ee1bcfa75 --- /dev/null +++ b/modules/ui/CtkTopBarView.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/CtkTrainUIView.py b/modules/ui/CtkTrainUIView.py new file mode 100644 index 000000000..41af55230 --- /dev/null +++ b/modules/ui/CtkTrainUIView.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) + + 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") + + 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.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.controller.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.controller.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.controller.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.controller.train_config), self.ui_state) + + def create_training_tab(self, master) -> CtkTrainingTabView: + return CtkTrainingTabView(master, TrainingTabController(self.controller.train_config), self.ui_state) + + def create_cloud_tab(self, master) -> CtkCloudTabView: + 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) + 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.controller.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.controller.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.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")) + + 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/CtkTrainingTabView.py b/modules/ui/CtkTrainingTabView.py new file mode 100644 index 000000000..a126b2fa4 --- /dev/null +++ b/modules/ui/CtkTrainingTabView.py @@ -0,0 +1,67 @@ + +from modules.ui.BaseTrainingTabView import BaseTrainingTabView +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) + + 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_scheduler_params(self): + self.master.wait_window(self.controller.open_scheduler_params_window(self.master, self.ui_state, CtkSchedulerParamsWindowView)) + + 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/CtkVideoToolUIView.py b/modules/ui/CtkVideoToolUIView.py new file mode 100644 index 000000000..d6190772f --- /dev/null +++ b/modules/ui/CtkVideoToolUIView.py @@ -0,0 +1,104 @@ +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 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/ui/GenerateCaptionsWindowController.py b/modules/ui/GenerateCaptionsWindowController.py new file mode 100644 index 000000000..2b2411b28 --- /dev/null +++ b/modules/ui/GenerateCaptionsWindowController.py @@ -0,0 +1,28 @@ +class GenerateCaptionsWindowController: + def __init__(self, parent): + self.parent = parent + self.view = None + + 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 + + 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", + }[mode_str] + + self.parent.captioning_model.caption_folder( + sample_dir=path, + initial_caption=initial_caption, + caption_prefix=caption_prefix, + caption_postfix=caption_postfix, + mode=mode, + progress_callback=self.view.set_progress, + include_subdirectories=include_subdirectories, + ) + self.parent.load_image() diff --git a/modules/ui/GenerateMasksWindowController.py b/modules/ui/GenerateMasksWindowController.py new file mode 100644 index 000000000..6c154ef46 --- /dev/null +++ b/modules/ui/GenerateMasksWindowController.py @@ -0,0 +1,32 @@ +class GenerateMasksWindowController: + def __init__(self, parent): + self.parent = parent + self.view = None + + 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 + + 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", + "Create if absent": "fill", + "Add to existing": "add", + "Subtract from existing": "subtract", + "Blend with existing": "blend", + }[mode_str] + + self.parent.masking_model.mask_folder( + sample_dir=path, + prompts=[prompt], + mode=mode, + 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() diff --git a/modules/ui/LoraTab.py b/modules/ui/LoraTab.py deleted file mode 100644 index 6e9975476..000000000 --- a/modules/ui/LoraTab.py +++ /dev/null @@ -1,215 +0,0 @@ - -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), - ("LoKr", PeftType.LOKR), - ], 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" - elif peft_type == PeftType.LOKR: - name = "LoKr" - 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) - - # 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, 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.") - 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") - - # 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/LoraTabController.py b/modules/ui/LoraTabController.py new file mode 100644 index 000000000..fe7964520 --- /dev/null +++ b/modules/ui/LoraTabController.py @@ -0,0 +1,23 @@ + +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), + ("LoKr", PeftType.LOKR), + ] + + def get_lora_weight_dtypes(self) -> list[tuple[str, DataType]]: + return [ + ("float32", DataType.FLOAT_32), + ("bfloat16", DataType.BFLOAT_16), + ] diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py deleted file mode 100644 index ff17ea3ba..000000000 --- a/modules/ui/ModelTab.py +++ /dev/null @@ -1,688 +0,0 @@ - -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/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/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/OptimizerParamsWindowController.py b/modules/ui/OptimizerParamsWindowController.py new file mode 100644 index 000000000..37f838907 --- /dev/null +++ b/modules/ui/OptimizerParamsWindowController.py @@ -0,0 +1,51 @@ + +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 ( + OPTIMIZER_DEFAULT_PARAMETERS, + change_optimizer, + load_optimizer_defaults, + update_optimizer_config, +) + + +class OptimizerParamsWindowController: + 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 load_defaults(self, ui_state): + optimizer_config = load_optimizer_defaults(self.config) + ui_state.get_var("optimizer").update(optimizer_config) + + 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.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 + + return adam_config, current_optimizer + + def save_muon_adam_config(self, adam_config: 'TrainOptimizerConfig'): + self.config.optimizer.muon_adam_config = adam_config.to_dict() diff --git a/modules/ui/ProfilingWindow.py b/modules/ui/ProfilingWindow.py deleted file mode 100644 index 8d298abe3..000000000 --- a/modules/ui/ProfilingWindow.py +++ /dev/null @@ -1,57 +0,0 @@ -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/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/PySide6AdditionalEmbeddingsTabView.py b/modules/ui/PySide6AdditionalEmbeddingsTabView.py new file mode 100644 index 000000000..196526138 --- /dev/null +++ b/modules/ui/PySide6AdditionalEmbeddingsTabView.py @@ -0,0 +1,54 @@ +from modules.ui.AdditionalEmbeddingsTabController import AdditionalEmbeddingsTabController +from modules.ui.BaseAdditionalEmbeddingsTabView import BaseAdditionalEmbeddingsTabView, BaseEmbeddingWidgetView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState + +from PySide6.QtWidgets import QWidget + + +class PySide6AdditionalEmbeddingsTabView(PySide6ConfigListView, BaseAdditionalEmbeddingsTabView): + + def __init__(self, master, controller: AdditionalEmbeddingsTabController, ui_state): + PySide6ConfigListView.__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 PySide6EmbeddingWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + + +class PySide6EmbeddingWidgetView(BaseEmbeddingWidgetView, QWidget): + + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command, controller): + QWidget.__init__(self, master) + BaseEmbeddingWidgetView.__init__(self, pyside6_components) + + self.element = element + ui_state = PySide6UIState(element) + + pyside6_components._layout(self).setColumnStretch(0, 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 = 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): + 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 new file mode 100644 index 000000000..620e4faab --- /dev/null +++ b/modules/ui/PySide6CaptionUIView.py @@ -0,0 +1,12 @@ +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.\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) diff --git a/modules/ui/PySide6CloudTabView.py b/modules/ui/PySide6CloudTabView.py new file mode 100644 index 000000000..ac4606ea4 --- /dev/null +++ b/modules/ui/PySide6CloudTabView.py @@ -0,0 +1,41 @@ +from modules.ui.BaseCloudTabView import BaseCloudTabView +from modules.ui.CloudTabController import CloudTabController +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QWidget + + +class PySide6CloudTabView(BaseCloudTabView, QWidget, metaclass=QtABCMeta): + + def __init__(self, master, controller: CloudTabController, ui_state): + QWidget.__init__(self, master) + BaseCloudTabView.__init__(self, pyside6_components, controller) + + self.ui_state = ui_state + + 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.clear() + self.gpu_types_menu.addItems(self.controller.get_gpu_types()) + + def _make_reattach_frame(self, frame): + 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 = 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 new file mode 100644 index 000000000..919450d01 --- /dev/null +++ b/modules/ui/PySide6ConceptTabView.py @@ -0,0 +1,175 @@ +from modules.ui.BaseConceptTabView import BaseConceptTabView, BaseConceptWidgetView +from modules.ui.ConceptTabController import ConceptTabController +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 + +from PIL.ImageQt import ImageQt +from PySide6.QtGui import QPixmap +from PySide6.QtWidgets import QCheckBox, QComboBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QWidget + + +class PySide6ConceptTabView(PySide6ConfigListView, BaseConceptTabView): + + def __init__(self, master, controller: ConceptTabController, ui_state): + # 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) + + PySide6ConfigListView.__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._add_search_bar() + + 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 PySide6ConceptWidgetView(master, element, i, open_command, remove_command, clone_command, save_command, self.controller) + + def _add_search_bar(self): + 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 _update_filters(self): + self._create_element_list(search=self.search_var.get(), + type=self.filter_var.get(), + show_disabled=self.show_disabled_var.get()) + + def _reset_filters(self): + 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() + + +class PySide6ConceptWidgetView(BaseConceptWidgetView, QWidget): + + def __init__(self, master, concept, i, open_command, remove_command, clone_command, save_command, controller): + QWidget.__init__(self, master) + BaseConceptWidgetView.__init__(self, pyside6_components, concept) + self.ui_state = PySide6UIState(concept) + self.image_ui_state = PySide6UIState(concept.image) + self.text_ui_state = PySide6UIState(concept.text) + self.i = i + + self.setFixedSize(160, 180) + + 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) + ) + + 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.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 + 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 new file mode 100644 index 000000000..87355a9a6 --- /dev/null +++ b/modules/ui/PySide6ConceptWindowView.py @@ -0,0 +1,207 @@ +import threading + +from modules.ui.BaseConceptWindowView import BaseConceptWindowView +from modules.ui.ConceptWindowController import ConceptWindowController +from modules.util.ui import pyside6_components + +from matplotlib import pyplot as plt +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, + controller: ConceptWindowController, + ui_state, + image_ui_state, + text_ui_state, + ): + QDialog.__init__(self, parent) + BaseConceptWindowView.__init__(self, pyside6_components) + + self.controller = controller + self.image_preview_file_index = 0 + self._preview_augmentations = True + self.bucket_fig = None + + 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 + ) + 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) + + 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) + 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_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() + + + 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) + 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 + ) + 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 + self.accept() diff --git a/modules/ui/PySide6ConfigListView.py b/modules/ui/PySide6ConfigListView.py new file mode 100644 index 000000000..6a369d6c7 --- /dev/null +++ b/modules/ui/PySide6ConfigListView.py @@ -0,0 +1,98 @@ +import contextlib +from abc import ABC + +from modules.ui.BaseConfigListView import BaseConfigListView +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QInputDialog, QWidget + + +class PySide6ConfigListView(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, pyside6_components) + + 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, + 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 = 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): + scroll, content = pyside6_components.scrollable_frame(master) + pyside6_components._layout(master).addWidget(scroll, 1, 0) + if self.is_full_width: + pyside6_components._layout(content).setColumnStretch(0, 1) + content._scroll_area = scroll + return content + + def _wait_for_window(self, window): + window.exec() + + def _remove_widget_from_layout(self, widget): + widget.hide() + + def _destroy_widget(self, widget): + with contextlib.suppress(RuntimeError, AttributeError): + widget.hide() + widget.deleteLater() + + def _destroy_frame(self, frame): + 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): + 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 new file mode 100644 index 000000000..959a2ea31 --- /dev/null +++ b/modules/ui/PySide6ConvertModelUIView.py @@ -0,0 +1,32 @@ +from modules.ui.BaseConvertModelUIView import BaseConvertModelUIView +from modules.ui.ConvertModelUIController import ConvertModelUIController +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState + +from PySide6.QtWidgets import QDialog, QGridLayout, QWidget + + +class PySide6ConvertModelUIView(BaseConvertModelUIView, QDialog): + def __init__(self, parent, controller: ConvertModelUIController): + QDialog.__init__(self, parent) + BaseConvertModelUIView.__init__(self, pyside6_components) + + ui_state = PySide6UIState(controller.convert_model_args) + + self.setWindowTitle("Convert models") + self.resize(600, 380) + + _pad = pyside6_components.PAD + outer = QGridLayout(self) + outer.setContentsMargins(_pad, _pad, _pad, _pad) + + frame = QWidget(self) + lo = pyside6_components._layout(frame) + lo.setColumnStretch(1, 1) + outer.addWidget(frame, 0, 0) + + self.build_content(frame, controller, ui_state) + lo.setRowStretch(lo.rowCount(), 1) + + def set_converting(self, active): + self.button.setEnabled(not active) 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..842017697 --- /dev/null +++ b/modules/ui/PySide6LoraTabView.py @@ -0,0 +1,47 @@ +from modules.ui.BaseLoraTabView import BaseLoraTabView +from modules.ui.LoraTabController import LoraTabController +from modules.util.enum.ModelType import PeftType +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QWidget + + +class PySide6LoraTabView(BaseLoraTabView, QWidget): + + def __init__(self, master, controller: LoraTabController, ui_state): + QWidget.__init__(self, master) + BaseLoraTabView.__init__(self, pyside6_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 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 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 new file mode 100644 index 000000000..0835a586f --- /dev/null +++ b/modules/ui/PySide6ModelTabView.py @@ -0,0 +1,41 @@ +from modules.ui.BaseModelTabView import BaseModelTabView +from modules.ui.ModelTabController import ModelTabController +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QWidget + + +class PySide6ModelTabView(BaseModelTabView, QWidget, metaclass=QtABCMeta): + + def __init__(self, master, controller: ModelTabController, ui_state): + QWidget.__init__(self, master) + BaseModelTabView.__init__(self, pyside6_components) + + self.master = master + self.controller = controller + self.ui_state = ui_state + self.scroll_frame = None + self.refresh_ui() + + def _make_svd_frames(self, parent, row: int): + 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 is not None: + self.scroll_frame.hide() + self.scroll_frame.deleteLater() + + scroll, frame = pyside6_components.scrollable_frame(self) + pyside6_components._layout(self).addWidget(scroll, 0, 0) + self.scroll_frame = scroll + + frame_lo = pyside6_components._layout(frame) + frame_lo.setColumnStretch(1, 10) + frame_lo.setColumnStretch(4, 1) + + self.build_content(frame, self.controller, self.ui_state) diff --git a/modules/ui/PySide6MuonAdamWindowView.py b/modules/ui/PySide6MuonAdamWindowView.py new file mode 100644 index 000000000..0deb4945a --- /dev/null +++ b/modules/ui/PySide6MuonAdamWindowView.py @@ -0,0 +1,29 @@ +from modules.ui.BaseMuonAdamWindowView import BaseMuonAdamWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton + + +class PySide6MuonAdamWindowView(BaseMuonAdamWindowView, QDialog): + def __init__(self, parent, controller: MuonAdamWindowController, ui_state): + QDialog.__init__(self, parent) + BaseMuonAdamWindowView.__init__(self, pyside6_components) + + self.setWindowTitle(controller.get_title()) + self.resize(800, 500) + + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + + 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) + + 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 new file mode 100644 index 000000000..bb6682f01 --- /dev/null +++ b/modules/ui/PySide6OptimizerParamsWindowView.py @@ -0,0 +1,79 @@ +from modules.ui.BaseOptimizerParamsWindowView import BaseOptimizerParamsWindowView +from modules.ui.MuonAdamWindowController import MuonAdamWindowController +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +from modules.ui.PySide6MuonAdamWindowView import PySide6MuonAdamWindowView +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState + +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QWidget + + +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._dynamic_frame = None + + self.setWindowTitle("Optimizer Settings") + self.resize(800, 500) + + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + + 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) + + ok = QPushButton("ok", self) + ok.clicked.connect(self._on_close) + outer.addWidget(ok, 1, 0) + + self.build_content(self._frame, controller, ui_state, self.optimizer_ui_state, + self.on_optimizer_change, self._load_defaults) + self._rebuild_dynamic_ui() + + + def _rebuild_dynamic_ui(self): + if self._dynamic_frame is not None: + self._dynamic_frame.setParent(None) + + self._dynamic_frame = QWidget(self._frame) + pyside6_components._layout(self._frame).addWidget(self._dynamic_frame, 1, 0, 1, 5) + + 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() + + 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_close(self): + self.controller.on_close() + self.accept() + + def toggle_muon_adam_button(self): + if self.muon_adam_button is not None: + muon_with_adam = self.optimizer_ui_state.get_var("MuonWithAuxAdam").get() + 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() + 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 new file mode 100644 index 000000000..b8c0a41e1 --- /dev/null +++ b/modules/ui/PySide6ProfilingWindowView.py @@ -0,0 +1,49 @@ +import contextlib + +from modules.ui.BaseProfilingWindowView import BaseProfilingWindowView +from modules.ui.ProfilingWindowController import ProfilingWindowController +from modules.util.ui import pyside6_components + +from PySide6.QtCore import Qt +from PySide6.QtWidgets import QGridLayout, QWidget + + +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.setWindowTitle("Profiling") + self.resize(512, 512) + + outer = QGridLayout(self) + outer.setRowStretch(2, 1) + + self._bottom_bar = QWidget(self) + QGridLayout(self._bottom_bar) + outer.addWidget(self._bottom_bar, 3, 0) + + self.build_content(self, self._bottom_bar, controller) + + def set_message(self, text: str): + self._message_label.setText(text) + + def set_profiling_active(self, active: bool): + if active: + 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.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 new file mode 100644 index 000000000..a936cccb3 --- /dev/null +++ b/modules/ui/PySide6SampleFrameView.py @@ -0,0 +1,46 @@ +from modules.ui.BaseSampleFrameView import BaseSampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QWidget + + +class PySide6SampleFrameView(BaseSampleFrameView, QWidget): + def __init__( + self, + parent: QWidget, + controller: SampleFrameController, + ui_state, + include_prompt: bool = True, + include_settings: bool = True, + ): + QWidget.__init__(self, parent) + BaseSampleFrameView.__init__(self, pyside6_components) + + 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 = 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 = 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 new file mode 100644 index 000000000..054712ccc --- /dev/null +++ b/modules/ui/PySide6SampleParamsWindowView.py @@ -0,0 +1,27 @@ +from modules.ui.BaseSampleParamsWindowView import BaseSampleParamsWindowView +from modules.ui.PySide6SampleFrameView import PySide6SampleFrameView +from modules.ui.SampleFrameController import SampleFrameController +from modules.ui.SampleParamsWindowController import SampleParamsWindowController +from modules.util.ui import pyside6_components + +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QWidget + + +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.setWindowTitle("Sample") + self.resize(800, 500) + + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + outer.setColumnStretch(0, 1) + + frame = PySide6SampleFrameView(self, SampleFrameController(controller.sample, controller.model_type), ui_state) + outer.addWidget(frame, 0, 0) + + 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 new file mode 100644 index 000000000..a41777f1d --- /dev/null +++ b/modules/ui/PySide6SampleWindowView.py @@ -0,0 +1,85 @@ +import threading + +from modules.modelSampler.BaseModelSampler import ( + ModelSamplerOutput, +) +from modules.ui.BaseSampleWindowView import BaseSampleWindowView +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 pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState + +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 PySide6SampleWindowView(BaseSampleWindowView, QDialog): + def __init__(self, parent, controller: SampleWindowController): + QDialog.__init__(self, parent) + BaseSampleWindowView.__init__(self, pyside6_components) + + self.setWindowTitle("Sample") + self.resize(1200, 800) + + 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) + + outer = QGridLayout(self) + outer.setRowStretch(1, 1) + outer.setColumnStretch(1, 1) + + 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) + + + 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.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 new file mode 100644 index 000000000..e83c4a727 --- /dev/null +++ b/modules/ui/PySide6SamplingTabView.py @@ -0,0 +1,66 @@ +from modules.ui.BaseSamplingTabView import BaseSampleWidgetView, BaseSamplingTabView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.ui.PySide6SampleParamsWindowView import PySide6SampleParamsWindowView +from modules.ui.SamplingTabController import SamplingTabController +from modules.util.ui import pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QWidget + + +class PySide6SamplingTabView(PySide6ConfigListView, BaseSamplingTabView): + + def __init__(self, master, controller: SamplingTabController, ui_state): + PySide6ConfigListView.__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): + 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 PySide6SampleWidgetView(master, element, i, open_command, remove_command, clone_command, save_command) + + +class PySide6SampleWidgetView(BaseSampleWidgetView, QWidget, metaclass=QtABCMeta): + + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + QWidget.__init__(self, master) + BaseSampleWidgetView.__init__(self, pyside6_components) + + from modules.util.ui.PySide6UIState import PySide6UIState + self.element = element + self.ui_state = PySide6UIState(element) + + 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.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): + 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 new file mode 100644 index 000000000..96deaae40 --- /dev/null +++ b/modules/ui/PySide6SchedulerParamsWindowView.py @@ -0,0 +1,93 @@ +from modules.ui.BaseSchedulerParamsWindowView import BaseKvParamsView, BaseSchedulerParamsWindowView +from modules.ui.PySide6ConfigListView import PySide6ConfigListView +from modules.ui.SchedulerParamsWindowController import KvParamsController, SchedulerParamsWindowController +from modules.util.ui import pyside6_components +from modules.util.ui.PySide6UIState import PySide6UIState + +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton, QScrollArea, QWidget + + +class PySide6KvParamsView(PySide6ConfigListView, BaseKvParamsView): + def __init__(self, master, controller: KvParamsController, ui_state): + 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, 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 PySide6KvWidget(master, element, i, open_command, remove_command, clone_command, save_command) + + +class PySide6KvWidget(QWidget): + def __init__(self, master, element, i, open_command, remove_command, clone_command, save_command): + super().__init__(master) + self.element = element + self.ui_state = PySide6UIState(element) + self.i = i + self.save_command = save_command + + lo = pyside6_components._layout(self) + lo.setColumnStretch(1, 1) + lo.setColumnStretch(2, 1) + + pyside6_components.colored_icon_button(self, 0, 0, "X", "#C00000", lambda: remove_command(self.i)) + + # Key + 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 + 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): + pyside6_components._layout(self.parent()).addWidget(self, getattr(self, 'visible_index', self.i), 0) + self.show() + + def destroy(self): + self.deleteLater() + + +class PySide6SchedulerParamsWindowView(BaseSchedulerParamsWindowView, QDialog): + def __init__(self, parent, controller: SchedulerParamsWindowController, ui_state): + QDialog.__init__(self, parent) + BaseSchedulerParamsWindowView.__init__(self, pyside6_components) + + self.setWindowTitle("Learning Rate Scheduler Settings") + self.resize(800, 500) + + outer = QGridLayout(self) + outer.setRowStretch(0, 1) + + scroll = QScrollArea(self) + scroll.setWidgetResizable(True) + inner = QWidget() + scroll.setWidget(inner) + inner_lo = pyside6_components._layout(inner) + inner_lo.setColumnStretch(1, 1) + + 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 new file mode 100644 index 000000000..3cd11d9ef --- /dev/null +++ b/modules/ui/PySide6TimestepDistributionWindowView.py @@ -0,0 +1,48 @@ +from modules.ui.BaseTimestepDistributionWindowView import BaseTimestepDistributionWindowView +from modules.ui.TimestepDistributionWindowController import TimestepDistributionWindowController +from modules.util.ui import pyside6_components + +from matplotlib import pyplot as plt +from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg +from PySide6.QtWidgets import QDialog, QGridLayout, QPushButton + + +class PySide6TimestepDistributionWindowView(BaseTimestepDistributionWindowView, QDialog): + def __init__(self, parent, controller: TimestepDistributionWindowController, ui_state): + QDialog.__init__(self, parent) + BaseTimestepDistributionWindowView.__init__(self, pyside6_components) + + self.setWindowTitle("Timestep Distribution") + self.resize(900, 600) + self._controller = controller + + outer = QGridLayout(self) + outer.setRowStretch(0, 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) + + fig, self._ax = plt.subplots() + self._canvas = FigureCanvasQTAgg(fig) + lo.addWidget(self._canvas, 0, 3, 8, 1) + self._update_preview() + + update_btn = QPushButton("Update Preview", frame) + update_btn.clicked.connect(self._update_preview) + lo.addWidget(update_btn, 8, 3) + + outer.addWidget(scroll, 0, 0) + + ok = QPushButton("ok", self) + ok.clicked.connect(self.accept) + outer.addWidget(ok, 1, 0) + + + 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..e852723c8 --- /dev/null +++ b/modules/ui/PySide6TopBarView.py @@ -0,0 +1,54 @@ +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 pyside6_components + +from PySide6.QtWidgets import QInputDialog, QWidget + + +class PySide6TopBarView(BaseTopBarView, QWidget): + + 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], + ): + 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(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): + from modules.util.ui.PySide6UIState import PySide6UIState + return PySide6UIState(data) + + def _get_dropdown_text(self, widget) -> str: + return widget.currentText() + + def _setup_frame_column_weight(self): + pyside6_components._layout(self.frame).setColumnStretch(5, 1) + + def _forget_dropdown(self): + 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): + 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 new file mode 100644 index 000000000..d189d05b9 --- /dev/null +++ b/modules/ui/PySide6TrainUIView.py @@ -0,0 +1,381 @@ +from collections.abc import Callable +from pathlib import Path + +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.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 +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 pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta +from modules.util.ui.PySide6UIState import PySide6UIState + +from PySide6.QtCore import QTimer +from PySide6.QtGui import QIcon +from PySide6.QtWidgets import QFileDialog, QGridLayout, QMainWindow, QMessageBox, QTabWidget, QWidget + + +class PySide6TrainView(BaseTrainUIView, QMainWindow, metaclass=QtABCMeta): + def __init__(self): + QMainWindow.__init__(self) + + train_config = TrainConfig.default_values() + ui_state = PySide6UIState(train_config) + controller = TrainUIController(train_config) + + BaseTrainUIView.__init__(self, pyside6_components, controller, ui_state) + self.controller.view = self + + self.setWindowTitle("OneTrainer") + self.setWindowIcon(QIcon("resources/icons/icon.png")) + self.resize(1100, 740) + + self.status_label = None + self.eta_label = None + self.training_button = None + self.export_button = 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 + + 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.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) + + bottom = self._build_bottom_bar(central) + central_lo.addWidget(bottom, 2, 0) + + self._create_tabs() + self.change_training_method(self.controller.train_config.training_method) + + self._profiling_controller = ProfilingWindowController() + self.profiling_window = PySide6ProfilingWindowView(self, self._profiling_controller) + + 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 closeEvent(self, event): + self.top_bar_component.save_default() + self.controller._stop_always_on_tensorboard() + 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): + # 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.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) + self.eta_label.setText(f"ETA: {eta_str}" if eta_str is not None else "") + + def schedule_on_main_thread(self, fn: Callable): + # 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 + + 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) + 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, PySide6CaptionUIView)) + + def open_video_tool(self): + 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, PySide6ConvertModelUIView)) + + def open_sampling_tool(self): + self.controller.open_sampling_tool(self, PySide6SampleWindowView) + + def open_manual_sample_window(self): + self.controller.open_manual_sample_window(self, PySide6SampleWindowView) + + def wait_window(self, window): + window.exec() + + def show_window(self, window): + window.show() + + def connect_window_closed(self, window, callback): + window.finished.connect(lambda _: callback()) + + # --- PySide6 layout builders --- + + def _build_top_bar(self, master): + return PySide6TopBarView( + master, + TopBarController(self.controller.train_config), + self.ui_state, + self.change_model_type, + self.change_training_method, + self.load_preset, + ) + + def _build_bottom_bar(self, parent): + frame = QWidget(parent) + lo = QGridLayout(frame) + lo.setColumnStretch(0, 1) + lo.setColumnStretch(2, 2) + + 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") + return frame + + 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) + 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): + self.build_tools_tab_content(frame, self.controller, self.ui_state) + pyside6_components._pack_form(frame) + + def _configure_embedding_frame(self, frame): + 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) + self.tabview.addTab(general_page, "general") + self._tab_widgets["general"] = general_page + + self.model_tab = PySide6ModelTabView(None, ModelTabController(self.controller.train_config), self.ui_state) + self.tabview.addTab(self.model_tab, "model") + self._tab_widgets["model"] = self.model_tab + + data_page = self._create_scrollable_tab(self._configure_data_frame) + self.tabview.addTab(data_page, "data") + self._tab_widgets["data"] = data_page + + concepts_page = QWidget() + self.concepts_tab = PySide6ConceptTabView(concepts_page, ConceptTabController(self.controller.train_config), self.ui_state) + self.tabview.addTab(concepts_page, "concepts") + self._tab_widgets["concepts"] = concepts_page + + self.training_tab = PySide6TrainingTabView(None, TrainingTabController(self.controller.train_config), self.ui_state) + self.tabview.addTab(self.training_tab, "training") + self._tab_widgets["training"] = self.training_tab + + sampling_page = self.create_sampling_tab() + self.tabview.addTab(sampling_page, "sampling") + self._tab_widgets["sampling"] = sampling_page + + backup_page = self._create_scrollable_tab(self._configure_backup_frame) + self.tabview.addTab(backup_page, "backup") + self._tab_widgets["backup"] = backup_page + + tools_page = self._create_scrollable_tab(self._configure_tools_frame) + self.tabview.addTab(tools_page, "tools") + self._tab_widgets["tools"] = tools_page + + additional_embeddings_page = QWidget() + self.additional_embeddings_tab = PySide6AdditionalEmbeddingsTabView( + additional_embeddings_page, + AdditionalEmbeddingsTabController(self.controller.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.controller.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) + + sampling_container = QWidget(tab_page) + tab_lo.addWidget(sampling_container, 1, 0) + self.sampling_tab = PySide6SamplingTabView( + sampling_container, SamplingTabController(self.controller.train_config), self.ui_state + ) + + return tab_page + + def open_profiling_tool(self): + 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() + + 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._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._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.controller.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 self.additional_embeddings_tab: + self.additional_embeddings_tab.refresh_ui() + + def _set_training_button_style(self, mode: str): + if not self.training_button: + return + 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, _ = 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): + dir_path = QFileDialog.getExistingDirectory(self, "Select Directory to Save Debug Package", ".") + if not dir_path: + return + self.controller.generate_debug_package(Path(dir_path) / "OneTrainer_debug_report.zip") diff --git a/modules/ui/PySide6TrainingTabView.py b/modules/ui/PySide6TrainingTabView.py new file mode 100644 index 000000000..d115c678d --- /dev/null +++ b/modules/ui/PySide6TrainingTabView.py @@ -0,0 +1,84 @@ +from modules.ui.BaseTrainingTabView import BaseTrainingTabView +from modules.ui.OptimizerParamsWindowController import OptimizerParamsWindowController +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 pyside6_components +from modules.util.ui.pyside6_abc import QtABCMeta + +from PySide6.QtWidgets import QScrollArea, QSizePolicy, QWidget + + +class PySide6TrainingTabView(BaseTrainingTabView, QWidget, metaclass=QtABCMeta): + + def __init__(self, master, controller: TrainingTabController, ui_state): + QWidget.__init__(self, master) + BaseTrainingTabView.__init__(self, pyside6_components) + + self.master = master + self.controller = controller + self.ui_state = ui_state + self.scroll_frame = None + self.refresh_ui() + + def refresh_ui(self): + 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) + + 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) + + self.build(column_0, column_1, column_2, self.controller, self.ui_state) + + for col_widget in (column_0, column_1, column_2): + lo = pyside6_components._layout(col_widget) + lo.setRowStretch(lo.rowCount(), 1) + + def restore_optimizer_config(self, variable: str): + self.controller.restore_optimizer_config(self.ui_state) + + def restore_scheduler(self, variable: str): + if not hasattr(self, 'lr_scheduler_adv_comp'): + return + self.lr_scheduler_adv_comp.setEnabled(self.controller.is_custom_scheduler_value(variable)) + + def open_optimizer_params(self): + PySide6OptimizerParamsWindowView(self, OptimizerParamsWindowController(self.controller.config), self.ui_state).exec() + + def open_scheduler_params(self): + PySide6SchedulerParamsWindowView(self, SchedulerParamsWindowController(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/PySide6VideoToolUIView.py b/modules/ui/PySide6VideoToolUIView.py new file mode 100644 index 000000000..88ac82f62 --- /dev/null +++ b/modules/ui/PySide6VideoToolUIView.py @@ -0,0 +1,126 @@ +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 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, + QGridLayout, + QLabel, + 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 + 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_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) + 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 update_status(self, status_text: str): + self._status_box.append(status_text) + + def clear_status(self): + self._status_box.clear() + + def update_preview(self, preview_image, label_text: str): + 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/ui/SampleFrame.py b/modules/ui/SampleFrame.py deleted file mode 100644 index 297caac29..000000000 --- a/modules/ui/SampleFrame.py +++ /dev/null @@ -1,134 +0,0 @@ -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/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/SampleParamsWindow.py b/modules/ui/SampleParamsWindow.py deleted file mode 100644 index 2b0b3f3f1..000000000 --- a/modules/ui/SampleParamsWindow.py +++ /dev/null @@ -1,39 +0,0 @@ -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/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/SampleWindow.py b/modules/ui/SampleWindowController.py similarity index 59% rename from modules/ui/SampleWindow.py rename to modules/ui/SampleWindowController.py index 0f91ad2fa..d1d24d643 100644 --- a/modules/ui/SampleWindow.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/SamplingTab.py b/modules/ui/SamplingTab.py deleted file mode 100644 index 5a3c44f08..000000000 --- a/modules/ui/SamplingTab.py +++ /dev/null @@ -1,124 +0,0 @@ -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/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/TimestepDistributionWindow.py b/modules/ui/TimestepDistributionWindow.py deleted file mode 100644 index 21e41ce3e..000000000 --- a/modules/ui/TimestepDistributionWindow.py +++ /dev/null @@ -1,186 +0,0 @@ - -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/TimestepDistributionWindowController.py b/modules/ui/TimestepDistributionWindowController.py new file mode 100644 index 000000000..682508599 --- /dev/null +++ b/modules/ui/TimestepDistributionWindowController.py @@ -0,0 +1,68 @@ + +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.TimestepDistribution import TimestepDistribution + +import torch +from torch import Tensor + + +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 TimestepDistributionWindowController: + def __init__(self, config: TrainConfig): + self.train_config = config + + def get_distribution_options(self) -> list[str]: + return [str(x) for x in list(TimestepDistribution)] + + def generate_preview_data(self) -> Tensor: + generator = TimestepGenerator( + 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, + ) + return generator.generate() diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py deleted file mode 100644 index 820fdb71a..000000000 --- a/modules/ui/TopBar.py +++ /dev/null @@ -1,260 +0,0 @@ -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/TopBarController.py b/modules/ui/TopBarController.py new file mode 100644 index 000000000..8a4946db4 --- /dev/null +++ b/modules/ui/TopBarController.py @@ -0,0 +1,80 @@ +import os +import webbrowser + +from modules.util import path_util +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.path_util import write_json_atomic + + +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), + ("Anima", ModelType.ANIMA), + ("Z-Image", ModelType.Z_IMAGE), + ("Ernie Image", ModelType.ERNIE), + ] + + def get_training_methods(self, model_type: ModelType) -> list[tuple[str, TrainingMethod]]: + labels = { + TrainingMethod.FINE_TUNE: "Fine Tune", + TrainingMethod.LORA: "LoRA", + TrainingMethod.EMBEDDING: "Embedding", + TrainingMethod.FINE_TUNE_VAE: "Fine Tune VAE", + } + return [(labels[m], m) for m in model_type.supported_training_methods()] + + 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(dir, path) + if path.endswith(".json") and os.path.isfile(path): + name = os.path.basename(path) + name = os.path.splitext(name)[0] + configs.append((name, path)) + configs.sort() + return configs + + 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_default(self): + self.save_to_file("#") + self.save_secrets("secrets.json") diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py deleted file mode 100644 index b9fa0c04a..000000000 --- a/modules/ui/TrainUI.py +++ /dev/null @@ -1,905 +0,0 @@ -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.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() - - 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, 2, 0, "Rolling Backup Count", - tooltip="Defines the number of backups to keep if rolling backups are enabled") - components.entry(frame, 2, 1, self.ui_state, "rolling_backup_count") - - # backup before save - components.label(frame, 3, 0, "Backup Before Save", - tooltip="Create a full backup before saving the final model") - components.switch(frame, 3, 1, self.ui_state, "backup_before_save") - - # save after - components.label(frame, 4, 0, "Save Every", - tooltip="The interval used when automatically saving the model during training") - components.time_entry(frame, 4, 1, self.ui_state, "save_every", "save_every_unit") - - # save now - components.button(frame, 4, 3, "save now", self.save_now) - - # skip save - components.label(frame, 5, 0, "Skip First", - tooltip="Start saving automatically after this interval has elapsed") - components.entry(frame, 5, 1, self.ui_state, "save_skip_first", width=50, sticky="nw") - - # 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, 6, 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: - assert self.start_time is not None and self.start_total_steps is not None - - spent_total = time.monotonic() - self.start_time - - # 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) - - 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) - 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): - # 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) - - 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) - - # Reset session tracking - actual values captured on first progress callback - self.start_total_steps = None - 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/TrainUIController.py b/modules/ui/TrainUIController.py new file mode 100644 index 000000000..a6bf4f21b --- /dev/null +++ b/modules/ui/TrainUIController.py @@ -0,0 +1,279 @@ +import datetime +import json +import os +import subprocess +import sys +import threading +import time +import traceback +import webbrowser +from pathlib import Path + +import scripts.generate_debug_report +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.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress +from modules.util.ui.validation import flush_and_validate_all + +import torch + + +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 = config.workspace_dir + self.start_time: float | None = None + self.start_total_steps: int | None = None + + 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 + 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 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: + assert self.start_time is not None and self.start_total_steps is not None + + spent_total = time.monotonic() - self.start_time + + # 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) + + 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) + 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 _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 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, parent, view_cls): + training_callbacks = self.training_callbacks + training_commands = self.training_commands + + if training_callbacks and training_commands: + controller = SampleWindowController( + self.train_config, + use_external_model=True, + callbacks=training_callbacks, + commands=training_commands, + ) + 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 + + 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.view.get_cloud_reattach()) + try: + trainer.start() + if self.train_config.cloud.enabled: + self.view.sync_cloud_secrets() + + # Reset session tracking - actual values captured on first progress callback + self.start_total_steps = None + self.start_time = time.monotonic() + trainer.train() + except Exception: + if self.train_config.cloud.enabled: + self.view.sync_cloud_secrets() + 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") + + # 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.view.schedule_on_main_thread(self._start_always_on_tensorboard) + + def start_training(self): + if self.training_thread is None: + self.view.save_default() + + errors = flush_and_validate_all() + if errors: + self.view.show_validation_errors(errors) + return + + 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() + + self.training_commands = TrainCommands() + torch_gc() + + self.training_thread = threading.Thread(target=self.__training_thread_function) + self.training_thread.start() + else: + self.view.on_training_stopping() + self.on_update_status("Stopping ...") + self.training_commands.stop() diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py deleted file mode 100644 index c839abf73..000000000 --- a/modules/ui/TrainingTab.py +++ /dev/null @@ -1,862 +0,0 @@ -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) - row = 0 - - if supports_training: - components.label(frame, row, 0, "Train Text Encoder", - tooltip="Enables training the text encoder model") - components.switch(frame, row, 1, self.ui_state, "text_encoder.train") - row += 1 - - # dropout - components.label(frame, row, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, row, 1, self.ui_state, "text_encoder.dropout_probability") - row += 1 - - if supports_training: - # train text encoder epochs - components.label(frame, row, 0, "Stop Training After", - tooltip="When to stop training the text encoder") - 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, row, 0, "Text Encoder Learning Rate", - tooltip="The learning rate of the text encoder. Overrides the base 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, row, 0, "Clip Skip", - tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, row, 1, self.ui_state, "text_encoder_layer_skip") - row += 1 - - 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/TrainingTabController.py b/modules/ui/TrainingTabController.py new file mode 100644 index 000000000..013c4d377 --- /dev/null +++ b/modules/ui/TrainingTabController.py @@ -0,0 +1,35 @@ + +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) diff --git a/modules/ui/VideoToolUI.py b/modules/ui/VideoToolUIController.py similarity index 51% rename from modules/ui/VideoToolUI.py rename to modules/ui/VideoToolUIController.py index c3291e6ea..294f458f4 100644 --- a/modules/ui/VideoToolUI.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/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/PrefetchIterator.py b/modules/util/PrefetchIterator.py new file mode 100644 index 000000000..aec97f3ee --- /dev/null +++ b/modules/util/PrefetchIterator.py @@ -0,0 +1,70 @@ +import queue +import threading +from collections.abc import Iterable, Iterator +from contextlib import nullcontext, suppress + +import torch + + +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. + + 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): + 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() + + 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. + while not stop_event.is_set(): + with suppress(queue.Full): + q.put(value, timeout=self._stop_poll_interval) + return True + return False + + def producer(): + 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() + + 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/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/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/config/TrainConfig.py b/modules/util/config/TrainConfig.py index f52988502..9af063c68 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 @@ -385,7 +401,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 @@ -404,7 +421,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 @@ -441,6 +459,7 @@ class TrainConfig(BaseConfig): timestep_distribution: TimestepDistribution min_noising_strength: float max_noising_strength: float + cep_gamma: float noising_weight: float noising_bias: float @@ -569,7 +588,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, @@ -581,6 +600,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, } ) @@ -727,12 +747,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 @@ -800,6 +822,50 @@ 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") + + 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 + + if "dataloader_threads" in migrated_data: + migrated_data["caching_threads"] = migrated_data.pop("dataloader_threads") + + 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, @@ -969,10 +1035,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)) @@ -980,7 +1043,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 @@ -997,7 +1061,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)) @@ -1033,6 +1098,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_gamma", 0.0, float, False)) # unet diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..b1100acf5 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -110,8 +110,12 @@ 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') + # 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.caching_threads > 1 and any(part.offload_fraction > 0 for part in config.model_part_configs()): + raise RuntimeError('layer offloading can not be activated if "caching_threads" > 1') if train_progress is None: train_progress = TrainProgress() @@ -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/dtype_util.py b/modules/util/dtype_util.py index d0df10c59..b6d6b2643 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. + 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: + 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( @@ -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)) 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/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 a3ad940ec..cb8c778ee 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' @@ -37,6 +39,8 @@ class ModelType(Enum): QWEN = 'QWEN' + ANIMA = 'ANIMA' + Z_IMAGE = 'Z_IMAGE' ERNIE = 'ERNIE' @@ -97,6 +101,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 +164,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() \ @@ -166,6 +174,60 @@ 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_anima() 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.ANIMA: ("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/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) 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/modules/util/path_util.py b/modules/util/path_util.py index 8fb80fdc1..796268efc 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 @@ -57,3 +58,8 @@ def is_supported_video_extension(extension: str) -> bool: def supported_caption_extensions() -> set[str]: return SUPPORTED_CAPTION_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/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) 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/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/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 87% rename from modules/util/ui/components.py rename to modules/util/ui/ctk_components.py index dd7b89719..ec5e0ecf9 100644 --- a/modules/util/ui/components.py +++ b/modules/util/ui/ctk_components.py @@ -7,10 +7,10 @@ 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 -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,25 +110,22 @@ 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, 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, + 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) @@ -183,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, @@ -216,7 +219,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 +242,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 +356,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 +377,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 +397,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 +429,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 +487,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 +560,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 new file mode 100644 index 000000000..5347ba6d4 --- /dev/null +++ b/modules/util/ui/ctk_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()) 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 new file mode 100644 index 000000000..c95add49a --- /dev/null +++ b/modules/util/ui/pyside6_components.py @@ -0,0 +1,734 @@ +import contextlib +from collections.abc import Callable +from pathlib import Path +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, supported_video_extensions +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 + + +# --------------------------------------------------------------------------- +# PySide6-only helpers +# --------------------------------------------------------------------------- + +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 + + return h | v + + +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 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) + + +# --------------------------------------------------------------------------- +# 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: + component.setToolTip(tooltip) + if underline: + 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: QWidget, + row: int, + column: int, + ui_state: BaseUIState, + 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[..., PySide6FieldValidator] | None = None, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, +) -> QLineEdit: + var = ui_state.get_var(var_name) + + if command: + ui_state.add_var_trace(var_name, command) + + component = QLineEdit(master) + component.setMinimumWidth(width) + _add(_layout(master), component, row, column, sticky=sticky) + + if tooltip: + component.setToolTip(tooltip) + + 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 = PySide6FieldValidator( + 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] + + return component + + +def path_entry( + master: QWidget, + row: int, + column: int, + ui_state: BaseUIState, + 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, + allow_video_files: bool = False, + command: Callable[[str], None] | None = None, + extra_validate: Callable[[str], str | None] | None = None, + required: bool = False, + columnspan: int = 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 PySide6PathValidator(comp, var, state, name, io_type=io_type, **kw) + + entry_component = entry( + frame, 0, 0, ui_state, var_name, + validator_factory=_path_validator_factory, + extra_validate=extra_validate, + required=required, + ) + + 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 _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(): + current_path_str = ui_state.get_var(var_name).get() or None + current_dir = "" + current_filename = "" + + 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) + + if mode == "dir": + chosen = QFileDialog.getExistingDirectory(frame, "", current_dir, QFileDialog.Option.ShowDirsOnly) + else: + filters = ["All Files (*.*)"] + if allow_model_files: + filters += [ + "Diffusers (model_index.json)", + "Checkpoint (*.ckpt *.pt *.bin)", + "Safetensors (*.safetensors)", + ] + if allow_image_files: + exts = " ".join(f"*.{x}" for x in supported_image_extensions()) + filters.append(f"Image ({exts})") + if allow_video_files: + exts = " ".join(f"*{e}" for e in supported_video_extensions()) + filters.append(f"Video ({exts})") + filter_str = ";;".join(filters) + init_path = str(Path(current_dir) / current_filename) if current_filename else current_dir + + if use_save_dialog: + chosen, _ = QFileDialog.getSaveFileName(frame, "", init_path, filter_str) + else: + 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) + + btn = QPushButton("...", frame) + btn.setFixedWidth(40) + btn.clicked.connect(_open_dialog) + frame_lo.addWidget(btn, 0, 1) + + return frame + + +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, 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()] + + options(frame, 0, 1, values, ui_state, unit_var_name) + + return frame + + +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 preset_set_layer_choice(selected: str): + if not selected or selected not in presets_list: + selected = presets_list[0] + + if selected == "custom": + layer_entry.setVisible(True) + layer_entry.setEnabled(True) + regex_label.setVisible(True) + regex_switch.setVisible(True) + else: + 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 + + 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.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, + ) + + 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.currentText()) + + return frame + + +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 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 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 + + +# --------------------------------------------------------------------------- +# Bound widgets +# --------------------------------------------------------------------------- + +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())) + + _updating = False + + 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()) + + return frame, {'component': combo, 'button_component': adv_btn} + + +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, _ in values] + str_values = [str(v) for _, v in values] + + if var.get() not in str_values and keys: + # store the str repr — UIState's enum trace looks up var_type[string] + var.set(str(values[0][1])) + + _updating = False + + def on_combo(key: str): + nonlocal _updating + if _updating: + return + _updating = True + for k, v in values: + if key == k: + var.set(str(v)) + if command: + command(v) + break + _updating = False + + 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 combo + + +def switch( + 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) + component = QCheckBox(text, master) + component.setChecked(bool(var.get())) + + if command: + ui_state.add_var_trace(var_name, command) + + _updating = False + + def on_toggle(checked: bool): + nonlocal _updating + if _updating: + return + _updating = True + var.set(checked) + _updating = False + + def on_var(value): + nonlocal _updating + if _updating: + return + _updating = True + component.setChecked(bool(value)) + _updating = False + + 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: 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: 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(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(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 + + +# --------------------------------------------------------------------------- +# Pure helper (toolkit-neutral) +# --------------------------------------------------------------------------- + +def set_widget_enabled(widget: QWidget, enabled: bool) -> None: + widget.setEnabled(enabled) + + +def set_label_text(label: QLabel, text: str) -> None: + label.setText(str(text)) + + +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 new file mode 100644 index 000000000..8e9e1e502 --- /dev/null +++ b/modules/util/ui/pyside6_validation.py @@ -0,0 +1,176 @@ +from collections.abc import Callable + +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, + BaseFieldValidator, + _validate_path_field, +) + +from PySide6.QtCore import QTimer +from PySide6.QtWidgets import QLineEdit + +_active_qt_validators: set["PySide6FieldValidator"] = set() + + +class PySide6FieldValidator(BaseFieldValidator): + def __init__( + self, + component: QLineEdit, + var: QtVar, + ui_state: BaseUIState, + 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 + self._original_style = component.styleSheet() + self._syncing = False + self._touched = False + self._var_trace_id: int | None = None + + self._debounce = QTimer(component) + self._debounce.setSingleShot(True) + self._debounce.setInterval(DEBOUNCE_TYPING_MS) + self._debounce.timeout.connect(self._on_debounce_fire) + + def _apply_error(self) -> None: + self.component.setStyleSheet(f"border: 1px solid {ERROR_BORDER_COLOR};") + + def _clear_error(self) -> None: + self.component.setStyleSheet(self._original_style) + + def attach(self) -> None: + self._syncing = True + self.component.setText(str(self.var.get())) + self._syncing = False + + 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_qt_validators.add(self) + + def detach(self) -> None: + if not self._bound: + return + self._bound = False + _active_qt_validators.discard(self) + self._debounce.stop() + self._commit() + 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: + val = self.component.text() + if val != str(self.var.get()): + self._syncing = True + self.var.set(val) + self._syncing = False + + def _on_text_changed(self, _text: str) -> None: + if self._syncing: + return + self._touched = True + self._debounce.start() + + def _on_debounce_fire(self) -> None: + val = self.component.text() + if self._validate_and_style(val): + self._commit() + + def _on_editing_finished(self) -> None: + self._debounce.stop() + if self._touched: + val = self.component.text() + if self._validate_and_style(val): + self._commit() + self._touched = False + + def _on_real_var_write(self, _0, _1, _2) -> None: + if self._syncing: + return + self._syncing = True + self.component.setText(str(self.var.get())) + self._syncing = False + self._validate_and_style(self.component.text()) + + def flush(self) -> str | None: + 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 PySide6PathValidator(PySide6FieldValidator): + def __init__( + self, + 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) + 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: + 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/modules/util/ui/theme.py b/modules/util/ui/theme.py new file mode 100644 index 000000000..c9dc7b1b7 --- /dev/null +++ b/modules/util/ui/theme.py @@ -0,0 +1,31 @@ +import platform + +from PySide6.QtCore import Qt +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; + } + QCheckBox::indicator { + width: 16px; + height: 16px; + } + QProgressBar { + background-color: #c8c8c8; + } +""" + +def apply_theme(app: QApplication) -> None: + 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) + 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) diff --git a/modules/util/ui/validation.py b/modules/util/ui/validation.py index 2f611bc80..91117fb42 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.""" @@ -345,12 +236,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: @@ -360,140 +245,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/requirements-global.txt b/requirements-global.txt index 299f765b0..84b9d660d 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 @@ -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/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/scripts/train_ui.py b/scripts/train_ui_ctk.py similarity index 63% rename from scripts/train_ui.py rename to scripts/train_ui_ctk.py index 46ee8f1e6..562c73feb 100644 --- a/scripts/train_ui.py +++ b/scripts/train_ui_ctk.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() diff --git a/scripts/train_ui_qt.py b/scripts/train_ui_qt.py new file mode 100644 index 000000000..1d49d3fd1 --- /dev/null +++ b/scripts/train_ui_qt.py @@ -0,0 +1,26 @@ +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 modules.util.ui.theme import apply_theme + +from PySide6.QtWidgets import QApplication + + +def main(): + app = QApplication(sys.argv) + apply_theme(app) + 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(): diff --git a/start-ui.bat b/start-ui.bat index 05881410b..c8ce4ae2d 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 ) @@ -52,7 +52,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 b2960c262..e2db340ab 100755 --- a/start-ui.sh +++ b/start-ui.sh @@ -6,4 +6,4 @@ source "${BASH_SOURCE[0]%/*}/lib.include.sh" prepare_runtime_environment -run_python_in_active_env "scripts/train_ui.py" "$@" +run_python_in_active_env "scripts/train_ui_qt.py" "$@" 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" + } +} 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/#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", 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,