diff --git a/conditioning.py b/conditioning.py index 0a17a0fa..080206ba 100644 --- a/conditioning.py +++ b/conditioning.py @@ -12666,7 +12666,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): guess_mode=guess_mode, guess_direction=guess_direction, guess_range=guess_range, guess_freeze_seed=guess_freeze_seed) return ( - self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan), + self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip), status, training_info, loss_graph, @@ -12684,7 +12684,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): repair_feedback, current_family_slot, _vf_for_memory, _concept_dir, _concept_strength, _current_final) return ( - self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan), + self._v2_finalize_conditioning(output_conditioning, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=_guess_spread_cap, temporal_style=temporal_style, temporal_fallback_text=prompt_to_encode, scene_refinement_keys=scene_refinement_keys, learning_profile=learning_profile, conditioning_plan=_conditioning_plan, clip=clip), status + enhancement_status, training_info, loss_graph, @@ -13223,11 +13223,84 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): out.append([cond, meta] if isinstance(entry, list) else (cond, meta)) return out + def _v2_token_seq_len(self, tokenized): + """Best-effort token-sequence length from clip.tokenize()'s output. Structure varies by + text-encoder (dict of encoder-name -> nested batch/chunk lists of (id, weight) tuples or + plain ids); we just need ANY one real token sequence's length, so recurse and return the + first list found whose elements look like tokens. Returns None on anything unexpected — + callers must treat that as "can't determine, skip".""" + def _walk(node): + if isinstance(node, dict): + for v in node.values(): + r = _walk(v) + if r is not None: + return r + return None + if isinstance(node, (list, tuple)): + if node and all(isinstance(e, (tuple, list)) and len(e) >= 1 for e in node): + return len(node) # this level IS the token list: [(id, weight), ...] + if node and all(isinstance(e, int) for e in node): + return len(node) + for e in node: + r = _walk(e) + if r is not None: + return r + return None + try: + return _walk(tokenized) + except Exception: + return None + + def _v2_bounded_attention_split_tokens(self, text, clip): + """EXPERIMENTAL: split a scene prompt into two halves by sentence count (subject 1's + description vs subject 2's), then tokenize each half SEPARATELY with the same CLIP/text + encoder used for the real encode, to find the approximate token-index boundary between + them within the full tokenized sequence. Returns the token count of half 1, or None if + the prompt can't be split (single sentence) or clip/tokenization isn't available. + + This is a heuristic, not exact: BPE/SentencePiece merges can shift a few tokens across the + boundary vs. tokenizing the two halves separately. Good enough for an experimental, + sentence-granularity attention split — not for precise word-level boundaries.""" + if clip is None: + return None + text = str(text or "").strip() + if not text: + return None + sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()] + if len(sentences) < 2: + return None + mid = (len(sentences) + 1) // 2 + span1 = " ".join(sentences[:mid]) + try: + n1 = self._v2_token_seq_len(clip.tokenize(span1)) + except Exception: + n1 = None + return int(n1) if n1 is not None and n1 > 0 else None + + def _v2_apply_bounded_attention(self, conditioning_list, clip, fallback_text=""): + """Always-on, cheap text-only step: tag each scene entry with an approximate subject-1/ + subject-2 token-split boundary (funpack_bound_split_tokens) for the EXPERIMENTAL Bounded + Attention sampler toggle to read. Pure metadata — no behavior change unless that sampler + toggle is enabled. No-op (and cheap) when the prompt has under two sentences or clip is + unavailable.""" + out = [] + for entry in conditioning_list or []: + if not (isinstance(entry, (list, tuple)) and len(entry) >= 2 and isinstance(entry[1], dict)): + out.append(entry) + continue + cond, meta = entry[0], dict(entry[1]) + text = str(meta.get("funpack_scene_text") or "").strip() or str(fallback_text or "") + n1 = self._v2_bounded_attention_split_tokens(text, clip) + if n1 is not None: + meta["funpack_bound_split_tokens"] = n1 + out.append([cond, meta] if isinstance(entry, list) else (cond, meta)) + return out + def _v2_finalize_conditioning(self, conditioning_list, refinement_key, value_guidance, steer_mode, absolute_strength, spread_cap=None, temporal_style="natural", temporal_fallback_text="", scene_refinement_keys=None, learning_profile=None, - conditioning_plan=None): + conditioning_plan=None, clip=None): """Single output hook for both steering modes. Relative = per-key VF ascend (current behaviour). Absolute = global taste pull. Both = layer them. Finally, if Interactive Guessing has learned a safe-spread ceiling, clamp the output conditioning's video-channel @@ -13239,6 +13312,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): mode = str(steer_mode or "relative").lower() out = self._v2_apply_pulse_temporal(conditioning_list, temporal_style) out = self._v2_apply_auto_temporal(out, temporal_style, fallback_text=temporal_fallback_text) + out = self._v2_apply_bounded_attention(out, clip, fallback_text=temporal_fallback_text) if mode in ("relative", "both"): out = self._v2_apply_scene_refinement_keys( out, scene_refinement_keys, refinement_key, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 303d08ae..5de8ce56 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -204,6 +204,7 @@ { name: "joyai_audio_memory", label: "Paired audio memory", kind: "bool", default: false, dependsOn: "joyai_memory" }, { name: "v2a_grad_scale", label: "Video→audio coupling", kind: "float", default: 1.0, min: 0.0, max: 4.0, step: 0.25, dependsOn: "joyai_audio_memory" }, { name: "alg_blur_guides", label: "Blur i2v guides and JoyAI memory", kind: "bool", default: false }, + { name: "bounded_attention_enabled", label: "Bounded attention (multi-subject)", kind: "bool", default: false }, ]; const CONTINUITY_DEFAULTS = { @@ -344,7 +345,7 @@ { id: "chain_cont", title: "Manual continuity", defaultOpen: false, knobs: ["carry_i2v_guides"] }, { id: "chain_guid", title: "Guidance", defaultOpen: false, knobs: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength"] }, { id: "chain_dec", title: "Decode", defaultOpen: false, knobs: ["decode_noise_scale", "decode_timestep", "decode_tile_size"] }, - { id: "chain_exp", title: "Experimental", defaultOpen: false, knobs: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides"] }, + { id: "chain_exp", title: "Experimental", defaultOpen: false, knobs: ["mid_scene_guide", "mid_scene_guide_strength", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "bounded_attention_enabled"] }, ]; const SAMPLER_KNOB_MAP = Object.fromEntries(SAMPLER_KNOBS.map((k) => [k.name, k])); diff --git a/samplers.py b/samplers.py index f73bd198..b6657e2f 100644 --- a/samplers.py +++ b/samplers.py @@ -2194,6 +2194,10 @@ class FunPackLTXAVSceneChainSampler: "default": False, "tooltip": "EXPERIMENTAL: extends ALG (see the sampler's alg_enabled) from just the i2v anchor to also blur newly-appended guide-attention frames this scene (mid_scene_guide / carry_i2v_guides-as-guide / configured per-scene guides / JoyAI memory), for the same early steps. No effect unless the chosen sampler has alg_enabled on, or no guide frames were appended this scene.", }), + "bounded_attention_enabled": ("BOOLEAN", { + "default": False, + "tooltip": "EXPERIMENTAL: Bounded Attention (arXiv:2403.16990-inspired). Masks text cross-attention so the left half of the frame can only attend to subject-1's prompt tokens and the right half only to subject-2's (split by sentence count, set automatically by Studio) — aims to stop attribute/anatomy bleed between two figures in one frame. No-op on single-sentence prompts or single-subject scenes. Works on any sampler (model-level hook, not sampler-specific).", + }), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -2578,7 +2582,8 @@ class FunPackLTXAVSceneChainSampler: return result def _sample_chunk(self, model, sampler, sigmas, seed, cfg, positive, negative, latent, - pbar=None, step_offset=0, alg_guide_tail_frames=0): + pbar=None, step_offset=0, alg_guide_tail_frames=0, + bounded_attention_enabled=False): if sampler is None: raise ValueError("sampler input is required.") if not isinstance(sigmas, torch.Tensor): @@ -2601,11 +2606,19 @@ class FunPackLTXAVSceneChainSampler: if isinstance(extra_options, dict) and sampler.sampler_function is sample_funpack_distilled_flow: extra_options["alg_guide_tail_frames"] = int(alg_guide_tail_frames) - sampled = comfy.sample.sample_custom( - model, noise, float(cfg), sampler, sigmas, positive, negative, samples, - noise_mask=latent.get("noise_mask"), seed=int(seed), - callback=_progress_cb if pbar is not None else None, - ) + # EXPERIMENTAL Bounded Attention: model-level attention hooks (sampler-agnostic, unlike + # the toggles above which only work on Distilled Flow), so install/remove here rather + # than via extra_options. Cheap to attempt (no-ops fast without the right metadata). + _ba_handles = self._install_bounded_attention(model, latent, positive) if bounded_attention_enabled else [] + + try: + sampled = comfy.sample.sample_custom( + model, noise, float(cfg), sampler, sigmas, positive, negative, samples, + noise_mask=latent.get("noise_mask"), seed=int(seed), + callback=_progress_cb if pbar is not None else None, + ) + finally: + self._remove_bounded_attention(_ba_handles) latent["samples"] = sampled latent.pop("noise_mask", None) return latent @@ -3613,6 +3626,83 @@ class FunPackLTXAVSceneChainSampler: except Exception: pass + def _bounded_attention_region_mask(self, t, h, w, device): + """[T*H*W] region id per video token: 0 = left half (by width), 1 = right half. + Assumes (t, h, w) row-major flattening — matches the packed-latent layout + _alg_blur_frames already relies on; LTX's patchify is 1:1 so the transformer's + token order should match. Best-effort for an experimental feature, not guaranteed + if that assumption ever changes upstream.""" + idx = torch.arange(t * h * w, device=device) + w_idx = idx % w + return (w_idx >= (w // 2)).long() + + def _install_bounded_attention(self, model, latent, positive): + """EXPERIMENTAL (arXiv:2403.16990-inspired, see [[project_bounded_attention]]): mask the + text cross-attention (attn2 only — never audio_attn2, never self-attention) so the LEFT + half of the video frame (by width) can only attend to subject-1's prompt tokens and the + RIGHT half only to subject-2's, per the sentence-count split Studio recorded as + funpack_bound_split_tokens on the scene's conditioning metadata. Pure attention masking, + no extra forward pass. Returns [] (no-op) without that metadata, a 5D video tensor, an + existing padding mask already in play, or recognizable transformer blocks.""" + try: + n1 = None + if isinstance(positive, list) and positive and isinstance(positive[0], (list, tuple)) and len(positive[0]) >= 2: + meta = positive[0][1] + if isinstance(meta, dict): + n1 = meta.get("funpack_bound_split_tokens") + if n1 is None: + return [] + n1 = int(n1) + tensors = self._latent_tensors(latent) + video = max(tensors, key=lambda v: v.dim()) + if video.dim() != 5: + return [] + _, _, t, h, w = video.shape + if w < 2: + return [] + blocks = model.model.diffusion_model.transformer_blocks + except Exception: + return [] + if not blocks: + return [] + region = self._bounded_attention_region_mask(t, h, w, video.device) # [Q], 0=left, 1=right + + def _hook(_module, args, kwargs): + try: + if kwargs.get("mask") is not None: + return args, kwargs # real padding mask already present -> skip, don't fight it + x = args[0] if args else None + context = kwargs.get("context") + if x is None or context is None: + return args, kwargs + q = int(x.shape[1]) + if q != region.shape[0]: + return args, kwargs # resolution mismatch -> no-op safely + k = int(context.shape[1]) + n1c = max(0, min(k, n1)) + key_idx = torch.arange(k, device=x.device) + allow0 = key_idx < n1c + mask = torch.where(region.unsqueeze(1) == 0, allow0.unsqueeze(0), (~allow0).unsqueeze(0)) + kwargs = dict(kwargs) + kwargs["mask"] = mask + return args, kwargs + except Exception: + return args, kwargs + + handles = [] + for blk in blocks: + sub = getattr(blk, "attn2", None) + if sub is not None: + handles.append(sub.register_forward_pre_hook(_hook, with_kwargs=True)) + return handles + + def _remove_bounded_attention(self, handles): + for h in handles or []: + try: + h.remove() + except Exception: + pass + def _vae_with_decode_noise(self, vae, timestep, scale, seed): """Return a shallow copy of the VAE stamped with LTX decode-time noise settings so its internal decoder restores fine detail/grain. Never mutates the shared input VAE. Mirrors @@ -3648,6 +3738,7 @@ class FunPackLTXAVSceneChainSampler: funpack_scene_anchors="", funpack_scene_media_refs="", alg_blur_guides=False, + bounded_attention_enabled=False, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -3689,6 +3780,7 @@ class FunPackLTXAVSceneChainSampler: joyai_memory_strength=joyai_memory_strength, joyai_audio_memory=joyai_audio_memory, v2a_grad_scale=v2a_grad_scale, alg_blur_guides=alg_blur_guides, + bounded_attention_enabled=bounded_attention_enabled, ) max_scene_count = max(1, int(max_scenes)) @@ -3900,6 +3992,7 @@ class FunPackLTXAVSceneChainSampler: model, sampler, sigmas, scene_seed, cfg, scene_positive, scene_negative, chunk, pbar=pbar, step_offset=scene_index * steps_per_scene, alg_guide_tail_frames=(guide_tail if (alg_blur_guides and guide_tail > 0) else 0), + bounded_attention_enabled=bounded_attention_enabled, ) finally: self._remove_v2a_scale(_v2a_handles) @@ -4129,7 +4222,7 @@ class FunPackLTXAVSceneChainSampler: joyai_memory=False, joyai_memory_size=7, joyai_fix_frames=3, joyai_frame_select="center", joyai_memory_strength=0.3, joyai_audio_memory=False, v2a_grad_scale=1.0, - alg_blur_guides=False): + alg_blur_guides=False, bounded_attention_enabled=False): """Sample one chain per Studio-packed variant entry (seed + index), persisting each result (latent + preview + per-entry cond + manifest) under ComfyUI temp for rating in Studio. Reuses sample() per entry with only the seed changed, so each entry is a clean generation.""" @@ -4176,6 +4269,7 @@ class FunPackLTXAVSceneChainSampler: transition_duration=transition_duration, decode_tile_size=decode_tile_size, refinement_key_input=key, alg_blur_guides=alg_blur_guides, + bounded_attention_enabled=bounded_attention_enabled, unique_id=None, prompt=None, ) last = out