diff --git a/config/examples/train_three_phase_trigger_binding_ideogram4.yaml b/config/examples/train_three_phase_trigger_binding_ideogram4.yaml new file mode 100644 index 00000000..9f58906b --- /dev/null +++ b/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: + 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: + 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 diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index e42a4a2c..aa957da1 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/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 diff --git a/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py b/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py index c080dc84..0594684f 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/pipeline.py +++ b/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 diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index aef5c7bb..991a0dac 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/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 diff --git a/extensions_built_in/sd_trainer/ThreePhaseTriggerTrainer.py b/extensions_built_in/sd_trainer/ThreePhaseTriggerTrainer.py new file mode 100644 index 00000000..447025a0 --- /dev/null +++ b/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') diff --git a/extensions_built_in/sd_trainer/__init__.py b/extensions_built_in/sd_trainer/__init__.py index 065ff818..26688412 100644 --- a/extensions_built_in/sd_trainer/__init__.py +++ b/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, ] diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 6ec2f070..9d9caf8b 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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) diff --git a/testing/test_ideogram4_trigger_activator.py b/testing/test_ideogram4_trigger_activator.py new file mode 100644 index 00000000..89f046f5 --- /dev/null +++ b/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() diff --git a/testing/test_sd_trainer_three_phase_runtime.py b/testing/test_sd_trainer_three_phase_runtime.py new file mode 100644 index 00000000..e4817175 --- /dev/null +++ b/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() diff --git a/testing/test_three_phase_trigger_training.py b/testing/test_three_phase_trigger_training.py new file mode 100644 index 00000000..8d05d6c5 --- /dev/null +++ b/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': '', + '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, '') + self.assertEqual(config.get_phase('b').steps, 20) + self.assertEqual(config.literal, '') + + def test_enabled_config_requires_native_placeholder(self): + raw = _three_phase_block() + raw['trigger']['placeholder'] = '' + with self.assertRaisesRegex(ValueError, 'native'): + validate_three_phase_trigger_training_config( + ThreePhaseTriggerTrainingConfig(**raw), + '', + ) + + 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), + '', + ) + + 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), + '', + ) + + +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': '', + '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() diff --git a/testing/test_trigger_binding.py b/testing/test_trigger_binding.py new file mode 100644 index 00000000..58408604 --- /dev/null +++ b/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="", 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 = "" if add_generation_prompt else "" + return f"{text}{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 = "" + 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 beta 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 ", 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("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() diff --git a/testing/test_trigger_binding_artifacts.py b/testing/test_trigger_binding_artifacts.py new file mode 100644 index 00000000..88403da0 --- /dev/null +++ b/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": ""}, + ) + 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() diff --git a/testing/test_trigger_binding_losses.py b/testing/test_trigger_binding_losses.py new file mode 100644 index 00000000..9290b1f9 --- /dev/null +++ b/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() diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 67d09131..3c3f09df 100644 --- a/toolkit/config_modules.py +++ b/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.") diff --git a/toolkit/models/ideogram4_trigger_activator.py b/toolkit/models/ideogram4_trigger_activator.py new file mode 100644 index 00000000..5c9188f8 --- /dev/null +++ b/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", +] diff --git a/toolkit/trigger_binding.py b/toolkit/trigger_binding.py new file mode 100644 index 00000000..b775dd75 --- /dev/null +++ b/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) diff --git a/toolkit/trigger_binding_artifacts.py b/toolkit/trigger_binding_artifacts.py new file mode 100644 index 00000000..bd116438 --- /dev/null +++ b/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) diff --git a/toolkit/trigger_binding_losses.py b/toolkit/trigger_binding_losses.py new file mode 100644 index 00000000..17a1d19a --- /dev/null +++ b/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, + )