Dark mode for Qt6

.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

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 <preset.json> [--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, "#<model> <method>[ <vram>].json"
+embedding_templates/ plain-text prompt files (`<embedding>` 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

CLAUDE.md

@@ -0,0 +1,3 @@
+# CLAUDE.md
+
+This project's contributor rules for AI coding agents live in **[AGENTS.md](AGENTS.md)**.

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.<NAME>:`.
+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).

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)

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)