Audio-safe conditioning steering on LTXAV (video-only channel-slice)

LTXAV splits one conditioning tensor into video|audio channel-slices
(av_model _prepare_context: torch.split(context,[v_dim,a_dim],-1)) and
cross-attends each stream to its own slice. Steering the whole tensor
corrupted the audio's text conditioning -> degraded audio (observed live).

Add protect_audio_channels: restore the audio channel-slice from the
unsteered conditioning so every edit moves video only. Applied at all
steering chokepoints:
- _v2_apply_conditioning_memory (per-prompt: liked/axis/concept/vf-grad)
- _v2_apply_absolute (global taste pull)
- _v2_ascend_conditioning (relative VF ascend+search)
- attn2 K/V direction patch + phrase emphasis
- FunPackConditioningAdjust (manual)
- sampler embed_guidance wrapper + per-scene VF ascend

Split auto-detected from width (7680->3840, 6144->4096); logs once.
Single-stream LTXV / unknown width = no-op (no regression).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+Made all conditioning steering **audio-safe on LTXAV**. LTXAV conditions video and audio from two separate text cross-attentions that the model carves out of one conditioning tensor by splitting its channel dim (`comfy/ldm/lightricks/av_model.py` `_prepare_context`: `torch.split(context, [v_context_dim, a_context_dim], -1)` — leading channels → video `attn2`, trailing → audio `audio_attn2`). Previously every steer (relative/absolute pull, value-function ascent + search, embed_guidance, the attn2 direction patch, and manual Conditioning Adjust) shifted the *whole* tensor, corrupting the audio's own text conditioning and degrading audio. Now a shared `protect_audio_channels` confines every edit to the video channel-slice and restores the audio slice from the unsteered conditioning — effectively "modified conditioning for video, original for audio", with no model patching or extra forward pass. The split is auto-detected from the channel width (7680→3840, 6144→4096) and logs the detected layout once; single-stream LTXV (unrecognised width) is a clean no-op, so video-only models are unaffected. Note: this protects the audio's *direct* text conditioning (the dominant lever); the per-block cross-modal `video_to_audio_attn` can still carry a weaker second-order influence from heavily-steered video.
+
 Added an **Absolute / Relative steer mode** to Refiner V2 and Studio. *Relative* (the default, unchanged behaviour) is per-prompt: it learns and applies the best conditioning for one specific prompt. *Absolute* is prompt-agnostic — it accumulates a single global "taste" prior across **every** rated generation and pulls conditioning toward it regardless of the prompt. In Absolute mode a rating means "this has (or lacks) details I like **in general**", not prompt adherence: a Perfect means "I love this, give me more of it everywhere", and a low rating means "this is missing the details I like in general". The pull layers two engines — a pooled liked/disliked direction (the learned, automated analogue of the manual Conditioning Adjust phrase shift) and the keyless value function on top — and is applied at the conditioning output (`_v2_finalize_conditioning`). *Both* keeps the per-prompt fit and layers the global prior under it; pure *Absolute* bypasses per-prompt memory entirely. The global store learns from rated runs even with no refinement key wired (it is keyless by design). The Scene Chain sampler's `embed_guidance` gains an `embed_guidance_source` (`relative` / `absolute`) so the per-step nudge can draw from the global taste direction too. Studio exposes `Steer mode` + `Absolute strength` in the refiner panel.
 
 Ported **velocity bias + reactive rescue** to the `FunPack Distilled Flow` sampler (`sample_funpack_distilled_flow`), reusing the same capture/memory/rescue machinery as the Hybrid Euler-2S and LTXAV/RF samplers. The sampler now exposes `velocity_bias_mode`, `velocity_bias_strength`, `velocity_bias_source`, `velocity_refinement_key`, `rescue_mode`, `rescue_threshold`, and `rescue_strength`, and these are wired into Studio's **Distilled Flow** sampler config panel (with the same blank/`default` velocity-key → wired-refinement-key fallback the Hybrid panel uses). Steering is applied magnitude-preservingly and, on packed LTXAV latents, confined to the video stream (audio-safe). Note: few-step distilled schedules may only land on one or two velocity targets, so apply/rescue fire less often than on an 8-step run; they no-op cleanly when no target matches. With everything off the sampler is byte-identical to the previous deterministic ODE.

conditioning.py

@@ -366,6 +366,53 @@ def update_refinement_sampler_context(refinement_key, sampler_context):
 FUNPACK_ABSOLUTE_KEY = "__funpack_absolute_global__"
 
 
+# --- Audio-safe conditioning steering -------------------------------------------------------
+# LTXAV conditions video and audio from two separate text contexts that the model carves out of
+# ONE conditioning tensor by splitting its channel dim: comfy/ldm/lightricks/av_model.py
+# `_prepare_context` does `torch.split(context, [v_context_dim, a_context_dim], -1)` — leading
+# channels drive video (self.attn2), trailing channels drive audio (self.audio_attn2). So any
+# conditioning steering that shifts the WHOLE tensor also corrupts the audio's text cross-
+# attention (observed live as degraded audio). Confining edits to the video channel-slice gives
+# "original conditioning for audio, modified for video" with no model patching.
+#   LTXv2  fully-processed context = caption_channels*2 = 7680           → video = 3840
+#   LTXv2.3 fully-processed context = cross_attn_dim + audio_cross_attn_dim = 4096+2048 = 6144 → video = 4096
+# Anything else (single-stream LTXV, unrecognised dim) returns the full width, so callers no-op.
+_FUNPACK_AV_SPLIT_TABLE = {7680: 3840, 6144: 4096}
+_FUNPACK_AV_LOGGED = set()
+
+
+def ltxav_video_channels(dim):
+    """Leading channel count that conditions VIDEO for an LTXAV concat-context of width `dim`.
+    Returns `dim` (no audio half) for single-stream / unrecognised layouts."""
+    dim = int(dim)
+    v = _FUNPACK_AV_SPLIT_TABLE.get(dim, dim)
+    if dim not in _FUNPACK_AV_LOGGED:
+        _FUNPACK_AV_LOGGED.add(dim)
+        if v < dim:
+            print(f"[FunPack AV] conditioning dim={dim}: steering confined to video [:{v}], audio [{v}:] protected")
+        else:
+            print(f"[FunPack AV] conditioning dim={dim}: single-stream (no audio half); steering applies to all channels")
+    return v
+
+
+def protect_audio_channels(steered, original):
+    """Restore the audio channel-slice of a steered conditioning tensor from `original`, so the
+    steering only moves the video stream. No-op when shapes mismatch or there is no audio half."""
+    try:
+        if not isinstance(steered, torch.Tensor) or not isinstance(original, torch.Tensor):
+            return steered
+        if tuple(steered.shape) != tuple(original.shape):
+            return steered
+        v = ltxav_video_channels(steered.shape[-1])
+        if v >= steered.shape[-1]:
+            return steered
+        out = steered.clone()
+        out[..., v:] = original[..., v:].to(device=out.device, dtype=out.dtype)
+        return out
+    except Exception:
+        return steered
+
+
 def refinement_state_path(refinement_key, mode, prefix="refine", extension="json"):
     base_dir = os.path.dirname(os.path.abspath(__file__))
     refinements_dir = os.path.join(base_dir, "refinements")
@@ -8917,6 +8964,10 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                             v_out[:, start:end] = v_out[:, start:end] * 1.25
                 except Exception:
                     pass
+            # If this fires on an AV-concat context, keep the audio channel-slice unmodified
+            # (no-op on already-projected / single-stream widths). Protects audio_attn2.
+            k_out = protect_audio_channels(k_out, k)
+            v_out = protect_audio_channels(v_out, v)
             return q, k_out, v_out
 
         return patch
@@ -9380,7 +9431,16 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
         except Exception:
             return None
 
-    def _v2_apply_conditioning_memory(self, conditioning, global_state, rating_profile, axis_feedback=None, intent_family_slot=None, vf=None, concept_delta_dir=None, concept_delta_strength=0.0, current_final_texts=None):
+    def _v2_apply_conditioning_memory(self, conditioning, *args, **kwargs):
+        """Audio-safe wrapper: confine all per-prompt conditioning steering to the video
+        channel-slice, restoring the audio channels from the unsteered input (LTXAV splits the
+        conditioning into video|audio channels; see protect_audio_channels)."""
+        result = self._v2_apply_conditioning_memory_impl(conditioning, *args, **kwargs)
+        if isinstance(result, tuple) and result and isinstance(result[0], torch.Tensor):
+            return (protect_audio_channels(result[0], conditioning),) + tuple(result[1:])
+        return result
+
+    def _v2_apply_conditioning_memory_impl(self, conditioning, global_state, rating_profile, axis_feedback=None, intent_family_slot=None, vf=None, concept_delta_dir=None, concept_delta_strength=0.0, current_final_texts=None):
         if not isinstance(conditioning, torch.Tensor):
             return conditioning, "Adaptation: unavailable."
         original = conditioning.clone()
@@ -12388,6 +12448,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 c = vf.ascend(tensor)
                 # Monte Carlo search around ascended point (explore neighbourhood)
                 c = vf.search(c)
+                c = protect_audio_channels(c, tensor)  # keep audio on the original conditioning
                 ascended.append((c, extra))
             print(f"[FunPackRefiner] Conditioning ascent+search applied ({vf.n_trained} samples, {len(ascended)} scene(s))")
             return ascended
@@ -12545,6 +12606,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                     mixed = gvf.ascend(mixed)
                 except Exception:
                     pass
+            mixed = protect_audio_channels(mixed, entry[0])  # audio rides the original conditioning
             out.append((mixed, extra))
             applied = True
         if applied:
@@ -13466,7 +13528,7 @@ class FunPackConditioningAdjust:
             shift = total_delta.to(cond_tensor.dtype)
             while shift.dim() < cond_tensor.dim():
                 shift = shift.unsqueeze(0)
-            new_tensor = cond_tensor + shift
+            new_tensor = protect_audio_channels(cond_tensor + shift, cond_tensor)
             new_conditioning.append((new_tensor, meta))
 
         status = "Conditioning adjust: " + (", ".join(log) if log else "nothing applied")

samplers.py

@@ -2177,6 +2177,18 @@ class FunPackLTXAVSceneChainSampler:
             from conditioning import FUNPACK_ABSOLUTE_KEY
         return FUNPACK_ABSOLUTE_KEY
 
+    def _protect_audio(self, steered, original):
+        """Confine a conditioning edit to LTXAV's video channel-slice; audio keeps the original
+        conditioning. No-op for single-stream LTXV. Mirrors the Refiner's protect_audio_channels."""
+        try:
+            try:
+                from .conditioning import protect_audio_channels
+            except ImportError:
+                from conditioning import protect_audio_channels
+            return protect_audio_channels(steered, original)
+        except Exception:
+            return steered
+
     def _load_value_function(self, refinement_key):
         """Load the online value function if trained and ready."""
         try:
@@ -2246,7 +2258,8 @@ class FunPackLTXAVSceneChainSampler:
                     else:
                         d = _fixed.to(cond.device, cond.dtype).expand_as(cond)
                     new_c = dict(c)
-                    new_c["c_crossattn"] = cond + (_s * scale) * d
+                    # Steer video channels only; audio keeps its original text conditioning.
+                    new_c["c_crossattn"] = self._protect_audio(cond + (_s * scale) * d, cond)
                     args = dict(args)
                     args["c"] = new_c
             if _ew is not None:
@@ -2426,7 +2439,7 @@ class FunPackLTXAVSceneChainSampler:
 
             if embed_guidance and _value_fn is not None and _value_fn.is_ready():
                 orig_cond, orig_extra = scene_positive[0][0], scene_positive[0][1]
-                ascended = _value_fn.ascend(orig_cond)
+                ascended = self._protect_audio(_value_fn.ascend(orig_cond), orig_cond)
                 scene_positive = [[ascended, orig_extra]] + list(scene_positive[1:])
             if embed_guidance and _liked_dir is not None:
                 _eg_old_wrapper = self._build_embed_guidance_wrapper(model, _liked_dir, embed_guidance_strength, value_fn=_value_fn)