From c2921dd5a90fea5656be2e30f3ce345b6aaa01da Mon Sep 17 00:00:00 2001 From: bghira Date: Mon, 31 Aug 2026 23:37:26 -0600 Subject: [PATCH 01/19] Round training metrics display precision --- simpletuner/static/js/training_metrics_chart.js | 4 ++-- tests/js/training_metrics_component.test.js | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/simpletuner/static/js/training_metrics_chart.js b/simpletuner/static/js/training_metrics_chart.js index 853e884f9..d2d85b93c 100644 --- a/simpletuner/static/js/training_metrics_chart.js +++ b/simpletuner/static/js/training_metrics_chart.js @@ -65,8 +65,8 @@ function formatMetricValue(value) { if (typeof value !== 'number' || !Number.isFinite(value)) return '—'; const absolute = Math.abs(value); - if (absolute !== 0 && (absolute >= 10000 || absolute < 0.001)) return value.toExponential(3); - return value.toLocaleString(undefined, { maximumFractionDigits: 6 }); + if (absolute !== 0 && (absolute >= 10000 || absolute < 0.001)) return value.toExponential(2); + return value.toLocaleString(undefined, { maximumFractionDigits: 2 }); } function timestampMs(record) { diff --git a/tests/js/training_metrics_component.test.js b/tests/js/training_metrics_component.test.js index 39e2e50ec..0b6dda323 100644 --- a/tests/js/training_metrics_component.test.js +++ b/tests/js/training_metrics_component.test.js @@ -117,6 +117,15 @@ describe('training metrics component', () => { expect(selected).toEqual(['train_loss']); }); + test('formats scalar metric values with two decimal places at most', () => { + const charts = window.TrainingMetricsCharts; + + expect(charts.formatMetricValue(1.23456)).toBe('1.23'); + expect(charts.formatMetricValue(1.2)).toBe('1.2'); + expect(charts.formatMetricValue(0.0001234)).toBe('1.23e-4'); + expect(charts.formatMetricValue(Number.NaN)).toBe('—'); + }); + test('limits the chart to eight selected metrics', () => { state.trainingCharts = [{ id: 'chart-a', kind: 'scalar', name: 'Loss', metrics: ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'], metricSearch: '' }]; state.selectedTrainingMetrics = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']; From d0ab1793cf622e649184b419b41fe9e0b5d5a966 Mon Sep 17 00:00:00 2001 From: bghira Date: Fri, 4 Sep 2026 09:21:45 -0600 Subject: [PATCH 02/19] Allow mixed ConvRot groups in MiniMax H3 checkpoints --- .../helpers/models/minimaxh3/transformer.py | 11 +++-- tests/test_minimaxh3.py | 41 +++++++++++++++++++ 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/simpletuner/helpers/models/minimaxh3/transformer.py b/simpletuner/helpers/models/minimaxh3/transformer.py index 43f53e617..f4ae6f670 100644 --- a/simpletuner/helpers/models/minimaxh3/transformer.py +++ b/simpletuner/helpers/models/minimaxh3/transformer.py @@ -1857,21 +1857,20 @@ def from_single_file( result_dtype=torch_dtype or torch.bfloat16, hadamard_group_size=hadamard_group_size, ) - if len(hadamard_group_sizes) != 1: - raise RuntimeError( - f"MiniMax-H3 ConvRot checkpoint uses multiple Hadamard group sizes: {sorted(hadamard_group_sizes)}" - ) - group_size = hadamard_group_sizes.pop() model.quantization_method = "minimax_h3_comfy_convrot_sdnq" model.quantization_config = { "quant_method": "sdnq_training", "weights_dtype": "int8", "quantized_matmul_dtype": "int8", "use_hadamard": True, - "hadamard_group_size": group_size, "group_size": -1, "source_format": "comfy_minimax_h3_convrot", } + sorted_group_sizes = sorted(hadamard_group_sizes) + if len(sorted_group_sizes) == 1: + model.quantization_config["hadamard_group_size"] = sorted_group_sizes[0] + else: + model.quantization_config["hadamard_group_sizes"] = sorted_group_sizes elif fp8_state_dict: model.quantization_method = "minimax_h3_comfy_fp8" model.quantization_config = { diff --git a/tests/test_minimaxh3.py b/tests/test_minimaxh3.py index f332ba925..50ab07348 100644 --- a/tests/test_minimaxh3.py +++ b/tests/test_minimaxh3.py @@ -3655,6 +3655,47 @@ def test_single_file_loader_accepts_abiray_convrot_metadata(self): self.assertEqual(wrap_convrot.call_args.args[1], "transformer_blocks.0.attn.to_out.0") self.assertEqual(wrap_convrot.call_args.kwargs["hadamard_group_size"], 256) + def test_single_file_loader_accepts_mixed_convrot_group_sizes(self): + model = tiny_h3_transformer(num_layers=1) + state_dict = dict(model.state_dict()) + + qkv_weights = [state_dict.pop(f"transformer_blocks.0.attn.to_{branch}.weight") for branch in ("q", "k", "v")] + qkv_source = "blocks.0.attn.qkv_proj" + qkv_weight = torch.cat(qkv_weights, dim=0) + state_dict[f"{qkv_source}.weight"] = torch.zeros(qkv_weight.shape, dtype=torch.int8) + state_dict[f"{qkv_source}.weight_scale"] = torch.ones(qkv_weight.shape[0], 1, dtype=torch.float32) + state_dict[f"{qkv_source}.comfy_quant"] = comfy_quant_metadata_tensor( + {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 64} + ) + + out_target = "transformer_blocks.0.attn.to_out.0.weight" + out_weight = state_dict.pop(out_target) + out_source = "blocks.0.attn.out_proj" + state_dict[f"{out_source}.weight"] = torch.zeros(out_weight.shape, dtype=torch.int8) + state_dict[f"{out_source}.weight_scale"] = torch.ones(out_weight.shape[0], 1, dtype=torch.float32) + state_dict[f"{out_source}.comfy_quant"] = comfy_quant_metadata_tensor( + {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256} + ) + + with tempfile.TemporaryDirectory() as tmpdir: + path = f"{tmpdir}/tiny-h3-mixed-convrot.safetensors" + save_file(state_dict, path) + with patch("simpletuner.helpers.models.z_image.quantized_loading._wrap_convrot_linear") as wrap_convrot: + loaded = MiniMaxH3Transformer3DModel.from_single_file(path, torch_dtype=torch.float32) + + group_sizes_by_module = {call.args[1]: call.kwargs["hadamard_group_size"] for call in wrap_convrot.call_args_list} + self.assertEqual( + group_sizes_by_module, + { + "transformer_blocks.0.attn.to_q": 64, + "transformer_blocks.0.attn.to_k": 64, + "transformer_blocks.0.attn.to_v": 64, + "transformer_blocks.0.attn.to_out.0": 256, + }, + ) + self.assertEqual(loaded.quantization_config["hadamard_group_sizes"], [64, 256]) + self.assertNotIn("hadamard_group_size", loaded.quantization_config) + def test_single_file_loader_accepts_comfy_fp8_scale_metadata(self): model = tiny_h3_transformer(num_layers=1) state_dict = dict(model.state_dict()) From 1f9681260a07d6536dac79868f6f4d7c6a11b28b Mon Sep 17 00:00:00 2001 From: bghira Date: Fri, 4 Sep 2026 09:29:24 -0600 Subject: [PATCH 03/19] Preserve activation dtype for ConvRot linears --- .../helpers/models/minimaxh3/transformer.py | 24 +++++++++++++------ .../models/z_image/quantized_loading.py | 4 +++- tests/test_minimaxh3.py | 12 ++++++++++ 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/simpletuner/helpers/models/minimaxh3/transformer.py b/simpletuner/helpers/models/minimaxh3/transformer.py index f4ae6f670..28680dd4a 100644 --- a/simpletuner/helpers/models/minimaxh3/transformer.py +++ b/simpletuner/helpers/models/minimaxh3/transformer.py @@ -69,6 +69,16 @@ _H3_MASKED_CONTEXT_PARALLEL_BACKENDS = frozenset({AttentionBackendName.NATIVE, AttentionBackendName._NATIVE_CUDNN}) +def _linear_compute_dtype(linear: nn.Module) -> torch.dtype: + compute_dtype = getattr(linear, "compute_dtype", None) + if isinstance(compute_dtype, torch.dtype): + return compute_dtype + weight = linear.weight + dequantizer = getattr(weight, "sdnq_dequantizer", None) + result_dtype = getattr(dequantizer, "result_dtype", None) + return result_dtype if isinstance(result_dtype, torch.dtype) else weight.dtype + + class _MiniMaxH3AllGather(torch.autograd.Function): """Gather sequence shards without PyTorch's unsupported NCCL coalesced path.""" @@ -631,7 +641,7 @@ def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: # The activation runs at `temb`'s own precision and only the projection input is aligned to the projection # weight. Every block reads the same `temb`, so early rounding biases every block's modulation coherently. temb = nn.functional.silu(temb) if self.apply_silu else temb - temb = self.linear(temb.to(self.linear.weight.dtype)) + temb = self.linear(temb.to(_linear_compute_dtype(self.linear))) temb = temb.view(-1, 6 * self.hidden_size) return temb.chunk(6, dim=-1) @@ -661,7 +671,7 @@ def forward( ) -> torch.Tensor: # As in `MiniMaxH3AdaLayerNormModulation`: activate at `temb`'s precision, cast to the projection's dtype after. temb = nn.functional.silu(temb) if self.apply_silu else temb - shift, scale = self.linear(temb.to(self.linear.weight.dtype)).chunk(2, dim=-1) + shift, scale = self.linear(temb.to(_linear_compute_dtype(self.linear))).chunk(2, dim=-1) activation_dtype = hidden_states.dtype hidden_states = self.norm(hidden_states) shift = _select_modulation(shift, timestep_indices).to(dtype=activation_dtype) @@ -1522,7 +1532,7 @@ def _time_embedding( temb = blend_flowmap_embeddings(temb, delta_temb, self.flowmap_delta_emb_gate) return temb - dtype = self.time_embedder.linear_1.weight.dtype + dtype = _linear_compute_dtype(self.time_embedder.linear_1) temb = flowmap_timestep_embedding( time_proj=self.time_proj, timestep_embedder=self.time_embedder, @@ -2021,9 +2031,9 @@ def forward( # mixed-precision (the two patch projections are float32 while `context_embedder` and the block stack are # bfloat16 — see `_keep_in_fp32_modules`), so every input is aligned with its projection's parameter dtype, # mirroring the reference's explicit casts. The text stream sets the dtype of the packed sequence. - video_embeds = self.proj_in(hidden_states.to(self.proj_in.weight.dtype)) - audio_embeds = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype)) - text_embeds = self.context_embedder(encoder_hidden_states.to(self.context_embedder.weight.dtype)) + video_embeds = self.proj_in(hidden_states.to(_linear_compute_dtype(self.proj_in))) + audio_embeds = self.audio_proj_in(audio_hidden_states.to(_linear_compute_dtype(self.audio_proj_in))) + text_embeds = self.context_embedder(encoder_hidden_states.to(_linear_compute_dtype(self.context_embedder))) self.token_refiner.gradient_checkpointing = self.gradient_checkpointing text_attention_mask = None if packed_valid_mask is not None: @@ -2374,7 +2384,7 @@ def run_checkpointed_block( # 5. Both heads run over every row, then the rows of each modality are selected. The heads are listed in # `_keep_in_fp32_modules`, so they stay float32 while the block stack runs in the requested `torch_dtype`; # align the activation with their parameter dtype. - hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(self.proj_out.weight.dtype) + hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(_linear_compute_dtype(self.proj_out)) video_output = _gather_h3_context_parallel_output(self.proj_out(hidden_states), cp_config, dim=1).index_select( 1, video_indices.to(hidden_states.device) ) diff --git a/simpletuner/helpers/models/z_image/quantized_loading.py b/simpletuner/helpers/models/z_image/quantized_loading.py index c04b53fee..98fee45f5 100644 --- a/simpletuner/helpers/models/z_image/quantized_loading.py +++ b/simpletuner/helpers/models/z_image/quantized_loading.py @@ -200,7 +200,9 @@ def _wrap_convrot_linear( True, -1, ) - _set_module(model, module_name, get_sdnq_wrapper_class(module, forward)) + wrapped_module = get_sdnq_wrapper_class(module, forward) + wrapped_module.compute_dtype = result_dtype + _set_module(model, module_name, wrapped_module) def _validate_quant_metadata(checkpoint, key: str) -> int: diff --git a/tests/test_minimaxh3.py b/tests/test_minimaxh3.py index 50ab07348..128eb609c 100644 --- a/tests/test_minimaxh3.py +++ b/tests/test_minimaxh3.py @@ -64,6 +64,7 @@ _convert_minimax_h3_native_swiglu_scale_to_diffusers, _convert_minimax_h3_native_swiglu_to_diffusers, _gather_h3_context_parallel_output, + _linear_compute_dtype, _pad_h3_context_parallel_layout, resolve_h3_reference_mode, ) @@ -768,6 +769,17 @@ def sample(self, generator=None): class MiniMaxH3Tests(unittest.TestCase): + def test_linear_compute_dtype_uses_quantized_result_dtype(self): + weight = SimpleNamespace( + dtype=torch.int8, + sdnq_dequantizer=SimpleNamespace(result_dtype=torch.bfloat16), + ) + self.assertEqual(_linear_compute_dtype(SimpleNamespace(weight=weight)), torch.bfloat16) + + def test_linear_compute_dtype_prefers_module_contract(self): + linear = SimpleNamespace(weight=SimpleNamespace(dtype=torch.int8), compute_dtype=torch.float16) + self.assertEqual(_linear_compute_dtype(linear), torch.float16) + def test_registry_metadata_resolves(self): model_cls = ModelRegistry.get("minimaxh3") self.assertEqual(model_cls.NAME, "MiniMax H3") From a734a6cd3408c436dac9e16b17eecdf03245f0a6 Mon Sep 17 00:00:00 2001 From: bghira Date: Fri, 4 Sep 2026 12:23:50 -0600 Subject: [PATCH 04/19] Support fused QKV config inference for MiniMax H3 --- simpletuner/helpers/models/minimaxh3/transformer.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/simpletuner/helpers/models/minimaxh3/transformer.py b/simpletuner/helpers/models/minimaxh3/transformer.py index 28680dd4a..f3d27a8ea 100644 --- a/simpletuner/helpers/models/minimaxh3/transformer.py +++ b/simpletuner/helpers/models/minimaxh3/transformer.py @@ -450,7 +450,13 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]: audio_weight = _get_checkpoint_tensor(checkpoint, "audio_proj_in.weight") context_weight = _get_checkpoint_tensor(checkpoint, "context_embedder.weight") q_norm_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.norm_q.weight") - q_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight") + if "transformer_blocks.0.attn.to_q.weight" in raw_keys: + q_output_dim = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight").shape[0] + else: + qkv_weight = _get_checkpoint_tensor(checkpoint, "blocks.0.attn.qkv_proj.weight") + if qkv_weight.shape[0] % 3 != 0: + raise RuntimeError("MiniMax-H3 fused QKV tensor blocks.0.attn.qkv_proj.weight cannot be split into q/k/v") + q_output_dim = qkv_weight.shape[0] // 3 ffn_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.ff.net.0.proj.weight") has_adaln_curve = "adaln_t_table" in raw_keys adaln_curve_table = _get_checkpoint_tensor(checkpoint, "adaln_t_table") if has_adaln_curve else None @@ -466,7 +472,7 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]: "audio_in_channels": audio_weight.shape[1], "text_dim": context_weight.shape[1], "attention_head_dim": q_norm_weight.shape[0], - "num_attention_heads": q_weight.shape[0] // q_norm_weight.shape[0], + "num_attention_heads": q_output_dim // q_norm_weight.shape[0], "freq_dim": time_in.shape[1] if time_in is not None else 256, "time_embed_hidden_dim": time_in.shape[0] if time_in is not None else 5376, "time_embed_dim": adaln_curve_table.shape[1] if has_adaln_curve else time_out.shape[0], From 8b7f0064d1f442a9b584798448ac1b6cac685ac9 Mon Sep 17 00:00:00 2001 From: bghira Date: Sun, 6 Sep 2026 11:49:22 -0600 Subject: [PATCH 05/19] webshart: custom caption key(s) for multicaption support --- documentation/DATALOADER.es.md | 1 + documentation/DATALOADER.hi.md | 1 + documentation/DATALOADER.ja.md | 1 + documentation/DATALOADER.md | 1 + documentation/DATALOADER.pt-BR.md | 1 + documentation/DATALOADER.zh.md | 1 + .../helpers/data_backend/builders/webshart.py | 1 + .../helpers/data_backend/config/image.py | 5 + .../helpers/data_backend/config/validators.py | 8 + simpletuner/helpers/data_backend/webshart.py | 32 ++++ .../helpers/metadata/backends/webshart.py | 13 +- .../static/js/dataloader-section-component.js | 10 + .../dataloader/sections/storage.html | 3 + .../dataloader/sections/storage_body.html | 5 +- .../dataloader/webshart_caption_key.html | 10 + simpletuner/templates/trainer_htmx.html | 3 + tests/test_webshart_backend.py | 173 +++++++++++++++++- tests/test_webui_e2e.py | 85 +++++++++ 18 files changed, 348 insertions(+), 6 deletions(-) create mode 100644 simpletuner/templates/components/dataloader/webshart_caption_key.html diff --git a/documentation/DATALOADER.es.md b/documentation/DATALOADER.es.md index 8bc0230b4..fdf97a3a9 100644 --- a/documentation/DATALOADER.es.md +++ b/documentation/DATALOADER.es.md @@ -1377,6 +1377,7 @@ Los datasets Webshart cargan shards tar estilo WebDataset mediante el paquete `w - `metadata` es opcional y puede apuntar a metadatos separados con captions. Para repositorios Hugging Face de metadata como `webshart/conceptual-captions-12m-webdataset-metadata`, pasa el repo id; Webshart sigue el layout de subcarpetas del source, como `data/`. - `metadata_backend` debe ser `webshart`; `caption_strategy` debe ser `webshart` o `instanceprompt`. - `webshart.cache_dir` almacena la metadata de SimpleTuner y las caches de Webshart. `shard_cache_gb` y `parallel_downloads` se pasan a la cache de shards de Webshart; define `shard_cache_gb` como `0` para desactivar la cache de shards completos y mantener lecturas por rango indexadas. +- `webshart.caption_key` permite seleccionar campos de caption: usa `"long_caption"` para una clave o `["long_caption", "short_caption"]` para recopilar varias en el orden indicado. Las claves son nombres literales que se buscan en los metadatos JSON de la muestra, después en los metadatos del índice y, finalmente, en las entradas con nombre de sus respectivos diccionarios `captions`; se usa la primera ubicación que contiene la clave. Los valores de texto y listas se convierten en variantes de caption, sin concatenarse en un único prompt. Se ignoran los valores ausentes o vacíos; con `caption_strategy: "webshart"`, se omiten las muestras sin captions seleccionadas aunque tengan captions predeterminadas o un sidecar `.txt`. Omitir la opción conserva la búsqueda predeterminada. Se leen los sidecars JSON cuando su contenido no está en el índice. Las cachés de captions y buckets se separan según el selector configurado. En los ajustes Webshart de la WebUI, introduce una clave por línea. - `webshart_optimize_captions` (grafía alternativa `webshart_optimise_captions`; también se acepta como `optimize_captions`/`optimise_captions` dentro del bloque `webshart`) sondea el layout de captions al arrancar y, cuando los captions residen en miembros tar sidecar `.txt`/`.json` en lugar del índice de metadata, los consolida una sola vez en la cache local de metadata de Webshart. Sin esta opción, los datasets con captions en sidecars (por ejemplo `laion/conceptual-captions-12m-webdataset`) pagan una lectura por rango por muestra cada vez que se enumeran los captions — al arrancar, al guardar checkpoints y al generar la model card. Los datasets cuya metadata ya incluye los captions omiten la consolidación automáticamente. #### Optimizar captions por adelantado diff --git a/documentation/DATALOADER.hi.md b/documentation/DATALOADER.hi.md index 999910cc9..5db48b099 100644 --- a/documentation/DATALOADER.hi.md +++ b/documentation/DATALOADER.hi.md @@ -1377,6 +1377,7 @@ Webshart datasets `webshart` package के जरिए WebDataset-style tar sh - `metadata` optional है और captions वाले separate metadata location को point कर सकता है। `webshart/conceptual-captions-12m-webdataset-metadata` जैसे Hugging Face metadata repos के लिए repo id दें; Webshart source shard के `data/` जैसे subfolder layout को follow करता है। - `metadata_backend` को `webshart` होना चाहिए; `caption_strategy` `webshart` या `instanceprompt` हो सकता है। - `webshart.cache_dir` SimpleTuner metadata और Webshart caches store करता है। `shard_cache_gb` और `parallel_downloads` Webshart shard cache को pass किए जाते हैं; whole-shard caching disable करने और indexed range reads बनाए रखने के लिए `shard_cache_gb` को `0` सेट करें। +- `webshart.caption_key` से caption फ़ील्ड चुन सकते हैं: एक कुंजी के लिए `"long_caption"` या कई कुंजियों को क्रम से लेने के लिए `["long_caption", "short_caption"]` दें। कुंजियाँ सीधे नाम के रूप में मिलाई जाती हैं: पहले sample के JSON metadata में, फिर indexed metadata में, और अंत में उनके `captions` dictionaries की नामित entries में। कुंजी जिस पहले स्थान पर मिलती है, वही उपयोग होता है। String और list के मान अलग caption विकल्प बनते हैं, एक prompt में जोड़े नहीं जाते। अनुपस्थित या खाली मान छोड़ दिए जाते हैं; `caption_strategy: "webshart"` के साथ चुने हुए captions न होने पर sample छोड़ दिया जाता है, भले ही उसमें default caption या `.txt` sidecar हो। यह विकल्प न देने पर default lookup बना रहता है। JSON सामग्री index में न होने पर JSON sidecar पढ़ा जाता है। Caption और bucket caches चुनी गई कुंजियों के अनुसार अलग रखे जाते हैं। WebUI की Webshart settings में हर पंक्ति पर एक कुंजी लिखें। - `webshart_optimize_captions` (alternate spelling `webshart_optimise_captions`; `webshart` block के अंदर `optimize_captions`/`optimise_captions` भी accepted हैं) startup पर caption layout probe करता है और, जब captions metadata index की बजाय `.txt`/`.json` sidecar tar members में हों, उन्हें एक बार local Webshart metadata cache में fold कर देता है। इसके बिना, sidecar-caption datasets (जैसे `laion/conceptual-captions-12m-webdataset`) को हर बार captions enumerate होने पर — startup, checkpointing और model card generation में — प्रति sample एक range read की कीमत चुकानी पड़ती है। जिन datasets की metadata में captions पहले से embedded हैं, वे coalescing अपने आप skip कर देते हैं। #### Captions को पहले से optimize करना diff --git a/documentation/DATALOADER.ja.md b/documentation/DATALOADER.ja.md index 5753389d1..a304ab320 100644 --- a/documentation/DATALOADER.ja.md +++ b/documentation/DATALOADER.ja.md @@ -1378,6 +1378,7 @@ Webshart データセットは `webshart` パッケージで WebDataset 形式 - `metadata` は任意で、captions を含む別 metadata location を指定できます。`webshart/conceptual-captions-12m-webdataset-metadata` のような Hugging Face metadata repo では repo id だけを渡します。Webshart は source shard の `data/` などのサブフォルダ構成に従います。 - `metadata_backend` は `webshart`、`caption_strategy` は `webshart` または `instanceprompt` にします。 - `webshart.cache_dir` は SimpleTuner metadata と Webshart caches を保存します。`shard_cache_gb` と `parallel_downloads` は Webshart の shard cache に渡されます。`shard_cache_gb` を `0` にすると、shard 全体の cache を無効にし、index 付き range read を維持します。 +- `webshart.caption_key` でキャプションのフィールドを選択できます。1つなら `"long_caption"`、複数なら `["long_caption", "short_caption"]` を指定すると、その順で収集します。キーはパスではなくそのままの名前として、サンプルの JSON メタデータ、インデックスのメタデータ、それぞれの `captions` 辞書内の名前付き項目の順に検索し、最初にキーが見つかった場所を使います。文字列やリストの値は1つのプロンプトに連結せず、キャプション候補になります。欠落した値や空の値は無視されます。`caption_strategy: "webshart"` では、選択したキャプションがないサンプルは、既定のキャプションや `.txt` サイドカーがあってもスキップされます。省略すると既定の取得方法を維持します。JSON の内容がインデックスにない場合はサイドカーを読み込みます。キャプションとバケットのキャッシュは設定したキーに応じて分離されます。WebUI の Webshart 設定では1行に1つのキーを入力します。 - `webshart_optimize_captions`（別綴り `webshart_optimise_captions`。`webshart` ブロック内では `optimize_captions`/`optimise_captions` も受け付けます）は起動時に caption layout を probe し、captions が metadata index ではなく `.txt`/`.json` の sidecar tar member にある場合、それらをローカルの Webshart metadata cache に一度だけ統合します。このオプションがないと、sidecar caption の dataset（たとえば `laion/conceptual-captions-12m-webdataset`）は captions を列挙するたび — 起動時、checkpoint 時、model card 生成時 — にサンプルごとに 1 回の range read が発生します。metadata に captions が既に埋め込まれている dataset では、統合は自動的にスキップされます。 #### captions を事前に最適化する diff --git a/documentation/DATALOADER.md b/documentation/DATALOADER.md index 87f43e5d4..fe7f51e73 100644 --- a/documentation/DATALOADER.md +++ b/documentation/DATALOADER.md @@ -1433,6 +1433,7 @@ Webshart datasets load WebDataset-style tar shards through the `webshart` packag - `metadata_backend` must be `webshart`; it reads dimensions and captions from Webshart metadata. - `caption_strategy` should be `webshart` to train from metadata captions, or `instanceprompt` to ignore stored captions. - `webshart.cache_dir` stores SimpleTuner metadata plus Webshart metadata and shard caches. `shard_cache_gb` and `parallel_downloads` are passed to Webshart's shard cache; set `shard_cache_gb` to `0` to disable whole-shard caching and retain indexed range reads. +- `webshart.caption_key` optionally selects caption fields: use `"long_caption"` for one key or `["long_caption", "short_caption"]` to collect multiple keys in order. Keys are literal names, checked in the sample’s JSON metadata, then its indexed metadata, then named entries inside their `captions` dictionaries; the first location containing a key wins. String and list values become caption variants, rather than being joined into one prompt. Missing or empty values are ignored; with `caption_strategy: "webshart"`, samples with no selected captions are skipped, even if they have default captions or a `.txt` sidecar. Omitting the option keeps the default caption lookup. JSON sidecars are read when their contents are absent from the index. Caption and bucket caches are separated by the configured selector. In the WebUI’s Webshart settings, enter one key per line. - `webshart_optimize_captions` (alt spelling `webshart_optimise_captions`; also accepted as `optimize_captions`/`optimise_captions` inside the `webshart` block) probes the caption layout at startup and, when captions live in `.txt`/`.json` sidecar tar members rather than the metadata index, folds them into the local Webshart metadata cache once. Without it, sidecar-caption datasets (for example `laion/conceptual-captions-12m-webdataset`) pay one range read per sample every time captions are enumerated — startup, checkpointing, and model card generation. Datasets whose metadata already embeds captions skip the coalescing automatically. #### Optimizing captions ahead of time diff --git a/documentation/DATALOADER.pt-BR.md b/documentation/DATALOADER.pt-BR.md index f9a77722a..06029b9cd 100644 --- a/documentation/DATALOADER.pt-BR.md +++ b/documentation/DATALOADER.pt-BR.md @@ -1377,6 +1377,7 @@ Datasets Webshart carregam shards tar no estilo WebDataset pelo pacote `webshart - `metadata` é opcional e pode apontar para metadados separados com captions. Para repos Hugging Face de metadata como `webshart/conceptual-captions-12m-webdataset-metadata`, passe o repo id; o Webshart segue o layout de subpastas do source, como `data/`. - `metadata_backend` deve ser `webshart`; `caption_strategy` deve ser `webshart` ou `instanceprompt`. - `webshart.cache_dir` armazena os metadados do SimpleTuner e os caches do Webshart. `shard_cache_gb` e `parallel_downloads` são passados ao cache de shards do Webshart; defina `shard_cache_gb` como `0` para desativar o cache de shards completos e manter leituras por intervalo indexadas. +- `webshart.caption_key` permite selecionar campos de caption: use `"long_caption"` para uma chave ou `["long_caption", "short_caption"]` para coletar várias na ordem indicada. As chaves são nomes literais, buscados nos metadados JSON da amostra, depois nos metadados do índice e, por fim, nas entradas nomeadas dos respectivos dicionários `captions`; vale o primeiro local que contém a chave. Valores de texto e listas tornam-se variantes de caption, sem serem concatenados em um único prompt. Valores ausentes ou vazios são ignorados; com `caption_strategy: "webshart"`, amostras sem captions selecionadas são ignoradas mesmo que tenham captions padrão ou um sidecar `.txt`. Omitir a opção mantém a busca padrão. Sidecars JSON são lidos quando seu conteúdo não está no índice. As caches de captions e buckets são separadas conforme o seletor configurado. Nas configurações Webshart da WebUI, insira uma chave por linha. - `webshart_optimize_captions` (grafia alternativa `webshart_optimise_captions`; também aceito como `optimize_captions`/`optimise_captions` dentro do bloco `webshart`) sonda o layout de captions na inicialização e, quando os captions ficam em membros tar sidecar `.txt`/`.json` em vez do índice de metadados, consolida-os uma única vez no cache local de metadados do Webshart. Sem essa opção, datasets com captions em sidecars (por exemplo `laion/conceptual-captions-12m-webdataset`) pagam uma leitura por intervalo por amostra sempre que os captions são enumerados — na inicialização, nos checkpoints e na geração do model card. Datasets cujos metadados já embutem os captions pulam a consolidação automaticamente. #### Otimizando captions com antecedência diff --git a/documentation/DATALOADER.zh.md b/documentation/DATALOADER.zh.md index f1d3d53f7..cfddca5e7 100644 --- a/documentation/DATALOADER.zh.md +++ b/documentation/DATALOADER.zh.md @@ -1377,6 +1377,7 @@ Webshart 数据集通过 `webshart` 包加载 WebDataset 风格的 tar shards。 - `metadata` 可选，可指向包含 captions 的独立 metadata location。对于 `webshart/conceptual-captions-12m-webdataset-metadata` 这样的 Hugging Face metadata repo，传 repo id 即可；Webshart 会跟随 source shard 的 `data/` 等子目录布局。 - `metadata_backend` 必须为 `webshart`；`caption_strategy` 应为 `webshart` 或 `instanceprompt`。 - `webshart.cache_dir` 存储 SimpleTuner metadata 与 Webshart caches。`shard_cache_gb` 和 `parallel_downloads` 会传给 Webshart 的 shard cache；将 `shard_cache_gb` 设为 `0` 可禁用整 shard cache，并保留基于索引的 range reads。 +- `webshart.caption_key` 可用于选择字幕字段：单个键使用 `"long_caption"`，多个键使用 `["long_caption", "short_caption"]`，按列表顺序收集。键按字面名称匹配，依次检查样本的 JSON 元数据、索引元数据及两者 `captions` 字典中的命名条目；采用第一个包含该键的位置。字符串和列表值作为字幕候选，不会拼接成一个提示词。缺失或空值会被忽略；使用 `caption_strategy: "webshart"` 时，没有选中字幕的样本将被跳过，即使它有默认字幕或 `.txt` 伴随文件。省略此选项将保留默认查找方式。如果索引未包含 JSON 内容，则读取 JSON 伴随文件。字幕和分桶缓存按配置的键分别保存。在 WebUI 的 Webshart 设置中，每行输入一个键。 - `webshart_optimize_captions`（另一拼写 `webshart_optimise_captions`；在 `webshart` 块内也接受 `optimize_captions`/`optimise_captions`）会在启动时探测 caption 布局，当 captions 位于 `.txt`/`.json` sidecar tar 成员中而不是 metadata 索引中时，将它们一次性合并进本地 Webshart metadata cache。若不启用，sidecar caption 数据集（例如 `laion/conceptual-captions-12m-webdataset`）在每次枚举 captions 时——启动、checkpoint、生成 model card——都要为每个样本付出一次 range read。metadata 中已内嵌 captions 的数据集会自动跳过合并。 #### 提前优化 captions diff --git a/simpletuner/helpers/data_backend/builders/webshart.py b/simpletuner/helpers/data_backend/builders/webshart.py index 8795606aa..020950893 100644 --- a/simpletuner/helpers/data_backend/builders/webshart.py +++ b/simpletuner/helpers/data_backend/builders/webshart.py @@ -55,6 +55,7 @@ def _create_backend(self, config: BaseBackendConfig) -> WebshartDataBackend: "compress_cache": self._get_compression_setting(config), "dataset_type": getattr(config, "dataset_type", "image"), "optimize_captions": bool(getattr(config, "webshart_optimize_captions", None) or False), + "caption_key": config.webshart_caption_key, } if is_mock_backend: backend_kwargs["identifier"] = config.id diff --git a/simpletuner/helpers/data_backend/config/image.py b/simpletuner/helpers/data_backend/config/image.py index 45d2e98a9..3bf17b2c3 100644 --- a/simpletuner/helpers/data_backend/config/image.py +++ b/simpletuner/helpers/data_backend/config/image.py @@ -63,6 +63,7 @@ class ImageBackendConfig(BaseBackendConfig): webshart_buffer_size: Optional[int] = None webshart_max_file_size: Optional[int] = None webshart_optimize_captions: Optional[bool] = None + webshart_caption_key: Optional[Union[str, List[str]]] = None vae_cache_clear_each_epoch: Optional[bool] = None probability: float = 1.0 @@ -187,6 +188,8 @@ def _get_arg(key: str, default: Any = None) -> Any: if config.backend_type == "webshart": webshart_block = backend_dict.get("webshart", {}) or {} config.webshart = webshart_block + config.webshart_caption_key = webshart_block.get("caption_key") + validators.validate_webshart_caption_key(config.webshart_caption_key) config.webshart_source = backend_dict.get("source", webshart_block.get("source")) config.webshart_metadata = backend_dict.get("metadata", webshart_block.get("metadata")) config.webshart_hf_token = backend_dict.get("hf_token", webshart_block.get("hf_token")) @@ -570,6 +573,8 @@ def to_dict(self) -> Dict[str, Any]: webshart_config["max_file_size"] = self.webshart_max_file_size if self.webshart_optimize_captions is not None: webshart_config["optimize_captions"] = self.webshart_optimize_captions + if self.webshart_caption_key is not None: + webshart_config["caption_key"] = self.webshart_caption_key if self.video is not None: config["video"] = self.video diff --git a/simpletuner/helpers/data_backend/config/validators.py b/simpletuner/helpers/data_backend/config/validators.py index 8fa9b8a0a..1f1f1a168 100644 --- a/simpletuner/helpers/data_backend/config/validators.py +++ b/simpletuner/helpers/data_backend/config/validators.py @@ -145,6 +145,14 @@ def validate_huggingface_backend_settings( return {"metadata_backend": metadata_backend or "huggingface", "caption_strategy": caption_strategy or "huggingface"} +def validate_webshart_caption_key(caption_key) -> None: + if caption_key is None: + return + keys = [caption_key] if isinstance(caption_key, str) else caption_key + if not isinstance(keys, list) or not keys or any(not isinstance(key, str) or not key.strip() for key in keys): + raise ValueError("webshart.caption_key must be a non-empty string or a non-empty list of non-empty strings.") + + def validate_webshart_backend_settings( backend_type: str, metadata_backend: Optional[str], caption_strategy: Optional[str], backend_id: str ) -> Dict[str, str]: diff --git a/simpletuner/helpers/data_backend/webshart.py b/simpletuner/helpers/data_backend/webshart.py index d1d8ab7be..f0685fd36 100644 --- a/simpletuner/helpers/data_backend/webshart.py +++ b/simpletuner/helpers/data_backend/webshart.py @@ -12,9 +12,11 @@ import torch from simpletuner.helpers.data_backend.base import BaseDataBackend +from simpletuner.helpers.data_backend.config.validators import validate_webshart_caption_key from simpletuner.helpers.data_backend.dataset_types import DatasetType, ensure_dataset_type from simpletuner.helpers.data_backend.filters import DatasetFilter from simpletuner.helpers.image_manipulation.load import load_image, load_video +from simpletuner.helpers.prompts import PromptHandler from simpletuner.helpers.training import video_file_extensions from simpletuner.helpers.training.multi_process import should_log @@ -55,9 +57,12 @@ def __init__( compress_cache: bool = False, dataset_type: Union[str, DatasetType] = DatasetType.IMAGE, optimize_captions: bool = False, + caption_key: Optional[Union[str, List[str]]] = None, ): if not source: raise ValueError("source is required for Webshart data backends.") + validate_webshart_caption_key(caption_key) + self.caption_key = caption_key try: import webshart @@ -345,6 +350,14 @@ def get_caption(self, image_path: str) -> Optional[Union[str, List[str], dict]]: sample_ref = self.parse_sample_id(image_path) sample_metadata = self.get_shard_metadata(sample_ref.shard_idx).get(sample_ref.filename, {}) or {} + if self.caption_key is not None: + if sample_metadata.get("json_metadata") is None and sample_metadata.get("json_path"): + reader = self.dataset.open_shard(sample_ref.shard_idx) + payload = reader.read_sample_json(sample_ref.sample_idx) + if payload is not None: + sample_metadata["json_metadata"] = json.loads(payload) + return self._select_caption_keys(sample_metadata) + caption = sample_metadata.get("captions") if caption: if isinstance(caption, dict): @@ -364,6 +377,23 @@ def get_caption(self, image_path: str) -> Optional[Union[str, List[str], dict]]: caption = caption.decode("utf-8") return str(caption).strip() + def _select_caption_keys(self, sample_metadata: dict) -> Optional[Union[str, List[str]]]: + json_metadata = sample_metadata.get("json_metadata") or {} + sources = [json_metadata, sample_metadata] + sources.extend(source.get("captions") for source in list(sources) if isinstance(source, dict)) + keys = [self.caption_key] if isinstance(self.caption_key, str) else self.caption_key + captions = [] + for key in keys: + for source in sources: + if isinstance(source, dict) and key in source: + captions.extend(PromptHandler._normalize_caption_payload(source[key])) + break + if not captions: + return None + if isinstance(self.caption_key, str) and len(captions) == 1: + return captions[0] + return captions + def read(self, identifier: Union[str, Path], as_byteIO: bool = False) -> Any: if self.is_sample_id(identifier): identifier = self.normalize_sample_id(identifier) @@ -501,6 +531,7 @@ def get_instance_representation(self) -> dict: "max_file_size": self.max_file_size, "compress_cache": self.compress_cache, "dataset_type": self.dataset_type.value, + "caption_key": self.caption_key, } @staticmethod @@ -523,6 +554,7 @@ def from_instance_representation(representation: dict) -> "WebshartDataBackend": max_file_size=representation.get("max_file_size", 500 * 1024 * 1024), compress_cache=representation.get("compress_cache", False), dataset_type=representation.get("dataset_type", DatasetType.IMAGE), + caption_key=representation.get("caption_key"), ) def num_shards(self) -> int: diff --git a/simpletuner/helpers/metadata/backends/webshart.py b/simpletuner/helpers/metadata/backends/webshart.py index 8471a2207..23414e72a 100644 --- a/simpletuner/helpers/metadata/backends/webshart.py +++ b/simpletuner/helpers/metadata/backends/webshart.py @@ -7,6 +7,7 @@ import time from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext +from hashlib import sha256 from pathlib import Path from typing import Any, Dict, List, Optional, Union @@ -63,6 +64,12 @@ def __init__( repeats: int = 0, max_num_samples: int = None, ): + if not isinstance(data_backend, WebshartDataBackend): + raise ValueError("WebshartMetadataBackend requires WebshartDataBackend") + if data_backend.caption_key is not None: + caption_digest = sha256(json.dumps(data_backend.caption_key).encode("utf-8")).hexdigest()[:16] + cache_file = f"{cache_file}_captions_{caption_digest}" + metadata_file = f"{metadata_file}_captions_{caption_digest}" super().__init__( id=id, instance_data_dir=instance_data_dir, @@ -86,8 +93,6 @@ def __init__( repeats=repeats, max_num_samples=max_num_samples, ) - if not isinstance(data_backend, WebshartDataBackend): - raise ValueError("WebshartMetadataBackend requires WebshartDataBackend") if self.dataset_type not in {DatasetType.IMAGE, DatasetType.VIDEO, DatasetType.CONDITIONING, DatasetType.EVAL}: raise ValueError("WebshartMetadataBackend supports image, video, conditioning, and eval datasets only.") @@ -321,6 +326,8 @@ def _metadata_for_entry(self, shard_metadata: dict, filename: str, entry: dict, metadata["original_size"] = (int(width), int(height)) if "captions" in file_metadata: metadata["captions"] = file_metadata["captions"] + if self.data_backend.caption_key is not None: + metadata["captions"] = self.data_backend.get_caption(sample_path) json_metadata = file_metadata.get("json_metadata") or {} if json_metadata: metadata["json_metadata"] = json_metadata @@ -526,7 +533,7 @@ def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = Fals # metadata (e.g. cc12m); get_caption() range-reads those at runtime. # get_shard_metadata returns a flat mapping keyed by member filename. caption_member = Path(str(entry["filename"])).with_suffix(".txt").name - if caption_member not in shard_metadata: + if self.data_backend.caption_key is not None or caption_member not in shard_metadata: statistics["skipped"]["caption_missing"] += 1 continue aspect_ratio_bucket_updates.setdefault(bucket_key, []).append(sample_path) diff --git a/simpletuner/static/js/dataloader-section-component.js b/simpletuner/static/js/dataloader-section-component.js index 094c5993d..eff3b09ca 100644 --- a/simpletuner/static/js/dataloader-section-component.js +++ b/simpletuner/static/js/dataloader-section-component.js @@ -1348,6 +1348,16 @@ function dataloaderSectionComponent() { trainer.markDatasetsDirty(); } }, + setWebshartCaptionKeys(dataset, value) { + const keys = value.split('\n').map(key => key.trim()).filter(Boolean); + if (keys.length) { + dataset.webshart = dataset.webshart || {}; + dataset.webshart.caption_key = keys.length === 1 ? keys[0] : keys; + } else if (dataset.webshart) { + delete dataset.webshart.caption_key; + } + this.markAsUnsaved(); + }, onStorageBackendChange(dataset) { if (!dataset) { return; diff --git a/simpletuner/templates/components/dataloader/sections/storage.html b/simpletuner/templates/components/dataloader/sections/storage.html index 8b7d7457a..021c3653e 100644 --- a/simpletuner/templates/components/dataloader/sections/storage.html +++ b/simpletuner/templates/components/dataloader/sections/storage.html @@ -259,6 +259,9 @@