diff --git a/TST_implementation.MD b/TST_implementation.MD new file mode 100644 index 00000000..54a65f47 --- /dev/null +++ b/TST_implementation.MD @@ -0,0 +1,672 @@ +# TRIGGER-SELECTIVE STYLE LORA TRAINING + +## Ideogram 4 + AI Toolkit + +### Implementation Specification - TST v1 + +Project trigger: + +*Dataset captions use AI Toolkit's native [trigger] placeholder* + +Version 1.0 | 13 August 2026 + +Purpose: controlled experiment for arbitrary-trigger binding without changing inference architecture + +# 1. Plain-language summary + +The existing V3 experiment already showed that ordinary transformer LoRA training can learn the desired visual style, but the arbitrary literal trigger does not reliably expose that learned effect. The objective of this method is therefore not to relearn the style with a new architecture. It is to train the same kind of LoRA while giving the optimizer three separate responsibilities. + +| Path | Simple job | Training behavior | +| --- | --- | --- | +| Path 1 | Learn the visual style | Normal dataset training with the literal trigger. This remains the dominant objective. | +| Path 2 | Do not leak | When [trigger] is replaced by an ordinary or competing style phrase, keep the LoRA close to the base model for that same condition. | +| Path 3 | Prefer the literal trigger | Measure how much the LoRA improves the dataset target under the literal trigger versus a matched non-trigger condition, and require a scheduled trigger advantage. | + +> **Core principle:** Early training mostly learns appearance. Later training gradually increases pressure on trigger selectivity. The final LoRA is still an ordinary Ideogram 4 transformer LoRA: no Qwen LoRA, no token-aware inference, and no custom ComfyUI runtime behavior. + +The method is intentionally designed as a falsifiable experiment. If the style learns again but the trigger gain does not separate from decoy gain, that is evidence that a static diffusion LoRA may not have enough conditional selectivity for this task. Only then should more invasive methods such as token-aware masking be justified. + +# 2. Evidence and design constraints + +- **Known positive result.** V3 used a rank-32 / alpha-16 Ideogram transformer LoRA and successfully carried the target visual style. + +- **Known binding failure.** The arbitrary literal trigger was weak, while semantic style wording could expose the V3 effect much more strongly. + +- **V3 already used DOP.** The V3 YAML had diff_output_preservation=true, multiplier=1, and class="painting". Therefore simply adding ordinary DOP again is not a new hypothesis. + +- **Keep content mode.** The learned visual behavior includes composition/content tendencies, so content_or_style remains "content". + +- **Minimize architecture changes.** The first test must stay diffusion-side only and preserve normal LoRA inference compatibility. + +- **Use early kill gates.** Do not spend 1000-2000 steps unless trigger-selective metrics start moving. + +> **Important correction to earlier plans:** The dataset does not literally contain . It contains [trigger] at deliberately chosen positions, and AI Toolkit replaces that placeholder with the configured trigger_word. TST must reuse that mechanism for every branch rather than manually searching for the literal trigger after conditioning. + +# 3. Final three-path training objective + +``` +ONE dataset sample / one noisy latent x_t / one timestep t + +Raw caption template contains: [trigger] + + AI Toolkit trigger resolver + | + +------------+-------------+ + | | + resolve as literal resolve as sampled + non-trigger style s + | | + c_trigger c_decoy + | | + +------+-------+ +----+-----+ + | | | | +LoRA ON LoRA OFF LoRA ON LoRA OFF + | | | | +Path 1 / 3 base trigger Path 2/3 base decoy +student reference student reference +``` + +All four predictions use the same current training image latent, the same sampled noise, the same timestep, and the same flow/diffusion target. Text conditioning is the manipulated variable. + +# 4. Trigger placeholder handling - reuse AI Toolkit, do not duplicate it + +AI Toolkit officially supports [trigger] in captions and replaces it with trigger_word. TST must operate from the raw caption template that still contains [trigger]. The new code should not perform a second ad-hoc literal-trigger search after the normal prompt has already been conditioned. + +## 4.1 Required behavior + +1. Load the raw caption template exactly as the existing dataset path does. It contains one or more [trigger] placeholders. + +1. Path 1 resolves every [trigger] using the existing configured trigger_word, exactly as ordinary AI Toolkit training already does. + +1. For a decoy branch, call the same trigger-placeholder resolver but provide the sampled non-trigger phrase as the effective replacement value for that branch. + +1. Path 3 reuses the already-resolved Path-1 trigger prompt and the already-resolved Path-2 decoy prompt. It must not create a third independent prompt transformation. + +1. Do not rewrite JSON structure, reorder fields, or reconstruct captions through JSON serialization. Only the [trigger] replacement value changes. + +## 4.2 Recommended small refactor + +If the current AI Toolkit trigger logic is embedded inside a broader conditioning function, extract the smallest reusable helper without changing standard behavior. Conceptually: + +``` +resolve_trigger_placeholder(raw_prompt, replacement) + +Path 1: replacement = configured trigger_word +Path 2: replacement = sampled decoy phrase +Path 3: reuse Path 1 and Path 2 outputs + +Standard AI Toolkit training when TST is disabled: unchanged +``` + +Because this dataset intentionally positions [trigger], TST should support a strict validation mode and default it on for this project: every training caption must contain at least one [trigger]. This avoids the normal automatic-prepend fallback creating different placement semantics between branches. + +``` +trigger_selective_training: + enabled: true + require_trigger_placeholder: true +``` + +# 5. Path 1 - ordinary style acquisition + +Path 1 is deliberately boring. It is the normal training forward that already made V3 a good style carrier. Do not create a parallel implementation. + +``` +raw caption + AI Toolkit [trigger] resolution + -> caption containing + -> Ideogram text conditioning + -> base + trainable transformer LoRA + -> normal flow/diffusion prediction + -> normal dataset target loss L1 +``` + +The existing network rank, alpha, optimizer, learning rate, sigmoid timestep distribution, content mode, EMA, Differential Guidance and other V3 diffusion-side controls continue to come from the ordinary YAML configuration. This keeps the style-learning mechanism causal and comparable to V3. + +> **Compatibility rule:** When trigger_selective_training.enabled=false, Path 1 must execute byte-for-byte-equivalent training logic to the current working trainer as far as the new patch is concerned. + +# 6. Non-trigger style bank and sampling + +Paths 2 and 3 use a rotating decoy condition. The user controls the number of categories, each category probability, and the phrases within each category from YAML. The first run should include neutral, hard/near semantic activators, and far competing styles. + +| Category | Initial probability | Purpose | Examples | +| --- | --- | --- | --- | +| neutral | 0.30 | Test the real deployment case: no alternate style phrase. | "" (empty replacement) | +| hard | 0.40 | Protect against semantic regions already known or suspected to expose the LoRA. | painting; illustration; anime; Ghibli anime; storybook illustration | +| far | 0.30 | Prevent broad leakage into strongly different style semantics. | line art; photorealistic photograph; 3D render; technical drawing; ink drawing | + +Sampling should be per training item by default. With batch_size=2, item A and item B may receive different decoys. Within each item, however, Path 2 and Path 3 must use the exact same sampled phrase. This improves diversity without contaminating the paired comparison. + +## 6.1 YAML schema + +``` +trigger_selective_training: + negative_styles: + expected_category_count: 3 + sample_scope: per_item + + categories: + - name: neutral + probability: 0.30 + phrases: + - "" + + - name: hard + probability: 0.40 + phrases: + - "painting" + - "illustration" + - "anime" + - "Ghibli anime" + - "storybook illustration" + + - name: far + probability: 0.30 + phrases: + - "line art" + - "photorealistic photograph" + - "3D render" + - "technical drawing" + - "ink drawing" +``` + +- expected_category_count is a validation value, not the source of truth. Actual count is len(categories). Mismatch is a startup error. + +- Category probabilities must sum to 1.0 within tolerance. + +- Phrases are uniformly sampled inside a category for v1. Phrase-level weighting can be added later without touching the loss design. + +- Reject any decoy phrase containing [trigger] or the configured literal trigger. + +# 7. Path 2 - conservative non-trigger preservation + +Path 2 is a generalized rotating DOP branch. It does not try to make a decoy prompt bad at the dataset style. It only asks the new LoRA not to materially change what base Ideogram would have done for that same decoy-conditioned caption. + +``` +c_decoy = resolve_trigger_placeholder(raw_caption, sampled_style_s) + +Teacher: base model, LoRA OFF, c_decoy, same x_t and t -> P_base_decoy +Student: base + LoRA, LoRA ON, c_decoy, same x_t and t -> P_lora_decoy + +L2 = MSE(P_lora_decoy, stopgrad(P_base_decoy)) +``` + +This starts close to zero when the effective LoRA is near zero and becomes a guardrail only when training causes the LoRA to leak into non-trigger conditions. + +## 7.1 Interaction with legacy AI Toolkit DOP + +Do not run legacy diff_output_preservation simultaneously with TST Path 2 in the first controlled experiment. V3 already used DOP with the class "painting". TST Path 2 replaces that single fixed preservation class with a sampled bank of conditions. + +``` +# Recommended validation when TST is enabled +train: + diff_output_preservation: false + +# TST Path 2 may reuse DOP's existing prior-prediction / LoRA-bypass plumbing internally. +``` + +If both legacy DOP and TST are enabled accidentally, the trainer should raise a clear configuration error unless a future explicit compatibility mode is added. Silent double-preservation would make the experiment uninterpretable. + +# 8. Path 3 - preferential trigger gain + +Path 2 alone can suppress leakage but does not guarantee that the literal trigger becomes a strong activator. Path 3 adds a relative objective: the LoRA should improve prediction of the dataset target more under the literal trigger than under the matched non-trigger decoy. + +The comparison must be baseline-normalized. Raw trigger loss and raw decoy loss cannot be directly ranked because base Ideogram may naturally be closer to the dataset under one phrase than another even before the LoRA learns anything. + +``` +For the same x_t, timestep t and target y: + +D_trigger = dataset_loss(LoRA(c_trigger), y) +B_trigger = dataset_loss(Base(c_trigger), y) + +D_decoy = dataset_loss(LoRA(c_decoy), y) +B_decoy = dataset_loss(Base(c_decoy), y) + +G_trigger = 1 - D_trigger / (B_trigger + eps) +G_decoy = 1 - D_decoy / (B_decoy + eps) +``` + +G is the normalized improvement attributable to the LoRA under that condition. Near initialization, both gains should be near zero even if the base model has very different absolute losses for "line art" and "Ghibli anime". + +## 8.1 Critical gradient rule + +Path 3 must not be allowed to satisfy the ranking by deliberately making the decoy branch worse. Therefore G_decoy is detached inside Path 3. Path 2 alone controls the decoy branch. Path 3 can only increase the trigger-conditioned advantage. + +``` +L3 = max(0, margin(step) - G_trigger + stopgrad(G_decoy)) +``` + +Hinge ranking is recommended for v1. Once the required advantage is satisfied, Path 3 stops pushing that item. This reduces the risk of an endlessly amplifying trigger objective. + +# 9. Scheduled Path-3 margin + +The margin must itself be scheduled. Early training should require only a small trigger advantage because the LoRA is still learning the visual effect. Later training should require a larger separation, matching the curriculum in which binding becomes more important after the style carrier is established. + +| Step | Recommended margin | Interpretation | +| --- | --- | --- | +| 0 | 0.02 | Very mild selectivity pressure while appearance learning dominates. | +| 1000 | 0.08 | Require a clearer normalized trigger advantage at mid-training. | +| 1500+ | 0.12 | Hold a stronger selectivity requirement through the final phase. | + +These values are starting defaults, not sacred constants. Margin is a normalized gain gap, not a direct percentage of visible style strength. + +## 9.1 YAML schema + +``` +trigger_selective_training: + path3: + loss_type: hinge + gain_epsilon: 1.0e-6 + + margin_schedule: + interpolation: smoothstep + keyframes: + - step: 0 + value: 0.02 + - step: 1000 + value: 0.08 + - step: 1500 + value: 0.12 +``` + +After the last margin keyframe, clamp to the last value. Before the first keyframe, clamp to the first. The same generic scheduler utility should support both linear and smoothstep interpolation. + +# 10. Scheduled three-path loss weights + +The user controls the nominal percentage assigned to each objective. The initial schedule intentionally keeps style acquisition dominant, then gradually shifts more coefficient weight toward trigger selectivity. + +| Step | Path 1 | Path 2 | Path 3 | Training emphasis | +| --- | --- | --- | --- | --- | +| 0 | 80% | 10% | 10% | Mostly learn the style; light leakage/selectivity constraints. | +| 1000 | 60% | 15% | 25% | Style should be established; binding pressure grows. | +| 1500+ | 50% | 15% | 35% | Final phase: preserve style while strongly optimizing trigger selectivity. | + +## 10.1 YAML schema + +``` +trigger_selective_training: + loss_schedule: + interpolation: smoothstep + normalize_weights: true + + keyframes: + - step: 0 + path1: 0.80 + path2: 0.10 + path3: 0.10 + + - step: 1000 + path1: 0.60 + path2: 0.15 + path3: 0.25 + + - step: 1500 + path1: 0.50 + path2: 0.15 + path3: 0.35 +``` + +> **Interpret percentages correctly:** These are normalized loss coefficients, not guaranteed percentages of the optimizer update. Different losses can have different scales and gradient norms. Log branch gradient norms during short diagnostics before treating 80/10/10 as literal update shares. + +# 11. Interpolation behavior + +Both the loss-weight schedule and the margin schedule use the same generic keyframe scheduler. Support two modes in v1: + +| Mode | Behavior | Recommended use | +| --- | --- | --- | +| linear | Straight interpolation between adjacent keyframes. | Useful for debugging because values are easy to predict. | +| smoothstep | Uses u^2(3-2u) inside each interval, giving smooth acceleration/deceleration. | Recommended first real run to avoid abrupt slope changes at keyframes. | + +Weights are normalized after interpolation when normalize_weights=true. Therefore 80/10/10 and 0.8/0.1/0.1 are equivalent inputs. Keyframe steps must be strictly increasing. + +# 12. Recommended first-run YAML block + +The following block is the proposed v1 interface. Existing V3 network/train/model settings remain outside it. + +``` +trigger_word: "" + +trigger_selective_training: + enabled: true + require_trigger_placeholder: true + + negative_styles: + expected_category_count: 3 + sample_scope: per_item + categories: + - name: neutral + probability: 0.30 + phrases: [""] + + - name: hard + probability: 0.40 + phrases: + - "painting" + - "illustration" + - "anime" + - "Ghibli anime" + - "storybook illustration" + + - name: far + probability: 0.30 + phrases: + - "line art" + - "photorealistic photograph" + - "3D render" + - "technical drawing" + - "ink drawing" + + path3: + loss_type: hinge + gain_epsilon: 1.0e-6 + margin_schedule: + interpolation: smoothstep + keyframes: + - {step: 0, value: 0.02} + - {step: 1000, value: 0.08} + - {step: 1500, value: 0.12} + + loss_schedule: + interpolation: smoothstep + normalize_weights: true + keyframes: + - {step: 0, path1: 0.80, path2: 0.10, path3: 0.10} + - {step: 1000, path1: 0.60, path2: 0.15, path3: 0.25} + - {step: 1500, path1: 0.50, path2: 0.15, path3: 0.35} + + logging: + log_every: 25 + log_category_stats: true + log_gain_stats: true + debug_gradient_contributions: false + +train: + # Path 2 supersedes legacy fixed-class DOP in this experiment. + diff_output_preservation: false +``` + +# 13. AI Toolkit integration plan + +## 13.1 Config layer + +- Add TriggerSelectiveTrainingConfig and nested config objects in toolkit/config_modules.py or the project-equivalent config module. + +- Validate category count, probability sum, non-empty categories, schedule ordering, non-negative loss weights and legal interpolation modes at startup. + +- When TST is enabled, reject legacy diff_output_preservation=true for the first controlled implementation. + +## 13.2 New isolated helper module + +Prefer a small new module such as toolkit/trigger_selective_training.py. It should contain logic that is independent of Ideogram model internals: + +- negative category/phrase sampling; + +- shared keyframe scheduler for scalar margin and vector loss weights; + +- [trigger] variant resolution wrapper that reuses the existing AI Toolkit trigger mechanism; + +- Path-3 gain and hinge-loss computation; + +- validation and lightweight metrics aggregation. + +## 13.3 SDTrainer integration + +Keep the main Path-1 training route intact. Around it, add the decoy and base-reference predictions. Reuse existing DOP prior-prediction / network-bypass mechanisms where they already provide the correct model state. + +``` +per training item/batch: + +1. sample x_t, t, target exactly once +2. sample decoy category + phrase s +3. build c_trigger and c_decoy from the same raw [trigger] caption template +4. compute base-decoy prediction (no grad, LoRA off) +5. compute student-decoy prediction (grad, LoRA on) +6. compute L2 and detached G_decoy +7. compute base-trigger prediction (no grad, LoRA off) +8. run the normal Path-1 student-trigger forward (grad, LoRA on) +9. compute L1, G_trigger and L3 using scheduled margin +10. get scheduled w1/w2/w3 +11. accumulate weighted gradients +12. optimizer step follows existing AI Toolkit cadence +``` + +# 14. Memory-conscious backward strategy + +A naive implementation holds both gradient-bearing student graphs at once. That may be unnecessarily expensive on a 24 GB card. Because Path 3 deliberately stop-grads G_decoy, the total gradient can be accumulated in two stages without changing the intended mathematics. + +``` +# Stage A: decoy branch +base_decoy = no_grad_base_forward(c_decoy) +student_decoy = grad_forward(c_decoy) +L2 = preservation_loss(student_decoy, base_decoy) +G_decoy = gain(student_decoy, base_decoy, target).detach() +backward(w2 * L2) +# release decoy student graph + +# Stage B: trigger branch +base_trigger = no_grad_base_forward(c_trigger) +student_trigger = normal_path1_grad_forward(c_trigger) +L1 = dataset_loss(student_trigger, target) +G_trigger = gain(student_trigger, base_trigger, target) +L3 = hinge(margin - G_trigger + G_decoy) +backward(w1 * L1 + w3 * L3) + +optimizer.step() # once, using accumulated gradients +``` + +This staged backward is mathematically equivalent to a single combined backward under the v1 detach rule, while allowing the decoy graph to be freed before the trigger graph is created. It must be tested carefully with AI Toolkit gradient accumulation and Accelerator semantics. If uncertain, implement the single-combined-loss version first for correctness, then switch to staged backward only after an equivalence test. + +> **Reuse rule:** If Differential Guidance or another existing AI Toolkit path already computes an exactly matching base-trigger prediction under the same x_t, timestep and conditioning, it may be reused. Do not reuse a superficially similar prior tensor unless the conditions are proven identical. + +# 15. Model-state and autograd safety + +- **Teacher/base forwards.** LoRA disabled, torch.no_grad(), detached outputs, base model frozen. + +- **Student forwards.** LoRA enabled and trainable; base Ideogram weights remain frozen. + +- **State restoration.** Use a scoped context manager or try/finally for LoRA enable/disable. Never scatter manual state toggles through the trainer. + +- **Same stochastic state.** All branches share the same sampled x_t, timestep and target. Do not resample noise per branch. + +- **Differential Guidance.** Keep the V3 setting unchanged for the first TST experiment. Do not simultaneously change the known-good diffusion recipe. + +# 16. Logging and diagnostics + +The experiment is only useful if it tells us why it succeeds or fails. Log raw losses, weighted losses, gains, schedule values and negative-category statistics separately. + +``` +loss/path1_raw +loss/path2_raw +loss/path3_raw +loss/path1_weighted +loss/path2_weighted +loss/path3_weighted +loss/total + +weight/path1 +weight/path2 +weight/path3 +path3/margin +path3/margin_satisfied + +gain/trigger +gain/decoy +gain/gap + +negative/category +negative/phrase +negative/neutral_count +negative/hard_count +negative/far_count +``` + +For running aggregate analysis, maintain separate EMAs or averages of gain/gap for neutral, hard and far categories. This can reveal, for example, that the trigger separates cleanly from neutral prompts while still failing against anime-like semantics. + +## 16.1 Sparse gradient diagnostics + +The 80/10/10 schedule is a coefficient schedule, not a guarantee of 80/10/10 gradient influence. During short diagnostic runs, optionally measure branch gradient norms at sparse steps such as 1, 10, 25, 50, 100, 250, 500, 1000 and 1500. Keep this behind a debug flag because it is expensive. + +# 17. Required tests before real training + +| Area | Required assertion | +| --- | --- | +| Disabled-mode regression | TST disabled -> existing trainer behavior unchanged. | +| Placeholder | Every raw caption contains [trigger] when strict mode is enabled. | +| Placeholder | Path 1 and Path 2 use the same native resolver with different replacement values. | +| Placeholder | All [trigger] occurrences in a caption are resolved; no manual literal-trigger search is required. | +| Sampling | Category probabilities validate and sampled phrase is shared by Path 2 and Path 3 for each item. | +| Scheduling | Exact keyframe values, linear interpolation and smoothstep interpolation all unit-test correctly. | +| Path 2 | Teacher is no-grad/LoRA-off; student is grad/LoRA-on. | +| Path 3 | G_decoy is detached; Path-3-only backward cannot change decoy-branch parameters through its graph. | +| Branch pairing | All predictions share the same x_t, timestep and target. | +| State | LoRA enabled/disabled state is restored after every teacher context. | +| Legacy DOP | TST + legacy diff_output_preservation raises a clear error in v1. | +| Backward | Staged and combined backward produce matching gradients within tolerance on a small deterministic test. | + +# 18. Step-zero sanity expectations + +| Metric | Expected at effective LoRA initialization | +| --- | --- | +| L1 | High relative to a trained checkpoint; normal dataset acquisition has not happened yet. | +| L2 | Near zero because LoRA student should initially approximate the base teacher on the decoy. | +| G_trigger | Near zero. | +| G_decoy | Near zero. | +| L3 | Usually active because the scheduled margin is positive and the gain gap is near zero. | + +If L2 is large at step zero, the teacher/student model state is wrong. If G_trigger and G_decoy are already very different, inspect base-loss normalization before training. + +# 19. First controlled experiment and kill gates + +Use a fresh LoRA initialization with the same known-good V3 diffusion-side configuration, except that legacy fixed-class DOP is disabled and TST supplies Path 2. Do not simultaneously change rank, timestep sampling, content mode, Differential Guidance, dataset or caption structure. + +| Checkpoint | What to inspect | +| --- | --- | +| 25 / 50 | Do L1 and gains move at all? Is L2 still small? Are schedules returning expected values? | +| 100 / 200 | Does G_trigger begin separating from G_decoy? Does visual trigger ON/OFF show any difference? | +| 500 | If no selectivity movement exists, stop rather than assuming 2000 steps will fix it. | +| 1000 | Confirm mid-training schedule transition: 60/15/25 and margin ~0.08. | +| 1500+ | Confirm final schedule: 50/15/35 and margin ~0.12; inspect whether style fidelity remains intact. | + +## 19.1 Primary evaluation matrix + +``` +Fixed seed / sampler / steps / LoRA strength: + +1. novel neutral prompt +2. same novel prompt + + +3. Ghibli anime prompt +4. same Ghibli anime prompt + + +5. line art prompt +6. same line art prompt + + +7. unseen style phrase +8. same unseen phrase + +``` + +The first success criterion is not perfect suppression of all no-trigger LoRA effects. It is a material, repeatable increase in target-style activation when the literal trigger is present, while Path 2 keeps non-trigger leakage conservative enough that base prompt behavior remains usable. + +# 20. Failure interpretation + +| Observed behavior | Likely interpretation / next action | +| --- | --- | +| L1 decreases, G_trigger and G_decoy stay similar | Style is relearning without conditional binding. Static LoRA selectivity may be insufficient or Path 3 pressure may be too weak. | +| L2 rises strongly while L1 improves | Leakage outpaces preservation. Inspect actual gradient norms before simply increasing Path-2 coefficient. | +| G_trigger rises, L2 stays low, images still do not change | Diffusion-prediction gain may not map to visible style activation strongly enough; inspect output residual direction and timestep dependence. | +| Path 3 remains unsatisfied late and L1 degrades | Margin and/or Path-3 weight may be too aggressive; reduce scheduled endpoint rather than training longer. | +| Neutral separation works, hard negatives fail | Binding exists but is not robust against pretrained style semantics. Increase hard-negative sampling or add a later simultaneous competitor+trigger extension. | + +# 21. Recommended code footprint + +| File / area | Change | +| --- | --- | +| toolkit/config_modules.py | Add TST config objects, validation and schedule schema. | +| toolkit/trigger_selective_training.py (new) | Sampling, schedule interpolation, [trigger] variant helper/wrapper, gain calculation, Path-3 hinge, metrics helpers. | +| extensions_built_in/sd_trainer/SDTrainer.py | Integrate decoy/base forwards and weighted losses around the existing Path-1 training route. | +| Existing trigger resolver location | Only if necessary, extract a reusable replacement helper. Preserve existing [trigger] behavior exactly for ordinary training. | +| Generic LoRA/network code | No change expected for v1. | +| Ideogram/Qwen model code | No change expected for v1. | +| ComfyUI loader | No change. Result remains a standard transformer LoRA. | + +# 22. Compact mathematical definition + +This section is intentionally small. It defines only the quantities the implementation needs. + +``` +Definitions +----------- +c_raw = raw dataset caption containing [trigger] +tau = configured literal trigger +s = sampled non-trigger style phrase +R(c,x) = AI Toolkit's native [trigger] resolution using replacement x + +c_tau = R(c_raw, tau) +c_s = R(c_raw, s) + +f_theta = base Ideogram + trainable LoRA +f_0 = frozen base Ideogram with trainable LoRA disabled +y = current flow/diffusion target + +Dataset losses +-------------- +D_tau = L(f_theta(x_t, c_tau, t), y) +B_tau = L(f_0 (x_t, c_tau, t), y) +D_s = L(f_theta(x_t, c_s, t), y) +B_s = L(f_0 (x_t, c_s, t), y) + +Normalized LoRA gains +--------------------- +G_tau = 1 - D_tau / (B_tau + eps) +G_s = 1 - D_s / (B_s + eps) + +Three losses +------------ +L1 = D_tau +L2 = MSE(f_theta(x_t,c_s,t), stopgrad(f_0(x_t,c_s,t))) +L3 = max(0, m(step) - G_tau + stopgrad(G_s)) + +Total objective +--------------- +L_total(step) = w1(step)*L1 + w2(step)*L2 + w3(step)*L3 + +Smoothstep interpolation +------------------------ +u in [0,1] between adjacent keyframes +s(u) = u^2 * (3 - 2u) +value = value_A + s(u) * (value_B - value_A) +``` + +# 23. Decisive implementation summary + +- **Do preserve:** the exact V3-style ordinary diffusion training path as Path 1. + +- **Do reuse:** AI Toolkit native [trigger] placeholder resolution for both literal-trigger and decoy prompt variants. + +- **Do replace:** legacy single-class DOP with Path 2 rotating preservation when TST is enabled. + +- **Do add:** Path 3 baseline-normalized trigger-vs-decoy gain ranking, with the decoy gain detached. + +- **Do schedule:** both three-path loss coefficients and the Path-3 margin through YAML keyframes using linear or smoothstep interpolation. + +- **Do keep:** same x_t, noise, timestep and target across all paired branches. + +- **Do expose:** negative categories, category probabilities, phrases, loss schedules, interpolation, margin schedule and diagnostics. + +- **Do not add yet:** Qwen LoRA, TARA token masking, text embeddings, custom inference or hard runtime trigger gating. + +> **Hypothesis under test:** A standard Ideogram transformer LoRA has enough conditional capacity to learn a strong style residual while making the arbitrary literal trigger preferentially expose that residual, provided training separately optimizes style acquisition, non-trigger preservation, and trigger-relative gain. + +# 24. Source basis + +Project evidence and current implementation details used to form this specification: + +- **Project V3 YAML.** 2026_08_11_ig4_r1X1dOn9mA2_v3.yaml. Relevant facts: trigger_word=, LoRA rank 32 / alpha 16, content mode, sigmoid timesteps, DOP enabled with class "painting", Differential Guidance scale 3. + +- **Project continuation handoff.** Dated 2026-08-12: IDEOGRAM 4 ARBITRARY-TRIGGER STYLE LORA PROJECT. Relevant facts: target objective, dataset structure, previous TI/Qwen-LoRA results, early kill-gate preference, and Ideogram/Qwen architecture notes. + +- **AI Toolkit official documentation.** ui/src/docs.tsx (checked 2026-08-13): [trigger] in captions is automatically replaced with the configured trigger_word. + +- **AI Toolkit official configuration code.** toolkit/config_modules.py (checked 2026-08-13): current DOP fields and semantics, including diff_output_preservation_class. + +- **AI Toolkit trainer source / project trainer snapshots.** The existing prior-prediction path already supports LoRA-bypassed teacher predictions and can be reused as plumbing for Path 2 where conditions match. + +This specification intentionally distinguishes established AI Toolkit behavior from the new TST hypothesis. TST itself is an experimental project design, not a claim of an existing published algorithm. diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index c1562c31..1c690b1a 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -41,6 +41,19 @@ from toolkit.unloader import unload_text_encoder from PIL import Image from torchvision.transforms import functional as TF from toolkit.basic import flush +from toolkit.trigger_selective_training import ( + TSTMetricsWriter, + apply_differential_guidance_target, + get_scheduled_loss_weights, + get_scheduled_margin, + network_disabled, + normalized_gain, + per_item_mse, + resolve_prompt_variants, + sample_negative_styles, + shared_loss_target, + trigger_advantage_hinge, +) adapter_transforms = transforms.Compose([ @@ -84,6 +97,13 @@ class SDTrainer(BaseSDTrainProcess): # fallback class-only embeds for when the text encoder is unloaded and # per item DOP embeds were not cached to disk self.cached_dop_class_embeds: Optional[PromptEmbeds] = None + self.tst_rng = random.Random(self.training_seed) + self.tst_metrics_writer: Optional[TSTMetricsWriter] = None + if self.trigger_selective_training.enabled: + self.tst_metrics_writer = TSTMetricsWriter( + self.save_root, + self.trigger_selective_training.logging.metrics_filename, + ) self.dfe: Optional[DiffusionFeatureExtractor] = None self.unconditional_embeds = None @@ -506,6 +526,7 @@ class SDTrainer(BaseSDTrainProcess): batch: 'DataLoaderBatchDTO', mask_multiplier: Union[torch.Tensor, float] = 1.0, prior_pred: Union[torch.Tensor, None] = None, + target_override: Union[torch.Tensor, None] = None, **kwargs ): loss_target = self.train_config.loss_target @@ -530,12 +551,12 @@ class SDTrainer(BaseSDTrainProcess): if self.train_config.pred_scaler != 1.0: noise_pred = noise_pred * self.train_config.pred_scaler - target = None + target = target_override - if self.train_config.target_noise_multiplier != 1.0: + if target is None and self.train_config.target_noise_multiplier != 1.0: noise = noise * self.train_config.target_noise_multiplier - if self.train_config.correct_pred_norm or (self.train_config.inverted_mask_prior and prior_pred is not None and has_mask): + if target is None and (self.train_config.correct_pred_norm or (self.train_config.inverted_mask_prior and prior_pred is not None and has_mask)): if self.train_config.correct_pred_norm and not is_reg: with torch.no_grad(): # this only works if doing a prior pred @@ -594,14 +615,14 @@ class SDTrainer(BaseSDTrainProcess): target = (noise - batch.latents).detach() else: target = noise - elif prior_pred is not None and not self.train_config.do_prior_divergence: + elif target is None and prior_pred is not None and not self.train_config.do_prior_divergence: assert not self.train_config.train_turbo # matching adapter prediction target = prior_pred - elif self.sd.prediction_type == 'v_prediction': + elif target is None and self.sd.prediction_type == 'v_prediction': # v-parameterization training target = self.sd.noise_scheduler.get_velocity(batch.tensor, noise, timesteps) - elif self.train_config.do_signal_amplification: + elif target is None and self.train_config.do_signal_amplification: if not self.sd.is_flow_matching: raise ValueError("Signal amplification is only supported for flow matching models") with torch.no_grad(): @@ -612,19 +633,19 @@ class SDTrainer(BaseSDTrainProcess): aug = batch.latents * nas target = noise - (batch.latents + aug) target = target.detach() - elif hasattr(self.sd, 'get_loss_target'): + elif target is None and hasattr(self.sd, 'get_loss_target'): target = self.sd.get_loss_target( noise=noise, batch=batch, timesteps=timesteps, ).detach() - elif self.sd.is_flow_matching: + elif target is None and self.sd.is_flow_matching: # forward ODE target = (noise - batch.latents).detach() # reverse ODE # target = (batch.latents - noise).detach() - else: + elif target is None: target = noise if self.dfe is not None: @@ -1354,6 +1375,199 @@ class SDTrainer(BaseSDTrainProcess): ) + 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: + prompt_kwargs['control_images'] = batch.control_tensor.to( + self.sd.device_torch, dtype=self.sd.torch_dtype + ) + with torch.no_grad(): + trigger_embeds = self.sd.encode_prompt( + trigger_prompts, + long_prompts=self.do_long_prompts, + **prompt_kwargs, + ).to(self.device_torch, dtype=dtype).detach() + decoy_embeds = self.sd.encode_prompt( + decoy_prompts, + long_prompts=self.do_long_prompts, + **prompt_kwargs, + ).to(self.device_torch, dtype=dtype).detach() + return trigger_embeds, decoy_embeds + + def _tst_trainable_params(self): + params = [] + for group in self.params: + if isinstance(group, dict): + params.extend(group.get('params', [])) + else: + params.append(group) + return [param for param in params if param.requires_grad] + + def _tst_gradient_norm(self, loss): + grads = torch.autograd.grad( + loss, + self._tst_trainable_params(), + retain_graph=True, + allow_unused=True, + ) + squared_norm = torch.zeros((), device=loss.device, dtype=torch.float32) + for grad in grads: + if grad is not None: + squared_norm = squared_norm + grad.detach().float().pow(2).sum() + return squared_norm.sqrt().item() + + def _calculate_tst_loss( + self, + noisy_latents, + noise, + timesteps, + batch, + conditioned_prompts, + prompt_2, + pred_kwargs, + mask_multiplier, + network, + dtype, + ): + raw_templates = [ + getattr(file_item, 'caption_template', None) or file_item.raw_caption + for file_item in batch.file_items + ] + negative_samples = sample_negative_styles( + self.trigger_selective_training, + len(raw_templates), + self.tst_rng, + ) + trigger_prompts, decoy_prompts = resolve_prompt_variants( + raw_templates, + self.trigger_word, + negative_samples, + self.trigger_selective_training.require_trigger_placeholder, + ) + trigger_embeds, decoy_embeds = self._encode_tst_prompt_variants( + batch, trigger_prompts, decoy_prompts, dtype + ) + + with network_disabled(network): + with torch.no_grad(): + base_decoy = self.predict_noise( + noisy_latents=noisy_latents, + timesteps=timesteps, + conditional_embeds=decoy_embeds, + unconditional_embeds=None, + batch=batch, + **pred_kwargs, + ).detach() + base_trigger = self.predict_noise( + noisy_latents=noisy_latents, + timesteps=timesteps, + conditional_embeds=trigger_embeds, + unconditional_embeds=None, + batch=batch, + **pred_kwargs, + ).detach() + + student_decoy = self.predict_noise( + noisy_latents=noisy_latents, + timesteps=timesteps, + conditional_embeds=decoy_embeds, + unconditional_embeds=None, + batch=batch, + **pred_kwargs, + ) + student_trigger = self.predict_noise( + noisy_latents=noisy_latents, + timesteps=timesteps, + conditional_embeds=trigger_embeds, + unconditional_embeds=None, + batch=batch, + is_primary_pred=True, + **pred_kwargs, + ) + + target = shared_loss_target(self, noise, batch, timesteps) + target = apply_differential_guidance_target(self, target, student_trigger) + path1 = self.calculate_loss( + noise_pred=student_trigger, + noise=noise, + noisy_latents=noisy_latents, + timesteps=timesteps, + batch=batch, + mask_multiplier=mask_multiplier, + target_override=target, + ) + path2_per_item = per_item_mse(student_decoy, base_decoy) + base_decoy_loss = per_item_mse(base_decoy, target) + student_decoy_loss = per_item_mse(student_decoy, target) + base_trigger_loss = per_item_mse(base_trigger, target) + student_trigger_loss = per_item_mse(student_trigger, target) + decoy_gain = normalized_gain(student_decoy_loss, base_decoy_loss, self.trigger_selective_training.path3.gain_epsilon).detach() + trigger_gain = normalized_gain(student_trigger_loss, base_trigger_loss, self.trigger_selective_training.path3.gain_epsilon) + margin = get_scheduled_margin(self.trigger_selective_training, self.step_num) + path3_per_item = trigger_advantage_hinge(trigger_gain, decoy_gain, margin) + weights = get_scheduled_loss_weights(self.trigger_selective_training, self.step_num) + path2 = path2_per_item.mean() + path3 = path3_per_item.mean() + weighted_path1 = weights['path1'] * path1 + weighted_path2 = weights['path2'] * path2 + weighted_path3 = weights['path3'] * path3 + loss = weighted_path1 + weighted_path2 + weighted_path3 + + gradient_logs = {} + logging_config = self.trigger_selective_training.logging + if ( + logging_config.debug_gradient_contributions + and self.step_num in logging_config.gradient_diagnostic_steps + ): + gradient_logs = { + 'grad_norm/path1': self._tst_gradient_norm(weighted_path1), + 'grad_norm/path2': self._tst_gradient_norm(weighted_path2), + 'grad_norm/path3': self._tst_gradient_norm(weighted_path3), + } + + category_counts = {} + category_gain_gaps = {} + for category, gap in zip( + [sample.category for sample in negative_samples], + (trigger_gain.detach() - decoy_gain).tolist(), + ): + category_counts[category] = category_counts.get(category, 0) + 1 + category_gain_gaps.setdefault(category, []).append(float(gap)) + self.additional_logs.update({ + 'loss/path1_raw': path1.detach().item(), + 'loss/path2_raw': path2.detach().item(), + 'loss/path3_raw': path3.detach().item(), + 'loss/path1_weighted': (weights['path1'] * path1.detach()).item(), + 'loss/path2_weighted': (weights['path2'] * path2.detach()).item(), + 'loss/path3_weighted': (weights['path3'] * path3.detach()).item(), + 'loss/total': loss.detach().item(), + 'weight/path1': weights['path1'], + 'weight/path2': weights['path2'], + 'weight/path3': weights['path3'], + 'path3/margin': margin, + 'path3/margin_satisfied': (path3_per_item <= 0).float().mean().item(), + 'gain/trigger': trigger_gain.detach().mean().item(), + 'gain/decoy': decoy_gain.mean().item(), + 'gain/gap': (trigger_gain.detach() - decoy_gain).mean().item(), + **{ + f'negative/{category}_count': count + for category, count in category_counts.items() + }, + **{ + f'gain/{category}_gap': sum(gaps) / len(gaps) + for category, gaps in category_gain_gaps.items() + }, + **gradient_logs, + }) + if self.tst_metrics_writer is not None and self.step_num % self.trigger_selective_training.logging.log_every == 0: + self.tst_metrics_writer.write({ + 'step': self.step_num, + **{key: float(value) for key, value in self.additional_logs.items() if isinstance(value, (int, float))}, + 'negative_category': [sample.category for sample in negative_samples], + 'negative_phrase': [sample.phrase for sample in negative_samples], + }) + return loss + def train_single_accumulation(self, batch: DataLoaderBatchDTO): with torch.no_grad(): self.timer.start('preprocess_batch') @@ -2126,36 +2340,53 @@ class SDTrainer(BaseSDTrainProcess): prior_pred=prior_pred, ) else: - with self.timer('predict_unet'): - noise_pred = self.predict_noise( - noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype), - timesteps=timesteps, - conditional_embeds=conditional_embeds.to(self.device_torch, dtype=dtype), - unconditional_embeds=unconditional_embeds, - batch=batch, - is_primary_pred=True, - **pred_kwargs - ) - self.after_unet_predict() + tst_loss = None + if 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), + noise=noise.to(self.device_torch, dtype=dtype).detach(), + timesteps=timesteps, + batch=batch, + conditioned_prompts=conditioned_prompts, + prompt_2=prompt_2, + pred_kwargs=pred_kwargs, + mask_multiplier=mask_multiplier, + network=network, + dtype=dtype, + ) + else: + with self.timer('predict_unet'): + noise_pred = self.predict_noise( + noisy_latents=noisy_latents.to(self.device_torch, dtype=dtype), + timesteps=timesteps, + conditional_embeds=conditional_embeds.to(self.device_torch, dtype=dtype), + unconditional_embeds=unconditional_embeds, + batch=batch, + is_primary_pred=True, + **pred_kwargs + ) + self.after_unet_predict() - with self.timer('calculate_loss'): - noise = noise.to(self.device_torch, dtype=dtype).detach() - prior_to_calculate_loss = prior_pred - # if we are doing diff_output_preservation and not noing inverted masked prior - # then we need to send none here so it will not target the prior - doing_preservation = self.train_config.diff_output_preservation or self.train_config.blank_prompt_preservation - if doing_preservation and not do_inverted_masked_prior: - prior_to_calculate_loss = None - - loss = self.calculate_loss( - noise_pred=noise_pred, - noise=noise, - noisy_latents=noisy_latents, - timesteps=timesteps, - batch=batch, - mask_multiplier=mask_multiplier, - prior_pred=prior_to_calculate_loss, - ) + with self.timer('calculate_loss'): + noise = noise.to(self.device_torch, dtype=dtype).detach() + prior_to_calculate_loss = prior_pred + # if we are doing diff_output_preservation and not noing inverted masked prior + # then we need to send none here so it will not target the prior + doing_preservation = self.train_config.diff_output_preservation or self.train_config.blank_prompt_preservation + if doing_preservation and not do_inverted_masked_prior: + prior_to_calculate_loss = None + + tst_loss = self.calculate_loss( + noise_pred=noise_pred, + noise=noise, + noisy_latents=noisy_latents, + timesteps=timesteps, + batch=batch, + mask_multiplier=mask_multiplier, + prior_pred=prior_to_calculate_loss, + ) + loss = tst_loss if self.train_config.diff_output_preservation or self.train_config.blank_prompt_preservation: with torch.no_grad(): diff --git a/handoff.txt b/handoff.txt new file mode 100644 index 00000000..d9821fad --- /dev/null +++ b/handoff.txt @@ -0,0 +1,1721 @@ +# IDEOGRAM 4 ARBITRARY-TRIGGER STYLE LORA PROJECT +# CONTINUATION HANDOFF — V7-B SELECTED +# +# Date: 2026-08-12 +# +# Main objective: +# Make the arbitrary literal trigger reliably activate the +# user's target visual style on novel prompts, including prompts containing +# competing style semantics. +# +# SELECTED NEXT EXPERIMENT: +# V7-B = TARA-like token-aware Qwen activator + semantic anchor alignment. +# +# IMPORTANT CHANGE: +# The semantic/class anchor is now: +# +# "painting" +# +# NOT "Ghibli anime". +# +# Do not silently substitute another anchor. +# +# ------------------------------------------------------------------------- +# 1. USER / PROJECT CONSTRAINTS +# ------------------------------------------------------------------------- +# +# The user is highly technical and wants controlled causal experiments. +# +# Important preferences: +# +# - Do not burn 1000–2000 steps before an early kill gate. +# - Separate: +# mechanical parameter movement +# gradient flow +# inference loading +# semantic success +# - Do not interpret "weights changed" as "binding succeeded". +# - Do not recommend changing: +# +# content_or_style: content +# +# to "style". +# +# The user deliberately wants content mode because the desired learned +# behavior includes composition/content tendencies as part of the visual style. +# +# - Do not go back to pure Textual Inversion as the primary solution. +# - Do not simply increase V6 text-encoder LR and retry the same formulation. +# - Preserve current working AI Toolkit infrastructure unless a change is +# specifically required for V7-B. +# +# +# ------------------------------------------------------------------------- +# 2. DATASET / TRIGGER +# ------------------------------------------------------------------------- +# +# Dataset: +# +# 40 curated target-style images +# +# Captions: +# +# detailed manually authored Ideogram 4 JSON captions +# +# Trigger: +# +# +# +# The trigger appears multiple times in the captions, generally around three +# occurrences in different style_description positions. +# +# Important: +# +# detect ALL actual occurrences dynamically; +# do not hardcode "exactly 3" unless dataset validation proves that. +# +# The target visual style itself is intentionally NOT described by ordinary +# style words. The random trigger is supposed to carry the binding. +# +# +# ------------------------------------------------------------------------- +# 3. IDEOGRAM 4 TEXT ARCHITECTURE ALREADY ESTABLISHED +# ------------------------------------------------------------------------- +# +# Current Ideogram4 AI Toolkit implementation uses: +# +# Qwen3-VL-8B-Instruct +# +# Text features are collected from 13 Qwen layers: +# +# (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35) +# +# Qwen hidden size: +# +# 4096 +# +# Ideogram concatenates those 13 tapped representations for conditioning. +# +# Qwen has 36 text decoder layers. +# +# Current AI Toolkit Ideogram prompt encoding processes prompts individually +# through Qwen: +# +# diffusion batch = 2 +# Qwen encode batch = 1 per caption +# +# +# ------------------------------------------------------------------------- +# 4. EXISTING V3 — GOOD STYLE CARRIER +# ------------------------------------------------------------------------- +# +# There is already a strong transformer-only LoRA called V3. +# +# V3: +# +# Ideogram transformer LoRA +# rank 32 +# good visual reproduction of the target style +# +# Main weakness: +# +# activation/binding is unreliable. +# +# Known observations: +# +# Base + semantic style phrase can move somewhat toward target. +# V3 + certain semantic style phrases activates target style strongly. +# Literal random trigger is unreliable. +# +# V3 therefore appears to contain the desired visual style reasonably well. +# +# Treat V3 as the STYLE CARRIER. +# +# For V7-B: +# +# V3 should be loaded ACTIVE + FROZEN. +# +# DO NOT make V3 trainable in the initial V7-B experiment. +# +# This is essential. +# +# One major problem with V6 was that the trainable transformer LoRA could solve +# diffusion reconstruction itself, giving the Qwen LoRA no incentive to become +# a meaningful trigger activator. +# +# V7-B intentionally removes that shortcut: +# +# style carrier = frozen V3 +# activator = trainable token-aware Qwen LoRA +# +# +# ------------------------------------------------------------------------- +# 5. PREVIOUS TI EXPERIMENTS — DO NOT REPEAT +# ------------------------------------------------------------------------- +# +# Several Textual Inversion experiments were already performed. +# +# Pure TI moved substantially in parameter space but did not provide robust +# semantic activation. +# +# A frozen-V3 + TI stress experiment also failed: +# +# random TI initialization was sometimes better than trained TI. +# +# Therefore: +# +# large embedding movement != useful semantic binding. +# +# Pure TI is not the selected path. +# +# +# ------------------------------------------------------------------------- +# 6. V6 — JOINT QWEN LORA + TRANSFORMER LORA +# ------------------------------------------------------------------------- +# +# V6 trained: +# +# Ideogram transformer LoRA +# + +# Qwen attention LoRA +# +# simultaneously under ordinary diffusion MSE. +# +# V6 Qwen target: +# +# Qwen3VLTextAttention +# +# which created: +# +# 36 layers +# × q/k/v/o +# = 144 Qwen LoRA modules +# +# = 288 A/B tensors when saved. +# +# Example V6 config: +# +# network: +# type: lora +# linear: 32 +# linear_alpha: 16 +# network_kwargs: +# target_te_modules: +# - Qwen3VLTextAttention +# +# train: +# batch_size: 2 +# train_unet: true +# train_text_encoder: false +# train_text_encoder_lora: true +# gradient_checkpointing: true +# timestep_type: sigmoid +# content_or_style: content +# cache_text_embeddings: false +# unload_text_encoder: false +# unet_lr: 8e-5 +# text_encoder_lr: 2e-5 +# dtype: bf16 +# do_differential_guidance: true +# differential_guidance_scale: 3.0 +# +# +# ------------------------------------------------------------------------- +# 7. IMPORTANT TRAINING INFRASTRUCTURE ALREADY FIXED +# ------------------------------------------------------------------------- +# +# Several plumbing bugs were solved. Preserve these. +# +# +# A. Separate flag: +# +# train_text_encoder_lora +# +# Base Qwen stays frozen. +# +# Qwen LoRA is trainable. +# +# +# B. Qwen autograd +# +# In Ideogram pipeline: +# +# get_qwen3_vl_features(...) +# +# MUST NOT have @torch.no_grad(). +# +# +# C. Qwen must execute under autograd during TE-LoRA training. +# +# SDTrainer currently uses: +# +# grad_on_text_encoder = ( +# train_text_encoder +# or train_text_encoder_lora +# ) +# +# +# D. Base Qwen must NOT be Accelerator-prepared as a fully trainable text model. +# +# Keep: +# +# train_text_encoder: false +# +# +# E. Critical batch/multiplier fix +# +# Ideogram4 Qwen encodes captions one at a time (B=1), while normal AI Toolkit +# LoRA network multiplier follows diffusion batch (e.g. [1,1]). +# +# This previously caused the apparent Qwen stall. +# +# During Qwen encoding, temporarily collapse a uniform diffusion-batch +# multiplier: +# +# [1.0, 1.0] +# +# to: +# +# 1.0 +# +# Restore the original multiplier before diffusion-transformer forward. +# +# Preserve strict validation: +# +# if batch weights differ, do NOT silently collapse them. +# +# +# F. Native Hugging Face Qwen gradient checkpointing +# +# Qwen needs to be temporarily in train() mode during the gradient-bearing +# student encode so HF GradientCheckpointingLayer actually activates. +# +# Base Qwen weights remain requires_grad=False. +# +# Restore the previous train/eval state after the encode. +# +# +# G. Do not restore old debugging hacks: +# +# no manual torch checkpoint wrapper +# no manual q_proj reconstruction +# no attention instrumentation +# no forced SDPA backend +# no gradient_checkpointing_disable() +# +# +# ------------------------------------------------------------------------- +# 8. V6 INFERENCE PLUMBING IS NOW WORKING +# ------------------------------------------------------------------------- +# +# ComfyUI initially rejected every AI Toolkit Qwen key: +# +# lora_te.language_model.layers.N.self_attn.q_proj.lora_A.weight +# ... +# +# A custom ComfyUI-gen2 compatibility loader was implemented. +# +# It now aliases AI Toolkit adapter BASE NAMES into ComfyUI's native CLIP +# key map and lets stock ComfyUI parse A/B tensors. +# +# Current successful log: +# +# AI-Toolkit Qwen TE keys +# detected=288 +# aliased=288 +# recognized=288 +# unrecognized=0 +# +# No lora-key-not-loaded warnings remain. +# +# Therefore do NOT treat V6's semantic failure as a simple key-loading failure. +# +# +# ------------------------------------------------------------------------- +# 9. V6 SEMANTIC RESULT — FAILURE +# ------------------------------------------------------------------------- +# +# The following fixed-style comparisons were performed visually. +# +# B: +# +# transformer/model LoRA = 1 +# Qwen/clip LoRA = 0 +# +# D: +# +# transformer/model LoRA = 1 +# Qwen/clip LoRA = 1 +# +# On a dataset prompt: +# +# B and D show no obvious visual difference. +# +# Also: +# +# model=0 +# clip=1 +# +# on a dataset prompt: +# +# no target style. +# +# And: +# +# model=1 +# clip=1 +# +# on novel prompts: +# +# no target style. +# +# Do not call V6 successful. +# +# Loader success != semantic success. +# +# +# ------------------------------------------------------------------------- +# 10. WHY V7-B EXISTS +# ------------------------------------------------------------------------- +# +# Ordinary V6 joint training has an identifiability / shortcut problem: +# +# diffusion target +# | +# +--> transformer LoRA can directly learn image appearance +# | +# +--> Qwen LoRA would need to learn a much harder semantic binding +# +# Nothing forces the optimizer to route style activation through the trigger. +# +# V7-B removes this competition. +# +# +# ------------------------------------------------------------------------- +# 11. V7-B HIGH-LEVEL DESIGN +# ------------------------------------------------------------------------- +# +# V7-B: +# +# FROZEN: +# base Qwen +# base Ideogram transformer +# V3 transformer style LoRA +# +# TRAINABLE: +# new Qwen token-aware LoRA only +# +# Target behavior: +# +# +# | +# v +# Qwen token-aware activator +# | +# v +# conditioning that behaves like a stable "painting" semantic anchor +# | +# v +# frozen V3 style carrier interprets it +# | +# v +# target style +# +# +# ------------------------------------------------------------------------- +# 12. IMPORTANT: ANCHOR IS "painting" +# ------------------------------------------------------------------------- +# +# The semantic/class anchor for V7-B is: +# +# painting +# +# Use lowercase literal: +# +# "painting" +# +# unless tokenizer/context inspection proves casing matters. +# +# Do NOT use: +# +# "Ghibli anime" +# +# anywhere in V7-B. +# +# The purpose of "painting" is NOT to describe the full target style. +# +# It is a pretrained semantic anchor / class-like reference that the random +# trigger can align toward. +# +# The actual target visual appearance continues to come primarily from: +# +# dataset images + frozen V3 style carrier. +# +# +# ------------------------------------------------------------------------- +# 13. V7-B IS TARA-LIKE, NOT A LITERAL TARA PORT +# ------------------------------------------------------------------------- +# +# Original TARA's useful principles are: +# +# TFM: +# LoRA residual is allowed only at the rare-token positions. +# +# TAL: +# rare-token representation is aligned toward a meaningful class-token +# representation. +# +# We adapt these ideas to Qwen3-VL. +# +# Ideogram4 does not expose the same SD/SDXL UNet cross-attention architecture, +# so do NOT blindly copy TARA's exact modules. +# +# +# ------------------------------------------------------------------------- +# 14. INITIAL V7-B QWEN TARGETS +# ------------------------------------------------------------------------- +# +# Initial experiment should target only: +# +# k_proj +# v_proj +# +# inside: +# +# Qwen3VLTextAttention +# +# NOT q_proj/o_proj initially. +# +# Reason: +# +# - closest analogue to TARA's K/V treatment; +# - lower-capacity, more controlled experiment; +# - trigger's K/V can influence subsequent causal Qwen tokens; +# - makes attribution easier. +# +# Expected module count: +# +# 36 × 2 = 72 LoRA modules +# +# Expected A/B tensor count: +# +# 72 × 2 = 144 tensors +# +# Keep rank: +# +# 32 +# +# initially, to avoid introducing rank as another experimental variable. +# +# +# ------------------------------------------------------------------------- +# 15. DO NOT USE GLOBAL TE `only_if_contains` IF IT ALSO AFFECTS TRANSFORMER +# ------------------------------------------------------------------------- +# +# Add a TEXT-ENCODER-SPECIFIC child-module filter, e.g.: +# +# target_te_child_modules: +# - k_proj +# - v_proj +# +# or an equivalent isolated mechanism. +# +# Do not accidentally filter the Ideogram transformer LoRA module discovery. +# +# +# ------------------------------------------------------------------------- +# 16. TOKEN FOCUS MASK — CORE V7-B MECHANISM +# ------------------------------------------------------------------------- +# +# Train-time Qwen LoRA residual: +# +# delta = B(A(x)) +# +# Instead of: +# +# y = W(x) + delta +# +# use: +# +# y = W(x) + M_trigger * delta +# +# where: +# +# M_trigger shape approximately [B, sequence_length, 1] +# +# and: +# +# 1 = token belongs to an occurrence of +# 0 = every other token +# +# Therefore: +# +# LoRA output on trigger tokens = trainable +# LoRA output on every other token = exactly zero +# +# This mask must apply to the LoRA RESIDUAL ONLY. +# +# Never mask: +# +# base Qwen output +# hidden states +# attention globally +# +# +# ------------------------------------------------------------------------- +# 17. MASK ALL SUBTOKENS OF THE TRIGGER +# ------------------------------------------------------------------------- +# +# is not assumed to be one tokenizer token. +# +# Determine the actual token span. +# +# All constituent Qwen tokens corresponding to each literal trigger occurrence +# should receive mask=1. +# +# If there are multiple trigger occurrences in the JSON caption, mask every +# occurrence. +# +# +# ------------------------------------------------------------------------- +# 18. TOKEN-SPAN DETECTION +# ------------------------------------------------------------------------- +# +# Preferred: +# +# tokenizer offset mappings +# +# Match literal character spans of: +# +# +# +# against token character offsets. +# +# A token receives mask=1 when its source-character range overlaps the trigger +# literal span. +# +# This is preferable to assuming that: +# +# tokenizer(trigger) +# +# always produces the exact same IDs in every surrounding punctuation context. +# +# If Qwen tokenizer path cannot provide offsets, build a tested fallback based +# on token-ID subsequence matching and validate it exhaustively over all 40 +# captions. +# +# +# ------------------------------------------------------------------------- +# 19. MASK VALIDATION MUST HAPPEN BEFORE TRAINING +# ------------------------------------------------------------------------- +# +# Build a dataset inspection utility. +# +# For every caption print/verify: +# +# raw trigger occurrence count +# tokenized trigger span count +# token indices for each occurrence +# number of masked tokens +# +# Assertions: +# +# every literal trigger occurrence maps to >=1 token +# no unrelated token is masked +# masked token count > 0 for every training caption +# +# Save a short summary such as: +# +# captions: 40 +# trigger occurrence distribution: ... +# unmapped occurrences: 0 +# +# Do not launch training if span detection is ambiguous. +# +# +# ------------------------------------------------------------------------- +# 20. IMPLEMENT TOKEN-AWARE LORA AS AN ISOLATED SUBCLASS +# ------------------------------------------------------------------------- +# +# Do not broadly modify generic ToolkitModuleMixin behavior. +# +# Preferred design: +# +# TokenAwareLoRAModule +# +# subclassing the existing LoRAModule or equivalent. +# +# It should only be instantiated for: +# +# Qwen TE k_proj/v_proj +# +# when V7-B token focus is enabled. +# +# Pseudocode: +# +# base = original_forward(x) +# delta = lora_forward(x) +# +# mask = network.current_te_token_mask +# +# if mask is required: +# delta = delta * broadcast(mask) +# +# return base + multiplier * scale * delta +# +# Match the installed AI Toolkit's existing LoRA scaling/multiplier semantics. +# +# Do NOT copy code from a different current-upstream version without checking +# the local repository. +# +# +# ------------------------------------------------------------------------- +# 21. RUNTIME MASK CONTEXT +# ------------------------------------------------------------------------- +# +# The network needs an ephemeral current token mask. +# +# Example concept: +# +# with network.use_te_token_mask(mask): +# encode_prompt(...) +# +# Restore previous mask in finally. +# +# Avoid stale masks crossing between prompts. +# +# Student trigger encode: +# +# mask = trigger mask +# +# Teacher/anchor encode: +# +# V7-B LoRA disabled +# +# or: +# +# zero mask +# +# but preferably explicit adapter-disabled mode. +# +# +# ------------------------------------------------------------------------- +# 22. KEY INVARIANT — NO TRIGGER MUST MEAN NO TE EFFECT +# ------------------------------------------------------------------------- +# +# This is one of V7-B's biggest advantages. +# +# If a prompt contains no trigger: +# +# M_trigger = all zeros +# +# therefore: +# +# Qwen_with_V7B(prompt) +# +# must be numerically equivalent to: +# +# base_Qwen(prompt) +# +# apart from floating-point noise. +# +# This becomes an extremely strong regression test. +# +# +# ------------------------------------------------------------------------- +# 23. ANCHOR PROMPT CONSTRUCTION +# ------------------------------------------------------------------------- +# +# For every actual training caption: +# +# Student prompt: +# +# original JSON caption containing +# +# Teacher/anchor prompt: +# +# exact same caption +# +# except each literal: +# +# +# +# is replaced in the SAME POSITION by: +# +# painting +# +# Do not otherwise rewrite or simplify the caption. +# +# This keeps: +# +# content +# composition +# object nouns +# surrounding JSON context +# +# matched between teacher and student. +# +# +# ------------------------------------------------------------------------- +# 24. PAIR TRIGGER OCCURRENCES WITH PAINTING OCCURRENCES +# ------------------------------------------------------------------------- +# +# Occurrence #1 of trigger: +# +# aligns to occurrence #1 of "painting" +# +# occurrence #2: +# +# aligns to occurrence #2 +# +# etc. +# +# Because trigger and painting may tokenize to different token counts, compare +# POOLED SPAN representations rather than individual token positions. +# +# For span S: +# +# span_repr = mean(features[:, S, :], token_dimension) +# +# This avoids requiring equal token counts. +# +# +# ------------------------------------------------------------------------- +# 25. V7-B ANCHOR ALIGNMENT LOSS +# ------------------------------------------------------------------------- +# +# Use a TAL-inspired alignment between: +# +# student trigger representation +# +# and: +# +# frozen teacher "painting" representation +# +# Preferred first implementation: +# +# align K-projection outputs. +# +# For selected Qwen layer l: +# +# K_anchor_l +# = effective BASE k_proj output for painting span +# +# K_trigger_l +# = effective base + token-aware-LoRA k_proj output +# for trigger span +# +# Teacher: +# +# no_grad +# V7-B adapter disabled +# +# Student: +# +# grad enabled +# trigger mask active +# +# Pool each span independently. +# +# Suggested loss: +# +# L_anchor = +# mean_l,occurrence( +# L1( +# normalize(K_trigger), +# normalize(K_anchor) +# ) +# ) +# +# or cosine-distance + a small magnitude-preserving L1 term. +# +# Start simple. +# +# Do NOT begin with many competing auxiliary objectives. +# +# +# ------------------------------------------------------------------------- +# 26. SAFE IMPLEMENTATION OF ANCHOR K TARGETS +# ------------------------------------------------------------------------- +# +# Qwen is quantized during training. +# +# Avoid manually doing: +# +# W_K @ hidden_state +# +# from raw quantized weights. +# +# Prefer capturing the REAL runtime k_proj output. +# +# Teacher forward: +# +# capture base k_proj outputs for painting spans +# detach() +# +# Student forward: +# +# capture effective masked-LoRA k_proj outputs +# +# This keeps quantization behavior identical to actual forward execution. +# +# +# ------------------------------------------------------------------------- +# 27. GRADIENT CHECKPOINTING CAUTION +# ------------------------------------------------------------------------- +# +# Hugging Face gradient checkpointing may recompute Qwen layers during backward. +# +# Therefore do not implement activation capture using naive append-only hooks +# that accumulate duplicate forward/recompute entries. +# +# The coding agent must explicitly design around checkpoint recomputation. +# +# Possible robust approaches: +# +# A. keyed activation slots overwritten by layer index; +# +# B. compute/return the needed representations through the existing Qwen +# feature-extraction path; +# +# C. dedicated context object keyed by: +# layer index +# occurrence +# current forward generation id +# +# Do not allow checkpoint recomputation to double-count L_anchor. +# +# +# ------------------------------------------------------------------------- +# 28. MAIN DIFFUSION LOSS REMAINS +# ------------------------------------------------------------------------- +# +# Student branch must still train against the actual dataset image target. +# +# Frozen: +# +# Ideogram base transformer +# V3 transformer LoRA +# +# Trainable: +# +# V7-B Qwen K/V token-aware LoRA +# +# Continue ordinary Ideogram diffusion loss: +# +# L_diff +# +# This teaches the activator to steer the frozen style carrier toward the actual +# training images. +# +# +# ------------------------------------------------------------------------- +# 29. TOTAL INITIAL LOSS +# ------------------------------------------------------------------------- +# +# First V7-B experiment: +# +# L_total = +# L_diff +# + +# lambda_anchor * L_anchor +# +# Start: +# +# lambda_anchor = 1.0 +# +# but expose it in YAML/UI/config. +# +# Do not hardcode. +# +# Log both losses independently: +# +# loss_diff +# loss_anchor +# loss_total +# +# +# ------------------------------------------------------------------------- +# 30. DO NOT ADD A TRAINABLE TRANSFORMER RESIDUAL YET +# ------------------------------------------------------------------------- +# +# Very important: +# +# train_unet: false +# +# for the NEW V7-B adapter. +# +# V3 is loaded frozen. +# +# If V7-B cannot alter behavior when it is the only trainable route, we learn +# something useful. +# +# If transformer LoRA is trainable again immediately, we recreate V6's +# shortcut problem. +# +# +# ------------------------------------------------------------------------- +# 31. INITIAL CONFIG DIRECTION +# ------------------------------------------------------------------------- +# +# Example conceptual config: +# +# trigger_word: "" +# +# network: +# type: lora +# linear: 32 +# linear_alpha: 16 +# network_kwargs: +# target_te_modules: +# - Qwen3VLTextAttention +# target_te_child_modules: +# - k_proj +# - v_proj +# te_token_focus: true +# +# train: +# batch_size: 2 +# train_unet: false +# train_text_encoder: false +# train_text_encoder_lora: true +# gradient_checkpointing: true +# cache_text_embeddings: false +# unload_text_encoder: false +# timestep_type: sigmoid +# content_or_style: content +# dtype: bf16 +# do_differential_guidance: true +# differential_guidance_scale: 3.0 +# +# text_encoder_lr: 5e-5 +# +# model: +# quantize: true +# qtype: qfloat8 +# quantize_te: true +# qtype_te: qfloat8 +# +# v7b: +# token_focus: true +# trigger_word: "" +# anchor_word: "painting" +# anchor_loss_weight: 1.0 +# +# These are STARTING values, not sacred constants. +# +# text_encoder_lr=5e-5 is a controlled initial increase from V6's 2e-5 now +# that Qwen is the only trainable branch. +# +# Do not immediately run 2000 steps. +# +# +# ------------------------------------------------------------------------- +# 32. LOAD V3 AS ACTIVE + FROZEN +# ------------------------------------------------------------------------- +# +# Earlier Ideogram4 code already acquired support for loading a frozen +# transformer LoRA via model_config.lora_path / load_frozen_lora(). +# +# Reuse this if still present. +# +# Required state: +# +# V3 parameters requires_grad=False +# V3 active during student diffusion forward +# V3 optimizer params = 0 +# +# Print at startup: +# +# V3 modules: +# V3 total params: +# V3 trainable params: 0 +# +# +# ------------------------------------------------------------------------- +# 33. PARAMETER-STATE PROOF BEFORE TRAINING +# ------------------------------------------------------------------------- +# +# Required: +# +# Qwen base: +# requires_grad = false +# +# Ideogram base transformer: +# requires_grad = false +# +# V3 transformer LoRA: +# requires_grad = false +# +# V7-B Qwen K/V LoRA: +# requires_grad = true +# +# There should be NO other trainable parameter group. +# +# +# ------------------------------------------------------------------------- +# 34. STEP-0 TOKEN-FOCUS PROOF +# ------------------------------------------------------------------------- +# +# For one real caption, inspect one Qwen layer. +# +# Compute LoRA residual norm: +# +# trigger_positions_norm +# non_trigger_positions_norm +# +# Requirement: +# +# trigger norm can be nonzero +# +# non-trigger norm must equal 0 +# except insignificant numerical noise if implementation requires casting +# +# Ideally: +# +# non_trigger_max_abs == 0 +# +# because masking is explicit multiplication. +# +# +# ------------------------------------------------------------------------- +# 35. GRADIENT PROOF +# ------------------------------------------------------------------------- +# +# After backward, before optimizer step: +# +# Verify: +# +# V7-B k_proj LoRA grad != None +# V7-B v_proj LoRA grad != None +# grad norm > 0 +# +# and: +# +# V3 grad is None +# Qwen base grad is None +# Ideogram base grad is None +# +# Do this once, then remove noisy diagnostics. +# +# +# ------------------------------------------------------------------------- +# 36. ANCHOR-LOSS PROOF +# ------------------------------------------------------------------------- +# +# Before real training: +# +# For a sample: +# +# L_anchor must be finite +# trigger spans detected > 0 +# painting spans detected == trigger occurrence count +# +# Check that random/init V7-B does NOT already give zero anchor loss. +# +# Otherwise the objective is ineffective. +# +# +# ------------------------------------------------------------------------- +# 37. SHORT TRAINING SCHEDULE / KILL GATES +# ------------------------------------------------------------------------- +# +# Suggested first run: +# +# save step 25 +# save step 50 +# save step 100 +# save step 200 +# +# Do NOT start with a blind 2000-step commitment. +# +# +# ------------------------------------------------------------------------- +# 38. PRIMARY SEMANTIC EVALUATION MATRIX +# ------------------------------------------------------------------------- +# +# At each checkpoint evaluate: +# +# 1. Dataset prompt, trigger ON +# +# 2. Same dataset content, trigger REMOVED +# +# 3. Novel neutral prompt, trigger ON +# +# 4. Same novel neutral prompt, trigger OFF +# +# 5. Novel prompt containing a competing style phrase, trigger ON +# +# 6. Same competing prompt, trigger OFF +# +# Use: +# +# identical seed +# identical sampler +# identical steps +# identical V3 strength +# +# +# ------------------------------------------------------------------------- +# 39. MOST IMPORTANT SUCCESS CONDITION +# ------------------------------------------------------------------------- +# +# V7-B is NOT successful merely because dataset reconstruction improves. +# +# Main success: +# +# novel prompt + trigger +# +# must show materially more target-style activation than: +# +# same novel prompt without trigger. +# +# Also test competing style semantics. +# +# +# ------------------------------------------------------------------------- +# 40. NO-TRIGGER EXACT-PRESERVATION TEST +# ------------------------------------------------------------------------- +# +# Because TFM zeros the V7-B residual when no trigger exists: +# +# compare Qwen conditioning: +# +# V3 + V7B loaded, prompt WITHOUT trigger +# +# versus: +# +# V3 + no V7B, same prompt +# +# Required: +# +# relative difference approximately zero. +# +# This should be evaluated numerically, not visually. +# +# +# ------------------------------------------------------------------------- +# 41. NUMERICAL CONDITIONING DIAGNOSTIC +# ------------------------------------------------------------------------- +# +# We did not yet have a convenient ComfyUI conditioning comparison tool. +# +# Build one during this project. +# +# For two conditioning tensors A/B report: +# +# L2 norm A +# L2 norm B +# delta L2 +# mean(abs(delta)) +# max(abs(delta)) +# relative L2: +# +# ||B-A|| / ||A|| +# +# cosine similarity +# +# This prevents future visual ambiguity. +# +# +# ------------------------------------------------------------------------- +# 42. ANCHOR SIMILARITY METRIC +# ------------------------------------------------------------------------- +# +# During training log: +# +# cosine(trigger_span_K, painting_span_K) +# +# before and after V7-B. +# +# We want the similarity to move in the intended direction. +# +# But do not accept anchor similarity alone as success. +# +# Final criterion is generated behavior on novel prompts. +# +# +# ------------------------------------------------------------------------- +# 43. TOKEN-FOCUS LEAKAGE METRIC +# ------------------------------------------------------------------------- +# +# Compute: +# +# sum(abs(LoRA residual on non-trigger tokens)) +# ------------------------------------------------ +# sum(abs(all LoRA residuals)) +# +# Expected: +# +# ~0 +# +# by construction. +# +# +# ------------------------------------------------------------------------- +# 44. IF STEP 100 SHOWS ZERO BEHAVIORAL DIFFERENCE +# ------------------------------------------------------------------------- +# +# First inspect: +# +# gradient norms +# anchor loss trajectory +# conditioning delta +# trigger K/painting K similarity +# +# Do NOT immediately train longer. +# +# Cases: +# +# A. +# Anchor loss not decreasing: +# +# optimization / implementation problem. +# +# B. +# Anchor loss decreases strongly but Qwen conditioning barely changes: +# +# loss may be aligned to a representation that Ideogram does not use +# strongly enough. +# +# C. +# Conditioning changes materially but output does not: +# +# V3 may not decode the chosen "painting" direction sufficiently. +# +# D. +# Dataset output changes but novel output does not: +# +# still overfitting / insufficient semantic generalization. +# +# +# ------------------------------------------------------------------------- +# 45. IMPORTANT FALLBACK, BUT DO NOT IMPLEMENT INITIALLY +# ------------------------------------------------------------------------- +# +# If TFM + K anchor alignment clearly optimizes but does not produce visual +# activation, the next controlled extension is: +# +# teacher diffusion-prediction distillation +# +# Teacher: +# +# frozen Qwen +# anchor prompt with "painting" +# frozen V3 +# +# Student: +# +# trigger prompt +# V7-B token-aware Qwen +# frozen V3 +# +# Same: +# +# latent +# noise +# timestep +# +# Additional loss: +# +# MSE( +# student_velocity, +# teacher_velocity.detach() +# ) +# +# Do NOT add this in the first V7-B implementation unless necessary. +# +# We want to know whether the simpler TFM + TAL-style anchor works first. +# +# +# ------------------------------------------------------------------------- +# 46. COMFYUI CONSEQUENCE — VERY IMPORTANT +# ------------------------------------------------------------------------- +# +# V7-B Qwen TE LoRA CANNOT be deployed as an ordinary static ComfyUI LoRA +# patch if token focus is part of the model function. +# +# Standard ModelPatcher weight LoRA implements effectively: +# +# W' = W + delta_W +# +# globally for every token. +# +# V7-B requires: +# +# output = +# W(x) +# + +# M_trigger(x_tokens) * delta_W(x) +# +# which is input/token dependent. +# +# Therefore inference also needs runtime token-aware behavior. +# +# Original TFM-style masking must be active at inference too. +# +# +# ------------------------------------------------------------------------- +# 47. CURRENT COMFYUI CUSTOM LOADER +# ------------------------------------------------------------------------- +# +# Existing custom node: +# +# ComfyUI-gen2/gen2_sampling/ideogram4_aitk_lora.py +# +# currently solves namespace compatibility for ordinary AI Toolkit Qwen LoRA. +# +# It aliases: +# +# lora_te.language_model.layers... +# +# into ComfyUI's Qwen CLIP targets. +# +# For V7-B this loader will need an additional TOKEN-AWARE TE path. +# +# +# ------------------------------------------------------------------------- +# 48. DO NOT THROW AWAY STOCK COMFYUI MODEL PATCHING FOR TRANSFORMER +# ------------------------------------------------------------------------- +# +# V3 / model-side LoRA can remain stock ComfyUI ModelPatcher behavior. +# +# Only the V7-B Qwen K/V adapter needs dynamic masking. +# +# +# ------------------------------------------------------------------------- +# 49. RECOMMENDED R&D ARTIFACT STRATEGY +# ------------------------------------------------------------------------- +# +# During V7-B development, keep: +# +# V3 transformer style carrier +# +# and: +# +# V7-B Qwen activator +# +# as SEPARATE files. +# +# This makes experiments clean: +# +# V3 only +# V7B only +# V3 + V7B +# +# It also prevents accidental retraining/merging confusion. +# +# Only package them together after semantic success has been demonstrated. +# +# +# ------------------------------------------------------------------------- +# 50. V7-B SAVE METADATA +# ------------------------------------------------------------------------- +# +# Add metadata if possible: +# +# token_aware_te_lora = true +# trigger_word = "" +# anchor_word = "painting" +# token_mask_mode = "literal_trigger_span" +# target_te_projections = "k_proj,v_proj" +# architecture = "ideogram4_qwen3vl" +# +# This lets the ComfyUI custom loader detect that this is NOT a normal static +# TE LoRA. +# +# +# ------------------------------------------------------------------------- +# 51. COMFYUI TOKEN-MASK IMPLEMENTATION +# ------------------------------------------------------------------------- +# +# Do not implement this before AI Toolkit V7-B shows semantic promise. +# +# Once training works: +# +# modify the custom ComfyUI loader so Qwen K/V LoRA residual is applied at +# runtime only to trigger token positions. +# +# Need to inspect current: +# +# comfy/text_encoders/ideogram4.py +# comfy/text_encoders/qwen3vl.py +# CLIP tokenization / encode_from_tokens path +# +# The runtime adapter must have access to the current token IDs / mask before +# Qwen forward. +# +# +# ------------------------------------------------------------------------- +# 52. DO NOT USE TRAIN-TIME MASK + UNMASKED INFERENCE +# ------------------------------------------------------------------------- +# +# This is not acceptable. +# +# If the LoRA is trained with: +# +# delta only on trigger tokens +# +# but inference applies: +# +# delta to every token +# +# then inference is a different function and reintroduces exactly the +# token-leakage problem V7-B is meant to remove. +# +# +# ------------------------------------------------------------------------- +# 53. IDEOGRAM CONDITIONAL / UNCONDITIONAL INFERENCE +# ------------------------------------------------------------------------- +# +# Existing practical setup: +# +# Conditional branch: +# +# V3 MODEL adaptation +# + +# V7-B token-aware Qwen CLIP adaptation +# +# Unconditional diffusion branch: +# +# V3 MODEL adaptation +# +# No separate Qwen encoding is required specifically for the unconditional +# diffusion model in the standard Ideogram4 ComfyUI workflow. +# +# Do not apply the V7-B Qwen adapter twice. +# +# +# ------------------------------------------------------------------------- +# 54. DO NOT CONFUSE THE OFFICIAL IDEOGRAM UNCONDITIONAL LORA +# ------------------------------------------------------------------------- +# +# Current AI Toolkit model config also uses: +# +# ostris/ideogram_4_unconditional_lora/... +# +# This is a separate Ideogram unconditional adapter. +# +# It is NOT V3 and NOT V7-B. +# +# Keep those concepts separate. +# +# +# ------------------------------------------------------------------------- +# 55. FIRST CODING PHASE +# ------------------------------------------------------------------------- +# +# Before editing anything, inspect current working files: +# +# toolkit/config_modules.py +# toolkit/lora_special.py +# toolkit/network_mixins.py +# jobs/process/BaseSDTrainProcess.py +# extensions_built_in/sd_trainer/SDTrainer.py +# extensions_built_in/diffusion_models/ideogram4/ideogram4.py +# extensions_built_in/diffusion_models/ideogram4/src/pipeline.py +# +# Identify the smallest integration points. +# +# Prefer ADDITIVE / isolated implementation. +# +# +# ------------------------------------------------------------------------- +# 56. SUGGESTED CODE FOOTPRINT +# ------------------------------------------------------------------------- +# +# Prefer something approximately: +# +# toolkit/config_modules.py +# + V7-B config fields +# +# toolkit/lora_special.py +# + text-encoder-specific child-module target support +# + instantiate token-aware TE LoRA class when requested +# +# new: +# toolkit/token_aware_lora.py +# + TokenAwareLoRAModule +# + runtime token-mask context +# +# Ideogram pipeline/tokenizer path +# + robust trigger-span mask generation +# + optional anchor prompt construction +# +# SDTrainer.py +# + teacher anchor encode +# + L_anchor +# + total-loss integration +# +# Avoid broad changes to: +# +# generic network_mixins.py +# +# if a subclass can do the job. +# +# +# ------------------------------------------------------------------------- +# 57. TESTS REQUIRED BEFORE REAL TRAINING +# ------------------------------------------------------------------------- +# +# Unit: +# +# [ ] literal trigger character span → correct token mask +# +# [ ] multiple trigger occurrences → all detected +# +# [ ] prompt without trigger → zero mask +# +# [ ] painting replacement preserves surrounding text +# +# [ ] trigger and painting span pairing correct +# +# [ ] TokenAwareLoRAModule masks residual only +# +# [ ] non-trigger residual exactly zero +# +# [ ] K/V only are adapted +# +# [ ] q_proj/o_proj untouched +# +# +# Integration: +# +# [ ] Qwen base frozen +# +# [ ] V3 frozen +# +# [ ] V7-B optimizer contains only 72 K/V LoRA modules +# +# [ ] student encode grad works +# +# [ ] teacher encode has no grad +# +# [ ] scalar TE multiplier fix still works with batch=2 +# +# [ ] Qwen checkpointing works +# +# [ ] anchor loss finite +# +# [ ] backward gives V7-B gradients +# +# +# ------------------------------------------------------------------------- +# 58. FIRST TRAINING KILL GATE +# ------------------------------------------------------------------------- +# +# Run only enough to obtain: +# +# step 25 / 50 / 100 +# +# At step 100, stop and evaluate. +# +# Continue only if at least one of these is clearly true: +# +# trigger materially changes conditioning; +# trigger materially changes V3 output; +# anchor similarity is moving meaningfully; +# novel-prompt target activation starts appearing. +# +# If literally nothing differs from V3-only: +# +# do not continue to 2000. +# +# +# ------------------------------------------------------------------------- +# 59. WHAT SUCCESS WOULD LOOK LIKE +# ------------------------------------------------------------------------- +# +# Best-case V7-B: +# +# WITHOUT trigger: +# +# normal V3/base behavior +# +# WITH : +# +# target style strongly activates +# +# ON novel content: +# +# target style persists +# +# WITH competing style words: +# +# trigger meaningfully resists competing semantics +# +# AND: +# +# content/prompt structure remains usable. +# +# +# ------------------------------------------------------------------------- +# 60. RESEARCH REFERENCE +# ------------------------------------------------------------------------- +# +# The conceptual inspiration is TARA: +# +# Token-Aware LoRA for Composable Personalization in Diffusion Models +# Peng et al. +# arXiv:2508.08812 / AAAI 2026 +# +# Relevant concepts: +# +# Token Focus Masking +# Token Alignment Loss +# +# But this project is an adaptation for: +# +# Ideogram4 +# Qwen3-VL +# single-stream diffusion conditioning +# style activation +# +# not a direct SD/SDXL reproduction. +# +# +# ------------------------------------------------------------------------- +# 61. FINAL DESIGN SUMMARY +# ------------------------------------------------------------------------- +# +# V3: +# +# frozen visual style carrier +# +# V7-B: +# +# Qwen K/V LoRA +# rank 32 +# trigger-token residual masking +# only active on token spans +# +# Anchor: +# +# "painting" +# +# Loss: +# +# actual dataset diffusion loss +# + +# TAL-like trigger → painting K alignment +# +# Training: +# +# no trainable transformer LoRA +# +# Inference: +# +# trigger-aware Qwen runtime masking REQUIRED +# +# Evaluation: +# +# novel prompts are the main criterion +# not training-image reconstruction +# +# Key philosophical goal: +# +# Do not merely teach the model the style again. +# Teach the random trigger to become the CONTROL SIGNAL for the style +# that V3 already knows how to render. +# +# END HANDOFF \ No newline at end of file diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index ef415dc5..864c96f8 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 + DecoratorConfig, TriggerSelectiveTrainingConfig from toolkit.logging_aitk import create_logger from diffusers import FluxTransformer2DModel from toolkit.accelerator import get_accelerator, unwrap_model @@ -110,6 +110,9 @@ class BaseSDTrainProcess(BaseTrainProcess): else: self.network_config = None self.train_config = TrainConfig(**self.get_conf('train', {})) + self.trigger_selective_training = TriggerSelectiveTrainingConfig( + **self.get_conf('trigger_selective_training', {}) + ) model_config = self.get_conf('model', {}) self.modules_being_trained: List[torch.nn.Module] = [] @@ -258,7 +261,15 @@ class BaseSDTrainProcess(BaseTrainProcess): self.snr_gos: Union[LearnableSNRGamma, None] = None self.ema: ExponentialMovingAverage = None - validate_configs(self.train_config, self.model_config, self.save_config, self.dataset_configs) + validate_configs( + self.train_config, + self.model_config, + self.save_config, + self.dataset_configs, + self.trigger_selective_training, + self.trigger_word, + self.network_config, + ) do_profiler = self.get_conf('torch_profiler', False) self.torch_profiler = None if not do_profiler else torch.profiler.profile( diff --git a/testing/test_trigger_selective_training.py b/testing/test_trigger_selective_training.py new file mode 100644 index 00000000..429ca9a2 --- /dev/null +++ b/testing/test_trigger_selective_training.py @@ -0,0 +1,125 @@ +import random +import unittest + +import torch + +from toolkit.config_modules import TriggerSelectiveTrainingConfig +from toolkit.trigger_selective_training import ( + apply_differential_guidance_target, + get_scheduled_loss_weights, + get_scheduled_margin, + network_disabled, + normalized_gain, + resolve_prompt_variants, + sample_negative_styles, + trigger_advantage_hinge, + validate_trigger_selective_config, +) + + +class _Network: + def __init__(self): + self.is_active = True + + +class TriggerSelectiveTrainingTest(unittest.TestCase): + def setUp(self): + self.config = TriggerSelectiveTrainingConfig( + enabled=True, + negative_styles={ + 'expected_category_count': 2, + 'categories': [ + {'name': 'neutral', 'probability': 0.5, 'phrases': ['']}, + {'name': 'hard', 'probability': 0.5, 'phrases': ['painting', 'illustration']}, + ], + }, + path3={ + 'margin_schedule': { + 'interpolation': 'linear', + 'keyframes': [{'step': 0, 'value': 0.02}, {'step': 100, 'value': 0.12}], + }, + }, + loss_schedule={ + 'interpolation': 'linear', + 'keyframes': [ + {'step': 0, 'path1': 0.8, 'path2': 0.1, 'path3': 0.1}, + {'step': 100, 'path1': 0.6, 'path2': 0.15, 'path3': 0.25}, + ], + }, + ) + + def test_validation_and_schedule_clamping(self): + validate_trigger_selective_config(self.config, '') + self.assertAlmostEqual(get_scheduled_margin(self.config, 0), 0.02) + self.assertAlmostEqual(get_scheduled_margin(self.config, 50), 0.07) + self.assertAlmostEqual(get_scheduled_margin(self.config, 1000), 0.12) + weights = get_scheduled_loss_weights(self.config, 50) + self.assertAlmostEqual(sum(weights.values()), 1.0) + self.assertAlmostEqual(weights['path1'], 0.7) + + def test_sampling_and_shared_placeholder_resolution(self): + samples = sample_negative_styles(self.config, 8, random.Random(3)) + trigger, decoy = resolve_prompt_variants( + ['a [trigger] portrait [trigger]'] * 8, + '', + samples, + ) + self.assertTrue(all(prompt.count('') == 2 for prompt in trigger)) + self.assertTrue(all('[trigger]' not in prompt for prompt in decoy)) + self.assertEqual(len(trigger), len(decoy)) + + def test_gain_and_hinge_stop_gradient_on_decoy_gain(self): + student_loss = torch.tensor([2.0], requires_grad=True) + base_loss = torch.tensor([4.0]) + decoy_gain = normalized_gain(student_loss, base_loss, 1.0e-6) + trigger_gain = torch.tensor([0.01], requires_grad=True) + loss = trigger_advantage_hinge(trigger_gain, decoy_gain, 0.1).sum() + loss.backward() + self.assertIsNotNone(trigger_gain.grad) + self.assertIsNone(student_loss.grad) + + def test_differential_guidance_target_is_detached_and_shared(self): + class _Config: + do_guidance_loss = True + do_differential_guidance = True + differential_guidance_scale = 3.0 + + class _Trainer: + train_config = _Config() + + target = torch.ones(1, 2) + prediction = torch.zeros(1, 2, requires_grad=True) + shared = apply_differential_guidance_target(_Trainer(), target, prediction) + self.assertFalse(shared.requires_grad) + self.assertTrue(torch.equal(shared, torch.full_like(target, 3.0))) + + def test_differential_guidance_preserves_effective_v3_behavior(self): + class _Config: + do_guidance_loss = False + do_differential_guidance = True + differential_guidance_scale = 3.0 + + class _Trainer: + train_config = _Config() + + target = torch.ones(1, 2) + prediction = torch.zeros(1, 2, requires_grad=True) + shared = apply_differential_guidance_target(_Trainer(), target, prediction) + self.assertIs(shared, target) + + def test_network_state_is_restored(self): + network = _Network() + with network_disabled(network): + self.assertFalse(network.is_active) + self.assertTrue(network.is_active) + network.is_active = False + with network_disabled(network): + self.assertFalse(network.is_active) + self.assertFalse(network.is_active) + + def test_disabled_config_does_not_require_tst_fields(self): + validate_trigger_selective_config(TriggerSelectiveTrainingConfig(enabled=False), None) + + +if __name__ == '__main__': + unittest.main() diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index a3a51138..8b0694f8 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -372,6 +372,71 @@ ContentOrStyleType = Literal['balanced', 'style', 'content'] LossTarget = Literal['noise', 'source', 'unaugmented', 'differential_noise'] +class TriggerSelectiveNegativeStyleCategoryConfig: + def __init__(self, **kwargs): + self.name: str = kwargs.get('name', '') + self.probability: float = float(kwargs.get('probability', 0.0)) + self.phrases: List[str] = kwargs.get('phrases', []) + + +class TriggerSelectiveNegativeStylesConfig: + def __init__(self, **kwargs): + self.expected_category_count: Optional[int] = kwargs.get('expected_category_count', None) + self.sample_scope: str = kwargs.get('sample_scope', 'per_item') + categories = kwargs.get('categories', []) + self.categories: List[TriggerSelectiveNegativeStyleCategoryConfig] = [ + TriggerSelectiveNegativeStyleCategoryConfig(**category) for category in categories + ] + + +class TriggerSelectiveMarginScheduleConfig: + def __init__(self, **kwargs): + self.interpolation: str = kwargs.get('interpolation', 'smoothstep') + self.keyframes: List[Dict] = kwargs.get('keyframes', []) + + +class TriggerSelectivePath3Config: + def __init__(self, **kwargs): + self.loss_type: str = kwargs.get('loss_type', 'hinge') + self.gain_epsilon: float = float(kwargs.get('gain_epsilon', 1.0e-6)) + self.margin_schedule = TriggerSelectiveMarginScheduleConfig( + **kwargs.get('margin_schedule', {}) + ) + + +class TriggerSelectiveLossScheduleConfig: + def __init__(self, **kwargs): + self.interpolation: str = kwargs.get('interpolation', 'smoothstep') + self.normalize_weights: bool = kwargs.get('normalize_weights', True) + self.keyframes: List[Dict] = kwargs.get('keyframes', []) + + +class TriggerSelectiveLoggingConfig: + def __init__(self, **kwargs): + self.log_every: int = int(kwargs.get('log_every', 25)) + self.log_category_stats: bool = kwargs.get('log_category_stats', True) + self.log_gain_stats: bool = kwargs.get('log_gain_stats', True) + self.debug_gradient_contributions: bool = kwargs.get('debug_gradient_contributions', False) + self.gradient_diagnostic_steps: List[int] = kwargs.get( + 'gradient_diagnostic_steps', [1, 10, 25, 50, 100, 250, 500, 1000, 1500] + ) + self.metrics_filename: str = kwargs.get('metrics_filename', 'tst_metrics.jsonl') + + +class TriggerSelectiveTrainingConfig: + def __init__(self, **kwargs): + self.enabled: bool = kwargs.get('enabled', False) + self.require_trigger_placeholder: bool = kwargs.get('require_trigger_placeholder', True) + self.negative_styles = TriggerSelectiveNegativeStylesConfig( + **kwargs.get('negative_styles', {}) + ) + self.path3 = TriggerSelectivePath3Config(**kwargs.get('path3', {})) + self.loss_schedule = TriggerSelectiveLossScheduleConfig( + **kwargs.get('loss_schedule', {}) + ) + self.logging = TriggerSelectiveLoggingConfig(**kwargs.get('logging', {})) + + class TrainConfig: def __init__(self, **kwargs): self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm') @@ -1472,7 +1537,10 @@ def validate_configs( train_config: TrainConfig, model_config: ModelConfig, save_config: SaveConfig, - dataset_configs: List[DatasetConfig] + dataset_configs: List[DatasetConfig], + trigger_selective_training: Optional[TriggerSelectiveTrainingConfig] = None, + trigger_word: Optional[str] = None, + network_config: Optional[NetworkConfig] = None, ): if model_config.is_flux: if save_config.save_format != 'diffusers': @@ -1505,6 +1573,30 @@ def validate_configs( if train_config.diff_output_preservation and train_config.blank_prompt_preservation: raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.") - + + if trigger_selective_training is not None and trigger_selective_training.enabled: + if train_config.diff_output_preservation: + raise ValueError( + 'trigger_selective_training cannot be combined with legacy diff_output_preservation in v1.' + ) + if model_config.arch != 'ideogram4': + raise ValueError('trigger_selective_training v1 is restricted to model.arch: ideogram4') + if network_config is None or network_config.type != 'lora': + raise ValueError('trigger_selective_training v1 requires network.type: lora') + if train_config.train_text_encoder: + raise ValueError('trigger_selective_training v1 does not support training the text encoder') + if train_config.cache_text_embeddings or any(dataset.cache_text_embeddings for dataset in dataset_configs): + raise ValueError('trigger_selective_training requires dynamic prompt encoding; disable text embedding caching') + if train_config.unload_text_encoder: + raise ValueError('trigger_selective_training requires the text encoder to remain loaded') + if train_config.do_guidance_loss or train_config.train_turbo or train_config.loss_target != 'noise': + raise ValueError( + 'trigger_selective_training v1 requires do_guidance_loss=false, train_turbo=false and loss_target=noise' + ) + if train_config.single_item_batching: + raise ValueError('trigger_selective_training v1 does not support single_item_batching') + from toolkit.trigger_selective_training import validate_trigger_selective_config + validate_trigger_selective_config(trigger_selective_training, trigger_word) + 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/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index c1a46451..638dec64 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -169,6 +169,7 @@ class FileItemDTO( # self.caption_path: str = kwargs.get('caption_path', None) self.raw_caption: str = kwargs.get("raw_caption", None) + self.caption_template: str = kwargs.get("caption_template", None) # we scale first, then crop self.scale_to_width: int = kwargs.get( "scale_to_width", int(self.width * self.dataset_config.scale) diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 2a58bba8..3d4dc2b8 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -321,6 +321,8 @@ class CaptionProcessingDTOMixin: self.raw_caption_short: str = None self.caption: str = None self.caption_short: str = None + self.caption_template: str = None + self.caption_short_template: str = None # caption with the trigger word replaced by the diff output preservation class self.caption_dop: str = None @@ -432,6 +434,10 @@ class CaptionProcessingDTOMixin: # join back together caption = ', '.join(token_list) + if short_caption: + self.caption_short_template = caption + else: + self.caption_template = caption caption = inject_trigger_into_prompt(caption, trigger, to_replace_list, add_if_not_present) if self.dataset_config.random_triggers: diff --git a/toolkit/trigger_selective_training.py b/toolkit/trigger_selective_training.py new file mode 100644 index 00000000..84c911a4 --- /dev/null +++ b/toolkit/trigger_selective_training.py @@ -0,0 +1,284 @@ +import json +import math +import os +import random +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import torch +import torch.nn.functional as F + +from toolkit.config_modules import TriggerSelectiveTrainingConfig +from toolkit.prompt_utils import inject_trigger_into_prompt + + +_ALLOWED_INTERPOLATIONS = {'linear', 'smoothstep'} + + +@dataclass(frozen=True) +class NegativeStyleSample: + category: str + phrase: str + + +def _validate_keyframes(keyframes: Sequence[Dict], value_keys: Sequence[str], name: str): + if not keyframes: + raise ValueError(f'{name} must contain at least one keyframe') + previous_step = None + for index, keyframe in enumerate(keyframes): + if 'step' not in keyframe: + raise ValueError(f'{name} keyframe {index} is missing step') + step = int(keyframe['step']) + if step < 0: + raise ValueError(f'{name} keyframe steps must be non-negative') + if previous_step is not None and step <= previous_step: + raise ValueError(f'{name} keyframe steps must be strictly increasing') + previous_step = step + for key in value_keys: + if key not in keyframe: + raise ValueError(f'{name} keyframe {index} is missing {key}') + value = float(keyframe[key]) + if not math.isfinite(value): + raise ValueError(f'{name} keyframe {index} has a non-finite {key}') + if value < 0: + raise ValueError(f'{name} keyframe {index} has a negative {key}') + + +def validate_trigger_selective_config( + config: TriggerSelectiveTrainingConfig, + trigger_word: Optional[str], +): + if not config.enabled: + return + if not trigger_word or not trigger_word.strip(): + raise ValueError('trigger_selective_training requires a non-empty trigger_word') + + negative_styles = config.negative_styles + if negative_styles.sample_scope != 'per_item': + raise ValueError("trigger_selective_training negative_styles.sample_scope must be 'per_item' in v1") + if not negative_styles.categories: + raise ValueError('trigger_selective_training requires at least one negative style category') + if ( + negative_styles.expected_category_count is not None + and negative_styles.expected_category_count != len(negative_styles.categories) + ): + raise ValueError( + 'trigger_selective_training negative_styles.expected_category_count does not match categories' + ) + + names = set() + probability_sum = 0.0 + for category in negative_styles.categories: + if not category.name or not category.name.strip(): + raise ValueError('trigger_selective_training category names must be non-empty') + if category.name in names: + raise ValueError(f'duplicate trigger_selective_training category name: {category.name}') + names.add(category.name) + if not math.isfinite(category.probability) or category.probability < 0: + raise ValueError(f'invalid probability for trigger_selective_training category {category.name}') + probability_sum += category.probability + if not category.phrases: + raise ValueError(f'trigger_selective_training category {category.name} has no phrases') + for phrase in category.phrases: + if not isinstance(phrase, str): + raise ValueError(f'trigger_selective_training phrases in {category.name} must be strings') + if '[trigger]' in phrase or '[name]' in phrase: + raise ValueError(f'trigger_selective_training phrase in {category.name} contains a placeholder') + if trigger_word in phrase: + raise ValueError(f'trigger_selective_training phrase in {category.name} contains trigger_word') + if not math.isclose(probability_sum, 1.0, rel_tol=0.0, abs_tol=1.0e-6): + raise ValueError('trigger_selective_training category probabilities must sum to 1.0') + + if config.path3.loss_type != 'hinge': + raise ValueError("trigger_selective_training path3.loss_type must be 'hinge' in v1") + if not math.isfinite(config.path3.gain_epsilon) or config.path3.gain_epsilon <= 0: + raise ValueError('trigger_selective_training path3.gain_epsilon must be positive') + + for interpolation, name in ( + (config.path3.margin_schedule.interpolation, 'path3.margin_schedule'), + (config.loss_schedule.interpolation, 'loss_schedule'), + ): + if interpolation not in _ALLOWED_INTERPOLATIONS: + raise ValueError(f'{name}.interpolation must be linear or smoothstep') + + _validate_keyframes( + config.path3.margin_schedule.keyframes, + ('value',), + 'trigger_selective_training path3.margin_schedule', + ) + _validate_keyframes( + config.loss_schedule.keyframes, + ('path1', 'path2', 'path3'), + 'trigger_selective_training loss_schedule', + ) + for keyframe in config.loss_schedule.keyframes: + if sum(float(keyframe[key]) for key in ('path1', 'path2', 'path3')) <= 0: + raise ValueError('trigger_selective_training loss schedule weights cannot all be zero') + + if config.logging.log_every <= 0: + raise ValueError('trigger_selective_training logging.log_every must be positive') + if not config.logging.metrics_filename or os.path.basename(config.logging.metrics_filename) != config.logging.metrics_filename: + raise ValueError('trigger_selective_training logging.metrics_filename must be a filename') + if any(step < 0 for step in config.logging.gradient_diagnostic_steps): + raise ValueError('trigger_selective_training gradient diagnostic steps must be non-negative') + + +def _interpolation_fraction(fraction: float, interpolation: str) -> float: + fraction = min(max(float(fraction), 0.0), 1.0) + if interpolation == 'linear': + return fraction + if interpolation == 'smoothstep': + return fraction * fraction * (3.0 - 2.0 * fraction) + raise ValueError(f'unsupported interpolation: {interpolation}') + + +def interpolate_keyframes( + keyframes: Sequence[Dict], + step: int, + value_keys: Sequence[str], + interpolation: str, +) -> Dict[str, float]: + if step <= int(keyframes[0]['step']): + return {key: float(keyframes[0][key]) for key in value_keys} + if step >= int(keyframes[-1]['step']): + return {key: float(keyframes[-1][key]) for key in value_keys} + + for left, right in zip(keyframes, keyframes[1:]): + left_step = int(left['step']) + right_step = int(right['step']) + if left_step <= step <= right_step: + fraction = (step - left_step) / (right_step - left_step) + fraction = _interpolation_fraction(fraction, interpolation) + return { + key: float(left[key]) + fraction * (float(right[key]) - float(left[key])) + for key in value_keys + } + raise RuntimeError('could not interpolate keyframes') + + +def get_scheduled_margin(config: TriggerSelectiveTrainingConfig, step: int) -> float: + return interpolate_keyframes( + config.path3.margin_schedule.keyframes, + step, + ('value',), + config.path3.margin_schedule.interpolation, + )['value'] + + +def get_scheduled_loss_weights(config: TriggerSelectiveTrainingConfig, step: int) -> Dict[str, float]: + weights = interpolate_keyframes( + config.loss_schedule.keyframes, + step, + ('path1', 'path2', 'path3'), + config.loss_schedule.interpolation, + ) + if config.loss_schedule.normalize_weights: + total = sum(weights.values()) + if total <= 0: + raise ValueError('trigger_selective_training interpolated loss weights sum to zero') + weights = {key: value / total for key, value in weights.items()} + return weights + + +def sample_negative_styles( + config: TriggerSelectiveTrainingConfig, + count: int, + rng: Optional[random.Random] = None, +) -> List[NegativeStyleSample]: + rng = rng or random + categories = config.negative_styles.categories + selected = rng.choices(categories, weights=[category.probability for category in categories], k=count) + return [NegativeStyleSample(category=category.name, phrase=rng.choice(category.phrases)) for category in selected] + + +def resolve_trigger_placeholder(raw_prompt: str, replacement: str, require_placeholder: bool = True) -> str: + if require_placeholder and '[trigger]' not in raw_prompt: + raise ValueError('TST strict placeholder validation failed: caption does not contain [trigger]') + return inject_trigger_into_prompt(raw_prompt, replacement, add_if_not_present=False) + + +def resolve_prompt_variants( + raw_prompts: Sequence[str], + trigger_word: str, + negative_samples: Sequence[NegativeStyleSample], + require_placeholder: bool = True, +) -> Tuple[List[str], List[str]]: + if len(raw_prompts) != len(negative_samples): + raise ValueError('raw prompt and negative sample counts must match') + trigger_prompts = [] + decoy_prompts = [] + for raw_prompt, sample in zip(raw_prompts, negative_samples): + trigger_prompts.append(resolve_trigger_placeholder(raw_prompt, trigger_word, require_placeholder)) + decoy_prompts.append(resolve_trigger_placeholder(raw_prompt, sample.phrase, require_placeholder)) + return trigger_prompts, decoy_prompts + + +def per_item_mse(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + if prediction.shape != target.shape: + raise ValueError('prediction and target shapes must match') + return F.mse_loss(prediction.float(), target.float(), reduction='none').flatten(1).mean(1) + + +def normalized_gain(student_loss: torch.Tensor, base_loss: torch.Tensor, epsilon: float) -> torch.Tensor: + return 1.0 - student_loss / (base_loss.detach() + epsilon) + + +def trigger_advantage_hinge( + trigger_gain: torch.Tensor, + decoy_gain: torch.Tensor, + margin: float, +) -> torch.Tensor: + return torch.relu(torch.as_tensor(margin, device=trigger_gain.device, dtype=trigger_gain.dtype) - trigger_gain + decoy_gain.detach()) + + +def shared_loss_target( + trainer, + noise: torch.Tensor, + batch, + timesteps: torch.Tensor, +) -> torch.Tensor: + if hasattr(trainer.sd, 'get_loss_target'): + target = trainer.sd.get_loss_target(noise=noise, batch=batch, timesteps=timesteps) + elif trainer.sd.is_flow_matching: + target = noise - batch.latents + else: + target = noise + return target.detach() + + +def apply_differential_guidance_target( + trainer, + target: torch.Tensor, + reference_prediction: torch.Tensor, +) -> torch.Tensor: + if not ( + trainer.train_config.do_guidance_loss + and trainer.train_config.do_differential_guidance + ): + return target + scale = trainer.train_config.differential_guidance_scale + return (reference_prediction.detach() + scale * (target - reference_prediction.detach())).detach() + + +@contextmanager +def network_disabled(network): + if network is None: + yield + return + previous = network.is_active + network.is_active = False + try: + yield + finally: + network.is_active = previous + + +class TSTMetricsWriter: + def __init__(self, output_dir: str, filename: str = 'tst_metrics.jsonl'): + self.path = os.path.join(output_dir, filename) + + def write(self, record: Dict): + os.makedirs(os.path.dirname(self.path), exist_ok=True) + with open(self.path, 'a', encoding='utf-8') as handle: + handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + '\n')