feature: TST v4 (#4)

Co-authored-by: James <1561862923@qq.com>

config/examples/train_three_phase_trigger_binding_ideogram4.yaml

@@ -0,0 +1,161 @@
+---
+job: extension
+config:
+  name: ig4_three_phase_trigger_binding
+  process:
+    - type: three_phase_trigger_trainer
+      name: ig4_three_phase_trigger_binding
+      training_folder: ./output
+      device: cuda
+      trigger_word: <r1X1dOn9mA2>
+      network:
+        type: lora
+        linear: 32
+        linear_alpha: 16
+      datasets:
+        - folder_path: /path/to/ideogram4/images
+          caption_ext: json
+          cache_text_embeddings: false
+          resolution: [256, 768, 1280]
+      train:
+        batch_size: 1
+        steps: 1
+        train_unet: true
+        train_text_encoder: false
+        cache_text_embeddings: false
+        unload_text_encoder: false
+        optimizer: adamw8bit
+        lr: 0.00008
+        dtype: bf16
+        gradient_checkpointing: true
+        disable_sampling: true
+      save:
+        dtype: bf16
+        save_every: 100
+        save_format: safetensors
+      model:
+        name_or_path: ideogram-ai/ideogram-4-fp8
+        arch: ideogram4
+        quantize: true
+        qtype: qfloat8
+        quantize_te: false
+        low_vram: false
+      sample:
+        samples: []
+      trigger_selective_training:
+        enabled: true
+        caption_sources:
+          enabled: true
+          sources:
+            - name: json
+              use_main_dataset: true
+              caption_ext: .json
+              format: text
+            - name: natural
+              path: /path/to/natural/captions
+              caption_ext: .txt
+              format: text
+          schedule:
+            keyframes:
+              - {step: 0, json: 1.0, natural: 0.0}
+              - {step: 1000, json: 0.5, natural: 0.5}
+        negative_styles:
+          expected_category_count: 3
+          categories:
+            - {name: neutral, probability: 0.3, phrases: [""]}
+            - {name: hard, probability: 0.4, phrases: [painting, illustration, anime]}
+            - {name: far, probability: 0.3, phrases: [line art, photorealistic photograph, 3D render]}
+        path3:
+          decoy_gain_mode: positive_clamped
+          margin_schedule:
+            keyframes: [{step: 0, value: 0.02}, {step: 1000, value: 0.1}]
+          gain_floor:
+            enabled: true
+            weight: 1.0
+            schedule:
+              keyframes: [{step: 0, value: 0.0}, {step: 1000, value: 0.08}]
+        loss_schedule:
+          keyframes:
+            - {step: 0, path1: 0.8, path2: 0.1, path3: 0.1}
+            - {step: 1000, path1: 0.5, path2: 0.15, path3: 0.35}
+      three_phase_trigger_training:
+        enabled: true
+        trigger:
+          placeholder: "[trigger]"
+          literal: <r1X1dOn9mA2>
+          span_detection: offsets
+          mask_all_occurrences: true
+          occurrence_mode: additive
+        text_activator:
+          embedding:
+            enabled: true
+            tokens: 1
+            init_mode: semantic
+            init_words: illustration
+            dtype: bf16
+          te_adapter:
+            enabled: false
+            rank: 4
+            alpha: 4
+            child_modules: [down_proj]
+          tap_adapters:
+            enabled: false
+            rank: 4
+            alpha: 4
+            tap_layers: [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35]
+        reachability_probe:
+          enabled: true
+          require_finite_gradients: true
+          require_mask_zero_outside_trigger: true
+        phase_a1:
+          enabled: true
+          steps: 400
+          optimizer: adamw
+          learning_rates: {embedding: 0.001}
+          train: {embedding: true}
+          caption_sources: {paired: [json, natural]}
+          save_steps: [100, 200, 400]
+          losses:
+            trigger_gain_floor:
+              enabled: true
+              weight: 1.0
+              schedule:
+                keyframes: [{step: 0, value: 0.0}, {step: 400, value: 0.05}]
+            context_consistency:
+              enabled: true
+              weight: 0.02
+              loss_type: cosine
+              warmup_steps: 100
+              min_delta_norm: 0.000001
+        phase_b:
+          enabled: true
+          steps: 2000
+          optimizer: adamw8bit
+          learning_rates: {diffusion_lora: 0.00008}
+          train: {diffusion_lora: true}
+          text_activator_source: {phase: a1, step: final}
+          save_steps: [500, 1000, 1500, 2000]
+        phase_a2:
+          enabled: true
+          steps: 400
+          optimizer: adamw
+          learning_rates: {embedding: 0.0001}
+          train: {embedding: true}
+          text_activator_source: {phase: a1, step: final}
+          diffusion_lora_source: {phase: b, step: final}
+          caption_sources: {paired: [json, natural]}
+          save_steps: [100, 200, 400]
+          losses:
+            activator_gain_floor:
+              enabled: true
+              weight: 1.0
+              schedule:
+                keyframes: [{step: 0, value: 0.02}, {step: 400, value: 0.08}]
+            context_consistency:
+              enabled: true
+              weight: 0.02
+              loss_type: cosine
+              warmup_steps: 100
+              min_delta_norm: 0.000001
+        artifacts:
+          output_root: ./output/ig4_three_phase_trigger_binding

extensions_built_in/diffusion_models/ideogram4/ideogram4.py

@@ -1,5 +1,7 @@
+import importlib
+import inspect
 import os
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import torch
 import yaml
@@ -25,7 +27,11 @@ import huggingface_hub
 from huggingface_hub.errors import EntryNotFoundError
 from transformers import AutoModel, AutoTokenizer
 
-from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel
+from .src.transformer import (
+    QWEN3_VL_ACTIVATION_LAYERS,
+    Ideogram4Config,
+    Ideogram4Transformer2DModel,
+)
 from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict
 from .src.latent_norm import get_latent_norm
 from .src.pipeline import (
@@ -194,6 +200,8 @@ class Ideogram4Model(BaseModel):
         # CFG pass. Loaded from model_config.unconditional_lora_path if set; stays
         # inactive everywhere else (training, conditional pass).
         self.unconditional_lora: Optional[LoRASpecialNetwork] = None
+        self.text_activator: Any = None
+        self.text_activator_runtime_mode: Optional[str] = None
 
     @property
     def text_embedding_space_version(self):
@@ -204,6 +212,46 @@ class Ideogram4Model(BaseModel):
     def get_train_scheduler():
         return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
 
+    def install_text_activator(self, text_activator, runtime_mode: Optional[str] = None):
+        """Install an optional trigger activator while keeping its API decoupled."""
+        activator = text_activator
+        for module_name in (
+            "toolkit.trigger_binding",
+            "toolkit.models.ideogram4_trigger_activator",
+        ):
+            try:
+                module = importlib.import_module(module_name)
+            except ImportError:
+                continue
+            for name in (
+                "adapt_ideogram4_text_activator",
+                "adapt_text_activator",
+                "ensure_ideogram4_text_activator",
+            ):
+                adapter = getattr(module, name, None)
+                if callable(adapter):
+                    adapted = adapter(activator)
+                    if adapted is not None:
+                        activator = adapted
+                    break
+        self.text_activator = activator
+        self.text_activator_runtime_mode = runtime_mode
+        return activator
+
+    def set_text_activator(self, text_activator, runtime_mode: Optional[str] = None):
+        return self.install_text_activator(text_activator, runtime_mode)
+
+    def clear_text_activator(self):
+        self.text_activator = None
+        self.text_activator_runtime_mode = None
+
+    def set_text_activator_runtime_mode(self, runtime_mode: Optional[str]):
+        self.text_activator_runtime_mode = runtime_mode
+        activator = self.text_activator
+        setter = getattr(activator, "set_runtime_mode", None)
+        if callable(setter):
+            setter(runtime_mode)
+
     def get_bucket_divisibility(self):
         # 8 for the VAE downsample, 2 for the patch size.
         return self.vae_scale_factor * self.patch_size
@@ -502,23 +550,98 @@ class Ideogram4Model(BaseModel):
         )
         return pred
 
-    def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
+    def _resolve_trigger_mask(
+        self,
+        prompt: str,
+        token_ids: torch.Tensor,
+        trigger_mask: Optional[torch.Tensor],
+        text_activator: Any,
+    ) -> Optional[torch.Tensor]:
+        if trigger_mask is not None:
+            mask = torch.as_tensor(trigger_mask, device=token_ids.device)
+            if mask.dim() == 1:
+                mask = mask.unsqueeze(0)
+            if mask.shape != token_ids.shape:
+                raise ValueError(
+                    f"trigger_mask shape {tuple(mask.shape)} does not match token ids "
+                    f"shape {tuple(token_ids.shape)}"
+                )
+            return mask.to(dtype=torch.bool)
+
+        candidates = [text_activator]
+        for module_name in (
+            "toolkit.trigger_binding",
+            "toolkit.models.ideogram4_trigger_activator",
+        ):
+            try:
+                candidates.append(importlib.import_module(module_name))
+            except ImportError:
+                pass
+        kwargs = {
+            "prompt": prompt,
+            "text": prompt,
+            "token_ids": token_ids,
+            "input_ids": token_ids,
+            "tokenizer": self.tokenizer,
+        }
+        for candidate in candidates:
+            if candidate is None:
+                continue
+            for name in (
+                "get_trigger_mask",
+                "build_trigger_mask",
+                "resolve_trigger_mask",
+                "detect_trigger_mask",
+            ):
+                method = getattr(candidate, name, None)
+                if not callable(method):
+                    continue
+                try:
+                    signature = inspect.signature(method)
+                    supported = {
+                        key: value
+                        for key, value in kwargs.items()
+                        if key in signature.parameters
+                    }
+                    mask = method(**supported)
+                except (TypeError, ValueError):
+                    continue
+                if mask is not None:
+                    return self._resolve_trigger_mask(
+                        prompt, token_ids, mask, text_activator=None
+                    )
+        return None
+
+    def get_prompt_embeds(
+        self,
+        prompt,
+        trigger_mask=None,
+        text_activator=None,
+        runtime_mode: Optional[str] = None,
+        return_taps: bool = False,
+        **kwargs,
+    ) -> AdvancedPromptEmbeds:
         if isinstance(prompt, str):
             prompt = [prompt]
 
         if self.text_encoder.device == torch.device("cpu"):
             self.text_encoder.to(self.device_torch)
         device = self.text_encoder.device
+        activator = self.text_activator if text_activator is None else text_activator
+        mode = self.text_activator_runtime_mode if runtime_mode is None else runtime_mode
+        requested_masks = trigger_mask
+        if requested_masks is not None and not isinstance(requested_masks, (list, tuple)):
+            requested_masks = [requested_masks]
+        if requested_masks is not None and len(requested_masks) != len(prompt):
+            if len(prompt) == 1:
+                requested_masks = [trigger_mask]
+            else:
+                raise ValueError("trigger_mask batch length must match prompt batch length")
 
-        # Encode each caption at its natural length (no cross-sample padding) and
-        # store one feature tensor per batch item. Padding to a common length is
-        # deferred to the model call, so caching a prompt only stores its real
-        # length -- important for the long structured (JSON) captions.
         features_list = []
-        for p in prompt:
-            # Digest the prompt: migrate any old-format Ideogram caption into the
-            # current schema and serialize it compact (the form the renderer wants).
-            # Plain-text prompts pass straight through unchanged.
+        trigger_masks = []
+        taps_list = []
+        for prompt_index, p in enumerate(prompt):
             p = digest_caption_string(p)
             messages = [{"role": "user", "content": [{"type": "text", "text": p}]}]
             text = self.tokenizer.apply_chat_template(
@@ -536,19 +659,63 @@ class Ideogram4Model(BaseModel):
             token_ids = torch.tensor([ids], dtype=torch.long, device=device)
             attention_mask = torch.ones_like(token_ids)
             pos_2d = (attention_mask.cumsum(dim=-1) - 1).clamp(min=0).to(torch.long)
+            explicit_mask = (
+                requested_masks[prompt_index] if requested_masks is not None else None
+            )
+            resolved_mask = self._resolve_trigger_mask(
+                p, token_ids, explicit_mask, activator
+            )
 
-            features = get_qwen3_vl_features(
-                self.text_encoder, token_ids, attention_mask, pos_2d
-            )  # (1, Lt, D)
+            result = get_qwen3_vl_features(
+                self.text_encoder,
+                token_ids,
+                attention_mask,
+                pos_2d,
+                trigger_mask=resolved_mask,
+                text_activator=activator,
+                runtime_mode=mode,
+                return_taps=return_taps,
+            )
+            if return_taps:
+                features, taps = result
+                taps_list.append(torch.stack([tap[0] for tap in taps], dim=0))
+            else:
+                features = result
             features_list.append(features[0].to(self.torch_dtype))
+            if resolved_mask is not None:
+                trigger_masks.append(resolved_mask[0])
+            elif activator is not None or trigger_mask is not None:
+                trigger_masks.append(torch.zeros_like(token_ids[0], dtype=torch.bool))
 
-        return AdvancedPromptEmbeds(text_embeds=features_list)
+        embeds_kwargs = {"text_embeds": features_list}
+        if trigger_masks:
+            embeds_kwargs["trigger_masks"] = trigger_masks
+        if return_taps:
+            embeds_kwargs["text_taps"] = taps_list
+            tap_layers = torch.tensor(
+                QWEN3_VL_ACTIVATION_LAYERS, dtype=torch.long, device=device
+            )
+            embeds_kwargs["tap_layers"] = [tap_layers for _ in prompt]
+        embeds = AdvancedPromptEmbeds(**embeds_kwargs)
+        embeds.frozen_dtype_keys = [
+            key for key in ("trigger_masks", "tap_layers") if key in embeds
+        ]
+        return embeds
 
     def get_model_has_grad(self):
         return False
 
     def get_te_has_grad(self):
-        return False
+        activator = self.text_activator
+        if activator is None:
+            return False
+        has_trainable = getattr(activator, "has_trainable_parameters", None)
+        if callable(has_trainable):
+            return bool(has_trainable())
+        parameters = getattr(activator, "parameters", None)
+        if callable(parameters):
+            return any(parameter.requires_grad for parameter in parameters())
+        return True
 
     # ------------------------------------------------------------------
     # VAE

extensions_built_in/diffusion_models/ideogram4/src/pipeline.py

@@ -7,8 +7,11 @@ sampling pipeline used to render preview images during training.
 
 from __future__ import annotations
 
+import contextlib
+import importlib
+import inspect
 import math
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import torch
 from PIL import Image
@@ -104,48 +107,240 @@ def unpatchify_latents(z: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
 # ---------------------------------------------------------------------------
 
 
-@torch.no_grad()
+def _call_activator(
+    activator: Any,
+    method_names: tuple[str, ...],
+    value: torch.Tensor,
+    **kwargs,
+) -> torch.Tensor:
+    """Call the first supported activator hook without imposing a hard API."""
+    for method_name in method_names:
+        method = getattr(activator, method_name, None)
+        if not callable(method):
+            continue
+        try:
+            signature = inspect.signature(method)
+        except (TypeError, ValueError):
+            signature = None
+        leading_args = (value,)
+        if method_name == "apply_tap" and "tap_layer" in kwargs:
+            leading_args = (kwargs["tap_layer"], value)
+        call_kwargs = dict(kwargs)
+        if method_name == "apply_tap":
+            call_kwargs.pop("tap_layer", None)
+        if signature is not None:
+            accepted = signature.parameters
+            call_kwargs = {
+                key: item for key, item in call_kwargs.items() if key in accepted
+            }
+        result = method(*leading_args, **call_kwargs)
+        if result is None:
+            return value
+        if isinstance(result, dict):
+            for key in ("hidden_states", "inputs_embeds", "embeddings", "output"):
+                if key in result:
+                    return result[key]
+        if isinstance(result, (tuple, list)):
+            return result[0]
+        return result
+    return value
+
+
+def _activator_runtime_context(
+    activator: Any,
+    runtime_mode: Optional[str],
+    trigger_mask: Optional[torch.Tensor],
+):
+    if activator is None:
+        return contextlib.nullcontext()
+
+    stack = contextlib.ExitStack()
+    runtime_module = None
+    try:
+        runtime_module = importlib.import_module("toolkit.trigger_binding")
+    except ImportError:
+        pass
+    activator_module = None
+    try:
+        activator_module = importlib.import_module(
+            "toolkit.models.ideogram4_trigger_activator"
+        )
+    except ImportError:
+        pass
+
+    trigger_runtime = getattr(activator_module, "trigger_runtime", None)
+    if callable(trigger_runtime):
+        stack.enter_context(
+            trigger_runtime(
+                {
+                    "token_mask": trigger_mask,
+                    "trigger_mask": trigger_mask,
+                    "runtime_mode": runtime_mode,
+                }
+            )
+        )
+    mode_context = getattr(runtime_module, "activator_runtime_mode", None)
+    if callable(mode_context) and runtime_mode is not None:
+        stack.enter_context(mode_context(activator, runtime_mode))
+    else:
+        for method_name in ("runtime", "runtime_context", "use_runtime_mode"):
+            method = getattr(activator, method_name, None)
+            if callable(method):
+                try:
+                    stack.enter_context(method(runtime_mode))
+                except TypeError:
+                    stack.enter_context(method(mode=runtime_mode))
+                break
+    return stack
+
+
+def _runtime_component_enabled(runtime_mode: Optional[str], component: str) -> bool:
+    if runtime_mode is None:
+        return True
+    try:
+        module = importlib.import_module("toolkit.trigger_binding")
+        state = module.get_activator_runtime_state(runtime_mode)
+    except (ImportError, AttributeError):
+        return runtime_mode not in ("activator_bypass", "stock_literal")
+    return bool(getattr(state, f"{component}_enabled", False))
+
+
+def _adapt_text_activator(activator: Any) -> Any:
+    """Give optional toolkit modules a chance to adapt their evolving runtime API."""
+    if activator is None:
+        return None
+    for module_name in (
+        "toolkit.trigger_binding",
+        "toolkit.models.ideogram4_trigger_activator",
+    ):
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError:
+            continue
+        for adapter_name in (
+            "adapt_ideogram4_text_activator",
+            "adapt_text_activator",
+            "ensure_ideogram4_text_activator",
+        ):
+            adapter = getattr(module, adapter_name, None)
+            if callable(adapter):
+                adapted = adapter(activator)
+                if adapted is not None:
+                    activator = adapted
+                break
+    return activator
+
+
 def get_qwen3_vl_features(
     text_encoder,
     token_ids: torch.Tensor,
     attention_mask: torch.Tensor,
     pos_2d: torch.Tensor,
-) -> torch.Tensor:
-    """Run Qwen3-VL and concat the hidden states from the activation layers.
+    trigger_mask: Optional[torch.Tensor] = None,
+    text_activator: Any = None,
+    runtime_mode: Optional[str] = None,
+    return_taps: bool = False,
+):
+    """Run Qwen3-VL and optionally apply trigger-selective text activation.
 
-    Returns a (B, L, hidden_size * num_layers) tensor (in the encoder's dtype),
-    zeroed at non-text (padding) positions.
+    Gradient ownership belongs to the caller. With no activator and
+    ``return_taps=False`` this returns the original concatenated feature tensor.
     """
     language_model = text_encoder.language_model
+    text_activator = _adapt_text_activator(text_activator)
+    activator_kwargs = {
+        "runtime_mode": runtime_mode,
+        "token_ids": token_ids,
+        "attention_mask": attention_mask,
+    }
 
-    inputs_embeds = language_model.embed_tokens(token_ids)
+    with _activator_runtime_context(text_activator, runtime_mode, trigger_mask):
+        lookup_ids = token_ids
+        atomic_token_id = getattr(text_activator, "atomic_token_id", None)
+        lookup_token_id = getattr(text_activator, "lookup_token_id", None)
+        if atomic_token_id is not None and lookup_token_id is not None:
+            lookup_ids = token_ids.masked_fill(
+                token_ids == int(atomic_token_id), int(lookup_token_id)
+            )
+        inputs_embeds = language_model.embed_tokens(lookup_ids)
+        if text_activator is not None and _runtime_component_enabled(
+            runtime_mode, "embedding"
+        ):
+            inputs_embeds = _call_activator(
+                text_activator,
+                (
+                    "apply_embedding",
+                    "override_embeddings",
+                    "apply_embedding_override",
+                    "apply_embeddings",
+                ),
+                inputs_embeds,
+                token_mask=trigger_mask,
+                **activator_kwargs,
+            )
 
-    position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
-    text_position_ids = position_ids_4d[0]
-    mrope_position_ids = position_ids_4d[1:]
+        position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
+        text_position_ids = position_ids_4d[0]
+        mrope_position_ids = position_ids_4d[1:]
 
-    causal_mask = create_causal_mask(
-        config=language_model.config,
-        inputs_embeds=inputs_embeds,
-        attention_mask=attention_mask,
-        past_key_values=None,
-        position_ids=text_position_ids,
-    )
-    position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
-
-    tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
-    captured: dict[int, torch.Tensor] = {}
-    hidden_states = inputs_embeds
-    for layer_idx, decoder_layer in enumerate(language_model.layers):
-        hidden_states = decoder_layer(
-            hidden_states,
-            attention_mask=causal_mask,
-            position_ids=text_position_ids,
+        causal_mask = create_causal_mask(
+            config=language_model.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
             past_key_values=None,
-            position_embeddings=position_embeddings,
+            position_ids=text_position_ids,
         )
-        if layer_idx in tap_set:
-            captured[layer_idx] = hidden_states
+        position_embeddings = language_model.rotary_emb(
+            inputs_embeds, mrope_position_ids
+        )
+
+        tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
+        captured: dict[int, torch.Tensor] = {}
+        hidden_states = inputs_embeds
+        for layer_idx, decoder_layer in enumerate(language_model.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=text_position_ids,
+                past_key_values=None,
+                position_embeddings=position_embeddings,
+            )
+            if text_activator is not None and _runtime_component_enabled(
+                runtime_mode, "internal"
+            ):
+                hidden_states = _call_activator(
+                    text_activator,
+                    (
+                        "apply_te_adapter",
+                        "apply_internal_adapter",
+                        "apply_hidden_states",
+                    ),
+                    hidden_states,
+                    token_mask=trigger_mask,
+                    layer_idx=layer_idx,
+                    **activator_kwargs,
+                )
+            if layer_idx in tap_set:
+                tap = hidden_states
+                if text_activator is not None and _runtime_component_enabled(
+                    runtime_mode, "tap"
+                ):
+                    tap = _call_activator(
+                        text_activator,
+                        (
+                            "apply_tap",
+                            "adapt_tap",
+                            "apply_tap_adapter",
+                            "apply_pre_concat",
+                        ),
+                        tap,
+                        tap_layer=layer_idx,
+                        token_mask=trigger_mask,
+                        layer_idx=layer_idx,
+                        tap_index=QWEN3_VL_ACTIVATION_LAYERS.index(layer_idx),
+                        **activator_kwargs,
+                    )
+                captured[layer_idx] = tap
 
     selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
     batch_size, seq_len = token_ids.shape
@@ -155,6 +350,8 @@ def get_qwen3_vl_features(
 
     text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
     stacked = stacked * text_mask
+    if return_taps:
+        return stacked, selected
     return stacked
 
 

extensions_built_in/sd_trainer/SDTrainer.py

@@ -1,6 +1,10 @@
+import contextlib
+import importlib
+import inspect
 import os
 import random
 from collections import OrderedDict
+from types import MethodType
 from typing import Union, Literal, List, Optional
 
 import numpy as np
@@ -108,7 +112,14 @@ class SDTrainer(BaseSDTrainProcess):
                 self.save_root,
                 self.trigger_selective_training.logging.metrics_filename,
             )
-        
+
+        self.runtime_phase = None
+        self.text_activator = None
+        self._trigger_binding_modules = {}
+        self._trigger_binding_initial_parameters = {}
+        self._trigger_binding_prompt_encoder = None
+        self._trigger_binding_last_metrics = {}
+
         self.dfe: Optional[DiffusionFeatureExtractor] = None
         self.unconditional_embeds = None
         
@@ -142,6 +153,308 @@ class SDTrainer(BaseSDTrainProcess):
             raise ValueError(f"Unknown guidance loss target type {type(self.train_config.guidance_loss_target)}")
 
 
+    @property
+    def three_phase_enabled(self):
+        return bool(getattr(self.three_phase_trigger_training, 'enabled', False))
+
+    def _load_trigger_binding_modules(self):
+        if self._trigger_binding_modules:
+            return self._trigger_binding_modules
+        required = {
+            'runtime': 'toolkit.trigger_binding',
+            'activator': 'toolkit.models.ideogram4_trigger_activator',
+            'losses': 'toolkit.trigger_binding_losses',
+            'artifacts': 'toolkit.trigger_binding_artifacts',
+        }
+        loaded = {}
+        missing = []
+        for key, module_name in required.items():
+            try:
+                loaded[key] = importlib.import_module(module_name)
+            except ImportError as exc:
+                missing.append(f'{module_name}: {exc}')
+        if missing:
+            raise RuntimeError(
+                'three_phase_trigger_training requires the trigger-binding runtime modules; '
+                + '; '.join(missing)
+            )
+        self._trigger_binding_modules = loaded
+        return loaded
+
+    @staticmethod
+    def _call_supported(callable_obj, **kwargs):
+        try:
+            signature = inspect.signature(callable_obj)
+        except (TypeError, ValueError):
+            return callable_obj(**kwargs)
+        if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()):
+            return callable_obj(**kwargs)
+        return callable_obj(**{key: value for key, value in kwargs.items() if key in signature.parameters})
+
+    @staticmethod
+    def _first_callable(module, names):
+        for name in names:
+            candidate = getattr(module, name, None)
+            if callable(candidate):
+                return candidate
+        return None
+
+    def _resolve_runtime_phase(self):
+        configured = getattr(self.three_phase_trigger_training, 'active_phase', None)
+        phase = configured or os.environ.get('AI_TOOLKIT_RUNTIME_PHASE') or os.environ.get('AITK_RUNTIME_PHASE')
+        if phase is None:
+            enabled = [
+                name for name in ('a1', 'b', 'a2')
+                if getattr(getattr(self.three_phase_trigger_training, f'phase_{name}'), 'enabled', False)
+            ]
+            if len(enabled) == 1:
+                phase = enabled[0]
+        phase = str(phase).lower() if phase is not None else None
+        if phase not in {'a1', 'b', 'a2'}:
+            raise ValueError(
+                'three_phase_trigger_training child runtime requires runtime_phase a1, b, or a2 '
+                '(set config runtime_phase or AI_TOOLKIT_RUNTIME_PHASE)'
+            )
+        phase_config = getattr(self.three_phase_trigger_training, f'phase_{phase}')
+        if not phase_config.enabled:
+            raise ValueError(f'three_phase_trigger_training runtime_phase {phase!r} is disabled')
+        self.runtime_phase = phase
+        return phase
+
+    def _phase_config(self):
+        return getattr(self.three_phase_trigger_training, f'phase_{self.runtime_phase}')
+
+    def _activator_component_flags(self):
+        train = self._phase_config().train
+        return {
+            'embedding': bool(train.get('embedding', train.get('trigger_embedding', False))),
+            'te_adapter': bool(train.get('internal', train.get('te_adapter', train.get('text_encoder_adapter', False)))),
+            'tap_adapters': bool(train.get('tap', train.get('tap_adapters', False))),
+        }
+
+    def _configure_phase_trainability(self):
+        flags = self._activator_component_flags()
+        phase_is_b = self.runtime_phase == 'b'
+        for component, trainable in flags.items():
+            setter = getattr(self.text_activator, 'set_component_mode', None)
+            if callable(setter):
+                setter(component, active=True, trainable=trainable and not phase_is_b)
+            else:
+                module = getattr(self.text_activator, component, None)
+                if module is not None:
+                    module.requires_grad_(trainable and not phase_is_b)
+        if phase_is_b:
+            self.text_activator.requires_grad_(False)
+        if self.network is not None:
+            network_trainable = bool(self._phase_config().train.get('diffusion_lora', phase_is_b))
+            self.network.requires_grad_(network_trainable)
+        return flags
+
+    def _build_text_activator(self, modules):
+        activator_module = modules['activator']
+        factory = self._first_callable(
+            activator_module,
+            ('create_text_activator', 'build_text_activator', 'create_ideogram4_text_activator'),
+        )
+        config = self.three_phase_trigger_training.text_activator
+        common = {
+            'config': config,
+            'text_activator_config': config,
+            'sd': self.sd,
+            'model': self.sd,
+            'text_encoder': self.sd.text_encoder,
+            'tokenizer': self.sd.tokenizer,
+            'device': self.device_torch,
+            'dtype': get_torch_dtype(config.embedding.dtype),
+            'phase': self.runtime_phase,
+            'phase_config': self._phase_config(),
+        }
+        if factory is not None:
+            return self._call_supported(factory, **common)
+
+        TextActivator = getattr(activator_module, 'TextActivator')
+        language_model = self.sd.text_encoder.language_model
+        hidden_size = int(getattr(language_model.config, 'hidden_size'))
+        literal = self.three_phase_trigger_training.literal
+        tokenizer = self.sd.tokenizer
+        tokenizer.add_tokens([literal], special_tokens=True)
+        atomic_ids = tokenizer(literal, add_special_tokens=False)['input_ids']
+        if len(atomic_ids) != 1:
+            raise RuntimeError(f'failed to register atomic trigger token {literal!r}: {atomic_ids}')
+        atomic_token_id = int(atomic_ids[0])
+        init_ids = tokenizer(config.embedding.init_words, add_special_tokens=False)['input_ids']
+        if not init_ids:
+            raise RuntimeError('text activator semantic initializer produced no tokens')
+        embedding_table = language_model.embed_tokens
+        safe_ids = [token_id for token_id in init_ids if token_id < embedding_table.num_embeddings]
+        if not safe_ids:
+            raise RuntimeError('text activator initializer token IDs are outside the Qwen embedding table')
+        initializer_ids = torch.tensor(safe_ids, device=embedding_table.weight.device, dtype=torch.long)
+        with torch.no_grad():
+            initializer = embedding_table(initializer_ids).float().mean(dim=0, keepdim=True)
+        te_config = config.te_adapter
+        te_adapter = None
+        if te_config.enabled:
+            adapter_class = getattr(activator_module, 'MaskedLowRankAdapter')
+            te_adapter = adapter_class(
+                hidden_size=hidden_size,
+                rank=te_config.rank,
+                alpha=te_config.alpha,
+                dropout=te_config.dropout,
+            )
+        tap_config = config.tap_adapters
+        activator = TextActivator(
+            embedding_dim=hidden_size,
+            hidden_size=hidden_size,
+            embedding_tokens=config.embedding.tokens,
+            initializer=initializer,
+            te_adapter=te_adapter,
+            tap_layers=tap_config.tap_layers,
+            tap_rank=tap_config.rank,
+            tap_alpha=tap_config.alpha,
+            tap_dropout=tap_config.dropout,
+            tap_learnable_scale=tap_config.learnable_scale,
+            tap_scale_init=tap_config.scale_init,
+            per_tap=tap_config.per_tap,
+        ).to(self.device_torch, dtype=get_torch_dtype(config.embedding.dtype))
+        activator.atomic_token_id = atomic_token_id
+        activator.lookup_token_id = int(safe_ids[0])
+        return activator
+
+    def _load_activator_source(self, modules):
+        source = self._phase_config().text_activator_source
+        checkpoint = source.path or self.three_phase_trigger_training.text_activator.embedding.checkpoint_path
+        if checkpoint is None:
+            return
+        loader = self._first_callable(
+            modules['artifacts'],
+            ('load_text_activator_artifacts', 'load_text_activator', 'load_activator_artifacts'),
+        )
+        if loader is None:
+            state = load_file(checkpoint)
+            self.text_activator.load_state_dict(state, strict=True)
+            return
+        self._call_supported(
+            loader,
+            activator=self.text_activator,
+            text_activator=self.text_activator,
+            source=source,
+            checkpoint=checkpoint,
+            phase=self.runtime_phase,
+            config=self.three_phase_trigger_training,
+        )
+
+    def hook_after_model_load(self):
+        if not self.three_phase_enabled:
+            return
+        self._resolve_runtime_phase()
+        modules = self._load_trigger_binding_modules()
+        self.text_activator = self._build_text_activator(modules)
+        if not hasattr(self.text_activator, 'apply_embeddings') and hasattr(self.text_activator, 'apply_embedding'):
+            self.text_activator.apply_embeddings = MethodType(
+                lambda activator, value, **kwargs: activator.apply_embedding(
+                    value, token_mask=kwargs.get('trigger_mask')
+                ),
+                self.text_activator,
+            )
+        if not hasattr(self.text_activator, 'apply_internal_adapter') and hasattr(self.text_activator, 'apply_te_adapter'):
+            self.text_activator.apply_internal_adapter = MethodType(
+                lambda activator, value, **kwargs: activator.apply_te_adapter(
+                    value, token_mask=kwargs.get('trigger_mask')
+                ),
+                self.text_activator,
+            )
+        if not hasattr(self.text_activator, 'apply_tap_adapter') and hasattr(self.text_activator, 'apply_tap'):
+            self.text_activator.apply_tap_adapter = MethodType(
+                lambda activator, value, **kwargs: activator.apply_tap(
+                    kwargs.get('layer_idx', kwargs.get('tap_index')), value,
+                    token_mask=kwargs.get('trigger_mask'),
+                ),
+                self.text_activator,
+            )
+        self._load_activator_source(modules)
+        installer = getattr(self.sd, 'install_text_activator', None) or getattr(self.sd, 'set_text_activator', None)
+        if not callable(installer):
+            raise RuntimeError('loaded StableDiffusion model does not expose install_text_activator')
+        installer(self.text_activator, runtime_mode='full')
+        self._trigger_binding_initial_parameters = {
+            name: parameter.detach().float().cpu().clone()
+            for name, parameter in self.text_activator.named_parameters()
+        }
+        self._install_trigger_binding_prompt_encoder(modules)
+
+    def _install_trigger_binding_prompt_encoder(self, modules):
+        original = self.sd.get_prompt_embeds
+        self._trigger_binding_prompt_encoder = original
+
+        def encode_with_binding(sd_model, prompt):
+            if not self.three_phase_enabled or self.runtime_phase == 'b':
+                return original(prompt)
+            prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+            runtime = modules['runtime']
+            batch = runtime.bind_trigger_batch(
+                sd_model.tokenizer,
+                prompts,
+                self.three_phase_trigger_training.literal,
+                placeholder=self.three_phase_trigger_training.placeholder,
+                max_length=getattr(sd_model, 'max_text_length', None),
+                require_placeholder=True,
+                mask_all_occurrences=self.three_phase_trigger_training.mask_all_occurrences,
+                require_atomic=True,
+                expected_token_id=getattr(self.text_activator, 'atomic_token_id', None),
+            ).to(sd_model.text_encoder.device)
+            pipeline = importlib.import_module(
+                'extensions_built_in.diffusion_models.ideogram4.src.pipeline'
+            )
+            features = pipeline.get_qwen3_vl_features(
+                sd_model.text_encoder,
+                batch.input_ids,
+                batch.attention_mask,
+                (batch.attention_mask.cumsum(dim=-1) - 1).clamp(min=0).long(),
+                trigger_mask=batch.trigger_mask,
+                text_activator=self.text_activator,
+                runtime_mode='full',
+            )
+            from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+            return AdvancedPromptEmbeds(
+                text_embeds=[features[index].to(sd_model.torch_dtype) for index in range(features.shape[0])]
+            )
+
+        self.sd.get_prompt_embeds = MethodType(encode_with_binding, self.sd)
+
+    def hook_add_extra_train_params(self, params):
+        if not self.three_phase_enabled:
+            return params
+        self._configure_phase_trainability()
+        phase_is_b = self.runtime_phase == 'b'
+        allowed_ids = set()
+        if phase_is_b:
+            if self.network is not None:
+                allowed_ids.update(id(parameter) for parameter in self.network.parameters() if parameter.requires_grad)
+        else:
+            allowed_ids.update(id(parameter) for parameter in self.text_activator.parameters() if parameter.requires_grad)
+        filtered = []
+        for group in params:
+            if isinstance(group, dict):
+                kept = [parameter for parameter in group.get('params', []) if id(parameter) in allowed_ids]
+                if kept:
+                    copied = dict(group)
+                    copied['params'] = kept
+                    filtered.append(copied)
+            elif id(group) in allowed_ids:
+                filtered.append(group)
+        if not phase_is_b:
+            learning_rates = self._phase_config().learning_rates
+            groups = self.text_activator.parameter_groups(learning_rates)
+            present = {id(parameter) for group in filtered for parameter in (group.get('params', []) if isinstance(group, dict) else [group])}
+            for group in groups:
+                group['params'] = [parameter for parameter in group['params'] if id(parameter) not in present]
+                if group['params']:
+                    filtered.append(group)
+        if not filtered:
+            raise RuntimeError(f'runtime_phase {self.runtime_phase} selected no trainable parameters')
+        return filtered
+
     def before_model_load(self):
         pass
 
@@ -295,6 +608,34 @@ class SDTrainer(BaseSDTrainProcess):
 
     def hook_before_train_loop(self):
         super().hook_before_train_loop()
+        if self.three_phase_enabled:
+            probe_config = self.three_phase_trigger_training.reachability_probe
+            if probe_config.get('enabled', True):
+                modules = self._load_trigger_binding_modules()
+                probe = self._first_callable(
+                    modules['runtime'],
+                    ('run_reachability_probe', 'probe_reachability', 'validate_reachability'),
+                ) or self._first_callable(
+                    modules['losses'],
+                    ('run_reachability_probe', 'probe_reachability', 'validate_reachability'),
+                )
+                if probe is not None:
+                    result = self._call_supported(
+                        probe,
+                        trainer=self,
+                        sd=self.sd,
+                        activator=self.text_activator,
+                        text_activator=self.text_activator,
+                        phase=self.runtime_phase,
+                        config=probe_config,
+                        phase_config=self._phase_config(),
+                    )
+                    if result is False or (isinstance(result, dict) and not result.get('reachable', result.get('ok', True))):
+                        raise RuntimeError(f'trigger activator reachability probe failed: {result}')
+                else:
+                    diagnostics = self.text_activator.probe_diagnostics()
+                    if not diagnostics.active or diagnostics.total_parameters == 0:
+                        raise RuntimeError(f'trigger activator reachability probe failed: {diagnostics}')
         if self.is_caching_text_embeddings:
             # make sure model is on cpu for this part so we don't oom.
             self.sd.unet.to('cpu')
@@ -1379,6 +1720,140 @@ class SDTrainer(BaseSDTrainProcess):
         )
     
 
+    def _activator_mode(self, mode):
+        if not self.three_phase_enabled or self.text_activator is None:
+            return contextlib.nullcontext()
+        setter = getattr(self.sd, 'set_text_activator_runtime_mode', None)
+        previous = getattr(self.sd, 'text_activator_runtime_mode', None)
+
+        @contextlib.contextmanager
+        def runtime_context():
+            if callable(setter):
+                setter(mode)
+            try:
+                yield
+            finally:
+                if callable(setter):
+                    setter(previous)
+        return runtime_context()
+
+    def _calculate_trigger_binding_loss(
+        self,
+        noisy_latents,
+        noise,
+        timesteps,
+        batch,
+        pred_kwargs,
+        mask_multiplier,
+        dtype,
+    ):
+        raw_prompts = [
+            getattr(item, 'caption_template', None) or item.raw_caption
+            for item in batch.file_items
+        ]
+        with self._activator_mode('full'):
+            active_embeds = self.sd.encode_prompt(raw_prompts, long_prompts=self.do_long_prompts)
+            active_pred = self.predict_noise(
+                noisy_latents=noisy_latents,
+                timesteps=timesteps,
+                conditional_embeds=active_embeds.to(self.device_torch, dtype=dtype),
+                unconditional_embeds=None,
+                batch=batch,
+                is_primary_pred=True,
+                **pred_kwargs,
+            )
+        with self._activator_mode('activator_bypass'):
+            bypass_embeds = self.sd.encode_prompt(raw_prompts, long_prompts=self.do_long_prompts)
+            with torch.no_grad():
+                bypass_pred = self.predict_noise(
+                    noisy_latents=noisy_latents,
+                    timesteps=timesteps,
+                    conditional_embeds=bypass_embeds.to(self.device_torch, dtype=dtype),
+                    unconditional_embeds=None,
+                    batch=batch,
+                    **pred_kwargs,
+                ).detach()
+        target = shared_loss_target(self, noise, batch, timesteps)
+        loss_module = self._load_trigger_binding_modules()['losses']
+        loss_fn = self._first_callable(
+            loss_module,
+            ('calculate_trigger_binding_losses', 'trigger_binding_losses', 'compute_trigger_binding_losses'),
+        )
+        if loss_fn is not None:
+            result = self._call_supported(
+                loss_fn,
+                trainer=self,
+                phase=self.runtime_phase,
+                phase_config=self._phase_config(),
+                config=self.three_phase_trigger_training,
+                active_prediction=active_pred,
+                activator_prediction=active_pred,
+                bypass_prediction=bypass_pred,
+                baseline_prediction=bypass_pred,
+                target=target,
+                noise=noise,
+                noisy_latents=noisy_latents,
+                timesteps=timesteps,
+                batch=batch,
+                mask_multiplier=mask_multiplier,
+                active_embeddings=active_embeds,
+                bypass_embeddings=bypass_embeds,
+            )
+        elif self.runtime_phase == 'a1' and callable(getattr(loss_module, 'compute_a1_loss', None)):
+            floor_config = self._phase_config().activator_gain_floor
+            floor = 0.0
+            if floor_config.enabled and floor_config.schedule.keyframes:
+                floor = loss_module.scheduled_gain_floor(
+                    self.step_num,
+                    floor_config.schedule.keyframes,
+                    floor_config.schedule.interpolation,
+                )
+            result = loss_module.compute_a1_loss(
+                active_pred,
+                target,
+                bypass_prediction=bypass_pred,
+                gain_floor=floor,
+                gain_floor_weight=floor_config.weight if floor_config.enabled else 0.0,
+            )
+        elif self.runtime_phase == 'a2' and callable(getattr(loss_module, 'compute_a2_loss', None)):
+            floor_config = self._phase_config().activator_gain_floor
+            floor = 0.0
+            if floor_config.enabled and floor_config.schedule.keyframes:
+                floor = loss_module.scheduled_gain_floor(
+                    self.step_num,
+                    floor_config.schedule.keyframes,
+                    floor_config.schedule.interpolation,
+                )
+            result = loss_module.compute_a2_loss(
+                active_pred,
+                bypass_pred,
+                target,
+                gain_floor=floor,
+                gain_floor_weight=floor_config.weight if floor_config.enabled else 0.0,
+            )
+        else:
+            raise RuntimeError('toolkit.trigger_binding_losses has no supported loss entry point')
+        metrics = {}
+        if torch.is_tensor(result):
+            loss = result
+        elif isinstance(result, tuple):
+            loss, metrics = result[0], result[1] if len(result) > 1 else {}
+        elif isinstance(result, dict):
+            loss = result.get('loss', result.get('total_loss'))
+            metrics = result.get('metrics', result.get('logs', {}))
+        else:
+            loss = getattr(result, 'loss', getattr(result, 'total_loss', None))
+            metrics = getattr(result, 'metrics', {})
+        if loss is None:
+            raise RuntimeError('trigger binding loss entry point did not return a loss tensor')
+        self._trigger_binding_last_metrics = dict(metrics or {})
+        self.additional_logs['runtime_phase'] = self.runtime_phase
+        for key, value in self._trigger_binding_last_metrics.items():
+            if isinstance(value, (int, float)):
+                self.additional_logs[f'phase/{self.runtime_phase}/{key}'] = float(value)
+        self.additional_logs[f'phase/{self.runtime_phase}/loss'] = float(loss.detach().item())
+        return loss
+
     def _encode_tst_prompt_variants(self, batch, trigger_prompts, decoy_prompts, dtype):
         prompt_kwargs = {}
         if self.sd.encode_control_in_text_embeddings and batch.control_tensor is not None:
@@ -2439,7 +2914,18 @@ class SDTrainer(BaseSDTrainProcess):
                     )
                 else:
                     tst_loss = None
-                    if self.trigger_selective_training.enabled:
+                    if self.three_phase_enabled and self.runtime_phase in {'a1', 'a2'}:
+                        with self.timer('trigger_binding_predict_and_loss'):
+                            tst_loss = self._calculate_trigger_binding_loss(
+                                noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype),
+                                noise=noise.to(self.device_torch, dtype=dtype).detach(),
+                                timesteps=timesteps,
+                                batch=batch,
+                                pred_kwargs=pred_kwargs,
+                                mask_multiplier=mask_multiplier,
+                                dtype=dtype,
+                            )
+                    elif self.trigger_selective_training.enabled:
                         with self.timer('tst_predict_and_loss'):
                             tst_loss = self._calculate_tst_loss(
                                 noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype),
@@ -2547,6 +3033,92 @@ class SDTrainer(BaseSDTrainProcess):
         return loss.detach()
         # flush()
 
+    def post_save_hook(self, save_path):
+        if not self.three_phase_enabled or self.text_activator is None:
+            return
+        modules = self._load_trigger_binding_modules()
+        proof = {}
+        for name, parameter in self.text_activator.named_parameters():
+            initial = self._trigger_binding_initial_parameters.get(name)
+            if initial is None:
+                continue
+            proof[name] = float((parameter.detach().float().cpu() - initial).norm().item())
+        writer = self._first_callable(
+            modules['artifacts'],
+            ('save_phase_artifacts', 'save_trigger_binding_artifacts', 'save_text_activator_artifacts'),
+        )
+        artifact_config = getattr(
+            self.three_phase_trigger_training.artifacts, f'phase_{self.runtime_phase}'
+        )
+        if writer is not None:
+            self._call_supported(
+                writer,
+                trainer=self,
+                activator=self.text_activator,
+                text_activator=self.text_activator,
+                phase=self.runtime_phase,
+                step=self.step_num,
+                save_path=save_path,
+                save_root=self.save_root,
+                config=self.three_phase_trigger_training,
+                phase_config=self._phase_config(),
+                artifact_config=artifact_config,
+                metrics=self._trigger_binding_last_metrics,
+                parameter_change_proof=proof,
+            )
+            return
+        save_artifact = getattr(modules['artifacts'], 'save_artifact', None)
+        if not callable(save_artifact):
+            raise RuntimeError('toolkit.trigger_binding_artifacts has no supported save entry point')
+        phase_root = os.path.join(
+            self.three_phase_trigger_training.run_root or self.save_root,
+            f'phase_{self.runtime_phase}',
+        )
+        output_dir = os.path.join(
+            phase_root,
+            artifact_config.final_dir if self.step_num >= self._phase_config().steps else artifact_config.checkpoint_dir,
+        )
+        if self.step_num < self._phase_config().steps:
+            output_dir = os.path.join(output_dir, str(self.step_num))
+        os.makedirs(output_dir, exist_ok=True)
+        component_specs = (
+            ('embedding', self.text_activator.embedding, artifact_config.embedding_filename),
+            ('te_adapter', self.text_activator.te_adapter, artifact_config.te_adapter_filename),
+            ('tap_adapter', self.text_activator.tap_adapters, artifact_config.tap_adapter_filename),
+        )
+        for artifact_type, component, filename in component_specs:
+            if component is None:
+                continue
+            tensors = {
+                name: tensor.detach().cpu()
+                for name, tensor in component.state_dict().items()
+            }
+            if tensors:
+                save_artifact(
+                    os.path.join(output_dir, filename),
+                    artifact_type,
+                    tensors,
+                    phase=self._phase_config(),
+                    source=self._phase_config().text_activator_source,
+                    config=self.three_phase_trigger_training.text_activator,
+                    extra={
+                        'runtime_phase': self.runtime_phase,
+                        'step': self.step_num,
+                        'metrics': self._trigger_binding_last_metrics,
+                        'parameter_change_proof': proof,
+                    },
+                )
+        metrics_path = os.path.join(phase_root, artifact_config.metrics_file)
+        os.makedirs(os.path.dirname(metrics_path), exist_ok=True)
+        with open(metrics_path, 'a', encoding='utf-8') as handle:
+            import json
+            handle.write(json.dumps({
+                'phase': self.runtime_phase,
+                'step': self.step_num,
+                'metrics': self._trigger_binding_last_metrics,
+                'parameter_change_proof': proof,
+            }, sort_keys=True) + '\n')
+
     def hook_train_loop(self, batch: Union[DataLoaderBatchDTO, List[DataLoaderBatchDTO]]):
         if isinstance(batch, list):
             batch_list = batch

extensions_built_in/sd_trainer/ThreePhaseTriggerTrainer.py

@@ -0,0 +1,271 @@
+import copy
+import json
+import os
+import subprocess
+import sys
+from collections import OrderedDict
+from datetime import datetime, timezone
+from typing import Dict, Optional
+
+import yaml
+
+from jobs.process import BaseExtensionProcess
+from toolkit.config_modules import (
+    ThreePhaseTriggerTrainingConfig,
+    TriggerBindingPhaseSourceConfig,
+    validate_three_phase_trigger_training_config,
+)
+from toolkit.paths import TOOLKIT_ROOT, get_path
+
+
+class ThreePhaseTriggerTrainer(BaseExtensionProcess):
+    PHASE_NAMES = ('a1', 'b', 'a2')
+
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.three_phase_config = ThreePhaseTriggerTrainingConfig(
+            **self.get_conf('three_phase_trigger_training', {})
+        )
+        if self.three_phase_config.literal is None:
+            self.three_phase_config.literal = self.get_conf('trigger_word', None)
+        validate_three_phase_trigger_training_config(
+            self.three_phase_config,
+            self.get_conf('trigger_word', None),
+        )
+        self.phase_snapshots: Dict[str, str] = {}
+        self.phase_contracts: Dict[str, str] = {}
+        self.active_phase: Optional[str] = None
+        if not self.three_phase_config.enabled:
+            return
+
+        self.training_folder = get_path(self.get_conf('training_folder', required=True))
+        configured_output_root = self.three_phase_config.artifacts.output_root
+        if configured_output_root is None:
+            configured_output_root = os.path.join(self.training_folder, self.name)
+        self.run_root = get_path(configured_output_root)
+        self.snapshot_root = os.path.join(self.run_root, 'phase_configs')
+        self.contract_root = os.path.join(self.run_root, 'contracts')
+        os.makedirs(self.snapshot_root, exist_ok=True)
+        os.makedirs(self.contract_root, exist_ok=True)
+
+    def _phase_root(self, phase_name: str) -> str:
+        return os.path.join(self.run_root, f'phase_{phase_name}')
+
+    def _phase_artifacts(self, phase_name: str):
+        return getattr(self.three_phase_config.artifacts, f'phase_{phase_name}')
+
+    def _source_path(
+        self,
+        source: TriggerBindingPhaseSourceConfig,
+        artifact_field: str,
+    ) -> Optional[str]:
+        if source.path is not None:
+            return get_path(source.path)
+        if source.phase is None:
+            return None
+
+        source_artifacts = self._phase_artifacts(source.phase)
+        artifact_name = getattr(source_artifacts, artifact_field)
+        source_root = self._phase_root(source.phase)
+        if source.step == 'final':
+            return os.path.join(source_root, source_artifacts.final_dir, artifact_name)
+        return os.path.join(
+            source_root,
+            source_artifacts.checkpoint_dir,
+            str(source.step),
+            artifact_name,
+        )
+
+    def _resolve_phase_sources(self, phase_name: str) -> Dict[str, Optional[str]]:
+        phase = self.three_phase_config.get_phase(phase_name)
+        return {
+            'embedding': self._source_path(
+                phase.text_activator_source,
+                'embedding_filename',
+            ),
+            'te_adapter': self._source_path(
+                phase.text_activator_source,
+                'te_adapter_filename',
+            ),
+            'tap_adapters': self._source_path(
+                phase.text_activator_source,
+                'tap_adapter_filename',
+            ),
+            'diffusion_lora': self._source_path(
+                phase.diffusion_lora_source,
+                'diffusion_lora_filename',
+            ),
+        }
+
+    def _build_child_process_config(self, phase_name: str) -> OrderedDict:
+        phase = self.three_phase_config.get_phase(phase_name)
+        child = copy.deepcopy(OrderedDict(self.raw_process_config))
+        child['type'] = 'sd_trainer'
+        child['name'] = f'phase_{phase_name}'
+        child['training_folder'] = self.run_root
+        child['train'] = copy.deepcopy(child.get('train', {}))
+        child['train']['steps'] = phase.steps
+        child['train']['optimizer'] = phase.optimizer
+        child['train']['optimizer_params'] = copy.deepcopy(phase.optimizer_params)
+        child['save'] = copy.deepcopy(child.get('save', {}))
+        if phase.save_steps:
+            child['save']['save_every'] = min(phase.save_steps)
+
+        learning_rate_aliases = {
+            'lr': 'lr',
+            'diffusion': 'lr',
+            'diffusion_lora': 'lr',
+            'unet': 'unet_lr',
+            'text_encoder': 'text_encoder_lr',
+            'embedding': 'embedding_lr',
+            'adapter': 'adapter_lr',
+        }
+        for key, value in phase.learning_rates.items():
+            child['train'][learning_rate_aliases.get(key, key)] = value
+
+        child['train'].update(copy.deepcopy(phase.train))
+        child['three_phase_trigger_training'] = copy.deepcopy(
+            child.get('three_phase_trigger_training', {})
+        )
+        runtime = child['three_phase_trigger_training'].setdefault('runtime', {})
+        runtime.update({
+            'active_phase': phase_name,
+            'orchestrated': True,
+            'run_root': self.run_root,
+            'config_snapshot': self.phase_snapshots.get(
+                phase_name,
+                os.path.join(self.snapshot_root, f'phase_{phase_name}.yaml'),
+            ),
+            'completion_contract': self.phase_contracts.get(
+                phase_name,
+                os.path.join(self.contract_root, f'phase_{phase_name}.json'),
+            ),
+        })
+        sources = self._resolve_phase_sources(phase_name)
+        child['three_phase_trigger_training']['phase_runtime'] = {
+            'caption_sources': copy.deepcopy(phase.caption_sources),
+            'losses': copy.deepcopy(phase.losses),
+            'save_steps': list(phase.save_steps),
+            'resume': {
+                'enabled': phase.resume.enabled,
+                'checkpoint': phase.resume.checkpoint,
+            },
+            'sources': sources,
+        }
+        if phase_name == 'a2' and sources['diffusion_lora'] is not None:
+            child['network'] = copy.deepcopy(child.get('network', {}))
+            child['network']['pretrained_lora_path'] = sources['diffusion_lora']
+        if phase_name == 'b':
+            child['trigger_selective_training'] = copy.deepcopy(child.get('trigger_selective_training', {}))
+            child['trigger_selective_training']['phase_local_step'] = True
+            child['trigger_selective_training']['source_artifact_hashes'] = sources
+        return child
+
+    def build_child_job_config(self, phase_name: str) -> OrderedDict:
+        child_process = self._build_child_process_config(phase_name)
+        child_job = OrderedDict({
+            'job': 'extension',
+            'config': OrderedDict({
+                'name': child_process['name'],
+                'process': [child_process],
+            }),
+            'meta': copy.deepcopy(self.job.meta),
+        })
+        return child_job
+
+    def _write_yaml_atomic(self, path: str, data: OrderedDict):
+        temp_path = path + '.tmp'
+        serializable = json.loads(json.dumps(data))
+        with open(temp_path, 'w', encoding='utf-8') as handle:
+            yaml.safe_dump(serializable, handle, sort_keys=False, allow_unicode=True)
+        os.replace(temp_path, path)
+
+    def _write_json_atomic(self, path: str, data: Dict):
+        temp_path = path + '.tmp'
+        with open(temp_path, 'w', encoding='utf-8') as handle:
+            json.dump(data, handle, indent=2, sort_keys=True)
+        os.replace(temp_path, path)
+
+    def write_phase_snapshot(self, phase_name: str) -> str:
+        snapshot_path = os.path.join(self.snapshot_root, f'phase_{phase_name}.yaml')
+        self.phase_snapshots[phase_name] = snapshot_path
+        self._write_yaml_atomic(snapshot_path, self.build_child_job_config(phase_name))
+        return snapshot_path
+
+    def completion_contract(self, phase_name: str, status: str, return_code: Optional[int] = None) -> Dict:
+        phase = self.three_phase_config.get_phase(phase_name)
+        phase_artifacts = self._phase_artifacts(phase_name)
+        phase_root = self._phase_root(phase_name)
+        return {
+            'schema_version': 1,
+            'phase': phase_name,
+            'status': status,
+            'return_code': return_code,
+            'completed_at': datetime.now(timezone.utc).isoformat() if status in ('completed', 'failed') else None,
+            'config_snapshot': self.phase_snapshots[phase_name],
+            'phase_root': phase_root,
+            'steps': phase.steps,
+            'inputs': self._resolve_phase_sources(phase_name),
+            'artifacts': {
+                'metrics_file': os.path.join(phase_root, phase_artifacts.metrics_file),
+                'console_log': os.path.join(phase_root, phase_artifacts.console_log),
+                'checkpoint_dir': os.path.join(phase_root, phase_artifacts.checkpoint_dir),
+                'final_dir': os.path.join(phase_root, phase_artifacts.final_dir),
+                'embedding': os.path.join(phase_root, phase_artifacts.final_dir, phase_artifacts.embedding_filename),
+                'te_adapter': os.path.join(phase_root, phase_artifacts.final_dir, phase_artifacts.te_adapter_filename),
+                'tap_adapters': os.path.join(phase_root, phase_artifacts.final_dir, phase_artifacts.tap_adapter_filename),
+                'diffusion_lora': os.path.join(phase_root, phase_artifacts.final_dir, phase_artifacts.diffusion_lora_filename),
+            },
+        }
+
+    def write_completion_contract(
+        self,
+        phase_name: str,
+        status: str,
+        return_code: Optional[int] = None,
+    ) -> str:
+        contract_path = os.path.join(self.contract_root, f'phase_{phase_name}.json')
+        self.phase_contracts[phase_name] = contract_path
+        self._write_json_atomic(
+            contract_path,
+            self.completion_contract(phase_name, status, return_code),
+        )
+        return contract_path
+
+    def run_phase(self, phase_name: str):
+        self.active_phase = phase_name
+        snapshot_path = self.write_phase_snapshot(phase_name)
+        self.write_completion_contract(phase_name, 'running')
+        command = [sys.executable, os.path.join(TOOLKIT_ROOT, 'run.py'), snapshot_path]
+        result = subprocess.run(command, cwd=TOOLKIT_ROOT, check=False)
+        status = 'completed' if result.returncode == 0 else 'failed'
+        self.write_completion_contract(phase_name, status, result.returncode)
+        if result.returncode != 0:
+            raise RuntimeError(
+                f'Three-phase trigger child phase {phase_name} failed with exit code {result.returncode}'
+            )
+
+    def _contract_is_verified(self, phase_name: str) -> bool:
+        path = os.path.join(self.contract_root, f'phase_{phase_name}.json')
+        if not os.path.isfile(path):
+            return False
+        try:
+            with open(path, 'r', encoding='utf-8') as handle:
+                contract = json.load(handle)
+        except (OSError, json.JSONDecodeError):
+            return False
+        return contract.get('phase') == phase_name and contract.get('status') == 'completed' and contract.get('return_code') == 0
+
+    def run(self):
+        super().run()
+        if not self.three_phase_config.enabled:
+            return
+        for phase_name in self.PHASE_NAMES:
+            if self._contract_is_verified(phase_name):
+                continue
+            self.run_phase(phase_name)
+        self.active_phase = None
+
+    def on_error(self, error: Exception):
+        if self.active_phase is not None and self.active_phase in self.phase_snapshots:
+            self.write_completion_contract(self.active_phase, 'failed')

extensions_built_in/sd_trainer/__init__.py

@@ -39,6 +39,17 @@ class UITrainerExtension(Extension):
 
 
 # This is a universal trainer that can be from ui or api
+class ThreePhaseTriggerTrainerExtension(Extension):
+    uid = "three_phase_trigger_trainer"
+    name = "Three-Phase Trigger Trainer"
+
+    @classmethod
+    def get_process(cls):
+        from .ThreePhaseTriggerTrainer import ThreePhaseTriggerTrainer
+
+        return ThreePhaseTriggerTrainer
+
+
 class DiffusionTrainerExtension(Extension):
     # uid must be unique, it is how the extension is identified
     uid = "diffusion_trainer"
@@ -67,4 +78,5 @@ AI_TOOLKIT_EXTENSIONS = [
     TextualInversionTrainer,
     UITrainerExtension,
     DiffusionTrainerExtension,
+    ThreePhaseTriggerTrainerExtension,
 ]

jobs/process/BaseSDTrainProcess.py

@@ -63,7 +63,7 @@ from tqdm import tqdm
 
 from toolkit.config_modules import SaveConfig, LoggingConfig, SampleConfig, NetworkConfig, TrainConfig, ModelConfig, \
     GenerateImageConfig, EmbeddingConfig, DatasetConfig, preprocess_dataset_raw_config, AdapterConfig, GuidanceConfig, validate_configs, \
-    DecoratorConfig, TriggerSelectiveTrainingConfig
+    DecoratorConfig, TriggerSelectiveTrainingConfig, ThreePhaseTriggerTrainingConfig
 from toolkit.logging_aitk import create_logger
 from diffusers import FluxTransformer2DModel
 from toolkit.accelerator import get_accelerator, unwrap_model
@@ -113,6 +113,11 @@ class BaseSDTrainProcess(BaseTrainProcess):
         self.trigger_selective_training = TriggerSelectiveTrainingConfig(
             **self.get_conf('trigger_selective_training', {})
         )
+        self.three_phase_trigger_training = ThreePhaseTriggerTrainingConfig(
+            **self.get_conf('three_phase_trigger_training', {})
+        )
+        if self.three_phase_trigger_training.literal is None:
+            self.three_phase_trigger_training.literal = self.get_conf('trigger_word', None)
         model_config = self.get_conf('model', {})
         self.modules_being_trained: List[torch.nn.Module] = []
 
@@ -271,6 +276,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
             self.trigger_selective_training,
             self.trigger_word,
             self.network_config,
+            self.three_phase_trigger_training,
         )
         
         do_profiler = self.get_conf('torch_profiler', False)

testing/test_ideogram4_trigger_activator.py

@@ -0,0 +1,152 @@
+import unittest
+
+import torch
+from torch import nn
+
+from toolkit.models.ideogram4_trigger_activator import (
+    AtomicLearnedEmbedding,
+    DEFAULT_TAP_LAYERS,
+    MaskedLowRankAdapter,
+    TextActivator,
+    trigger_runtime,
+)
+
+
+class _FakeQwen(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.blocks = nn.ModuleList([nn.Linear(hidden_size, hidden_size, bias=False) for _ in range(4)])
+        for block in self.blocks:
+            nn.init.eye_(block.weight)
+
+    def forward(self, hidden_states):
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        return hidden_states
+
+
+class Ideogram4TriggerActivatorTest(unittest.TestCase):
+    def test_atomic_embedding_learned_frozen_and_bypass(self):
+        initializer = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+        embedding = AtomicLearnedEmbedding(4, initializer=initializer)
+        hidden = torch.zeros(1, 3, 4)
+        mask = torch.tensor([[0, 1, 0]], dtype=torch.bool)
+
+        embedding.weight.data.fill_(9.0)
+        learned = embedding(hidden, mask, mode="learned")
+        frozen = embedding(hidden, mask, mode="frozen")
+        bypass = embedding(hidden, mask, mode="bypass")
+
+        self.assertTrue(torch.equal(learned[0, 1], torch.full((4,), 9.0)))
+        self.assertTrue(torch.equal(frozen[0, 1], initializer[0]))
+        self.assertTrue(torch.equal(frozen[0, 0], hidden[0, 0]))
+        self.assertIs(bypass, hidden)
+        self.assertFalse(embedding.frozen_initializer.requires_grad)
+
+    def test_masked_low_rank_adapter_only_changes_trigger_span(self):
+        adapter = MaskedLowRankAdapter(4, rank=2, alpha=2)
+        nn.init.ones_(adapter.down.weight)
+        nn.init.ones_(adapter.up.weight)
+        hidden = torch.ones(1, 3, 4)
+        mask = torch.tensor([[0, 1, 0]], dtype=torch.bool)
+
+        output = adapter(hidden, mask)
+
+        self.assertTrue(torch.equal(output[:, 0], hidden[:, 0]))
+        self.assertFalse(torch.equal(output[:, 1], hidden[:, 1]))
+        self.assertTrue(torch.equal(output[:, 2], hidden[:, 2]))
+
+    def test_runtime_context_supplies_mask_without_hard_dependency(self):
+        adapter = MaskedLowRankAdapter(2, rank=1)
+        nn.init.ones_(adapter.down.weight)
+        nn.init.ones_(adapter.up.weight)
+        hidden = torch.ones(1, 2, 2)
+        with trigger_runtime({"token_mask": torch.tensor([[1, 0]])}):
+            output = adapter(hidden)
+        self.assertFalse(torch.equal(output[:, 0], hidden[:, 0]))
+        self.assertTrue(torch.equal(output[:, 1], hidden[:, 1]))
+
+    def test_exactly_thirteen_taps_are_keyed_by_actual_layer(self):
+        activator = TextActivator(4, tap_layers=DEFAULT_TAP_LAYERS)
+        self.assertEqual(activator.tap_layers, DEFAULT_TAP_LAYERS)
+        self.assertEqual(len(activator.tap_adapters), 13)
+        with self.assertRaises(KeyError):
+            activator.apply_tap(1, torch.zeros(1, 1, 4), torch.ones(1, 1))
+        with self.assertRaises(ValueError):
+            TextActivator(4, tap_layers=DEFAULT_TAP_LAYERS[:-1])
+
+    def test_component_active_trainable_and_parameter_groups(self):
+        te_adapter = MaskedLowRankAdapter(4, rank=1)
+        activator = TextActivator(4, te_adapter=te_adapter)
+        activator.set_component_mode("embedding", active=False, trainable=False)
+        activator.set_component_mode("te_adapter", active=True, trainable=True)
+        groups = activator.parameter_groups({"te_adapter": 2.0e-4, "tap_adapters": 1.0e-4})
+
+        self.assertFalse(activator.component_active["embedding"])
+        self.assertTrue(all(not parameter.requires_grad for parameter in activator.embedding.parameters()))
+        self.assertEqual(
+            [group["name"] for group in groups],
+            ["text_activator.te_adapter", "text_activator.tap_adapters"],
+        )
+        self.assertEqual([group["lr"] for group in groups], [2.0e-4, 1.0e-4])
+
+    def test_qwen_internal_hooks_apply_and_are_removable(self):
+        qwen = _FakeQwen(4)
+        activator = TextActivator(4)
+        tap = activator.tap_adapters[str(DEFAULT_TAP_LAYERS[0])]
+        nn.init.ones_(tap.down.weight)
+        nn.init.ones_(tap.up.weight)
+        hidden = torch.ones(1, 2, 4)
+        mask = torch.tensor([[1, 0]], dtype=torch.bool)
+
+        with trigger_runtime({"token_mask": mask}):
+            baseline = qwen(hidden)
+            activator.install_qwen_hooks(
+                qwen,
+                tap_module_names={DEFAULT_TAP_LAYERS[0]: "blocks.1"},
+            )
+            activated = qwen(hidden)
+            diagnostics = activator.probe_diagnostics()
+            activator.remove_qwen_hooks()
+            restored = qwen(hidden)
+
+        self.assertFalse(torch.equal(activated[:, 0], baseline[:, 0]))
+        self.assertTrue(torch.equal(activated[:, 1], baseline[:, 1]))
+        self.assertTrue(torch.equal(restored, baseline))
+        self.assertEqual(diagnostics.hook_count, 1)
+        self.assertGreater(diagnostics.adapter_update_norms[str(DEFAULT_TAP_LAYERS[0])], 0.0)
+
+    def test_wrapper_mode_restores_original_module(self):
+        qwen = _FakeQwen(4)
+        original = qwen.blocks[2]
+        activator = TextActivator(4, te_adapter=MaskedLowRankAdapter(4, rank=1))
+        activator.install_qwen_hooks(qwen, te_module_names=["blocks.2"], use_wrappers=True)
+        self.assertIsNot(qwen.blocks[2], original)
+        activator.remove_qwen_hooks()
+        self.assertIs(qwen.blocks[2], original)
+
+    def test_namespaced_state_dict_and_strict_validation(self):
+        activator = TextActivator(4, te_adapter=MaskedLowRankAdapter(4, rank=1))
+        state = activator.state_dict()
+        self.assertTrue(state)
+        self.assertTrue(all(key.startswith("ideogram4_text_activator.") for key in state))
+
+        clone = TextActivator(4, te_adapter=MaskedLowRankAdapter(4, rank=1))
+        clone.load_state_dict(state, strict=True)
+        missing = dict(state)
+        missing.pop(next(iter(missing)))
+        with self.assertRaises(RuntimeError):
+            clone.load_state_dict(missing, strict=True)
+        foreign = dict(state)
+        foreign["unrelated.weight"] = torch.ones(1)
+        with self.assertRaises(RuntimeError):
+            clone.load_state_dict(foreign, strict=True)
+        wrong_shape = dict(state)
+        first_key = next(iter(wrong_shape))
+        wrong_shape[first_key] = torch.ones(999)
+        with self.assertRaises(RuntimeError):
+            clone.load_state_dict(wrong_shape, strict=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

testing/test_sd_trainer_three_phase_runtime.py

@@ -0,0 +1,151 @@
+import ast
+import contextlib
+import importlib
+import inspect
+import os
+import types
+import unittest
+from pathlib import Path
+from types import MethodType, SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+
+def _load_runtime_methods():
+    source_path = Path(__file__).parents[1] / 'extensions_built_in' / 'sd_trainer' / 'SDTrainer.py'
+    tree = ast.parse(source_path.read_text(encoding='utf-8'))
+    class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == 'SDTrainer')
+    names = {
+        'three_phase_enabled', '_load_trigger_binding_modules', '_call_supported', '_first_callable',
+        '_phase_config', '_activator_component_flags', '_configure_phase_trainability',
+        'hook_add_extra_train_params', '_activator_mode', '_calculate_trigger_binding_loss',
+    }
+    selected = [node for node in class_node.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names]
+    module = ast.Module(body=[ast.ClassDef(name='SDTrainerRuntimeHarness', bases=[], keywords=[], body=selected, decorator_list=[])], type_ignores=[])
+    ast.fix_missing_locations(module)
+    namespace = {
+        'contextlib': contextlib,
+        'importlib': importlib,
+        'inspect': inspect,
+        'os': os,
+        'MethodType': MethodType,
+        'torch': torch,
+        'get_torch_dtype': lambda _dtype: torch.float32,
+        'shared_loss_target': lambda trainer, noise, batch, timesteps: trainer.sd.get_loss_target(
+            noise=noise, batch=batch, timesteps=timesteps
+        ).detach(),
+    }
+    exec(compile(module, str(source_path), 'exec'), namespace)
+    return namespace['SDTrainerRuntimeHarness']
+
+
+SDTrainer = _load_runtime_methods()
+
+
+class _FakeActivator(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = torch.nn.Linear(2, 2, bias=False)
+        self.te_adapter = torch.nn.Linear(2, 2, bias=False)
+        self.tap_adapters = torch.nn.Linear(2, 2, bias=False)
+        self.component_active = {}
+
+    def set_component_mode(self, component, active=None, trainable=None):
+        self.component_active[component] = active
+        getattr(self, component).requires_grad_(trainable)
+
+    def parameter_groups(self, learning_rates=None):
+        learning_rates = learning_rates or {}
+        groups = []
+        for name in ('embedding', 'te_adapter', 'tap_adapters'):
+            params = [parameter for parameter in getattr(self, name).parameters() if parameter.requires_grad]
+            if params:
+                group = {'params': params, 'name': name}
+                if name in learning_rates:
+                    group['lr'] = learning_rates[name]
+                groups.append(group)
+        return groups
+
+
+class ThreePhaseRuntimeTest(unittest.TestCase):
+    def _trainer(self, phase):
+        trainer = SDTrainer.__new__(SDTrainer)
+        trainer.runtime_phase = phase
+        trainer.text_activator = _FakeActivator()
+        trainer.network = torch.nn.Linear(2, 2, bias=False)
+        phase_config = SimpleNamespace(
+            train={
+                'embedding': phase != 'b',
+                'internal': False,
+                'tap': phase == 'a2',
+                'diffusion_lora': phase == 'b',
+            },
+            learning_rates={'embedding': 1e-3, 'tap_adapters': 2e-3},
+        )
+        trainer.three_phase_trigger_training = SimpleNamespace(
+            enabled=True,
+            phase_a1=phase_config,
+            phase_b=phase_config,
+            phase_a2=phase_config,
+        )
+        return trainer
+
+    def test_phase_whitelist_freezes_non_targets(self):
+        trainer = self._trainer('a2')
+        params = [{'params': list(trainer.network.parameters())}]
+        filtered = trainer.hook_add_extra_train_params(params)
+        selected = {id(parameter) for group in filtered for parameter in group['params']}
+        self.assertTrue(all(id(parameter) not in selected for parameter in trainer.network.parameters()))
+        self.assertTrue(all(parameter.requires_grad for parameter in trainer.text_activator.embedding.parameters()))
+        self.assertTrue(all(parameter.requires_grad for parameter in trainer.text_activator.tap_adapters.parameters()))
+        self.assertTrue(all(not parameter.requires_grad for parameter in trainer.text_activator.te_adapter.parameters()))
+
+    def test_b_phase_keeps_only_diffusion_lora_and_freezes_activator(self):
+        trainer = self._trainer('b')
+        params = [{'params': list(trainer.network.parameters()) + list(trainer.text_activator.parameters())}]
+        filtered = trainer.hook_add_extra_train_params(params)
+        selected = {id(parameter) for group in filtered for parameter in group['params']}
+        self.assertEqual(selected, {id(parameter) for parameter in trainer.network.parameters()})
+        self.assertTrue(all(not parameter.requires_grad for parameter in trainer.text_activator.parameters()))
+
+    def test_a_phase_loss_receives_shared_latent_noise_timestep_and_target(self):
+        trainer = self._trainer('a1')
+        trainer.device_torch = torch.device('cpu')
+        trainer.do_long_prompts = False
+        trainer.additional_logs = {}
+        trainer.sd = SimpleNamespace(
+            encode_prompt=lambda prompts, **kwargs: torch.ones(len(prompts), 1, 2),
+            get_loss_target=lambda noise, batch, timesteps: noise + 1,
+        )
+        trainer.predict_noise = lambda noisy_latents, **kwargs: noisy_latents * trainer.text_activator.embedding.weight.mean()
+        trainer._activator_mode = lambda mode: patch.object(trainer, '_mode', mode, create=True)
+        batch = SimpleNamespace(
+            file_items=[SimpleNamespace(caption_template='x [trigger]', raw_caption='unused')],
+            latents=torch.zeros(1, 2),
+        )
+        noisy = torch.randn(1, 2)
+        noise = torch.randn(1, 2)
+        timesteps = torch.tensor([10])
+        captured = {}
+
+        def fake_losses(**kwargs):
+            captured.update(kwargs)
+            return {'loss': (kwargs['active_prediction'] - kwargs['target']).pow(2).mean(), 'metrics': {'paired': 1}}
+
+        trainer._trigger_binding_modules = {
+            'losses': types.SimpleNamespace(calculate_trigger_binding_losses=fake_losses)
+        }
+        loss = trainer._calculate_trigger_binding_loss(
+            noisy, noise, timesteps, batch, {}, 1.0, torch.float32
+        )
+        self.assertTrue(loss.requires_grad)
+        self.assertIs(captured['noisy_latents'], noisy)
+        self.assertIs(captured['noise'], noise)
+        self.assertIs(captured['timesteps'], timesteps)
+        self.assertTrue(torch.equal(captured['target'], noise + 1))
+        self.assertEqual(trainer.additional_logs['phase/a1/paired'], 1.0)
+
+
+if __name__ == '__main__':
+    unittest.main()

testing/test_three_phase_trigger_training.py

@@ -0,0 +1,212 @@
+import json
+import os
+import tempfile
+import unittest
+from collections import OrderedDict
+
+import yaml
+
+from extensions_built_in.sd_trainer.ThreePhaseTriggerTrainer import ThreePhaseTriggerTrainer
+from toolkit.config_modules import (
+    ThreePhaseTriggerTrainingConfig,
+    validate_three_phase_trigger_training_config,
+)
+
+
+def _three_phase_block(enabled=True):
+    return {
+        'enabled': enabled,
+        'trigger': {
+            'placeholder': '[trigger]',
+            'literal': '<r1X1dOn9mA2>',
+            'span_detection': 'offsets',
+            'mask_all_occurrences': True,
+            'occurrence_mode': 'additive',
+        },
+        'text_activator': {
+            'embedding': {
+                'enabled': True,
+                'tokens': 1,
+                'init_mode': 'semantic',
+                'init_words': 'illustration',
+            },
+            'te_adapter': {'enabled': False},
+            'tap_adapters': {'enabled': False},
+        },
+        'reachability_probe': {'enabled': True},
+        'phase_a1': {
+            'enabled': True,
+            'steps': 10,
+            'optimizer': 'adamw',
+            'learning_rates': {'embedding': 0.001},
+            'train': {'train_embedding': True},
+            'save_steps': [5, 10],
+        },
+        'phase_b': {
+            'enabled': True,
+            'steps': 20,
+            'optimizer': 'adamw8bit',
+            'optimizer_params': {'weight_decay': 0.00005},
+            'learning_rates': {'diffusion_lora': 0.00008},
+            'train': {'train_unet': True, 'train_text_encoder': False},
+            'text_activator_source': {'phase': 'a1', 'step': 'final'},
+            'save_steps': [10, 20],
+        },
+        'phase_a2': {
+            'enabled': True,
+            'steps': 8,
+            'optimizer': 'adamw',
+            'learning_rates': {'embedding': 0.0001},
+            'train': {'train_embedding': True},
+            'text_activator_source': {'phase': 'a1', 'step': 'final'},
+            'diffusion_lora_source': {'phase': 'b', 'step': 'final'},
+            'save_steps': [4, 8],
+            'losses': {
+                'context_consistency': {
+                    'enabled': True,
+                    'weight': 0.05,
+                    'loss_type': 'cosine',
+                    'warmup_steps': 2,
+                    'min_delta_norm': 1.0e-6,
+                },
+                'activator_gain_floor': {
+                    'enabled': True,
+                    'weight': 1.0,
+                    'schedule': {
+                        'keyframes': [
+                            {'step': 0, 'value': 0.0},
+                            {'step': 8, 'value': 0.1},
+                        ],
+                    },
+                },
+            },
+        },
+    }
+
+
+class _Job:
+    def __init__(self, process_config):
+        self.name = 'parent_job'
+        self.meta = OrderedDict({'purpose': 'unit-test'})
+        self.raw_config = OrderedDict({
+            'job': 'extension',
+            'config': OrderedDict({
+                'name': self.name,
+                'process': [process_config],
+            }),
+            'meta': self.meta,
+        })
+
+
+class ThreePhaseTriggerTrainingConfigTest(unittest.TestCase):
+    def test_disabled_config_keeps_legacy_path_compatible(self):
+        validate_three_phase_trigger_training_config(
+            ThreePhaseTriggerTrainingConfig(enabled=False),
+            None,
+        )
+
+    def test_valid_three_phase_config(self):
+        config = ThreePhaseTriggerTrainingConfig(**_three_phase_block())
+        validate_three_phase_trigger_training_config(config, '<r1X1dOn9mA2>')
+        self.assertEqual(config.get_phase('b').steps, 20)
+        self.assertEqual(config.literal, '<r1X1dOn9mA2>')
+
+    def test_enabled_config_requires_native_placeholder(self):
+        raw = _three_phase_block()
+        raw['trigger']['placeholder'] = '<trigger>'
+        with self.assertRaisesRegex(ValueError, 'native'):
+            validate_three_phase_trigger_training_config(
+                ThreePhaseTriggerTrainingConfig(**raw),
+                '<r1X1dOn9mA2>',
+            )
+
+    def test_enabled_config_rejects_invalid_phase_dependency(self):
+        raw = _three_phase_block()
+        raw['phase_a2']['diffusion_lora_source']['phase'] = 'a1'
+        with self.assertRaisesRegex(ValueError, 'phase b'):
+            validate_three_phase_trigger_training_config(
+                ThreePhaseTriggerTrainingConfig(**raw),
+                '<r1X1dOn9mA2>',
+            )
+
+    def test_enabled_config_rejects_missing_trainable_component(self):
+        raw = _three_phase_block()
+        raw['phase_b']['train'] = {'train_unet': False}
+        with self.assertRaisesRegex(ValueError, 'trainable component'):
+            validate_three_phase_trigger_training_config(
+                ThreePhaseTriggerTrainingConfig(**raw),
+                '<r1X1dOn9mA2>',
+            )
+
+
+class ThreePhaseTriggerTrainerTest(unittest.TestCase):
+    def test_disabled_orchestrator_is_noop_compatible(self):
+        process_config = OrderedDict({
+            'type': 'three_phase_trigger_trainer',
+            'name': 'disabled_binding_run',
+            'three_phase_trigger_training': {'enabled': False},
+        })
+        process = ThreePhaseTriggerTrainer(0, _Job(process_config), process_config)
+        process.run()
+        self.assertFalse(hasattr(process, 'run_root'))
+
+    def _make_process(self, temp_dir):
+        process_config = OrderedDict({
+            'type': 'three_phase_trigger_trainer',
+            'name': 'binding_run',
+            'training_folder': temp_dir,
+            'trigger_word': '<r1X1dOn9mA2>',
+            'network': {'type': 'lora', 'linear': 32},
+            'train': {'dtype': 'bf16', 'steps': 999, 'optimizer': 'adamw'},
+            'model': {'name_or_path': 'test/model', 'arch': 'ideogram4'},
+            'datasets': [{'folder_path': 'dataset'}],
+            'save': {'save_every': 100},
+            'sample': {'samples': []},
+            'three_phase_trigger_training': _three_phase_block(),
+        })
+        return ThreePhaseTriggerTrainer(0, _Job(process_config), process_config)
+
+    def test_build_child_config_maps_phase_and_sources(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            process = self._make_process(temp_dir)
+            child_job = process.build_child_job_config('a2')
+            child = child_job['config']['process'][0]
+            self.assertEqual(child['type'], 'sd_trainer')
+            self.assertEqual(child['name'], 'phase_a2')
+            self.assertEqual(child['train']['steps'], 8)
+            self.assertEqual(child['train']['optimizer'], 'adamw')
+            self.assertAlmostEqual(child['train']['embedding_lr'], 0.0001)
+            self.assertTrue(child['train']['train_embedding'])
+            runtime = child['three_phase_trigger_training']['runtime']
+            self.assertEqual(runtime['active_phase'], 'a2')
+            self.assertTrue(runtime['orchestrated'])
+            sources = child['three_phase_trigger_training']['phase_runtime']['sources']
+            self.assertEqual(
+                sources['embedding'],
+                os.path.join(process.run_root, 'phase_a1', 'final', 'trigger_embedding.safetensors'),
+            )
+            self.assertEqual(
+                sources['diffusion_lora'],
+                os.path.join(process.run_root, 'phase_b', 'final', 'diffusion_lora.safetensors'),
+            )
+
+    def test_snapshot_and_completion_contract_are_written(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            process = self._make_process(temp_dir)
+            snapshot_path = process.write_phase_snapshot('a1')
+            contract_path = process.write_completion_contract('a1', 'completed', 0)
+            self.assertTrue(os.path.isfile(snapshot_path))
+            self.assertTrue(os.path.isfile(contract_path))
+            with open(snapshot_path, 'r', encoding='utf-8') as handle:
+                snapshot = yaml.safe_load(handle)
+            self.assertEqual(snapshot['config']['process'][0]['type'], 'sd_trainer')
+            with open(contract_path, 'r', encoding='utf-8') as handle:
+                contract = json.load(handle)
+            self.assertEqual(contract['status'], 'completed')
+            self.assertEqual(contract['return_code'], 0)
+            self.assertEqual(contract['phase'], 'a1')
+            self.assertTrue(contract['artifacts']['embedding'].endswith('trigger_embedding.safetensors'))
+
+
+if __name__ == '__main__':
+    unittest.main()

testing/test_trigger_binding.py

@@ -0,0 +1,211 @@
+import unittest
+
+import torch
+
+from toolkit.trigger_binding import (
+    ActivatorModeError,
+    TriggerAtomicityError,
+    TriggerConflictError,
+    TriggerPlaceholderError,
+    TriggerTokenizerError,
+    TriggerTruncationError,
+    activator_runtime_mode,
+    bind_trigger_batch,
+    bind_trigger_prompt,
+    get_activator_runtime_state,
+    resolve_trigger_literal,
+    validate_atomic_token_id,
+)
+
+
+class _FastTokenizer:
+    is_fast = True
+    pad_token_id = 0
+    eos_token_id = 2
+    unk_token_id = 1
+
+    def __init__(self, literal="<trigger>", atomic=True):
+        self.literal = literal
+        self.atomic = atomic
+        self.literal_id = 700
+
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
+        text = messages[0]["content"][0]["text"]
+        suffix = "<assistant>" if add_generation_prompt else ""
+        return f"<user>{text}</user>{suffix}"
+
+    def __call__(
+        self,
+        text,
+        add_special_tokens=False,
+        return_offsets_mapping=False,
+        truncation=False,
+        max_length=None,
+    ):
+        input_ids = []
+        offsets = []
+        index = 0
+        while index < len(text):
+            if self.atomic and text.startswith(self.literal, index):
+                input_ids.append(self.literal_id)
+                offsets.append((index, index + len(self.literal)))
+                index += len(self.literal)
+            else:
+                input_ids.append(100 + (ord(text[index]) % 500))
+                offsets.append((index, index + 1))
+                index += 1
+        if truncation and max_length is not None:
+            input_ids = input_ids[:max_length]
+            offsets = offsets[:max_length]
+        result = {"input_ids": input_ids, "attention_mask": [1] * len(input_ids)}
+        if return_offsets_mapping:
+            result["offset_mapping"] = offsets
+        return result
+
+
+class _SlowTokenizer(_FastTokenizer):
+    is_fast = False
+
+
+class _DuplicatingTemplateTokenizer(_FastTokenizer):
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
+        text = messages[0]["content"][0]["text"]
+        return f"{text} {self.literal}"
+
+
+class _Runtime:
+    def __init__(self, runtime_mode="full"):
+        self.runtime_mode = runtime_mode
+
+
+class TriggerBindingTest(unittest.TestCase):
+    def setUp(self):
+        self.literal = "<trigger>"
+        self.tokenizer = _FastTokenizer(self.literal)
+
+    def test_resolver_replaces_all_occurrences_and_records_spans(self):
+        resolved = resolve_trigger_literal(
+            "alpha [trigger] beta [trigger] omega",
+            self.literal,
+        )
+        self.assertEqual(resolved.text, "alpha <trigger> beta <trigger> omega")
+        self.assertEqual(resolved.occurrence_count, 2)
+        self.assertEqual(
+            tuple(resolved.text[start:end] for start, end in resolved.spans),
+            (self.literal, self.literal),
+        )
+
+    def test_resolver_rejects_missing_placeholder_and_literal_conflict(self):
+        with self.assertRaises(TriggerPlaceholderError):
+            resolve_trigger_literal("plain caption", self.literal)
+        with self.assertRaises(TriggerConflictError):
+            resolve_trigger_literal("[trigger] plus <trigger>", self.literal)
+
+    def test_chat_template_offsets_create_all_occurrence_mask(self):
+        binding = bind_trigger_prompt(
+            self.tokenizer,
+            "alpha [trigger] beta [trigger]",
+            self.literal,
+            require_atomic=True,
+            expected_token_id=self.tokenizer.literal_id,
+        )
+        self.assertTrue(binding.rendered_text.startswith("<user>alpha"))
+        self.assertEqual(binding.occurrence_count, 2)
+        self.assertEqual(len(binding.token_indices), 2)
+        self.assertEqual(sum(binding.trigger_mask), 2)
+        self.assertTrue(all(binding.input_ids[index] == self.tokenizer.literal_id for index in binding.token_indices))
+        self.assertEqual(
+            tuple(binding.rendered_text[start:end] for start, end in binding.character_spans),
+            (self.literal, self.literal),
+        )
+
+    def test_first_occurrence_mode_leaves_other_occurrences_unmasked(self):
+        binding = bind_trigger_prompt(
+            self.tokenizer,
+            "[trigger] then [trigger]",
+            self.literal,
+            mask_all_occurrences=False,
+        )
+        self.assertEqual(binding.occurrence_count, 2)
+        self.assertEqual(len(binding.token_indices), 1)
+        self.assertEqual(sum(binding.trigger_mask), 1)
+
+    def test_truncation_is_detected_instead_of_silently_dropping_trigger(self):
+        with self.assertRaises(TriggerTruncationError):
+            bind_trigger_prompt(
+                self.tokenizer,
+                "a long prefix [trigger]",
+                self.literal,
+                max_length=5,
+            )
+
+    def test_chat_template_literal_duplication_is_rejected(self):
+        tokenizer = _DuplicatingTemplateTokenizer(self.literal)
+        with self.assertRaises(TriggerConflictError):
+            bind_trigger_prompt(tokenizer, "[trigger]", self.literal)
+
+    def test_fast_tokenizer_and_atomic_id_validation(self):
+        with self.assertRaises(TriggerTokenizerError):
+            bind_trigger_prompt(_SlowTokenizer(self.literal), "[trigger]", self.literal)
+        self.assertEqual(validate_atomic_token_id(self.tokenizer, self.literal), self.tokenizer.literal_id)
+        with self.assertRaises(TriggerAtomicityError):
+            validate_atomic_token_id(self.tokenizer, self.literal, expected_token_id=999)
+        with self.assertRaises(TriggerAtomicityError):
+            validate_atomic_token_id(_FastTokenizer(self.literal, atomic=False), self.literal)
+
+    def test_batch_padding_masks_and_metadata(self):
+        batch = bind_trigger_batch(
+            self.tokenizer,
+            ["[trigger]", "longer [trigger] and [trigger]"],
+            self.literal,
+            require_atomic=True,
+            metadata={"phase": "a1"},
+        )
+        self.assertEqual(batch.input_ids.shape, batch.attention_mask.shape)
+        self.assertEqual(batch.input_ids.shape, batch.trigger_mask.shape)
+        self.assertEqual(batch.input_ids.dtype, torch.long)
+        self.assertEqual(batch.trigger_mask.dtype, torch.bool)
+        self.assertEqual(batch.trigger_mask.sum(dim=1).tolist(), [1, 2])
+        self.assertEqual(batch.metadata["batch_size"], 2)
+        self.assertEqual(batch.metadata["occurrence_counts"], (1, 2))
+        self.assertEqual(batch.metadata["phase"], "a1")
+
+    def test_all_runtime_modes_expose_expected_flags(self):
+        full = get_activator_runtime_state("full")
+        self.assertTrue(full.embedding_enabled and full.internal_enabled and full.tap_enabled)
+        self.assertTrue(get_activator_runtime_state("embedding_only").embedding_enabled)
+        self.assertTrue(get_activator_runtime_state("tap_only").tap_enabled)
+        self.assertTrue(get_activator_runtime_state("internal_only").internal_enabled)
+        self.assertTrue(get_activator_runtime_state("activator_bypass").activator_bypassed)
+        self.assertTrue(get_activator_runtime_state("stock_literal").stock_literal)
+        with self.assertRaises(ActivatorModeError):
+            get_activator_runtime_state("unknown")
+
+    def test_runtime_context_is_nested_and_exception_safe_for_objects(self):
+        runtime = _Runtime("full")
+        with activator_runtime_mode(runtime, "embedding_only"):
+            self.assertEqual(runtime.runtime_mode, "embedding_only")
+            with self.assertRaisesRegex(RuntimeError, "boom"):
+                with activator_runtime_mode(runtime, "tap_only"):
+                    self.assertEqual(runtime.runtime_mode, "tap_only")
+                    raise RuntimeError("boom")
+            self.assertEqual(runtime.runtime_mode, "embedding_only")
+        self.assertEqual(runtime.runtime_mode, "full")
+
+    def test_runtime_context_restores_mapping_and_removes_new_attribute(self):
+        runtime = {"runtime_mode": "stock_literal"}
+        with activator_runtime_mode(runtime, "activator_bypass"):
+            self.assertEqual(runtime["runtime_mode"], "activator_bypass")
+        self.assertEqual(runtime["runtime_mode"], "stock_literal")
+
+        class Empty:
+            pass
+
+        empty = Empty()
+        with activator_runtime_mode(empty, "internal_only"):
+            self.assertEqual(empty.runtime_mode, "internal_only")
+        self.assertFalse(hasattr(empty, "runtime_mode"))
+
+
+if __name__ == "__main__":
+    unittest.main()

testing/test_trigger_binding_artifacts.py

@@ -0,0 +1,210 @@
+import json
+import os
+import random
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import torch
+from safetensors.torch import save_file
+
+from toolkit.trigger_binding_artifacts import (
+    ARTIFACT_TYPES,
+    ArtifactValidationError,
+    artifact_reference,
+    build_checkpoint_manifest,
+    canonical_json_dumps,
+    config_fingerprint,
+    decode_rng_state,
+    encode_rng_state,
+    load_artifact,
+    load_checkpoint_manifest,
+    phase_fingerprint,
+    save_artifact,
+    save_checkpoint_manifest,
+    sha256_bytes,
+    source_fingerprint,
+)
+
+
+class TriggerBindingArtifactsTest(unittest.TestCase):
+    def setUp(self):
+        self.phase = {"name": "a1", "step": 12}
+        self.source = {"model": "ideogram-4", "revision": "abc123"}
+        self.config = {"rank": 1, "dtype": "bf16"}
+        self.tensors = {
+            "adapter.down.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3),
+            "adapter.up.weight": torch.ones(3, 2, dtype=torch.bfloat16),
+        }
+
+    def _save(self, directory, artifact_type="te_adapter", tensors=None):
+        path = Path(directory) / f"{artifact_type}.safetensors"
+        manifest = save_artifact(
+            path,
+            artifact_type,
+            tensors or self.tensors,
+            phase=self.phase,
+            source=self.source,
+            config=self.config,
+            extra={"trigger": "<literal>"},
+        )
+        return path, manifest
+
+    def test_all_artifact_types_round_trip_with_metadata_and_hashes(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            for artifact_type in sorted(ARTIFACT_TYPES):
+                with self.subTest(artifact_type=artifact_type):
+                    path, manifest = self._save(temp_dir, artifact_type)
+                    loaded, loaded_manifest = load_artifact(
+                        path,
+                        expected_type=artifact_type,
+                        expected_keys=self.tensors.keys(),
+                        expected_shapes={key: tensor.shape for key, tensor in self.tensors.items()},
+                        expected_phase_fingerprint=phase_fingerprint(self.phase),
+                        expected_source_fingerprint=source_fingerprint(self.source),
+                        expected_config_fingerprint=config_fingerprint(self.config),
+                    )
+                    self.assertEqual(manifest, loaded_manifest)
+                    for key, tensor in self.tensors.items():
+                        self.assertTrue(torch.equal(loaded[key], tensor))
+                        self.assertEqual(
+                            len(loaded_manifest["tensors"][key]["sha256"]),
+                            64,
+                        )
+
+    def test_fingerprints_are_canonical_and_order_independent(self):
+        self.assertEqual(config_fingerprint({"a": 1, "b": 2}), config_fingerprint({"b": 2, "a": 1}))
+        self.assertNotEqual(config_fingerprint({"a": 1}), config_fingerprint({"a": 2}))
+
+    def test_load_fails_closed_on_expected_key_shape_and_fingerprint_mismatch(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path, _ = self._save(temp_dir)
+            cases = [
+                {"expected_keys": ["adapter.down.weight"]},
+                {"expected_shapes": {key: [99] for key in self.tensors}},
+                {"expected_phase_fingerprint": "0" * 64},
+                {"expected_source_fingerprint": "1" * 64},
+                {"expected_config_fingerprint": "2" * 64},
+                {"expected_file_sha256": "3" * 64},
+                {"expected_type": "embedding"},
+            ]
+            for kwargs in cases:
+                with self.subTest(kwargs=kwargs), self.assertRaises(ArtifactValidationError):
+                    load_artifact(path, **kwargs)
+
+    def test_load_rejects_unmanaged_safetensors_without_required_metadata(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path = Path(temp_dir) / "bare.safetensors"
+            save_file({"weight": torch.ones(1)}, str(path))
+            with self.assertRaises(ArtifactValidationError):
+                load_artifact(path)
+
+    def test_load_rejects_tampered_tensor_manifest_hash(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path, manifest = self._save(temp_dir)
+            tampered = json.loads(json.dumps(manifest))
+            tampered["tensors"]["adapter.down.weight"]["sha256"] = "0" * 64
+            manifest_json = canonical_json_dumps(tampered)
+            metadata = {
+                "trigger_binding.schema": tampered["schema"],
+                "trigger_binding.schema_version": str(tampered["schema_version"]),
+                "trigger_binding.artifact_type": tampered["artifact_type"],
+                "trigger_binding.manifest": manifest_json,
+                "trigger_binding.manifest_sha256": sha256_bytes(manifest_json.encode("utf-8")),
+            }
+            save_file(self.tensors, str(path), metadata=metadata)
+            with self.assertRaisesRegex(ArtifactValidationError, "SHA-256 mismatch"):
+                load_artifact(path)
+
+    def test_atomic_save_preserves_previous_destination_on_replace_failure(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path = Path(temp_dir) / "te_adapter.safetensors"
+            path.write_bytes(b"previous-good-file")
+            with mock.patch("toolkit.trigger_binding_artifacts.os.replace", side_effect=OSError("replace failed")):
+                with self.assertRaises(OSError):
+                    save_artifact(
+                        path,
+                        "te_adapter",
+                        self.tensors,
+                        phase=self.phase,
+                        source=self.source,
+                        config=self.config,
+                    )
+            self.assertEqual(path.read_bytes(), b"previous-good-file")
+            leftovers = [item for item in os.listdir(temp_dir) if item != path.name]
+            self.assertEqual(leftovers, [])
+
+    def test_rng_state_json_round_trip_supports_python_numpy_and_torch(self):
+        state = {
+            "python": random.Random(7).getstate(),
+            "torch": torch.get_rng_state(),
+            "cuda": [torch.arange(8, dtype=torch.uint8)],
+        }
+        try:
+            import numpy as np
+
+            state["numpy"] = np.random.RandomState(11).get_state()
+        except ImportError:
+            np = None
+
+        encoded = encode_rng_state(state)
+        json.dumps(encoded, allow_nan=False)
+        decoded = decode_rng_state(encoded)
+        self.assertEqual(decoded["python"], state["python"])
+        self.assertTrue(torch.equal(decoded["torch"], state["torch"]))
+        self.assertTrue(torch.equal(decoded["cuda"][0], state["cuda"][0]))
+        if np is not None:
+            self.assertEqual(decoded["numpy"][0], state["numpy"][0])
+            self.assertTrue(np.array_equal(decoded["numpy"][1], state["numpy"][1]))
+            self.assertEqual(decoded["numpy"][2:], state["numpy"][2:])
+
+    def test_checkpoint_manifest_round_trip_and_artifact_verification(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            artifact_path, _ = self._save(temp_dir, "embedding")
+            reference = artifact_reference(artifact_path, relative_to=temp_dir)
+            manifest = build_checkpoint_manifest(
+                phase=self.phase,
+                source=self.source,
+                config=self.config,
+                step=42,
+                artifacts={"embedding": reference},
+                rng_state={"python": random.Random(3).getstate(), "torch": torch.get_rng_state()},
+                extra={"optimizer": "adamw"},
+            )
+            manifest_path = Path(temp_dir) / "checkpoint_manifest.json"
+            save_checkpoint_manifest(manifest_path, manifest)
+            loaded = load_checkpoint_manifest(
+                manifest_path,
+                expected_phase_fingerprint=phase_fingerprint(self.phase),
+                expected_source_fingerprint=source_fingerprint(self.source),
+                expected_config_fingerprint=config_fingerprint(self.config),
+            )
+            self.assertEqual(loaded, manifest)
+            decoded_rng = decode_rng_state(loaded["rng_state"])
+            self.assertTrue(torch.equal(decoded_rng["torch"], torch.get_rng_state()))
+
+            with open(artifact_path, "ab") as handle:
+                handle.write(b"tamper")
+            with self.assertRaises(ArtifactValidationError):
+                load_checkpoint_manifest(manifest_path)
+
+    def test_checkpoint_manifest_rejects_unknown_schema_keys(self):
+        manifest = build_checkpoint_manifest(
+            phase=self.phase,
+            source=self.source,
+            config=self.config,
+            step=0,
+            artifacts={},
+            rng_state=None,
+        )
+        manifest["unexpected"] = True
+        with tempfile.TemporaryDirectory() as temp_dir:
+            path = Path(temp_dir) / "checkpoint.json"
+            path.write_text(json.dumps(manifest), encoding="utf-8")
+            with self.assertRaises(ArtifactValidationError):
+                load_checkpoint_manifest(path, verify_artifacts=False)
+
+
+if __name__ == "__main__":
+    unittest.main()

testing/test_trigger_binding_losses.py

@@ -0,0 +1,183 @@
+import unittest
+
+import torch
+
+from toolkit.trigger_binding_losses import (
+    activator_gain_floor_hinge,
+    aggregate_paired_source_losses,
+    compute_a1_loss,
+    compute_a2_loss,
+    delta_context_consistency,
+    normalized_activator_gain,
+    per_item_diffusion_mse,
+    scheduled_gain_floor,
+)
+
+
+class TriggerBindingLossesTest(unittest.TestCase):
+    def test_per_item_diffusion_mse(self):
+        prediction = torch.tensor([[[1.0, 3.0]], [[2.0, 4.0]]])
+        target = torch.tensor([[[0.0, 1.0]], [[2.0, 2.0]]])
+        result = per_item_diffusion_mse(prediction, target)
+        torch.testing.assert_close(result, torch.tensor([2.5, 2.0]))
+
+    def test_normalized_gain_detaches_denominator(self):
+        activator_loss = torch.tensor([2.0], requires_grad=True)
+        bypass_loss = torch.tensor([4.0], requires_grad=True)
+        gain = normalized_activator_gain(activator_loss, bypass_loss)
+        gain.sum().backward()
+        self.assertAlmostEqual(gain.item(), 0.5000001, places=5)
+        self.assertLess(activator_loss.grad.item(), 0.0)
+        self.assertIsNone(bypass_loss.grad)
+
+    def test_gain_floor_schedule_and_hinge(self):
+        keyframes = [
+            {'step': 0, 'value': 0.0},
+            {'step': 100, 'value': 0.2},
+        ]
+        self.assertEqual(scheduled_gain_floor(-1, keyframes, 'linear'), 0.0)
+        self.assertAlmostEqual(scheduled_gain_floor(50, keyframes, 'linear'), 0.1)
+        self.assertEqual(scheduled_gain_floor(200, keyframes, 'smoothstep'), 0.2)
+        gain = torch.tensor([0.05, 0.15], requires_grad=True)
+        loss = activator_gain_floor_hinge(gain, 0.1)
+        torch.testing.assert_close(loss, torch.tensor([0.05, 0.0]))
+        loss.sum().backward()
+        torch.testing.assert_close(gain.grad, torch.tensor([-1.0, 0.0]))
+
+    def test_context_consistency_cosine_magnitude_mask_gate_and_warmup(self):
+        shape = (1, 13, 3, 2)
+        bypass = torch.zeros(shape)
+        reference_bypass = torch.zeros(shape)
+        reference_on = torch.zeros(shape)
+        on = torch.zeros(shape, requires_grad=True)
+        reference_on[:, :, 1] = torch.tensor([1.0, 0.0])
+        reference_on[:, :, 2] = torch.tensor([2.0, 0.0])
+        with torch.no_grad():
+            on[:, :, 1] = torch.tensor([0.0, 1.0])
+            on[:, :, 2] = torch.tensor([1.0, 0.0])
+        trigger_mask = torch.tensor([[True, False, False]])
+        token_mask = torch.tensor([[True, True, True]])
+        result = delta_context_consistency(
+            on,
+            bypass,
+            reference_on,
+            reference_bypass,
+            token_mask=token_mask,
+            trigger_mask=trigger_mask,
+            mask_mode='nontrigger',
+            cosine_weight=1.0,
+            magnitude_weight=0.5,
+            min_delta_norm=0.5,
+            step=5,
+            warmup_steps=10,
+        )
+        self.assertAlmostEqual(result.warmup_scale, 0.5)
+        self.assertEqual(result.valid_taps_per_item.item(), 26.0)
+        self.assertAlmostEqual(result.cosine_per_item.item(), 0.5, places=6)
+        self.assertAlmostEqual(result.magnitude_per_item.item(), 0.25, places=6)
+        self.assertAlmostEqual(result.per_item.item(), 0.3125, places=6)
+        result.loss.backward()
+        self.assertIsNotNone(on.grad)
+        self.assertIsNone(bypass.grad)
+
+    def test_context_trigger_mask_only_selects_trigger_tokens(self):
+        reference_on = torch.zeros(1, 13, 2, 2)
+        reference_on[:, :, 0, 0] = 1.0
+        reference_on[:, :, 1, 0] = 1.0
+        on = reference_on.clone()
+        on[:, :, 1] = torch.tensor([0.0, 1.0])
+        result = delta_context_consistency(
+            on,
+            torch.zeros_like(on),
+            reference_on,
+            torch.zeros_like(reference_on),
+            trigger_mask=torch.tensor([[True, False]]),
+            mask_mode='trigger',
+        )
+        torch.testing.assert_close(result.per_item, torch.zeros(1))
+        self.assertEqual(result.valid_taps_per_item.item(), 13.0)
+
+    def test_context_rejects_non_13_tap_input(self):
+        taps = torch.zeros(1, 12, 2, 3)
+        with self.assertRaisesRegex(ValueError, 'expected 13'):
+            delta_context_consistency(taps, taps, taps, taps)
+
+    def test_paired_source_aggregation(self):
+        aggregate, weighted, weights = aggregate_paired_source_losses(
+            {
+                'json': torch.tensor([1.0, 3.0]),
+                'natural': torch.tensor([5.0, 1.0]),
+            },
+            {'json': 3.0, 'natural': 1.0},
+        )
+        torch.testing.assert_close(aggregate, torch.tensor([2.0, 2.5]))
+        torch.testing.assert_close(weighted['json'], torch.tensor([0.75, 2.25]))
+        self.assertEqual(weights, {'json': 0.75, 'natural': 0.25})
+
+    def test_a1_result_has_detailed_metrics(self):
+        prediction = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
+        target = torch.zeros_like(prediction)
+        result = compute_a1_loss(prediction, target)
+        self.assertEqual(result.per_item.shape, (2,))
+        self.assertIn('a1/diffusion_mse', result.metrics)
+        self.assertIn('a1/source/primary', result.metrics)
+        result.loss.backward()
+        self.assertIsNotNone(prediction.grad)
+
+    def test_a1_gain_floor_cannot_worsen_bypass(self):
+        activator_prediction = torch.tensor([[1.0, 1.0]], requires_grad=True)
+        bypass_prediction = torch.tensor([[2.0, 2.0]], requires_grad=True)
+        target = torch.zeros_like(activator_prediction)
+        result = compute_a1_loss(
+            activator_prediction,
+            target,
+            bypass_prediction=bypass_prediction,
+            gain_floor=0.8,
+            diffusion_weight=0.0,
+            gain_floor_weight=1.0,
+        )
+        result.loss.backward()
+        self.assertGreater(activator_prediction.grad.abs().sum().item(), 0.0)
+        self.assertIsNone(bypass_prediction.grad)
+        self.assertIn('a1/activator_gain', result.metrics)
+
+    def test_a2_cannot_get_gradient_by_worsening_bypass(self):
+        activator_prediction = torch.tensor([[1.0, 1.0]], requires_grad=True)
+        bypass_prediction = torch.tensor([[2.0, 2.0]], requires_grad=True)
+        target = torch.zeros_like(activator_prediction)
+        result = compute_a2_loss(
+            activator_prediction,
+            bypass_prediction,
+            target,
+            gain_floor=0.8,
+            diffusion_weight=0.0,
+            gain_floor_weight=1.0,
+        )
+        self.assertGreater(result.gain_floor_per_item.item(), 0.0)
+        result.loss.backward()
+        self.assertIsNotNone(activator_prediction.grad)
+        self.assertGreater(activator_prediction.grad.abs().sum().item(), 0.0)
+        self.assertIsNone(bypass_prediction.grad)
+
+    def test_a2_combined_objective_still_never_updates_bypass(self):
+        activator_prediction = torch.tensor([[1.0, -1.0]], requires_grad=True)
+        bypass_prediction = torch.tensor([[0.5, -0.5]], requires_grad=True)
+        target = torch.zeros_like(activator_prediction)
+        result = compute_a2_loss(
+            activator_prediction,
+            bypass_prediction,
+            target,
+            gain_floor=0.5,
+            diffusion_weight=1.0,
+            gain_floor_weight=1.0,
+        )
+        result.loss.backward()
+        self.assertGreater(activator_prediction.grad.abs().sum().item(), 0.0)
+        self.assertIsNone(bypass_prediction.grad)
+        self.assertIn('a2/activator_gain', result.metrics)
+        self.assertIn('a2/gain_floor_satisfied', result.metrics)
+        self.assertIn('a2/bypass_diffusion_mse', result.metrics)
+
+
+if __name__ == '__main__':
+    unittest.main()

toolkit/config_modules.py

@@ -476,6 +476,377 @@ class TriggerSelectiveTrainingConfig:
         self.logging = TriggerSelectiveLoggingConfig(**kwargs.get('logging', {}))
 
 
+class TriggerBindingEmbeddingConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', True)
+        self.tokens: int = int(kwargs.get('tokens', 1))
+        self.init_mode: str = kwargs.get('init_mode', 'semantic')
+        self.init_words: str = kwargs.get('init_words', 'illustration')
+        self.checkpoint_path: Optional[str] = kwargs.get('checkpoint_path', None)
+        self.dtype: str = kwargs.get('dtype', 'bf16')
+        self.filename: str = kwargs.get('filename', 'trigger_embedding.safetensors')
+
+
+class TriggerBindingTEAdapterConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.type: str = kwargs.get('type', 'lora')
+        self.rank: int = int(kwargs.get('rank', 1))
+        self.alpha: float = float(kwargs.get('alpha', self.rank))
+        self.dropout: float = float(kwargs.get('dropout', 0.0))
+        self.parent_modules: List[str] = kwargs.get('parent_modules', ['Qwen3VLTextMLP'])
+        self.child_modules: List[str] = kwargs.get('child_modules', ['down_proj'])
+        self.layers = kwargs.get('layers', 'all')
+        self.token_mask_mode: str = kwargs.get('token_mask_mode', 'trigger_span')
+        self.filename: str = kwargs.get('filename', 'te_adapter.safetensors')
+
+
+class TriggerBindingTapAdapterConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.type: str = kwargs.get('type', 'lora')
+        self.rank: int = int(kwargs.get('rank', 1))
+        self.alpha: float = float(kwargs.get('alpha', self.rank))
+        self.dropout: float = float(kwargs.get('dropout', 0.0))
+        self.tap_layers: List[int] = [int(layer) for layer in kwargs.get(
+            'tap_layers', [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35]
+        )]
+        self.share_weights: bool = kwargs.get('share_weights', False)
+        self.token_mask_mode: str = kwargs.get('token_mask_mode', 'trigger_span')
+        self.learnable_scale: bool = kwargs.get('learnable_scale', False)
+        self.scale_init: float = float(kwargs.get('scale_init', 1.0))
+        self.per_tap: Dict = kwargs.get('per_tap', {})
+        self.filename: str = kwargs.get('filename', 'tap_adapters.safetensors')
+
+
+class TriggerBindingTextActivatorConfig:
+    def __init__(self, **kwargs):
+        self.embedding = TriggerBindingEmbeddingConfig(**kwargs.get('embedding', {}))
+        self.te_adapter = TriggerBindingTEAdapterConfig(**kwargs.get('te_adapter', {}))
+        self.tap_adapters = TriggerBindingTapAdapterConfig(**kwargs.get('tap_adapters', {}))
+
+
+class TriggerBindingContextConsistencyConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.weight: float = float(kwargs.get('weight', 0.0))
+        self.loss_type: str = kwargs.get('loss_type', kwargs.get('type', 'cosine'))
+        self.magnitude_weight: float = float(kwargs.get('magnitude_weight', 0.0))
+        self.warmup_steps: int = int(kwargs.get('warmup_steps', 0))
+        self.min_delta_norm: float = float(kwargs.get('min_delta_norm', 1.0e-6))
+        self.tap_layers: Optional[List[int]] = kwargs.get('tap_layers', None)
+
+
+class TriggerBindingScheduledFloorConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.weight: float = float(kwargs.get('weight', 1.0))
+        self.schedule = TriggerSelectiveScheduleConfig(**kwargs.get('schedule', {}))
+
+
+class TriggerBindingPhaseSourceConfig:
+    def __init__(self, **kwargs):
+        self.phase: Optional[str] = kwargs.get('phase', None)
+        self.step = kwargs.get('step', 'final')
+        self.path: Optional[str] = kwargs.get('path', None)
+
+
+class TriggerBindingResumeConfig:
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.checkpoint: Optional[str] = kwargs.get('checkpoint', None)
+
+
+class TriggerBindingPhaseConfig:
+    def __init__(self, phase_name: str, **kwargs):
+        self.phase_name = phase_name
+        self.enabled: bool = kwargs.get('enabled', False)
+        self.steps: int = int(kwargs.get('steps', 0))
+        self.optimizer: str = kwargs.get('optimizer', 'adamw')
+        self.optimizer_params: Dict = kwargs.get('optimizer_params', {})
+        self.learning_rates: Dict[str, float] = {
+            key: float(value) for key, value in kwargs.get('learning_rates', {}).items()
+        }
+        self.train: Dict[str, bool] = kwargs.get('train', {})
+        self.caption_sources: Dict = kwargs.get('caption_sources', {})
+        self.losses: Dict = kwargs.get('losses', {})
+        self.save_steps: List[int] = [int(step) for step in kwargs.get('save_steps', [])]
+        self.resume = TriggerBindingResumeConfig(**kwargs.get('resume', {}))
+        self.text_activator_source = TriggerBindingPhaseSourceConfig(
+            **kwargs.get('text_activator_source', kwargs.get('text_activator_init', {}))
+        )
+        self.diffusion_lora_source = TriggerBindingPhaseSourceConfig(
+            **kwargs.get('diffusion_lora_source', {})
+        )
+        self.context_consistency = TriggerBindingContextConsistencyConfig(
+            **kwargs.get('losses', {}).get('context_consistency', {})
+        )
+        floor_key = 'activator_gain_floor' if phase_name == 'a2' else 'trigger_gain_floor'
+        floor_kwargs = kwargs.get('losses', {}).get(floor_key, {})
+        if not floor_kwargs and phase_name != 'a2':
+            floor_kwargs = kwargs.get('losses', {}).get('activator_gain_floor', {})
+        self.activator_gain_floor = TriggerBindingScheduledFloorConfig(**floor_kwargs)
+
+
+class TriggerBindingPhaseArtifactConfig:
+    def __init__(self, **kwargs):
+        self.metrics_file: str = kwargs.get('metrics_file', 'metrics.jsonl')
+        self.console_log: str = kwargs.get('console_log', 'console.log')
+        self.checkpoint_dir: str = kwargs.get('checkpoint_dir', 'checkpoints')
+        self.final_dir: str = kwargs.get('final_dir', 'final')
+        self.embedding_filename: str = kwargs.get('embedding_filename', 'trigger_embedding.safetensors')
+        self.te_adapter_filename: str = kwargs.get('te_adapter_filename', 'te_adapter.safetensors')
+        self.tap_adapter_filename: str = kwargs.get('tap_adapter_filename', 'tap_adapters.safetensors')
+        self.diffusion_lora_filename: str = kwargs.get('diffusion_lora_filename', 'diffusion_lora.safetensors')
+
+
+class TriggerBindingArtifactConfig:
+    def __init__(self, **kwargs):
+        self.output_root: Optional[str] = kwargs.get('output_root', None)
+        self.phase_a1 = TriggerBindingPhaseArtifactConfig(**kwargs.get('phase_a1', {}))
+        self.phase_b = TriggerBindingPhaseArtifactConfig(**kwargs.get('phase_b', {}))
+        self.phase_a2 = TriggerBindingPhaseArtifactConfig(**kwargs.get('phase_a2', {}))
+
+
+class ThreePhaseTriggerTrainingConfig:
+    PHASE_NAMES = ('a1', 'b', 'a2')
+
+    def __init__(self, **kwargs):
+        self.enabled: bool = kwargs.get('enabled', False)
+        trigger = kwargs.get('trigger', {})
+        self.placeholder: str = trigger.get('placeholder', '[trigger]')
+        self.literal: Optional[str] = trigger.get('literal', None)
+        self.span_detection: str = trigger.get('span_detection', 'offsets')
+        self.mask_all_occurrences: bool = trigger.get('mask_all_occurrences', True)
+        self.occurrence_mode: str = trigger.get('occurrence_mode', 'additive')
+        self.text_activator = TriggerBindingTextActivatorConfig(**kwargs.get('text_activator', {}))
+        self.reachability_probe: Dict = kwargs.get('reachability_probe', {'enabled': True})
+        self.phase_a1 = TriggerBindingPhaseConfig('a1', **kwargs.get('phase_a1', {}))
+        self.phase_b = TriggerBindingPhaseConfig('b', **kwargs.get('phase_b', {}))
+        self.phase_a2 = TriggerBindingPhaseConfig('a2', **kwargs.get('phase_a2', {}))
+        self.artifacts = TriggerBindingArtifactConfig(**kwargs.get('artifacts', {}))
+        runtime = kwargs.get('runtime', {})
+        self.active_phase: Optional[str] = runtime.get('active_phase', None)
+        self.orchestrated: bool = runtime.get('orchestrated', False)
+        self.run_root: Optional[str] = runtime.get('run_root', None)
+        self.config_snapshot: Optional[str] = runtime.get('config_snapshot', None)
+        self.completion_contract: Optional[str] = runtime.get('completion_contract', None)
+
+    def get_phase(self, phase_name: str) -> TriggerBindingPhaseConfig:
+        if phase_name not in self.PHASE_NAMES:
+            raise ValueError(f'Unknown three-phase trigger phase: {phase_name}')
+        return getattr(self, f'phase_{phase_name}')
+
+
+def _validate_non_empty_string(value, field_name: str):
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f'{field_name} must be a non-empty string')
+
+
+def _validate_trigger_binding_source(
+    source: TriggerBindingPhaseSourceConfig,
+    field_name: str,
+    enabled_phases: Dict[str, bool],
+):
+    if source.phase is not None and source.path is not None:
+        raise ValueError(f'{field_name} must specify either phase or path, not both')
+    if source.phase is not None:
+        if source.phase not in ThreePhaseTriggerTrainingConfig.PHASE_NAMES:
+            raise ValueError(f'{field_name}.phase must be one of a1, b or a2')
+        if not enabled_phases[source.phase]:
+            raise ValueError(f'{field_name} references disabled phase {source.phase}')
+    if source.path is not None:
+        _validate_non_empty_string(source.path, f'{field_name}.path')
+    if source.step != 'final':
+        try:
+            if int(source.step) < 0:
+                raise ValueError
+        except (TypeError, ValueError):
+            raise ValueError(f'{field_name}.step must be final or a non-negative integer')
+
+
+def validate_three_phase_trigger_training_config(
+    config: Optional[ThreePhaseTriggerTrainingConfig],
+    trigger_word: Optional[str] = None,
+):
+    if config is None or not config.enabled:
+        return
+
+    _validate_non_empty_string(config.placeholder, 'three_phase_trigger_training.trigger.placeholder')
+    if config.placeholder != '[trigger]':
+        raise ValueError('three_phase_trigger_training must use the native [trigger] placeholder')
+    literal = config.literal if config.literal is not None else trigger_word
+    _validate_non_empty_string(literal, 'three_phase_trigger_training.trigger.literal')
+    if trigger_word is not None and literal != trigger_word:
+        raise ValueError('three_phase_trigger_training.trigger.literal must match trigger_word')
+    config.literal = literal
+
+    if config.span_detection != 'offsets':
+        raise ValueError('three_phase_trigger_training.trigger.span_detection must be offsets')
+    if not config.mask_all_occurrences:
+        raise ValueError('three_phase_trigger_training requires mask_all_occurrences=true')
+    if config.occurrence_mode != 'additive':
+        raise ValueError('three_phase_trigger_training.trigger.occurrence_mode must be additive')
+
+    embedding = config.text_activator.embedding
+    if embedding.enabled:
+        if embedding.tokens <= 0:
+            raise ValueError('three_phase_trigger_training.text_activator.embedding.tokens must be positive')
+        if embedding.init_mode not in ('semantic', 'random', 'checkpoint'):
+            raise ValueError('embedding.init_mode must be semantic, random or checkpoint')
+        if embedding.init_mode == 'semantic':
+            _validate_non_empty_string(embedding.init_words, 'embedding.init_words')
+        if embedding.init_mode == 'checkpoint':
+            _validate_non_empty_string(embedding.checkpoint_path, 'embedding.checkpoint_path')
+        if embedding.dtype not in ('bf16', 'bfloat16', 'fp16', 'float16', 'fp32', 'float32'):
+            raise ValueError('embedding.dtype must be a supported floating-point dtype')
+        _validate_non_empty_string(embedding.filename, 'embedding.filename')
+
+    te_adapter = config.text_activator.te_adapter
+    if te_adapter.enabled:
+        if te_adapter.type != 'lora':
+            raise ValueError('text_activator.te_adapter.type must be lora')
+        if te_adapter.rank <= 0 or te_adapter.alpha <= 0:
+            raise ValueError('text_activator.te_adapter rank and alpha must be positive')
+        if not 0.0 <= te_adapter.dropout < 1.0:
+            raise ValueError('text_activator.te_adapter.dropout must be in [0, 1)')
+        if te_adapter.token_mask_mode != 'trigger_span':
+            raise ValueError('text_activator.te_adapter.token_mask_mode must be trigger_span')
+        if not te_adapter.parent_modules or not te_adapter.child_modules:
+            raise ValueError('text_activator.te_adapter module filters cannot be empty')
+
+    tap_adapters = config.text_activator.tap_adapters
+    if tap_adapters.enabled:
+        if tap_adapters.type != 'lora':
+            raise ValueError('text_activator.tap_adapters.type must be lora')
+        if tap_adapters.rank <= 0 or tap_adapters.alpha <= 0:
+            raise ValueError('text_activator.tap_adapters rank and alpha must be positive')
+        if not 0.0 <= tap_adapters.dropout < 1.0:
+            raise ValueError('text_activator.tap_adapters.dropout must be in [0, 1)')
+        if tap_adapters.token_mask_mode != 'trigger_span':
+            raise ValueError('text_activator.tap_adapters.token_mask_mode must be trigger_span')
+        expected_tap_layers = [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35]
+        if tap_adapters.tap_layers != expected_tap_layers:
+            raise ValueError(
+                'text_activator.tap_adapters.tap_layers must exactly match Ideogram4 taps '
+                '0,3,6,9,12,15,18,21,24,27,30,33,35'
+            )
+        unknown_per_tap = set(str(key) for key in tap_adapters.per_tap) - set(str(layer) for layer in tap_adapters.tap_layers)
+        if unknown_per_tap:
+            raise ValueError(f'text_activator.tap_adapters.per_tap contains unknown layers: {sorted(unknown_per_tap)}')
+
+    if not any((embedding.enabled, te_adapter.enabled, tap_adapters.enabled)):
+        raise ValueError('three_phase_trigger_training requires at least one text activator')
+
+    if not isinstance(config.reachability_probe, dict):
+        raise ValueError('three_phase_trigger_training.reachability_probe must be a mapping')
+    if config.reachability_probe.get('enabled', True) is not True:
+        raise ValueError('three_phase_trigger_training requires reachability_probe.enabled=true')
+
+    phases = {name: config.get_phase(name) for name in config.PHASE_NAMES}
+    enabled_phases = {name: phase.enabled for name, phase in phases.items()}
+    if not all(enabled_phases.values()):
+        raise ValueError('three_phase_trigger_training requires phase_a1, phase_b and phase_a2 to be enabled')
+
+    for phase_name, phase in phases.items():
+        prefix = f'three_phase_trigger_training.phase_{phase_name}'
+        if phase.steps <= 0:
+            raise ValueError(f'{prefix}.steps must be positive')
+        _validate_non_empty_string(phase.optimizer, f'{prefix}.optimizer')
+        if not isinstance(phase.optimizer_params, dict):
+            raise ValueError(f'{prefix}.optimizer_params must be a mapping')
+        if not phase.learning_rates or any(value < 0 for value in phase.learning_rates.values()):
+            raise ValueError(f'{prefix}.learning_rates must contain non-negative values')
+        if not phase.train or not any(bool(value) for value in phase.train.values()):
+            raise ValueError(f'{prefix}.train must enable at least one trainable component')
+        aliases = {
+            'embedding': ('embedding', 'train_embedding', 'trigger_embedding'),
+            'te_adapter': ('te_adapter', 'internal', 'train_internal', 'text_encoder_adapter'),
+            'tap_adapters': ('tap_adapters', 'tap', 'train_tap'),
+            'diffusion_lora': ('diffusion_lora', 'train_unet', 'unet'),
+        }
+        for component, keys in aliases.items():
+            train_enabled = any(bool(phase.train.get(key, False)) for key in keys)
+            lr = next((phase.learning_rates[key] for key in keys if key in phase.learning_rates), 0.0)
+            if train_enabled and lr <= 0:
+                raise ValueError(f'{prefix} enables {component} training but its learning rate is not positive')
+            if not train_enabled and lr > 0:
+                raise ValueError(f'{prefix} configures a positive {component} learning rate while training is disabled')
+        if any(step <= 0 or step > phase.steps for step in phase.save_steps):
+            raise ValueError(f'{prefix}.save_steps must be within 1..steps')
+        if len(set(phase.save_steps)) != len(phase.save_steps):
+            raise ValueError(f'{prefix}.save_steps must be unique')
+        if phase.resume.enabled:
+            _validate_non_empty_string(phase.resume.checkpoint, f'{prefix}.resume.checkpoint')
+
+        _validate_trigger_binding_source(
+            phase.text_activator_source,
+            f'{prefix}.text_activator_source',
+            enabled_phases,
+        )
+        _validate_trigger_binding_source(
+            phase.diffusion_lora_source,
+            f'{prefix}.diffusion_lora_source',
+            enabled_phases,
+        )
+
+        consistency = phase.context_consistency
+        if consistency.enabled:
+            if consistency.weight < 0:
+                raise ValueError(f'{prefix}.losses.context_consistency.weight must be non-negative')
+            if consistency.loss_type not in ('cosine', 'mse'):
+                raise ValueError(f'{prefix}.losses.context_consistency.loss_type must be cosine or mse')
+            if consistency.warmup_steps < 0 or consistency.warmup_steps > phase.steps:
+                raise ValueError(f'{prefix}.losses.context_consistency.warmup_steps must be within phase steps')
+            if consistency.min_delta_norm <= 0:
+                raise ValueError(f'{prefix}.losses.context_consistency.min_delta_norm must be positive')
+
+        floor = phase.activator_gain_floor
+        if floor.enabled:
+            if floor.weight < 0:
+                raise ValueError(f'{prefix}.losses gain floor weight must be non-negative')
+            keyframes = floor.schedule.keyframes
+            if not keyframes:
+                raise ValueError(f'{prefix}.losses gain floor schedule requires keyframes')
+            previous_step = -1
+            for keyframe in keyframes:
+                step = int(keyframe.get('step', -1))
+                if step <= previous_step or step > phase.steps:
+                    raise ValueError(f'{prefix}.losses gain floor keyframe steps must increase within phase steps')
+                if float(keyframe.get('value', -1.0)) < 0:
+                    raise ValueError(f'{prefix}.losses gain floor values must be non-negative')
+                previous_step = step
+
+    if phases['a1'].text_activator_source.phase is not None:
+        raise ValueError('phase_a1 cannot source a text activator from another phase')
+    if phases['a1'].diffusion_lora_source.phase is not None:
+        raise ValueError('phase_a1 cannot source a diffusion LoRA from another phase')
+    if phases['b'].text_activator_source.phase not in (None, 'a1'):
+        raise ValueError('phase_b.text_activator_source may only reference phase a1')
+    if not any(bool(phases['b'].train.get(key, False)) for key in ('diffusion_lora', 'train_unet', 'unet')):
+        raise ValueError('phase_b must train the diffusion LoRA')
+    if phases['a2'].text_activator_source.phase not in (None, 'a1'):
+        raise ValueError('phase_a2.text_activator_source may only reference phase a1')
+    if phases['a2'].diffusion_lora_source.phase not in (None, 'b'):
+        raise ValueError('phase_a2.diffusion_lora_source may only reference phase b')
+
+    artifacts = config.artifacts
+    if artifacts.output_root is not None:
+        _validate_non_empty_string(artifacts.output_root, 'three_phase_trigger_training.artifacts.output_root')
+    for phase_name in config.PHASE_NAMES:
+        phase_artifacts = getattr(artifacts, f'phase_{phase_name}')
+        for field_name, value in vars(phase_artifacts).items():
+            _validate_non_empty_string(value, f'artifacts.phase_{phase_name}.{field_name}')
+
+    if config.active_phase is not None:
+        if config.active_phase not in config.PHASE_NAMES:
+            raise ValueError('three_phase_trigger_training.runtime.active_phase must be a1, b or a2')
+        if not config.orchestrated:
+            raise ValueError('three_phase_trigger_training.runtime.active_phase requires orchestrated=true')
+        for field_name in ('run_root', 'config_snapshot', 'completion_contract'):
+            _validate_non_empty_string(getattr(config, field_name), f'three_phase_trigger_training.runtime.{field_name}')
+
+
 class TrainConfig:
     def __init__(self, **kwargs):
         self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
@@ -1581,6 +1952,7 @@ def validate_configs(
     trigger_selective_training: Optional[TriggerSelectiveTrainingConfig] = None,
     trigger_word: Optional[str] = None,
     network_config: Optional[NetworkConfig] = None,
+    three_phase_trigger_training: Optional[ThreePhaseTriggerTrainingConfig] = None,
 ):
     if model_config.is_flux:
         if save_config.save_format != 'diffusers':
@@ -1638,5 +2010,25 @@ def validate_configs(
         from toolkit.trigger_selective_training import validate_trigger_selective_config
         validate_trigger_selective_config(trigger_selective_training, trigger_word)
 
+    validate_three_phase_trigger_training_config(
+        three_phase_trigger_training,
+        trigger_word,
+    )
+    if three_phase_trigger_training is not None and three_phase_trigger_training.enabled:
+        if model_config.arch != 'ideogram4':
+            raise ValueError('three_phase_trigger_training is restricted to model.arch: ideogram4')
+        active_phase = three_phase_trigger_training.active_phase
+        text_side_phase = active_phase in (None, 'a1', 'a2')
+        if text_side_phase:
+            if train_config.cache_text_embeddings or any(dataset.cache_text_embeddings for dataset in dataset_configs):
+                raise ValueError('three_phase_trigger_training A1/A2 require dynamic text encoding; disable caches')
+            if train_config.unload_text_encoder:
+                raise ValueError('three_phase_trigger_training A1/A2 require the text encoder to remain loaded')
+        if active_phase == 'b':
+            if network_config is None or network_config.type != 'lora':
+                raise ValueError('three_phase_trigger_training Phase B requires network.type: lora')
+            if not train_config.train_unet:
+                raise ValueError('three_phase_trigger_training Phase B must train the diffusion LoRA')
+
     if train_config.batch_size > 1 and any(dataset_config.auto_frame_count for dataset_config in dataset_configs):
         raise ValueError("Cannot use batch size greater than 1 with auto_frame_count. Please set batch_size to 1 or auto_frame_count to False.")

toolkit/models/ideogram4_trigger_activator.py

@@ -0,0 +1,520 @@
+"""Isolated text activator components for Ideogram 4 trigger binding.
+
+This module deliberately has no dependency on the Ideogram pipeline, model, or
+trainer. Runtime trigger masks can be passed explicitly or obtained lazily from
+``toolkit.trigger_binding`` when that integration module is available.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import contextvars
+import importlib
+import inspect
+import math
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, Tuple
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+DEFAULT_TAP_LAYERS: Tuple[int, ...] = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
+_STATE_PREFIX = "ideogram4_text_activator."
+_RUNTIME_CONTEXT: contextvars.ContextVar[Any] = contextvars.ContextVar(
+    "ideogram4_trigger_activator_runtime", default=None
+)
+
+
+def _runtime_module() -> Any:
+    try:
+        return importlib.import_module("toolkit.trigger_binding")
+    except ImportError:
+        return None
+
+
+def _runtime_value(name: str, default: Any = None) -> Any:
+    runtime = _RUNTIME_CONTEXT.get()
+    if runtime is not None:
+        if isinstance(runtime, Mapping) and name in runtime:
+            return runtime[name]
+        value = getattr(runtime, name, None)
+        if value is not None:
+            return value() if callable(value) else value
+    module = _runtime_module()
+    if module is None:
+        return default
+    for accessor in (f"get_current_{name}", f"get_{name}"):
+        value = getattr(module, accessor, None)
+        if callable(value):
+            return value()
+    value = getattr(module, name, default)
+    return value() if callable(value) else value
+
+
+@contextlib.contextmanager
+def trigger_runtime(runtime: Any) -> Iterator[Any]:
+    """Provide a local runtime compatible with the future trigger binding API."""
+
+    token = _RUNTIME_CONTEXT.set(runtime)
+    try:
+        yield runtime
+    finally:
+        _RUNTIME_CONTEXT.reset(token)
+
+
+def _extract_tensor(output: Any) -> Tuple[Optional[Tensor], Any]:
+    if torch.is_tensor(output):
+        return output, lambda value: value
+    if isinstance(output, tuple) and output and torch.is_tensor(output[0]):
+        return output[0], lambda value: (value,) + output[1:]
+    if isinstance(output, list) and output and torch.is_tensor(output[0]):
+        return output[0], lambda value: [value] + output[1:]
+    if isinstance(output, Mapping):
+        for key in ("hidden_states", "last_hidden_state"):
+            if key in output and torch.is_tensor(output[key]):
+                def rebuild(value: Tensor, key: str = key) -> Any:
+                    copied = output.copy()
+                    copied[key] = value
+                    return copied
+                return output[key], rebuild
+    return None, lambda value: output
+
+
+def _normalize_mask(mask: Optional[Tensor], reference: Tensor) -> Optional[Tensor]:
+    if mask is None:
+        return None
+    mask = torch.as_tensor(mask, device=reference.device)
+    if mask.ndim == reference.ndim - 1:
+        mask = mask.unsqueeze(-1)
+    while mask.ndim < reference.ndim:
+        mask = mask.unsqueeze(-1)
+    if mask.ndim != reference.ndim:
+        raise ValueError(f"mask rank {mask.ndim} is incompatible with tensor rank {reference.ndim}")
+    try:
+        torch.broadcast_shapes(mask.shape, reference.shape)
+    except RuntimeError as exc:
+        raise ValueError(
+            f"mask shape {tuple(mask.shape)} cannot broadcast to {tuple(reference.shape)}"
+        ) from exc
+    return mask.to(dtype=reference.dtype)
+
+
+class AtomicLearnedEmbedding(nn.Module):
+    """A standalone embedding parameter that never mutates Qwen's token table.
+
+    ``initializer`` is copied into a frozen buffer. Learned mode uses only the
+    independent parameter; frozen mode is useful for probes, and bypass mode
+    leaves the original hidden states unchanged.
+    """
+
+    MODES = {"learned", "frozen", "bypass"}
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        tokens: int = 1,
+        initializer: Optional[Tensor] = None,
+        init_std: float = 0.02,
+    ) -> None:
+        super().__init__()
+        if embedding_dim <= 0 or tokens <= 0:
+            raise ValueError("embedding_dim and tokens must be positive")
+        shape = (int(tokens), int(embedding_dim))
+        if initializer is None:
+            frozen = torch.empty(shape).normal_(mean=0.0, std=float(init_std))
+        else:
+            frozen = torch.as_tensor(initializer).detach().clone()
+            if frozen.ndim == 1:
+                frozen = frozen.unsqueeze(0)
+            if frozen.shape == (1, embedding_dim) and tokens > 1:
+                frozen = frozen.expand(tokens, -1).clone()
+            if tuple(frozen.shape) != shape:
+                raise ValueError(f"initializer shape must be {shape}, got {tuple(frozen.shape)}")
+        self.weight = nn.Parameter(frozen.clone())
+        self.register_buffer("frozen_initializer", frozen, persistent=True)
+        self.mode = "learned"
+        self.active = True
+
+    def set_mode(self, mode: str) -> None:
+        if mode not in self.MODES:
+            raise ValueError(f"unsupported embedding mode: {mode}")
+        self.mode = mode
+
+    def vectors(self, mode: Optional[str] = None) -> Optional[Tensor]:
+        mode = mode or self.mode
+        if mode == "learned":
+            return self.weight
+        if mode == "frozen":
+            return self.frozen_initializer.detach()
+        if mode == "bypass":
+            return None
+        raise ValueError(f"unsupported embedding mode: {mode}")
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        token_mask: Optional[Tensor] = None,
+        token_indices: Optional[Tensor] = None,
+        mode: Optional[str] = None,
+    ) -> Tensor:
+        vectors = self.vectors(mode)
+        if not self.active or vectors is None:
+            return hidden_states
+        mask = _normalize_mask(token_mask if token_mask is not None else _runtime_value("token_mask"), hidden_states)
+        if mask is None:
+            return hidden_states
+        if vectors.shape[0] == 1:
+            replacement = vectors[0]
+        else:
+            if token_indices is None:
+                token_indices = _runtime_value("token_indices")
+            if token_indices is None:
+                raise ValueError("token_indices are required for a multi-token atomic embedding")
+            indices = torch.as_tensor(token_indices, device=hidden_states.device, dtype=torch.long)
+            indices = indices.clamp(min=0, max=vectors.shape[0] - 1)
+            replacement = F.embedding(indices, vectors)
+        replacement = replacement.to(device=hidden_states.device, dtype=hidden_states.dtype)
+        while replacement.ndim < hidden_states.ndim:
+            replacement = replacement.unsqueeze(0)
+        return torch.lerp(hidden_states, replacement, mask)
+
+
+class MaskedLowRankAdapter(nn.Module):
+    """Low-rank residual adapter whose update is restricted to a token mask."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        rank: int = 1,
+        alpha: Optional[float] = None,
+        dropout: float = 0.0,
+        learnable_scale: bool = False,
+        scale_init: float = 1.0,
+    ) -> None:
+        super().__init__()
+        if hidden_size <= 0 or rank <= 0:
+            raise ValueError("hidden_size and rank must be positive")
+        self.hidden_size = int(hidden_size)
+        self.rank = int(rank)
+        self.alpha = float(rank if alpha is None else alpha)
+        self.down = nn.Linear(hidden_size, rank, bias=False)
+        self.up = nn.Linear(rank, hidden_size, bias=False)
+        self.dropout = nn.Dropout(float(dropout))
+        if learnable_scale:
+            self.scale = nn.Parameter(torch.tensor(float(scale_init)))
+        else:
+            self.register_buffer("scale", torch.tensor(float(scale_init)), persistent=True)
+        self.active = True
+        nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, hidden_states: Tensor, token_mask: Optional[Tensor] = None) -> Tensor:
+        if not self.active:
+            return hidden_states
+        mask = _normalize_mask(token_mask if token_mask is not None else _runtime_value("token_mask"), hidden_states)
+        if mask is None or not bool(torch.any(mask != 0)):
+            return hidden_states
+        update = self.up(self.down(self.dropout(hidden_states)))
+        scale = self.scale.to(device=hidden_states.device, dtype=hidden_states.dtype)
+        return hidden_states + update * mask * scale * (self.alpha / self.rank)
+
+
+class MaskedTapAdapter(MaskedLowRankAdapter):
+    def __init__(self, tap_layer: int, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.tap_layer = int(tap_layer)
+
+
+@dataclass
+class ProbeDiagnostics:
+    active: bool
+    trainable_parameters: int
+    total_parameters: int
+    embedding_delta_norm: float
+    adapter_update_norms: Dict[str, float]
+    tap_layers: Tuple[int, ...]
+    hook_count: int
+
+
+class _AdapterWrapper(nn.Module):
+    def __init__(self, wrapped: nn.Module, adapter: MaskedLowRankAdapter) -> None:
+        super().__init__()
+        self.wrapped = wrapped
+        self.adapter = adapter
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        output = self.wrapped(*args, **kwargs)
+        hidden, rebuild = _extract_tensor(output)
+        return rebuild(self.adapter(hidden)) if hidden is not None else output
+
+
+class TextActivator(nn.Module):
+    """Unified atomic embedding, Qwen adapter, and 13-tap aggregator."""
+
+    COMPONENTS = ("embedding", "te_adapter", "tap_adapters")
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        hidden_size: Optional[int] = None,
+        embedding_tokens: int = 1,
+        initializer: Optional[Tensor] = None,
+        te_adapter: Optional[MaskedLowRankAdapter] = None,
+        tap_layers: Sequence[int] = DEFAULT_TAP_LAYERS,
+        tap_rank: int = 1,
+        tap_alpha: Optional[float] = None,
+        tap_dropout: float = 0.0,
+        tap_learnable_scale: bool = False,
+        tap_scale_init: float = 1.0,
+        per_tap: Optional[Mapping[Any, Mapping[str, Any]]] = None,
+    ) -> None:
+        super().__init__()
+        hidden_size = int(hidden_size or embedding_dim)
+        layers = tuple(int(layer) for layer in tap_layers)
+        if len(layers) != 13 or len(set(layers)) != 13:
+            raise ValueError("Ideogram 4 requires exactly 13 unique tap layer keys")
+        self.embedding = AtomicLearnedEmbedding(embedding_dim, embedding_tokens, initializer)
+        self.te_adapter = te_adapter
+        per_tap = per_tap or {}
+        adapters: Dict[str, MaskedTapAdapter] = {}
+        for layer in layers:
+            overrides = dict(per_tap.get(layer, per_tap.get(str(layer), {})))
+            adapters[str(layer)] = MaskedTapAdapter(
+                tap_layer=layer,
+                hidden_size=hidden_size,
+                rank=int(overrides.pop("rank", tap_rank)),
+                alpha=overrides.pop("alpha", tap_alpha),
+                dropout=float(overrides.pop("dropout", tap_dropout)),
+                learnable_scale=bool(overrides.pop("learnable_scale", tap_learnable_scale)),
+                scale_init=float(overrides.pop("scale_init", tap_scale_init)),
+            )
+            if overrides:
+                raise ValueError(f"unknown per_tap options for layer {layer}: {sorted(overrides)}")
+        self.tap_adapters = nn.ModuleDict(adapters)
+        self.component_active = {
+            "embedding": True,
+            "te_adapter": te_adapter is not None,
+            "tap_adapters": True,
+        }
+        self._hooks: List[Any] = []
+        self._wrapped: List[Tuple[nn.Module, str, nn.Module]] = []
+        self._probe_inputs: Dict[str, Tensor] = {}
+        self._probe_updates: Dict[str, float] = {}
+
+    @property
+    def tap_layers(self) -> Tuple[int, ...]:
+        return tuple(int(key) for key in self.tap_adapters.keys())
+
+    def set_runtime_mode(self, mode: Optional[str]) -> None:
+        if mode is None:
+            return
+        try:
+            runtime = importlib.import_module("toolkit.trigger_binding")
+            state = runtime.get_activator_runtime_state(mode)
+        except (ImportError, AttributeError):
+            enabled = mode not in {"activator_bypass", "stock_literal"}
+            state = type("RuntimeState", (), {
+                "embedding_enabled": enabled,
+                "internal_enabled": enabled,
+                "tap_enabled": enabled,
+            })()
+        self.set_component_mode("embedding", active=state.embedding_enabled)
+        self.set_component_mode("te_adapter", active=state.internal_enabled)
+        self.set_component_mode("tap_adapters", active=state.tap_enabled)
+        self.embedding.set_mode("frozen" if not state.embedding_enabled else "learned")
+
+    def has_trainable_parameters(self) -> bool:
+        return any(parameter.requires_grad for parameter in self.parameters())
+
+    def set_component_mode(
+        self, component: str, *, active: Optional[bool] = None, trainable: Optional[bool] = None
+    ) -> None:
+        if component not in self.COMPONENTS:
+            raise KeyError(f"unknown component: {component}")
+        module = getattr(self, component)
+        if active is not None:
+            self.component_active[component] = bool(active)
+            if module is not None and hasattr(module, "active"):
+                module.active = bool(active)
+            if component == "tap_adapters":
+                for adapter in self.tap_adapters.values():
+                    adapter.active = bool(active)
+        if trainable is not None and module is not None:
+            module.requires_grad_(bool(trainable))
+
+    def apply_embedding(self, hidden_states: Tensor, **kwargs: Any) -> Tensor:
+        if not self.component_active["embedding"]:
+            return hidden_states
+        return self.embedding(hidden_states, **kwargs)
+
+    def apply_te_adapter(self, hidden_states: Tensor, token_mask: Optional[Tensor] = None) -> Tensor:
+        if not self.component_active["te_adapter"] or self.te_adapter is None:
+            return hidden_states
+        return self.te_adapter(hidden_states, token_mask)
+
+    def apply_tap(self, tap_layer: int, hidden_states: Tensor, token_mask: Optional[Tensor] = None) -> Tensor:
+        key = str(int(tap_layer))
+        if key not in self.tap_adapters:
+            raise KeyError(f"unconfigured tap layer: {tap_layer}")
+        if not self.component_active["tap_adapters"]:
+            return hidden_states
+        before = hidden_states
+        output = self.tap_adapters[key](hidden_states, token_mask)
+        if torch.is_grad_enabled():
+            self._probe_updates[key] = float((output.detach() - before.detach()).float().norm().item())
+        return output
+
+    def parameter_groups(self, learning_rates: Optional[Mapping[str, float]] = None) -> List[Dict[str, Any]]:
+        learning_rates = learning_rates or {}
+        groups: List[Dict[str, Any]] = []
+        for name in self.COMPONENTS:
+            module = getattr(self, name)
+            if module is None:
+                continue
+            params = [parameter for parameter in module.parameters() if parameter.requires_grad]
+            if params:
+                group: Dict[str, Any] = {"params": params, "name": f"text_activator.{name}"}
+                if name in learning_rates:
+                    group["lr"] = float(learning_rates[name])
+                groups.append(group)
+        return groups
+
+    def install_qwen_hooks(
+        self,
+        qwen: nn.Module,
+        te_module_names: Iterable[str] = (),
+        tap_module_names: Optional[Mapping[int, str]] = None,
+        use_wrappers: bool = False,
+    ) -> None:
+        self.remove_qwen_hooks()
+        named = dict(qwen.named_modules())
+        tap_module_names = tap_module_names or {}
+        for name in te_module_names:
+            if name not in named:
+                raise KeyError(f"Qwen module not found: {name}")
+            self._attach(named[name], name, self.apply_te_adapter, use_wrappers, qwen)
+        for layer, name in tap_module_names.items():
+            if str(int(layer)) not in self.tap_adapters:
+                raise KeyError(f"unconfigured tap layer: {layer}")
+            if name not in named:
+                raise KeyError(f"Qwen tap module not found: {name}")
+            callback = lambda hidden, layer=int(layer): self.apply_tap(layer, hidden)
+            self._attach(named[name], name, callback, use_wrappers, qwen)
+
+    def _attach(self, module: nn.Module, name: str, callback: Any, wrapper: bool, root: nn.Module) -> None:
+        if wrapper:
+            parent, attr = self._resolve_parent(root, name)
+            wrapped = _AdapterWrapper(module, _CallbackAdapter(callback))
+            setattr(parent, attr, wrapped)
+            self._wrapped.append((parent, attr, module))
+            return
+
+        def hook(_module: nn.Module, _inputs: Tuple[Any, ...], output: Any) -> Any:
+            hidden, rebuild = _extract_tensor(output)
+            return rebuild(callback(hidden)) if hidden is not None else output
+
+        self._hooks.append(module.register_forward_hook(hook))
+
+    @staticmethod
+    def _resolve_parent(root: nn.Module, name: str) -> Tuple[nn.Module, str]:
+        if not name:
+            raise ValueError("the root Qwen module cannot be replaced by a wrapper")
+        parts = name.split(".")
+        parent = root
+        for part in parts[:-1]:
+            parent = getattr(parent, part) if not part.isdigit() else parent[int(part)]
+        return parent, parts[-1]
+
+    def remove_qwen_hooks(self) -> None:
+        for handle in self._hooks:
+            handle.remove()
+        self._hooks.clear()
+        for parent, attr, original in reversed(self._wrapped):
+            setattr(parent, attr, original)
+        self._wrapped.clear()
+
+    def probe_diagnostics(self) -> ProbeDiagnostics:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        trainable = sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
+        embedding_delta = (self.embedding.weight.detach() - self.embedding.frozen_initializer).float().norm()
+        return ProbeDiagnostics(
+            active=any(self.component_active.values()),
+            trainable_parameters=trainable,
+            total_parameters=total,
+            embedding_delta_norm=float(embedding_delta.item()),
+            adapter_update_norms=dict(self._probe_updates),
+            tap_layers=self.tap_layers,
+            hook_count=len(self._hooks) + len(self._wrapped),
+        )
+
+    def activator_state_dict(self) -> "OrderedDict[str, Tensor]":
+        return OrderedDict((_STATE_PREFIX + key, value) for key, value in super().state_dict().items())
+
+    def load_activator_state_dict(
+        self, state_dict: Mapping[str, Tensor], strict: bool = True
+    ) -> torch.nn.modules.module._IncompatibleKeys:
+        foreign = sorted(key for key in state_dict if not key.startswith(_STATE_PREFIX))
+        if foreign and strict:
+            raise RuntimeError(f"foreign state dict keys: {foreign}")
+        stripped = OrderedDict(
+            (key[len(_STATE_PREFIX):], value)
+            for key, value in state_dict.items()
+            if key.startswith(_STATE_PREFIX)
+        )
+        expected = super().state_dict()
+        if strict:
+            missing = sorted(set(expected) - set(stripped))
+            unexpected = sorted(set(stripped) - set(expected))
+            shape_errors = sorted(
+                key for key in set(expected) & set(stripped)
+                if tuple(expected[key].shape) != tuple(stripped[key].shape)
+            )
+            if missing or unexpected or shape_errors:
+                raise RuntimeError(
+                    f"invalid activator state dict; missing={missing}, unexpected={unexpected}, "
+                    f"shape_mismatch={shape_errors}"
+                )
+        return super().load_state_dict(stripped, strict=strict)
+
+    def state_dict(self, *args: Any, **kwargs: Any) -> "OrderedDict[str, Tensor]":
+        destination = kwargs.get("destination")
+        if destination is not None or args:
+            return super().state_dict(*args, **kwargs)
+        return self.activator_state_dict()
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Tensor], strict: bool = True, assign: bool = False
+    ) -> torch.nn.modules.module._IncompatibleKeys:
+        if assign and "assign" not in inspect.signature(super().load_state_dict).parameters:
+            raise TypeError("this torch version does not support assign=True")
+        if assign:
+            stripped = OrderedDict(
+                (key[len(_STATE_PREFIX):], value)
+                for key, value in state_dict.items()
+                if key.startswith(_STATE_PREFIX)
+            )
+            return super().load_state_dict(stripped, strict=strict, assign=True)
+        return self.load_activator_state_dict(state_dict, strict=strict)
+
+
+class _CallbackAdapter(nn.Module):
+    def __init__(self, callback: Any) -> None:
+        super().__init__()
+        self.callback = callback
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        return self.callback(hidden_states)
+
+
+__all__ = [
+    "AtomicLearnedEmbedding",
+    "DEFAULT_TAP_LAYERS",
+    "MaskedLowRankAdapter",
+    "MaskedTapAdapter",
+    "ProbeDiagnostics",
+    "TextActivator",
+    "trigger_runtime",
+]

toolkit/trigger_binding.py

@@ -0,0 +1,469 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterator, List, Mapping, MutableMapping, Optional, Sequence, Tuple
+
+import torch
+
+
+SUPPORTED_ACTIVATOR_MODES = frozenset(
+    {
+        "full",
+        "embedding_only",
+        "tap_only",
+        "internal_only",
+        "activator_bypass",
+        "stock_literal",
+    }
+)
+
+
+class TriggerBindingError(ValueError):
+    pass
+
+
+class TriggerPlaceholderError(TriggerBindingError):
+    pass
+
+
+class TriggerConflictError(TriggerBindingError):
+    pass
+
+
+class TriggerTokenizerError(TriggerBindingError):
+    pass
+
+
+class TriggerTruncationError(TriggerBindingError):
+    pass
+
+
+class TriggerAtomicityError(TriggerBindingError):
+    pass
+
+
+class ActivatorModeError(TriggerBindingError):
+    pass
+
+
+@dataclass(frozen=True)
+class ResolvedTriggerText:
+    raw_text: str
+    text: str
+    placeholder: str
+    literal: str
+    spans: Tuple[Tuple[int, int], ...]
+
+    @property
+    def occurrence_count(self) -> int:
+        return len(self.spans)
+
+
+@dataclass(frozen=True)
+class TriggerBindingMetadata:
+    raw_text: str
+    resolved_text: str
+    rendered_text: str
+    literal: str
+    character_spans: Tuple[Tuple[int, int], ...]
+    token_spans: Tuple[Tuple[int, int], ...]
+    token_indices: Tuple[int, ...]
+    input_ids: Tuple[int, ...]
+    attention_mask: Tuple[int, ...]
+    trigger_mask: Tuple[int, ...]
+    atomic_token_id: Optional[int] = None
+
+    @property
+    def occurrence_count(self) -> int:
+        return len(self.character_spans)
+
+
+@dataclass
+class TriggerBindingBatch:
+    items: Tuple[TriggerBindingMetadata, ...]
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    trigger_mask: torch.Tensor
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to(self, *args, **kwargs) -> "TriggerBindingBatch":
+        self.input_ids = self.input_ids.to(*args, **kwargs)
+        self.attention_mask = self.attention_mask.to(*args, **kwargs)
+        self.trigger_mask = self.trigger_mask.to(*args, **kwargs)
+        return self
+
+
+@dataclass(frozen=True)
+class ActivatorRuntimeState:
+    mode: str
+    embedding_enabled: bool
+    internal_enabled: bool
+    tap_enabled: bool
+    activator_bypassed: bool
+    stock_literal: bool
+
+
+_MODE_STATES = {
+    "full": ActivatorRuntimeState("full", True, True, True, False, False),
+    "embedding_only": ActivatorRuntimeState("embedding_only", True, False, False, False, False),
+    "tap_only": ActivatorRuntimeState("tap_only", False, False, True, False, False),
+    "internal_only": ActivatorRuntimeState("internal_only", False, True, False, False, False),
+    "activator_bypass": ActivatorRuntimeState("activator_bypass", False, False, False, True, False),
+    "stock_literal": ActivatorRuntimeState("stock_literal", False, False, False, False, True),
+}
+
+
+def resolve_trigger_literal(
+    raw_text: str,
+    literal: str,
+    *,
+    placeholder: str = "[trigger]",
+    require_placeholder: bool = True,
+    reject_literal_conflicts: bool = True,
+) -> ResolvedTriggerText:
+    if not isinstance(raw_text, str):
+        raise TriggerPlaceholderError("raw trigger text must be a string")
+    if not isinstance(placeholder, str) or not placeholder:
+        raise TriggerPlaceholderError("trigger placeholder must be a non-empty string")
+    if not isinstance(literal, str) or not literal:
+        raise TriggerPlaceholderError("trigger literal must be a non-empty string")
+
+    placeholder_count = raw_text.count(placeholder)
+    if require_placeholder and placeholder_count == 0:
+        raise TriggerPlaceholderError(f"caption does not contain required placeholder {placeholder!r}")
+    if reject_literal_conflicts and literal in raw_text:
+        raise TriggerConflictError("raw caption already contains the literal trigger; binding would be ambiguous")
+
+    parts = raw_text.split(placeholder)
+    resolved_parts: List[str] = []
+    spans: List[Tuple[int, int]] = []
+    cursor = 0
+    for index, part in enumerate(parts):
+        resolved_parts.append(part)
+        cursor += len(part)
+        if index < len(parts) - 1:
+            spans.append((cursor, cursor + len(literal)))
+            resolved_parts.append(literal)
+            cursor += len(literal)
+
+    return ResolvedTriggerText(
+        raw_text=raw_text,
+        text="".join(resolved_parts),
+        placeholder=placeholder,
+        literal=literal,
+        spans=tuple(spans),
+    )
+
+
+def find_literal_spans(text: str, literal: str) -> Tuple[Tuple[int, int], ...]:
+    if not literal:
+        raise TriggerPlaceholderError("trigger literal must be non-empty")
+    spans: List[Tuple[int, int]] = []
+    start = 0
+    while True:
+        index = text.find(literal, start)
+        if index < 0:
+            return tuple(spans)
+        spans.append((index, index + len(literal)))
+        start = index + len(literal)
+
+
+def render_chat_prompt(tokenizer: Any, text: str, *, add_generation_prompt: bool = True) -> str:
+    messages = [{"role": "user", "content": [{"type": "text", "text": text}]}]
+    try:
+        rendered = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=False,
+        )
+    except (AttributeError, TypeError) as exc:
+        raise TriggerTokenizerError("tokenizer must support apply_chat_template(..., tokenize=False)") from exc
+    if not isinstance(rendered, str):
+        raise TriggerTokenizerError("chat template must return rendered text when tokenize=False")
+    return rendered
+
+
+def _as_flat_list(value: Any, name: str) -> List[Any]:
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().tolist()
+    if isinstance(value, tuple):
+        value = list(value)
+    if not isinstance(value, list):
+        raise TriggerTokenizerError(f"tokenizer output {name!r} must be a sequence")
+    if value and isinstance(value[0], (list, tuple)):
+        if len(value) != 1:
+            raise TriggerTokenizerError("single-prompt tokenization unexpectedly returned a batch")
+        value = list(value[0])
+    return value
+
+
+def _tokenize_with_offsets(tokenizer: Any, rendered_text: str, max_length: Optional[int]) -> Tuple[List[int], List[int], List[Tuple[int, int]]]:
+    if getattr(tokenizer, "is_fast", None) is False:
+        raise TriggerTokenizerError("trigger binding requires a fast tokenizer with offset mapping")
+    kwargs: Dict[str, Any] = {
+        "add_special_tokens": False,
+        "return_offsets_mapping": True,
+        "truncation": max_length is not None,
+    }
+    if max_length is not None:
+        if max_length <= 0:
+            raise TriggerTokenizerError("max_length must be positive")
+        kwargs["max_length"] = int(max_length)
+    try:
+        encoded = tokenizer(rendered_text, **kwargs)
+    except (TypeError, NotImplementedError) as exc:
+        raise TriggerTokenizerError("tokenizer does not provide fast offset mapping") from exc
+    if "input_ids" not in encoded or "offset_mapping" not in encoded:
+        raise TriggerTokenizerError("tokenizer output must contain input_ids and offset_mapping")
+
+    input_ids = [int(value) for value in _as_flat_list(encoded["input_ids"], "input_ids")]
+    raw_offsets = encoded["offset_mapping"]
+    if isinstance(raw_offsets, torch.Tensor):
+        raw_offsets = raw_offsets.detach().cpu().tolist()
+    if isinstance(raw_offsets, tuple):
+        raw_offsets = list(raw_offsets)
+    if (
+        isinstance(raw_offsets, list)
+        and len(raw_offsets) == 1
+        and raw_offsets
+        and isinstance(raw_offsets[0], list)
+        and (not raw_offsets[0] or isinstance(raw_offsets[0][0], (list, tuple)))
+    ):
+        raw_offsets = raw_offsets[0]
+    if not isinstance(raw_offsets, list):
+        raise TriggerTokenizerError("tokenizer output 'offset_mapping' must be a sequence")
+    offsets = [(int(pair[0]), int(pair[1])) for pair in raw_offsets]
+    if len(input_ids) != len(offsets):
+        raise TriggerTokenizerError("input_ids and offset_mapping lengths differ")
+    if "attention_mask" in encoded:
+        attention_mask = [int(value) for value in _as_flat_list(encoded["attention_mask"], "attention_mask")]
+        if len(attention_mask) != len(input_ids):
+            raise TriggerTokenizerError("input_ids and attention_mask lengths differ")
+    else:
+        attention_mask = [1] * len(input_ids)
+    return input_ids, attention_mask, offsets
+
+
+def validate_atomic_token_id(
+    tokenizer: Any,
+    literal: str,
+    *,
+    expected_token_id: Optional[int] = None,
+) -> int:
+    try:
+        encoded = tokenizer(literal, add_special_tokens=False, truncation=False)
+    except TypeError as exc:
+        raise TriggerAtomicityError("tokenizer cannot encode literal trigger") from exc
+    if "input_ids" not in encoded:
+        raise TriggerAtomicityError("tokenizer output does not contain input_ids")
+    token_ids = [int(value) for value in _as_flat_list(encoded["input_ids"], "input_ids")]
+    if len(token_ids) != 1:
+        raise TriggerAtomicityError(
+            f"literal trigger must map to exactly one token ID, got {len(token_ids)}: {token_ids}"
+        )
+    token_id = token_ids[0]
+    unk_token_id = getattr(tokenizer, "unk_token_id", None)
+    if unk_token_id is not None and token_id == int(unk_token_id):
+        raise TriggerAtomicityError("literal trigger maps to tokenizer unknown token ID")
+    if expected_token_id is not None and token_id != int(expected_token_id):
+        raise TriggerAtomicityError(
+            f"literal trigger token ID {token_id} does not match expected ID {int(expected_token_id)}"
+        )
+    return token_id
+
+
+def map_trigger_offsets(
+    offsets: Sequence[Tuple[int, int]],
+    character_spans: Sequence[Tuple[int, int]],
+    *,
+    mask_all_occurrences: bool = True,
+) -> Tuple[Tuple[int, ...], Tuple[Tuple[int, int], ...]]:
+    spans = tuple(character_spans if mask_all_occurrences else character_spans[:1])
+    token_indices: List[int] = []
+    token_spans: List[Tuple[int, int]] = []
+    for char_start, char_end in spans:
+        overlapping = [
+            index
+            for index, (token_start, token_end) in enumerate(offsets)
+            if token_end > token_start and token_start < char_end and token_end > char_start
+        ]
+        if not overlapping:
+            raise TriggerTruncationError(
+                f"trigger occurrence at character span ({char_start}, {char_end}) was truncated or not tokenized"
+            )
+        covered_start = min(offsets[index][0] for index in overlapping)
+        covered_end = max(offsets[index][1] for index in overlapping)
+        if covered_start > char_start or covered_end < char_end:
+            raise TriggerTruncationError(
+                f"trigger occurrence at character span ({char_start}, {char_end}) is only partially tokenized"
+            )
+        token_indices.extend(overlapping)
+        token_spans.append((overlapping[0], overlapping[-1] + 1))
+    return tuple(sorted(set(token_indices))), tuple(token_spans)
+
+
+def bind_trigger_prompt(
+    tokenizer: Any,
+    raw_text: str,
+    literal: str,
+    *,
+    placeholder: str = "[trigger]",
+    max_length: Optional[int] = None,
+    require_placeholder: bool = True,
+    reject_literal_conflicts: bool = True,
+    mask_all_occurrences: bool = True,
+    require_atomic: bool = False,
+    expected_token_id: Optional[int] = None,
+    add_generation_prompt: bool = True,
+) -> TriggerBindingMetadata:
+    resolved = resolve_trigger_literal(
+        raw_text,
+        literal,
+        placeholder=placeholder,
+        require_placeholder=require_placeholder,
+        reject_literal_conflicts=reject_literal_conflicts,
+    )
+    rendered = render_chat_prompt(tokenizer, resolved.text, add_generation_prompt=add_generation_prompt)
+    rendered_spans = find_literal_spans(rendered, literal)
+    if len(rendered_spans) != resolved.occurrence_count:
+        raise TriggerConflictError(
+            "chat template changed or duplicated literal-trigger occurrences; offset mapping is ambiguous"
+        )
+
+    input_ids, attention_mask, offsets = _tokenize_with_offsets(tokenizer, rendered, max_length)
+    token_indices, token_spans = map_trigger_offsets(
+        offsets,
+        rendered_spans,
+        mask_all_occurrences=mask_all_occurrences,
+    )
+    trigger_mask = [0] * len(input_ids)
+    for index in token_indices:
+        trigger_mask[index] = 1
+
+    atomic_token_id = None
+    if require_atomic or expected_token_id is not None:
+        atomic_token_id = validate_atomic_token_id(
+            tokenizer,
+            literal,
+            expected_token_id=expected_token_id,
+        )
+        for index in token_indices:
+            if input_ids[index] != atomic_token_id:
+                raise TriggerAtomicityError(
+                    "literal is atomic in isolation but has a different contextual token ID"
+                )
+        expected_occurrences = resolved.occurrence_count if mask_all_occurrences else min(1, resolved.occurrence_count)
+        if len(token_indices) != expected_occurrences:
+            raise TriggerAtomicityError("each masked trigger occurrence must map to exactly one token")
+
+    return TriggerBindingMetadata(
+        raw_text=raw_text,
+        resolved_text=resolved.text,
+        rendered_text=rendered,
+        literal=literal,
+        character_spans=rendered_spans,
+        token_spans=token_spans,
+        token_indices=token_indices,
+        input_ids=tuple(input_ids),
+        attention_mask=tuple(attention_mask),
+        trigger_mask=tuple(trigger_mask),
+        atomic_token_id=atomic_token_id,
+    )
+
+
+def bind_trigger_batch(
+    tokenizer: Any,
+    raw_texts: Sequence[str],
+    literal: str,
+    *,
+    pad_token_id: Optional[int] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    **binding_kwargs: Any,
+) -> TriggerBindingBatch:
+    if not raw_texts:
+        raise TriggerBindingError("trigger binding batch must not be empty")
+    items = tuple(bind_trigger_prompt(tokenizer, text, literal, **binding_kwargs) for text in raw_texts)
+    max_tokens = max(len(item.input_ids) for item in items)
+    if pad_token_id is None:
+        pad_token_id = getattr(tokenizer, "pad_token_id", None)
+    if pad_token_id is None:
+        pad_token_id = getattr(tokenizer, "eos_token_id", 0)
+    if pad_token_id is None:
+        pad_token_id = 0
+
+    ids: List[List[int]] = []
+    attention: List[List[int]] = []
+    masks: List[List[int]] = []
+    for item in items:
+        padding = max_tokens - len(item.input_ids)
+        ids.append(list(item.input_ids) + [int(pad_token_id)] * padding)
+        attention.append(list(item.attention_mask) + [0] * padding)
+        masks.append(list(item.trigger_mask) + [0] * padding)
+
+    batch_metadata = dict(metadata or {})
+    batch_metadata.update(
+        {
+            "batch_size": len(items),
+            "sequence_length": max_tokens,
+            "occurrence_counts": tuple(item.occurrence_count for item in items),
+            "token_indices": tuple(item.token_indices for item in items),
+            "character_spans": tuple(item.character_spans for item in items),
+            "literal": literal,
+        }
+    )
+    return TriggerBindingBatch(
+        items=items,
+        input_ids=torch.tensor(ids, dtype=torch.long),
+        attention_mask=torch.tensor(attention, dtype=torch.long),
+        trigger_mask=torch.tensor(masks, dtype=torch.bool),
+        metadata=batch_metadata,
+    )
+
+
+def get_activator_runtime_state(mode: str) -> ActivatorRuntimeState:
+    try:
+        return _MODE_STATES[mode]
+    except KeyError as exc:
+        supported = ", ".join(sorted(SUPPORTED_ACTIVATOR_MODES))
+        raise ActivatorModeError(f"unsupported activator runtime mode {mode!r}; expected one of: {supported}") from exc
+
+
+def _get_runtime_mode(target: Any) -> Any:
+    if isinstance(target, Mapping):
+        return target.get("runtime_mode")
+    return getattr(target, "runtime_mode", None)
+
+
+def _set_runtime_mode(target: Any, mode: Any) -> None:
+    if isinstance(target, MutableMapping):
+        target["runtime_mode"] = mode
+        return
+    setattr(target, "runtime_mode", mode)
+
+
+def _delete_runtime_mode(target: Any) -> None:
+    if isinstance(target, MutableMapping):
+        target.pop("runtime_mode", None)
+        return
+    try:
+        delattr(target, "runtime_mode")
+    except AttributeError:
+        pass
+
+
+@contextmanager
+def activator_runtime_mode(target: Any, mode: str) -> Iterator[ActivatorRuntimeState]:
+    state = get_activator_runtime_state(mode)
+    if target is None:
+        raise ActivatorModeError("activator runtime target must not be None")
+    existed = "runtime_mode" in target if isinstance(target, Mapping) else hasattr(target, "runtime_mode")
+    previous = _get_runtime_mode(target)
+    _set_runtime_mode(target, mode)
+    try:
+        yield state
+    finally:
+        if existed:
+            _set_runtime_mode(target, previous)
+        else:
+            _delete_runtime_mode(target)

toolkit/trigger_binding_artifacts.py

@@ -0,0 +1,594 @@
+"""Strict artifact persistence for three-phase trigger-binding training.
+
+This module is intentionally independent from trainer and model internals. It stores
+one artifact class per safetensors file and treats every persisted manifest as an
+untrusted input when loading.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import math
+import os
+import tempfile
+from collections.abc import Mapping, Sequence
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+ARTIFACT_SCHEMA = "ai-toolkit.trigger-binding-artifact"
+ARTIFACT_SCHEMA_VERSION = 1
+CHECKPOINT_SCHEMA = "ai-toolkit.trigger-binding-checkpoint"
+CHECKPOINT_SCHEMA_VERSION = 1
+ARTIFACT_TYPES = frozenset({"embedding", "te_adapter", "tap_adapter", "diffusion_lora"})
+ARTIFACT_MANIFEST_SCHEMAS = {
+    artifact_type: f"{ARTIFACT_SCHEMA}.{artifact_type}.v{ARTIFACT_SCHEMA_VERSION}"
+    for artifact_type in ARTIFACT_TYPES
+}
+
+_METADATA_SCHEMA = "trigger_binding.schema"
+_METADATA_VERSION = "trigger_binding.schema_version"
+_METADATA_TYPE = "trigger_binding.artifact_type"
+_METADATA_MANIFEST = "trigger_binding.manifest"
+_METADATA_MANIFEST_SHA256 = "trigger_binding.manifest_sha256"
+_REQUIRED_METADATA_KEYS = frozenset(
+    {
+        _METADATA_SCHEMA,
+        _METADATA_VERSION,
+        _METADATA_TYPE,
+        _METADATA_MANIFEST,
+        _METADATA_MANIFEST_SHA256,
+    }
+)
+
+PathLike = Union[str, os.PathLike]
+
+
+class ArtifactValidationError(ValueError):
+    """Raised when an artifact or checkpoint fails closed validation."""
+
+
+def canonical_json_dumps(value: Any) -> str:
+    """Serialize JSON deterministically for hashing and metadata."""
+    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_file(path: PathLike, chunk_size: int = 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while True:
+            chunk = handle.read(chunk_size)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def fingerprint(value: Any) -> str:
+    return sha256_bytes(canonical_json_dumps(to_json_compatible(value)).encode("utf-8"))
+
+
+def phase_fingerprint(phase: Any) -> str:
+    return fingerprint(phase)
+
+
+def source_fingerprint(source: Any) -> str:
+    return fingerprint(source)
+
+
+def config_fingerprint(config: Any) -> str:
+    return fingerprint(config)
+
+
+def _qualified_name(value: Any) -> str:
+    cls = value if isinstance(value, type) else type(value)
+    return f"{cls.__module__}.{cls.__qualname__}"
+
+
+def _object_config(value: Any) -> Dict[str, Any]:
+    fields = {
+        key: field_value
+        for key, field_value in vars(value).items()
+        if not key.startswith("_") and not callable(field_value)
+    }
+    return {"__type__": "object", "class": _qualified_name(value), "fields": to_json_compatible(fields)}
+
+
+def to_json_compatible(value: Any) -> Any:
+    """Encode checkpoint state, including Python/NumPy/Torch RNG state, as JSON values."""
+    if value is None or isinstance(value, (str, bool, int)):
+        return value
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            raise TypeError("Non-finite floats are not JSON compatible")
+        return value
+    if isinstance(value, bytes):
+        return {"__type__": "bytes", "base64": base64.b64encode(value).decode("ascii")}
+    if isinstance(value, bytearray):
+        return {"__type__": "bytearray", "base64": base64.b64encode(bytes(value)).decode("ascii")}
+    if isinstance(value, tuple):
+        return {"__type__": "tuple", "items": [to_json_compatible(item) for item in value]}
+    if isinstance(value, list):
+        return [to_json_compatible(item) for item in value]
+    if isinstance(value, Mapping):
+        if not all(isinstance(key, str) for key in value):
+            raise TypeError("JSON-compatible mappings require string keys")
+        return {key: to_json_compatible(item) for key, item in value.items()}
+    if isinstance(value, torch.dtype):
+        return {"__type__": "torch.dtype", "value": str(value).removeprefix("torch.")}
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach().cpu().contiguous()
+        raw = tensor.view(torch.uint8).numpy().tobytes()
+        return {
+            "__type__": "torch.Tensor",
+            "dtype": str(tensor.dtype).removeprefix("torch."),
+            "shape": list(tensor.shape),
+            "base64": base64.b64encode(raw).decode("ascii"),
+        }
+    module = type(value).__module__
+    if module.startswith("numpy"):
+        import numpy as np
+
+        if isinstance(value, np.ndarray):
+            array = np.ascontiguousarray(value)
+            return {
+                "__type__": "numpy.ndarray",
+                "dtype": array.dtype.str,
+                "shape": list(array.shape),
+                "base64": base64.b64encode(array.tobytes()).decode("ascii"),
+            }
+        if isinstance(value, np.generic):
+            return to_json_compatible(value.item())
+    if hasattr(value, "__dict__"):
+        return _object_config(value)
+    raise TypeError(f"Unsupported JSON-compatible value: {_qualified_name(value)}")
+
+
+def from_json_compatible(value: Any) -> Any:
+    """Decode values produced by :func:`to_json_compatible`."""
+    if isinstance(value, list):
+        return [from_json_compatible(item) for item in value]
+    if not isinstance(value, dict):
+        return value
+    tag = value.get("__type__")
+    if tag is None:
+        return {key: from_json_compatible(item) for key, item in value.items()}
+    if tag in {"bytes", "bytearray"}:
+        decoded = base64.b64decode(value["base64"], validate=True)
+        return decoded if tag == "bytes" else bytearray(decoded)
+    if tag == "tuple":
+        return tuple(from_json_compatible(item) for item in value["items"])
+    if tag == "torch.dtype":
+        dtype = getattr(torch, value["value"], None)
+        if not isinstance(dtype, torch.dtype):
+            raise ArtifactValidationError(f"Unknown torch dtype: {value['value']}")
+        return dtype
+    if tag == "torch.Tensor":
+        dtype = getattr(torch, value["dtype"], None)
+        if not isinstance(dtype, torch.dtype):
+            raise ArtifactValidationError(f"Unknown torch dtype: {value['dtype']}")
+        raw = base64.b64decode(value["base64"], validate=True)
+        tensor = torch.frombuffer(bytearray(raw), dtype=dtype).clone()
+        expected_numel = math.prod(value["shape"])
+        if tensor.numel() != expected_numel:
+            raise ArtifactValidationError("Encoded tensor byte count does not match its shape")
+        return tensor.reshape(value["shape"])
+    if tag == "numpy.ndarray":
+        import numpy as np
+
+        raw = base64.b64decode(value["base64"], validate=True)
+        array = np.frombuffer(raw, dtype=np.dtype(value["dtype"])).copy()
+        expected_size = math.prod(value["shape"])
+        if array.size != expected_size:
+            raise ArtifactValidationError("Encoded NumPy byte count does not match its shape")
+        return array.reshape(value["shape"])
+    if tag == "object":
+        return {"class": value["class"], "fields": from_json_compatible(value["fields"])}
+    raise ArtifactValidationError(f"Unknown JSON compatibility tag: {tag}")
+
+
+def encode_rng_state(rng_state: Any) -> Any:
+    return to_json_compatible(rng_state)
+
+
+def decode_rng_state(encoded: Any) -> Any:
+    return from_json_compatible(encoded)
+
+
+def _tensor_bytes(tensor: torch.Tensor) -> bytes:
+    return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
+
+
+def tensor_sha256(tensor: torch.Tensor) -> str:
+    return sha256_bytes(_tensor_bytes(tensor))
+
+
+def _validate_tensor_mapping(tensors: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    if not isinstance(tensors, Mapping) or not tensors:
+        raise ArtifactValidationError("Artifact tensors must be a non-empty mapping")
+    normalized: Dict[str, torch.Tensor] = {}
+    for key, tensor in tensors.items():
+        if not isinstance(key, str) or not key:
+            raise ArtifactValidationError("Artifact tensor keys must be non-empty strings")
+        if key in normalized:
+            raise ArtifactValidationError(f"Duplicate tensor key: {key}")
+        if not isinstance(tensor, torch.Tensor):
+            raise ArtifactValidationError(f"Artifact value for {key!r} is not a tensor")
+        if tensor.layout != torch.strided:
+            raise ArtifactValidationError(f"Artifact tensor {key!r} must use strided layout")
+        normalized[key] = tensor.detach().cpu().contiguous()
+    return normalized
+
+
+def build_artifact_manifest(
+    artifact_type: str,
+    tensors: Mapping[str, torch.Tensor],
+    *,
+    phase: Any,
+    source: Any,
+    config: Any,
+    extra: Optional[Mapping[str, Any]] = None,
+) -> Dict[str, Any]:
+    if artifact_type not in ARTIFACT_TYPES:
+        raise ArtifactValidationError(f"Unsupported artifact type: {artifact_type!r}")
+    normalized = _validate_tensor_mapping(tensors)
+    specs = {
+        key: {
+            "shape": list(tensor.shape),
+            "dtype": str(tensor.dtype).removeprefix("torch."),
+            "sha256": tensor_sha256(tensor),
+        }
+        for key, tensor in sorted(normalized.items())
+    }
+    manifest = {
+        "schema": ARTIFACT_SCHEMA,
+        "schema_version": ARTIFACT_SCHEMA_VERSION,
+        "artifact_type": artifact_type,
+        "artifact_schema": ARTIFACT_MANIFEST_SCHEMAS[artifact_type],
+        "phase_fingerprint": phase_fingerprint(phase),
+        "source_fingerprint": source_fingerprint(source),
+        "config_fingerprint": config_fingerprint(config),
+        "tensors": specs,
+    }
+    if extra is not None:
+        manifest["extra"] = to_json_compatible(dict(extra))
+    return manifest
+
+
+def _artifact_metadata(manifest: Mapping[str, Any]) -> Dict[str, str]:
+    manifest_json = canonical_json_dumps(manifest)
+    return {
+        _METADATA_SCHEMA: ARTIFACT_SCHEMA,
+        _METADATA_VERSION: str(ARTIFACT_SCHEMA_VERSION),
+        _METADATA_TYPE: str(manifest["artifact_type"]),
+        _METADATA_MANIFEST: manifest_json,
+        _METADATA_MANIFEST_SHA256: sha256_bytes(manifest_json.encode("utf-8")),
+    }
+
+
+def _atomic_replace(temp_path: str, destination: Path) -> None:
+    try:
+        with open(temp_path, "r+b") as handle:
+            handle.flush()
+            os.fsync(handle.fileno())
+    except OSError:
+        # Some Windows filesystems reject fsync on safetensors temp handles.
+        # os.replace remains atomic on the same volume.
+        pass
+    os.replace(temp_path, destination)
+    try:
+        directory_fd = os.open(str(destination.parent), os.O_RDONLY)
+    except (AttributeError, OSError):
+        return
+    try:
+        os.fsync(directory_fd)
+    finally:
+        os.close(directory_fd)
+
+
+def save_artifact(
+    path: PathLike,
+    artifact_type: str,
+    tensors: Mapping[str, torch.Tensor],
+    *,
+    phase: Any,
+    source: Any,
+    config: Any,
+    extra: Optional[Mapping[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Atomically save a typed safetensors artifact and return its manifest."""
+    destination = Path(path)
+    if destination.suffix.lower() != ".safetensors":
+        raise ArtifactValidationError("Artifact path must end with .safetensors")
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    normalized = _validate_tensor_mapping(tensors)
+    manifest = build_artifact_manifest(
+        artifact_type,
+        normalized,
+        phase=phase,
+        source=source,
+        config=config,
+        extra=extra,
+    )
+    file_descriptor, temp_path = tempfile.mkstemp(
+        prefix=f".{destination.name}.", suffix=".tmp", dir=str(destination.parent)
+    )
+    os.close(file_descriptor)
+    try:
+        save_file(normalized, temp_path, metadata=_artifact_metadata(manifest))
+        _atomic_replace(temp_path, destination)
+    except BaseException:
+        try:
+            os.unlink(temp_path)
+        except FileNotFoundError:
+            pass
+        raise
+    return manifest
+
+
+def _parse_artifact_manifest(metadata: Optional[Mapping[str, str]]) -> Dict[str, Any]:
+    if metadata is None or not _REQUIRED_METADATA_KEYS.issubset(metadata):
+        missing = sorted(_REQUIRED_METADATA_KEYS.difference(metadata or {}))
+        raise ArtifactValidationError(f"Missing required safetensors metadata: {missing}")
+    if metadata[_METADATA_SCHEMA] != ARTIFACT_SCHEMA:
+        raise ArtifactValidationError("Artifact metadata schema mismatch")
+    if metadata[_METADATA_VERSION] != str(ARTIFACT_SCHEMA_VERSION):
+        raise ArtifactValidationError("Unsupported artifact metadata schema version")
+    manifest_json = metadata[_METADATA_MANIFEST]
+    if sha256_bytes(manifest_json.encode("utf-8")) != metadata[_METADATA_MANIFEST_SHA256]:
+        raise ArtifactValidationError("Artifact manifest metadata hash mismatch")
+    try:
+        manifest = json.loads(manifest_json)
+    except (TypeError, json.JSONDecodeError) as exc:
+        raise ArtifactValidationError("Artifact manifest is not valid JSON") from exc
+    if canonical_json_dumps(manifest) != manifest_json:
+        raise ArtifactValidationError("Artifact manifest is not canonically encoded")
+    required = {
+        "schema",
+        "schema_version",
+        "artifact_type",
+        "artifact_schema",
+        "phase_fingerprint",
+        "source_fingerprint",
+        "config_fingerprint",
+        "tensors",
+    }
+    if set(manifest).difference(required | {"extra"}) or not required.issubset(manifest):
+        raise ArtifactValidationError("Artifact manifest has missing or unknown top-level keys")
+    if manifest["schema"] != ARTIFACT_SCHEMA or manifest["schema_version"] != ARTIFACT_SCHEMA_VERSION:
+        raise ArtifactValidationError("Artifact manifest schema mismatch")
+    if manifest["artifact_type"] not in ARTIFACT_TYPES:
+        raise ArtifactValidationError("Artifact manifest type is unsupported")
+    if manifest["artifact_schema"] != ARTIFACT_MANIFEST_SCHEMAS[manifest["artifact_type"]]:
+        raise ArtifactValidationError("Artifact-specific manifest schema mismatch")
+    if metadata[_METADATA_TYPE] != manifest["artifact_type"]:
+        raise ArtifactValidationError("Artifact type metadata disagrees with manifest")
+    for name in ("phase_fingerprint", "source_fingerprint", "config_fingerprint"):
+        value = manifest[name]
+        if not isinstance(value, str) or len(value) != 64:
+            raise ArtifactValidationError(f"Invalid {name}")
+    if not isinstance(manifest["tensors"], dict) or not manifest["tensors"]:
+        raise ArtifactValidationError("Artifact manifest tensor map is empty")
+    return manifest
+
+
+def load_artifact(
+    path: PathLike,
+    *,
+    expected_type: Optional[str] = None,
+    expected_keys: Optional[Sequence[str]] = None,
+    expected_shapes: Optional[Mapping[str, Sequence[int]]] = None,
+    expected_phase_fingerprint: Optional[str] = None,
+    expected_source_fingerprint: Optional[str] = None,
+    expected_config_fingerprint: Optional[str] = None,
+    expected_file_sha256: Optional[str] = None,
+    device: Union[str, torch.device] = "cpu",
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
+    """Load an artifact only after exact key, shape, dtype and hash validation."""
+    artifact_path = Path(path)
+    if not artifact_path.is_file():
+        raise ArtifactValidationError(f"Artifact file does not exist: {artifact_path}")
+    if expected_file_sha256 is not None and sha256_file(artifact_path) != expected_file_sha256:
+        raise ArtifactValidationError("Artifact file SHA-256 mismatch")
+    try:
+        with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
+            manifest = _parse_artifact_manifest(handle.metadata())
+            file_keys = set(handle.keys())
+            manifest_keys = set(manifest["tensors"])
+            if file_keys != manifest_keys:
+                raise ArtifactValidationError(
+                    f"Artifact tensor keys mismatch: file={sorted(file_keys)}, manifest={sorted(manifest_keys)}"
+                )
+            tensors = {key: handle.get_tensor(key) for key in sorted(file_keys)}
+    except ArtifactValidationError:
+        raise
+    except Exception as exc:
+        raise ArtifactValidationError(f"Unable to read safetensors artifact: {artifact_path}") from exc
+
+    if expected_type is not None and manifest["artifact_type"] != expected_type:
+        raise ArtifactValidationError("Artifact type does not match the expected type")
+    expected_fingerprints = {
+        "phase_fingerprint": expected_phase_fingerprint,
+        "source_fingerprint": expected_source_fingerprint,
+        "config_fingerprint": expected_config_fingerprint,
+    }
+    for name, expected in expected_fingerprints.items():
+        if expected is not None and manifest[name] != expected:
+            raise ArtifactValidationError(f"Artifact {name} mismatch")
+    if expected_keys is not None and set(expected_keys) != set(tensors):
+        raise ArtifactValidationError("Artifact keys do not match expected keys")
+    if expected_shapes is not None:
+        if set(expected_shapes) != set(tensors):
+            raise ArtifactValidationError("Expected shape keys do not match artifact keys")
+        for key, shape in expected_shapes.items():
+            if list(tensors[key].shape) != list(shape):
+                raise ArtifactValidationError(f"Artifact tensor {key!r} has unexpected shape")
+
+    for key, tensor in tensors.items():
+        spec = manifest["tensors"].get(key)
+        if not isinstance(spec, dict) or set(spec) != {"shape", "dtype", "sha256"}:
+            raise ArtifactValidationError(f"Invalid tensor manifest entry for {key!r}")
+        if list(tensor.shape) != spec["shape"]:
+            raise ArtifactValidationError(f"Artifact tensor {key!r} shape mismatch")
+        if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]:
+            raise ArtifactValidationError(f"Artifact tensor {key!r} dtype mismatch")
+        if tensor_sha256(tensor) != spec["sha256"]:
+            raise ArtifactValidationError(f"Artifact tensor {key!r} SHA-256 mismatch")
+    return {key: tensor.to(device) for key, tensor in tensors.items()}, manifest
+
+
+def artifact_reference(path: PathLike, *, relative_to: Optional[PathLike] = None) -> Dict[str, Any]:
+    artifact_path = Path(path)
+    with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
+        manifest = _parse_artifact_manifest(handle.metadata())
+    reference_path = os.path.relpath(artifact_path, relative_to) if relative_to is not None else str(artifact_path)
+    return {
+        "path": reference_path.replace(os.sep, "/"),
+        "artifact_type": manifest["artifact_type"],
+        "sha256": sha256_file(artifact_path),
+        "manifest_sha256": sha256_bytes(canonical_json_dumps(manifest).encode("utf-8")),
+    }
+
+
+def build_checkpoint_manifest(
+    *,
+    phase: Any,
+    source: Any,
+    config: Any,
+    step: int,
+    artifacts: Mapping[str, Mapping[str, Any]],
+    rng_state: Any = None,
+    extra: Optional[Mapping[str, Any]] = None,
+) -> Dict[str, Any]:
+    if not isinstance(step, int) or isinstance(step, bool) or step < 0:
+        raise ArtifactValidationError("Checkpoint step must be a non-negative integer")
+    if not isinstance(artifacts, Mapping):
+        raise ArtifactValidationError("Checkpoint artifacts must be a mapping")
+    manifest = {
+        "schema": CHECKPOINT_SCHEMA,
+        "schema_version": CHECKPOINT_SCHEMA_VERSION,
+        "step": step,
+        "phase_fingerprint": phase_fingerprint(phase),
+        "source_fingerprint": source_fingerprint(source),
+        "config_fingerprint": config_fingerprint(config),
+        "artifacts": to_json_compatible(dict(artifacts)),
+        "rng_state": encode_rng_state(rng_state),
+    }
+    if extra is not None:
+        manifest["extra"] = to_json_compatible(dict(extra))
+    validate_checkpoint_manifest(manifest)
+    return manifest
+
+
+def validate_checkpoint_manifest(manifest: Mapping[str, Any]) -> None:
+    required = {
+        "schema",
+        "schema_version",
+        "step",
+        "phase_fingerprint",
+        "source_fingerprint",
+        "config_fingerprint",
+        "artifacts",
+        "rng_state",
+    }
+    if not isinstance(manifest, Mapping):
+        raise ArtifactValidationError("Checkpoint manifest must be a mapping")
+    if set(manifest).difference(required | {"extra"}) or not required.issubset(manifest):
+        raise ArtifactValidationError("Checkpoint manifest has missing or unknown top-level keys")
+    if manifest["schema"] != CHECKPOINT_SCHEMA or manifest["schema_version"] != CHECKPOINT_SCHEMA_VERSION:
+        raise ArtifactValidationError("Checkpoint manifest schema mismatch")
+    if not isinstance(manifest["step"], int) or isinstance(manifest["step"], bool) or manifest["step"] < 0:
+        raise ArtifactValidationError("Checkpoint manifest step is invalid")
+    for name in ("phase_fingerprint", "source_fingerprint", "config_fingerprint"):
+        value = manifest[name]
+        if not isinstance(value, str) or len(value) != 64:
+            raise ArtifactValidationError(f"Checkpoint {name} is invalid")
+    if not isinstance(manifest["artifacts"], dict):
+        raise ArtifactValidationError("Checkpoint artifact references must be a mapping")
+    for name, reference in manifest["artifacts"].items():
+        if not isinstance(name, str) or not isinstance(reference, dict):
+            raise ArtifactValidationError("Checkpoint artifact reference is invalid")
+        if set(reference) != {"path", "artifact_type", "sha256", "manifest_sha256"}:
+            raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid keys")
+        if reference["artifact_type"] not in ARTIFACT_TYPES:
+            raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid type")
+        if not all(isinstance(reference[key], str) for key in reference):
+            raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} contains non-string values")
+        if len(reference["sha256"]) != 64 or len(reference["manifest_sha256"]) != 64:
+            raise ArtifactValidationError(f"Checkpoint artifact reference {name!r} has invalid hash")
+    try:
+        decode_rng_state(manifest["rng_state"])
+    except Exception as exc:
+        raise ArtifactValidationError("Checkpoint RNG state encoding is invalid") from exc
+
+
+def save_checkpoint_manifest(path: PathLike, manifest: Mapping[str, Any]) -> None:
+    validate_checkpoint_manifest(manifest)
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    payload = (canonical_json_dumps(manifest) + "\n").encode("utf-8")
+    file_descriptor, temp_path = tempfile.mkstemp(
+        prefix=f".{destination.name}.", suffix=".tmp", dir=str(destination.parent)
+    )
+    try:
+        with os.fdopen(file_descriptor, "wb") as handle:
+            handle.write(payload)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temp_path, destination)
+    except BaseException:
+        try:
+            os.unlink(temp_path)
+        except FileNotFoundError:
+            pass
+        raise
+
+
+def load_checkpoint_manifest(
+    path: PathLike,
+    *,
+    verify_artifacts: bool = True,
+    expected_phase_fingerprint: Optional[str] = None,
+    expected_source_fingerprint: Optional[str] = None,
+    expected_config_fingerprint: Optional[str] = None,
+) -> Dict[str, Any]:
+    manifest_path = Path(path)
+    try:
+        with open(manifest_path, "r", encoding="utf-8") as handle:
+            manifest = json.load(handle)
+    except Exception as exc:
+        raise ArtifactValidationError(f"Unable to read checkpoint manifest: {manifest_path}") from exc
+    validate_checkpoint_manifest(manifest)
+    expected_fingerprints = {
+        "phase_fingerprint": expected_phase_fingerprint,
+        "source_fingerprint": expected_source_fingerprint,
+        "config_fingerprint": expected_config_fingerprint,
+    }
+    for name, expected in expected_fingerprints.items():
+        if expected is not None and manifest[name] != expected:
+            raise ArtifactValidationError(f"Checkpoint {name} mismatch")
+    if verify_artifacts:
+        for name, reference in manifest["artifacts"].items():
+            artifact_path = Path(reference["path"])
+            if not artifact_path.is_absolute():
+                artifact_path = manifest_path.parent / artifact_path
+            _, artifact_manifest = load_artifact(
+                artifact_path,
+                expected_type=reference["artifact_type"],
+                expected_file_sha256=reference["sha256"],
+            )
+            actual_manifest_hash = sha256_bytes(canonical_json_dumps(artifact_manifest).encode("utf-8"))
+            if actual_manifest_hash != reference["manifest_sha256"]:
+                raise ArtifactValidationError(f"Checkpoint artifact {name!r} manifest hash mismatch")
+    return deepcopy(manifest)

toolkit/trigger_binding_losses.py

@@ -0,0 +1,394 @@
+from dataclasses import dataclass
+from typing import Dict, Mapping, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+
+
+TensorOrFloat = Union[torch.Tensor, float]
+TapTensors = Union[torch.Tensor, Sequence[torch.Tensor]]
+
+
+@dataclass(frozen=True)
+class ContextConsistencyResult:
+    per_item: torch.Tensor
+    cosine_per_item: torch.Tensor
+    magnitude_per_item: torch.Tensor
+    valid_taps_per_item: torch.Tensor
+    warmup_scale: float
+
+    @property
+    def loss(self) -> torch.Tensor:
+        return self.per_item.mean()
+
+
+@dataclass(frozen=True)
+class A1LossResult:
+    loss: torch.Tensor
+    per_item: torch.Tensor
+    diffusion_per_item: torch.Tensor
+    bypass_diffusion_per_item: torch.Tensor
+    activator_gain_per_item: torch.Tensor
+    gain_floor_per_item: torch.Tensor
+    context_per_item: torch.Tensor
+    source_per_item: Dict[str, torch.Tensor]
+    metrics: Dict[str, float]
+
+
+@dataclass(frozen=True)
+class A2LossResult:
+    loss: torch.Tensor
+    per_item: torch.Tensor
+    diffusion_per_item: torch.Tensor
+    bypass_diffusion_per_item: torch.Tensor
+    activator_gain_per_item: torch.Tensor
+    gain_floor_per_item: torch.Tensor
+    context_per_item: torch.Tensor
+    source_per_item: Dict[str, torch.Tensor]
+    metrics: Dict[str, float]
+
+
+def _validate_per_item(tensor: torch.Tensor, name: str) -> torch.Tensor:
+    if tensor.ndim != 1:
+        raise ValueError(f'{name} must be a one-dimensional per-item tensor')
+    return tensor
+
+
+def _scalar_like(value: TensorOrFloat, reference: torch.Tensor) -> torch.Tensor:
+    return torch.as_tensor(value, device=reference.device, dtype=reference.dtype)
+
+
+def per_item_diffusion_mse(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    if prediction.shape != target.shape:
+        raise ValueError('prediction and target shapes must match')
+    if prediction.ndim < 2:
+        raise ValueError('prediction and target must include batch and feature dimensions')
+    return F.mse_loss(prediction.float(), target.float(), reduction='none').flatten(1).mean(1)
+
+
+def normalized_activator_gain(
+    activator_loss: torch.Tensor,
+    bypass_loss: torch.Tensor,
+    epsilon: float = 1.0e-6,
+) -> torch.Tensor:
+    _validate_per_item(activator_loss, 'activator_loss')
+    _validate_per_item(bypass_loss, 'bypass_loss')
+    if activator_loss.shape != bypass_loss.shape:
+        raise ValueError('activator_loss and bypass_loss shapes must match')
+    if epsilon <= 0:
+        raise ValueError('epsilon must be positive')
+    return 1.0 - activator_loss / (bypass_loss.detach() + epsilon)
+
+
+def interpolate_schedule(
+    keyframes: Sequence[Mapping[str, float]],
+    step: int,
+    interpolation: str = 'smoothstep',
+) -> float:
+    if not keyframes:
+        raise ValueError('schedule must contain at least one keyframe')
+    if interpolation not in {'linear', 'smoothstep'}:
+        raise ValueError("interpolation must be 'linear' or 'smoothstep'")
+
+    normalized = []
+    previous_step = None
+    for keyframe in keyframes:
+        if 'step' not in keyframe or 'value' not in keyframe:
+            raise ValueError('each keyframe must contain step and value')
+        frame_step = int(keyframe['step'])
+        value = float(keyframe['value'])
+        if previous_step is not None and frame_step <= previous_step:
+            raise ValueError('keyframe steps must be strictly increasing')
+        previous_step = frame_step
+        normalized.append((frame_step, value))
+
+    if step <= normalized[0][0]:
+        return normalized[0][1]
+    if step >= normalized[-1][0]:
+        return normalized[-1][1]
+
+    for (left_step, left_value), (right_step, right_value) in zip(normalized, normalized[1:]):
+        if left_step <= step <= right_step:
+            fraction = (step - left_step) / (right_step - left_step)
+            if interpolation == 'smoothstep':
+                fraction = fraction * fraction * (3.0 - 2.0 * fraction)
+            return left_value + fraction * (right_value - left_value)
+    raise RuntimeError('could not interpolate schedule')
+
+
+def scheduled_gain_floor(
+    step: int,
+    keyframes: Sequence[Mapping[str, float]],
+    interpolation: str = 'smoothstep',
+) -> float:
+    return interpolate_schedule(keyframes, step, interpolation)
+
+
+def activator_gain_floor_hinge(
+    activator_gain: torch.Tensor,
+    floor: TensorOrFloat,
+) -> torch.Tensor:
+    _validate_per_item(activator_gain, 'activator_gain')
+    return torch.relu(_scalar_like(floor, activator_gain) - activator_gain)
+
+
+def _stack_taps(taps: TapTensors, name: str) -> torch.Tensor:
+    if isinstance(taps, torch.Tensor):
+        if taps.ndim < 4:
+            raise ValueError(f'{name} tensor must have shape [batch, taps, tokens, ...]')
+        return taps
+    taps = tuple(taps)
+    if not taps:
+        raise ValueError(f'{name} must contain at least one tap')
+    first_shape = taps[0].shape
+    if len(first_shape) < 3:
+        raise ValueError(f'{name} tap tensors must have shape [batch, tokens, ...]')
+    if any(tap.shape != first_shape for tap in taps):
+        raise ValueError(f'all {name} tap tensors must have matching shapes')
+    return torch.stack(taps, dim=1)
+
+
+def _prepare_context_mask(
+    token_mask: Optional[torch.Tensor],
+    trigger_mask: Optional[torch.Tensor],
+    mask_mode: str,
+    reference: torch.Tensor,
+) -> torch.Tensor:
+    batch, tap_count, token_count = reference.shape[:3]
+    if mask_mode not in {'all', 'trigger', 'nontrigger'}:
+        raise ValueError("mask_mode must be 'all', 'trigger', or 'nontrigger'")
+    selected = token_mask
+    if mask_mode in {'trigger', 'nontrigger'}:
+        if trigger_mask is None:
+            raise ValueError(f'trigger_mask is required for mask_mode={mask_mode}')
+        branch_mask = trigger_mask if mask_mode == 'trigger' else ~trigger_mask.bool()
+        selected = branch_mask if selected is None else selected.bool() & branch_mask.bool()
+    if selected is None:
+        selected = torch.ones((batch, token_count), device=reference.device, dtype=torch.bool)
+    selected = selected.to(device=reference.device, dtype=torch.bool)
+    if selected.ndim == 2:
+        if selected.shape != (batch, token_count):
+            raise ValueError('token masks must have shape [batch, tokens]')
+        selected = selected[:, None, :].expand(batch, tap_count, token_count)
+    elif selected.ndim == 3:
+        if selected.shape != (batch, tap_count, token_count):
+            raise ValueError('per-tap token masks must have shape [batch, taps, tokens]')
+    else:
+        raise ValueError('token masks must have two or three dimensions')
+    return selected
+
+
+def delta_context_consistency(
+    on_taps: TapTensors,
+    bypass_taps: TapTensors,
+    reference_on_taps: TapTensors,
+    reference_bypass_taps: TapTensors,
+    *,
+    token_mask: Optional[torch.Tensor] = None,
+    trigger_mask: Optional[torch.Tensor] = None,
+    mask_mode: str = 'nontrigger',
+    cosine_weight: float = 1.0,
+    magnitude_weight: float = 0.0,
+    min_delta_norm: float = 1.0e-6,
+    step: int = 0,
+    warmup_steps: int = 0,
+    expected_taps: int = 13,
+    epsilon: float = 1.0e-8,
+) -> ContextConsistencyResult:
+    on = _stack_taps(on_taps, 'on_taps').float()
+    bypass = _stack_taps(bypass_taps, 'bypass_taps').float()
+    reference_on = _stack_taps(reference_on_taps, 'reference_on_taps').float().detach()
+    reference_bypass = _stack_taps(reference_bypass_taps, 'reference_bypass_taps').float().detach()
+    if on.shape != bypass.shape or on.shape != reference_on.shape or on.shape != reference_bypass.shape:
+        raise ValueError('all context tap collections must have matching shapes')
+    if on.shape[1] != expected_taps:
+        raise ValueError(f'expected {expected_taps} context taps, got {on.shape[1]}')
+    if cosine_weight < 0 or magnitude_weight < 0:
+        raise ValueError('context consistency weights must be non-negative')
+    if min_delta_norm < 0 or epsilon <= 0:
+        raise ValueError('norm thresholds must be non-negative and epsilon must be positive')
+
+    on_delta = (on - bypass.detach()).flatten(start_dim=3)
+    reference_delta = (reference_on - reference_bypass).flatten(start_dim=3)
+    on_norm = torch.linalg.vector_norm(on_delta, dim=-1)
+    reference_norm = torch.linalg.vector_norm(reference_delta, dim=-1)
+    cosine = F.cosine_similarity(on_delta, reference_delta, dim=-1, eps=epsilon)
+    cosine_loss = 1.0 - cosine
+    magnitude_loss = torch.abs(on_norm - reference_norm) / (reference_norm.detach() + epsilon)
+
+    selected = _prepare_context_mask(token_mask, trigger_mask, mask_mode, on_delta)
+    norm_gate = (on_norm.detach() >= min_delta_norm) & (reference_norm.detach() >= min_delta_norm)
+    valid = selected & norm_gate
+    valid_float = valid.to(on_delta.dtype)
+    counts = valid_float.sum(dim=(1, 2))
+    denominator = counts.clamp_min(1.0)
+    cosine_per_item = (cosine_loss * valid_float).sum(dim=(1, 2)) / denominator
+    magnitude_per_item = (magnitude_loss * valid_float).sum(dim=(1, 2)) / denominator
+
+    if warmup_steps < 0:
+        raise ValueError('warmup_steps must be non-negative')
+    warmup_scale = 1.0 if warmup_steps == 0 else min(max(float(step) / warmup_steps, 0.0), 1.0)
+    per_item = warmup_scale * (
+        float(cosine_weight) * cosine_per_item + float(magnitude_weight) * magnitude_per_item
+    )
+    return ContextConsistencyResult(
+        per_item=per_item,
+        cosine_per_item=cosine_per_item,
+        magnitude_per_item=magnitude_per_item,
+        valid_taps_per_item=counts,
+        warmup_scale=warmup_scale,
+    )
+
+
+def aggregate_paired_source_losses(
+    losses_by_source: Mapping[str, torch.Tensor],
+    source_weights: Optional[Mapping[str, float]] = None,
+    *,
+    normalize_weights: bool = True,
+) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, float]]:
+    if not losses_by_source:
+        raise ValueError('losses_by_source must not be empty')
+    names = tuple(losses_by_source.keys())
+    first = _validate_per_item(losses_by_source[names[0]], f'losses_by_source[{names[0]}]')
+    for name in names[1:]:
+        current = _validate_per_item(losses_by_source[name], f'losses_by_source[{name}]')
+        if current.shape != first.shape:
+            raise ValueError('all paired source losses must have matching per-item shapes')
+
+    weights = {name: float(source_weights[name]) if source_weights is not None else 1.0 for name in names}
+    if source_weights is not None and set(source_weights) != set(names):
+        raise ValueError('source_weights keys must exactly match losses_by_source keys')
+    if any(weight < 0 for weight in weights.values()):
+        raise ValueError('source weights must be non-negative')
+    total_weight = sum(weights.values())
+    if total_weight <= 0:
+        raise ValueError('source weights must not all be zero')
+    if normalize_weights:
+        weights = {name: weight / total_weight for name, weight in weights.items()}
+
+    weighted = {name: losses_by_source[name] * weights[name] for name in names}
+    aggregate = torch.stack(tuple(weighted.values()), dim=0).sum(dim=0)
+    return aggregate, weighted, weights
+
+
+def _mean_metric(tensor: torch.Tensor) -> float:
+    return float(tensor.detach().float().mean().item())
+
+
+def _source_metrics(prefix: str, source_per_item: Mapping[str, torch.Tensor]) -> Dict[str, float]:
+    return {f'{prefix}/source/{name}': _mean_metric(value) for name, value in source_per_item.items()}
+
+
+def compute_a1_loss(
+    prediction: torch.Tensor,
+    target: torch.Tensor,
+    *,
+    bypass_prediction: Optional[torch.Tensor] = None,
+    gain_floor: TensorOrFloat = 0.0,
+    gain_epsilon: float = 1.0e-6,
+    diffusion_weight: float = 1.0,
+    gain_floor_weight: float = 0.0,
+    context: Optional[ContextConsistencyResult] = None,
+    context_weight: float = 1.0,
+    paired_source_losses: Optional[Mapping[str, torch.Tensor]] = None,
+    source_weights: Optional[Mapping[str, float]] = None,
+) -> A1LossResult:
+    diffusion = per_item_diffusion_mse(prediction, target)
+    if bypass_prediction is None:
+        bypass_diffusion = diffusion.detach()
+        gain = torch.zeros_like(diffusion)
+        floor_per_item = torch.zeros_like(diffusion)
+    else:
+        bypass_diffusion = per_item_diffusion_mse(bypass_prediction, target)
+        gain = normalized_activator_gain(diffusion, bypass_diffusion, gain_epsilon)
+        floor_per_item = activator_gain_floor_hinge(gain, gain_floor)
+    context_per_item = torch.zeros_like(diffusion) if context is None else context.per_item
+    if context_per_item.shape != diffusion.shape:
+        raise ValueError('context loss must match diffusion batch shape')
+    base = (
+        float(diffusion_weight) * diffusion
+        + float(gain_floor_weight) * floor_per_item
+        + float(context_weight) * context_per_item
+    )
+    source_inputs = {'primary': base} if paired_source_losses is None else dict(paired_source_losses)
+    per_item, source_per_item, effective_weights = aggregate_paired_source_losses(
+        source_inputs, source_weights
+    )
+    metrics = {
+        'a1/loss': _mean_metric(per_item),
+        'a1/diffusion_mse': _mean_metric(diffusion),
+        'a1/bypass_diffusion_mse': _mean_metric(bypass_diffusion),
+        'a1/activator_gain': _mean_metric(gain),
+        'a1/gain_floor_loss': _mean_metric(floor_per_item),
+        'a1/context': _mean_metric(context_per_item),
+        'a1/context_cosine': 0.0 if context is None else _mean_metric(context.cosine_per_item),
+        'a1/context_magnitude': 0.0 if context is None else _mean_metric(context.magnitude_per_item),
+        'a1/context_valid_taps': 0.0 if context is None else _mean_metric(context.valid_taps_per_item),
+        'a1/context_warmup_scale': 0.0 if context is None else context.warmup_scale,
+    }
+    metrics.update({f'a1/source_weight/{name}': weight for name, weight in effective_weights.items()})
+    metrics.update(_source_metrics('a1', source_per_item))
+    return A1LossResult(
+        per_item.mean(), per_item, diffusion, bypass_diffusion, gain,
+        floor_per_item, context_per_item, source_per_item, metrics
+    )
+
+
+def compute_a2_loss(
+    activator_prediction: torch.Tensor,
+    bypass_prediction: torch.Tensor,
+    target: torch.Tensor,
+    *,
+    gain_floor: TensorOrFloat,
+    gain_epsilon: float = 1.0e-6,
+    diffusion_weight: float = 1.0,
+    gain_floor_weight: float = 1.0,
+    context: Optional[ContextConsistencyResult] = None,
+    context_weight: float = 1.0,
+    paired_source_losses: Optional[Mapping[str, torch.Tensor]] = None,
+    source_weights: Optional[Mapping[str, float]] = None,
+) -> A2LossResult:
+    diffusion = per_item_diffusion_mse(activator_prediction, target)
+    bypass_diffusion = per_item_diffusion_mse(bypass_prediction, target)
+    gain = normalized_activator_gain(diffusion, bypass_diffusion, gain_epsilon)
+    floor_per_item = activator_gain_floor_hinge(gain, gain_floor)
+    context_per_item = torch.zeros_like(diffusion) if context is None else context.per_item
+    if context_per_item.shape != diffusion.shape:
+        raise ValueError('context loss must match diffusion batch shape')
+    base = (
+        float(diffusion_weight) * diffusion
+        + float(gain_floor_weight) * floor_per_item
+        + float(context_weight) * context_per_item
+    )
+    source_inputs = {'primary': base} if paired_source_losses is None else dict(paired_source_losses)
+    per_item, source_per_item, effective_weights = aggregate_paired_source_losses(
+        source_inputs, source_weights
+    )
+    floor_tensor = _scalar_like(gain_floor, gain)
+    metrics = {
+        'a2/loss': _mean_metric(per_item),
+        'a2/diffusion_mse': _mean_metric(diffusion),
+        'a2/bypass_diffusion_mse': _mean_metric(bypass_diffusion),
+        'a2/activator_gain': _mean_metric(gain),
+        'a2/gain_floor': float(floor_tensor.detach().float().mean().item()),
+        'a2/gain_floor_loss': _mean_metric(floor_per_item),
+        'a2/gain_floor_satisfied': _mean_metric((gain.detach() >= floor_tensor).float()),
+        'a2/context': _mean_metric(context_per_item),
+        'a2/context_cosine': 0.0 if context is None else _mean_metric(context.cosine_per_item),
+        'a2/context_magnitude': 0.0 if context is None else _mean_metric(context.magnitude_per_item),
+        'a2/context_valid_taps': 0.0 if context is None else _mean_metric(context.valid_taps_per_item),
+        'a2/context_warmup_scale': 0.0 if context is None else context.warmup_scale,
+    }
+    metrics.update({f'a2/source_weight/{name}': weight for name, weight in effective_weights.items()})
+    metrics.update(_source_metrics('a2', source_per_item))
+    return A2LossResult(
+        per_item.mean(),
+        per_item,
+        diffusion,
+        bypass_diffusion,
+        gain,
+        floor_per_item,
+        context_per_item,
+        source_per_item,
+        metrics,
+    )