From a6e4a934f132cbfac472347d2fa82c2fb2c4b50e Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:10:10 +0300 Subject: [PATCH] Segmented detailing: CLIPSeg-targeted latent crop-refine on Scene Chain ADetailer-for-video probe core. After a scene denoises, CLIPSeg locates user-named regions (hands, feet, ...) on decoded keyframes; the region is cropped as a spatiotemporal tube, 2x-upsampled through Lightricks' trained LatentUpsampler (comfy core, no port needed), re-denoised over the official stage-2 tail (0.85/0.725/0.4219/0), downscaled back and pasted through the feathered CLIPSeg silhouette. Final resolution unchanged; audio protected by construction (refined audio discarded); tubes >35% frame refused; empty targets/no-detection = bit-identical no-op. UNVALIDATED LIVE. Co-Authored-By: Claude Fable 5 --- detailing.py | 349 ++++++++++++++++++++++++++++++++++++++++ samplers.py | 50 ++++++ tests/test_detailing.py | 248 ++++++++++++++++++++++++++++ 3 files changed, 647 insertions(+) create mode 100644 detailing.py create mode 100644 tests/test_detailing.py diff --git a/detailing.py b/detailing.py new file mode 100644 index 00000000..832226f5 --- /dev/null +++ b/detailing.py @@ -0,0 +1,349 @@ +"""Segmented detailing: ADetailer-style region refine for the Scene Chain sampler. + +After a scene finishes denoising, CLIPSeg (text-prompted segmentation, the same +model family face-swap stacks use for occlusion masking) locates user-named +regions ("hands", "feet", ...) on a few decoded keyframes. The matched region is +cut out of the video latent as a spatiotemporal tube, pushed through Lightricks' +own trained LatentUpsampler (2x spatial, latent space - the official two-stage +pipeline's stage-2 model), re-noised to the official stage-2 re-entry sigma and +re-denoised for the tail schedule, then downscaled back to its original latent +size and pasted through the feathered CLIPSeg soft mask. Final resolution never +changes: the upsample exists only so the model can resolve structure (fingers, +edges) at a higher working resolution before the crop returns to the grid. + +Boundary law: this is a visual-behavior op, so it lives on the Chain Sampler. +Audio is never modified - the refine pass carries the scene's audio stream so +joint attention sees a legal AV latent, but its audio output is discarded. + +No-op guarantee: empty targets, no upsampler, or no detection above threshold +returns the input latent object untouched (bit-identical run). +""" + +import math + +import torch +import torch.nn.functional as F + +# Official LTX-2.3 two-stage distilled workflow, stage-2 ManualSigmas: re-noise the +# upsampled latent to 0.85 and run the tail of the 8-step distilled schedule. Using +# Lightricks' published re-entry point instead of inventing one: too much noise and +# the region redraws (drifts from the surrounding frame), too little and no detail +# is gained. +STAGE2_SIGMAS = (0.85, 0.725, 0.421875, 0.0) + +# CLIPSeg heat above this (post-sigmoid) counts as "the named thing is here". +DEFAULT_THRESHOLD = 0.35 + +# Tubes larger than this fraction of the frame are refused: at that size the pass +# stops being region detailing and becomes a second full render (cost ~= 4x area +# fraction x tail steps), which belongs to a different feature. +MAX_TUBE_AREA = 0.35 + +# Latent cells of context padding around the thresholded region, and the minimum +# tube edge. Padding gives the refine pass surrounding context to blend against; +# the minimum keeps the crop a legal latent for conv stacks and attention. +TUBE_PAD = 2 +MIN_TUBE_EDGE = 4 + +# Keyframes decoded per scene for detection. Single latent frames decode cheaply +# through the causal VAE (treated as stills) - enough for a temporal-union tube. +DETECTION_FRAMES = 3 + +# Cond keys that tie a scene's conditioning to the full-frame latent layout +# (i2v anchors / guide keyframes). The crop is its own small latent with no +# appended guide frames, so these must not travel into the refine pass. +_LAYOUT_COND_KEYS = ("keyframe_idxs", "guiding_latent", "concat_latent_image", "concat_mask") + +_CLIPSEG = None # (processor, model), lazy - loaded on first enabled run only + + +def _log(debug, msg): + if debug: + print(f"[FunPackDetail] {msg}") + + +# --------------------------------------------------------------------------- +# Models +# --------------------------------------------------------------------------- + +def load_latent_upsampler(model_name): + """Load Lightricks' LatentUpsampler from models/latent_upscale_models. + + Mirrors the LTX branch of comfy's LatentUpscaleModelLoader (safe_load + + config from safetensors metadata); the architecture itself ships in comfy + core (comfy.ldm.lightricks.latent_upsampler), so there is nothing to port. + """ + import json + + import comfy.model_management + import comfy.utils + import folder_paths + from comfy.ldm.lightricks.latent_upsampler import LatentUpsampler + + model_path = folder_paths.get_full_path_or_raise("latent_upscale_models", model_name) + sd, metadata = comfy.utils.load_torch_file(model_path, safe_load=True, return_metadata=True) + if "post_upsample_res_blocks.0.conv2.bias" not in sd: + raise ValueError( + f"{model_name} is not an LTX latent upsampler (expected post_upsample_res_blocks keys).") + config = json.loads(metadata["config"]) + model = LatentUpsampler.from_config(config).to( + dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) + model.load_state_dict(sd) + model.eval() + return model + + +def _get_clipseg(): + """CLIPSeg processor+model, loaded once and kept on CPU. + + CPU on purpose: ~150M params, a handful of 352px frames per scene - the cost + is milliseconds against a video denoise, and it never competes with the + diffusion model for VRAM. Loaded straight from transformers (which comfy + already ships); no GPL face-swap code involved. + """ + global _CLIPSEG + if _CLIPSEG is None: + from transformers import CLIPSegForImageSegmentation, CLIPSegProcessor + + name = "CIDAS/clipseg-rd64-refined" + processor = CLIPSegProcessor.from_pretrained(name) + model = CLIPSegForImageSegmentation.from_pretrained(name) + model.eval() + _CLIPSEG = (processor, model) + return _CLIPSEG + + +# --------------------------------------------------------------------------- +# Detection +# --------------------------------------------------------------------------- + +def parse_targets(detail_targets): + """Comma-separated user text -> clean prompt list ('' -> []).""" + if not detail_targets: + return [] + return [t.strip() for t in str(detail_targets).split(",") if t.strip()] + + +def _decode_detection_frames(vae, video, count, debug=False): + """Decode `count` evenly spaced single latent frames to [H, W, C] float images. + + Each latent frame is decoded alone as [B, C, 1, h, w]: the causal VAE treats + a length-1 clip as a still, which is exactly good enough for segmentation and + avoids decoding the whole scene just to find a region. + """ + frames = [] + f = int(video.shape[2]) + idxs = sorted({min(f - 1, round(i * (f - 1) / max(1, count - 1))) for i in range(max(1, count))}) + for i in idxs: + try: + decoded = vae.decode(video[:, :, i:i + 1]) + except Exception as exc: + _log(debug, f"keyframe {i} decode failed ({exc}); skipping") + continue + img = decoded + if img.dim() == 5: # [B, T, H, W, C] + img = img[0, 0] + elif img.dim() == 4: # [B, H, W, C] + img = img[0] + frames.append((i, img.detach().float().clamp(0.0, 1.0).cpu())) + return frames + + +def _clipseg_heat(frames, prompts, debug=False): + """Max-over-frames, max-over-prompts CLIPSeg heat map. + + Returns [H, W] in [0, 1] at the first frame's pixel size, or None when + nothing was decodable. Max (not mean) across frames on purpose: the tube + must cover the region wherever it appears during the scene. + """ + import numpy as np + from PIL import Image + + if not frames or not prompts: + return None + processor, model = _get_clipseg() + h, w = frames[0][1].shape[0], frames[0][1].shape[1] + heat = None + with torch.no_grad(): + for _, img in frames: + pil = Image.fromarray((img.numpy() * 255.0).astype(np.uint8)) + inputs = processor( + text=list(prompts), images=[pil] * len(prompts), return_tensors="pt") + logits = model(**inputs).logits # [P, 352, 352] (or [352, 352] for one prompt) + if logits.dim() == 2: + logits = logits[None] + frame_heat = torch.sigmoid(logits).amax(dim=0)[None, None] # union over prompts + frame_heat = F.interpolate(frame_heat, size=(h, w), mode="bilinear", align_corners=False)[0, 0] + heat = frame_heat if heat is None else torch.maximum(heat, frame_heat) + return heat + + +def find_tube(heat, latent_h, latent_w, threshold=DEFAULT_THRESHOLD, + area_cap=MAX_TUBE_AREA, debug=False): + """Threshold pixel-space heat into a latent-space tube box + soft paste mask. + + Returns (y0, y1, x0, x1, mask) in latent coords - mask is [1, 1, 1, th, tw] + matching the tube, feathered from the CLIPSeg heat itself so the paste + follows the actual silhouette instead of the rectangle. None when nothing + crossed the threshold or the tube breached the area cap. + """ + if heat is None: + return None + # Heat -> latent grid first; box math happens where the crop happens. + lat_heat = F.interpolate(heat[None, None], size=(latent_h, latent_w), + mode="bilinear", align_corners=False)[0, 0] + hot = lat_heat >= threshold + if not bool(hot.any()): + _log(debug, f"no region above threshold {threshold}") + return None + ys, xs = torch.where(hot) + y0 = max(0, int(ys.min()) - TUBE_PAD) + y1 = min(latent_h, int(ys.max()) + 1 + TUBE_PAD) + x0 = max(0, int(xs.min()) - TUBE_PAD) + x1 = min(latent_w, int(xs.max()) + 1 + TUBE_PAD) + # Enforce a minimum edge (grow symmetrically, clamped to the frame). + while (y1 - y0) < min(MIN_TUBE_EDGE, latent_h): + y0, y1 = max(0, y0 - 1), min(latent_h, y1 + 1) + while (x1 - x0) < min(MIN_TUBE_EDGE, latent_w): + x0, x1 = max(0, x0 - 1), min(latent_w, x1 + 1) + area = (y1 - y0) * (x1 - x0) / float(latent_h * latent_w) + if area > area_cap: + _log(debug, f"tube area {area:.0%} exceeds cap {area_cap:.0%}; refusing " + "(a region that large is a re-render, not a detail pass)") + return None + # Soft mask: the heat inside the tube, floored at 0 outside the threshold and + # lightly blurred so the seam never lands on a hard edge. + m = lat_heat[y0:y1, x0:x1].clone() + m = ((m - threshold) / max(1e-6, 1.0 - threshold)).clamp(0.0, 1.0) + m = F.avg_pool2d(m[None, None], kernel_size=3, stride=1, padding=1)[0, 0] + return y0, y1, x0, x1, m[None, None, None] + + +# --------------------------------------------------------------------------- +# Refine +# --------------------------------------------------------------------------- + +def _strip_layout_conds(conds): + """Copy conditioning minus full-frame layout keys (anchors/guides). + + The crop refine is its own small denoise: text still applies, but keyframe + indices and guide latents describe positions in the full-frame latent and + would fight (or crash) the crop-sized pass. + """ + out = [] + for entry in conds or []: + if isinstance(entry, (list, tuple)) and len(entry) == 2 and isinstance(entry[1], dict): + meta = {k: v for k, v in entry[1].items() if k not in _LAYOUT_COND_KEYS} + out.append([entry[0], meta]) + else: + out.append(entry) + return out + + +def _per_channel_stats(vae): + return getattr(getattr(vae, "first_stage_model", None), "per_channel_statistics", None) + + +def _run_upsampler(upsampler, video_crop, vae, debug=False): + """un_normalize -> LatentUpsampler(2x spatial) -> normalize, on the crop. + + Same recipe as comfy's LTXVLatentUpsampler node: the model was trained on + raw (denormalized) VAE latents, so the trip through per_channel_statistics + is mandatory, not cosmetic. + """ + import comfy.model_management + + device = comfy.model_management.get_torch_device() + model_dtype = next(upsampler.parameters()).dtype + in_dtype, in_device = video_crop.dtype, video_crop.device + stats = _per_channel_stats(vae) + try: + upsampler.to(device) + x = video_crop.to(device=device, dtype=model_dtype) + if stats is not None: + x = stats.un_normalize(x) + with torch.no_grad(): + x = upsampler(x) + if stats is not None: + x = stats.normalize(x) + finally: + upsampler.cpu() + return x.to(device=in_device, dtype=in_dtype) + + +def _downscale_to(video, h, w): + """Area-downscale [B, C, F, H, W] spatially to (h, w); frames independent.""" + b, c, f, hh, ww = video.shape + x = video.permute(0, 2, 1, 3, 4).reshape(b * f, c, hh, ww) + x = F.interpolate(x.float(), size=(h, w), mode="area") + return x.reshape(b, f, c, h, w).permute(0, 2, 1, 3, 4).to(dtype=video.dtype) + + +def detail_refine_scene(chain, model, vae, sampler, positive, negative, latent, + detail_targets, upsampler, seed, cfg, + threshold=DEFAULT_THRESHOLD, strength=1.0, debug=False): + """Run the segmented-detail pass on one finished scene latent. + + Returns (latent, note): the refined latent dict plus a run_mechanisms note, + or (latent, None) untouched - same object - when disabled or nothing found. + `chain` is the FunPackLTXAVSceneChainSampler instance (reuses _latent_tensors + and _sample_chunk so the refine denoise goes through the exact same path as + the scene it is polishing). + """ + targets = parse_targets(detail_targets) + if not targets or upsampler is None or strength <= 0: + return latent, None + + tensors = chain._latent_tensors(latent) + video = tensors[0] + if video.dim() != 5: + _log(debug, f"unexpected video latent rank {video.dim()}; skipping") + return latent, None + b, c, f, lat_h, lat_w = video.shape + + # 1) Detect on a few decoded keyframes. + frames = _decode_detection_frames(vae, video, DETECTION_FRAMES, debug=debug) + heat = _clipseg_heat(frames, targets, debug=debug) + tube = find_tube(heat, lat_h, lat_w, threshold=threshold, debug=debug) + if tube is None: + return latent, None + y0, y1, x0, x1, mask = tube + _log(debug, f"tube y[{y0}:{y1}] x[{x0}:{x1}] of {lat_h}x{lat_w} for {targets}") + + # 2) Crop -> 2x latent upsample (Lightricks stage-2 model). + crop = video[:, :, :, y0:y1, x0:x1] + crop_up = _run_upsampler(upsampler, crop, vae, debug=debug) + + # 3) Re-denoise the crop over the official stage-2 tail. The clean crop goes + # in as the latent; CONST noise_scaling (sigma*noise + (1-sigma)*x) re-noises + # it to sigma=0.85 exactly like the official two-stage workflow's stage 2. + # The scene's audio stream rides along so joint attention sees a legal AV + # latent, but its refined output is discarded below - audio is protected by + # construction, not by masking. + import comfy.nested_tensor + + crop_latent = {"samples": crop_up} + if len(tensors) > 1: + crop_latent["samples"] = comfy.nested_tensor.NestedTensor( + [crop_up] + [t.detach().clone() for t in tensors[1:]]) + tail_sigmas = torch.tensor(STAGE2_SIGMAS, dtype=torch.float32) + refined = chain._sample_chunk( + model, sampler, tail_sigmas, int(seed) + 7777, cfg, + _strip_layout_conds(positive), _strip_layout_conds(negative), crop_latent) + refined_crop = chain._latent_tensors(refined)[0] + + # 4) Back to grid size, feathered paste through the CLIPSeg silhouette. + refined_down = _downscale_to(refined_crop, y1 - y0, x1 - x0) + m = (mask.to(device=video.device, dtype=video.dtype) * float(strength)).clamp(0.0, 1.0) + out_video = video.clone() + region = out_video[:, :, :, y0:y1, x0:x1] + out_video[:, :, :, y0:y1, x0:x1] = region + (refined_down.to(region.device) - region) * m + + out = dict(latent) + if len(tensors) > 1: + out["samples"] = comfy.nested_tensor.NestedTensor([out_video] + list(tensors[1:])) + else: + out["samples"] = out_video + note = (f"segmented_detail({','.join(targets)}, tube={y1 - y0}x{x1 - x0}" + f"/{lat_h}x{lat_w}, s={strength})") + return out, note diff --git a/samplers.py b/samplers.py index 49ddad3d..ce6242b7 100644 --- a/samplers.py +++ b/samplers.py @@ -2450,6 +2450,22 @@ class FunPackLTXAVSceneChainSampler: "default": False, "tooltip": "EXPERIMENTAL: source embed_guidance / score_slider from the taste direction learned on the prompts NEAREST this scene's prompt, instead of the single global liked-direction average. On every liked rating the Refiner records (prompt fingerprint -> that run's liked direction); with this on, each scene retrieves the similarity-weighted direction of its closest matches (a forest prompt pulls what worked on forests, not the mean across all prompts). Non-parametric retrieval — no extra model forward, just a cosine lookup + vector mean, and it can't collapse into a spurious attractor the way a value function can. Falls back to the global liked direction when no rated prompt is close enough (or the index is empty). Only affects embed_guidance / score_slider; needs refinement_key_input (or embed_guidance_source=absolute). UNVALIDATED LIVE.", }), + "segmented_detailing": ("BOOLEAN", { + "default": False, + "tooltip": "EXPERIMENTAL ADetailer-for-video: after each scene finishes denoising, CLIPSeg (text-prompted segmentation) locates the regions named in detail_targets on a few decoded keyframes; the matched region is cut out of the video latent as a spatiotemporal tube, pushed through Lightricks' trained latent upsampler (2x spatial — the official two-stage pipeline's stage-2 model), re-noised to the official stage-2 sigma (0.85) and re-denoised for the 3-step tail schedule, then downscaled back to its ORIGINAL latent size and pasted through the feathered CLIPSeg silhouette. Final resolution never changes — the upsample only lets the model resolve structure (fingers, small faces) at a higher working resolution. Cost ~= 4 x tube area fraction x 3 steps (hands ~+15%); tubes over 35% of the frame are refused (that's a re-render, not a detail pass). Audio untouched by construction. Requires detail_upsampler; silent no-op when nothing matches. UNVALIDATED LIVE.", + }), + "detail_targets": ("STRING", { + "default": "hands", + "tooltip": "Comma-separated regions to detail, in plain words ('hands', 'hands, feet', 'face'). Each becomes a CLIPSeg text query; matched regions merge into one tube per scene. CLIPSeg matches broad CLIP semantics, so malformed anatomy still lights up for its name.", + }), + "detail_upsampler": (cls._detail_upsampler_choices(), { + "default": "None", + "tooltip": "Latent upsampler checkpoint from models/latent_upscale_models (the LTX 2.3 spatial upsampler used by the official two-stage workflows). Required — segmented_detailing no-ops while this is None.", + }), + "detail_strength": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "Blend of the refined region into the frame at paste-back (through the feathered CLIPSeg mask). 1.0 = full replacement inside the silhouette; 0 disables the pass entirely.", + }), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -2474,6 +2490,14 @@ class FunPackLTXAVSceneChainSampler: except Exception: return ["None"] + @classmethod + def _detail_upsampler_choices(cls): + try: + import folder_paths + return ["None"] + folder_paths.get_filename_list("latent_upscale_models") + except Exception: + return ["None"] + def _is_nested(self, samples): return bool(getattr(samples, "is_nested", False)) @@ -4637,6 +4661,8 @@ class FunPackLTXAVSceneChainSampler: carry_overlap_through_anchor=False, plateau_cache=False, plateau_cache_threshold=0.975, taste_nearest_prompt=False, + segmented_detailing=False, detail_targets="hands", + detail_upsampler="None", detail_strength=1.0, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -4825,6 +4851,7 @@ class FunPackLTXAVSceneChainSampler: scene_runs: list = [] joyai_bank = _JoyAIMemoryBank(joyai_memory_size, joyai_fix_frames) if joyai_memory else None _identity_overlap_state: dict = {} + _detail_upsampler_model = None # lazy: loaded at the first detailed scene for scene_index, scene_cond in enumerate(scene_conditionings): scene_positive = [scene_cond] @@ -5083,6 +5110,29 @@ class FunPackLTXAVSceneChainSampler: sampled = self._crop_video_tail(sampled, guide_tail) if audio_tail > 0: sampled = self._crop_audio_tail(sampled, audio_tail) + # Segmented detailing runs on the clean, fully-cropped scene (guide/audio + # tails gone, overlap still present — the tube is spatial, so carried head + # frames detail together with the rest) and BEFORE the JoyAI harvest, so + # cross-shot memory banks the improved frames. + if segmented_detailing and detail_strength > 0 and detail_upsampler != "None": + try: + try: + from . import detailing as _detailing + except ImportError: + import detailing as _detailing + if _detail_upsampler_model is None: + _detail_upsampler_model = _detailing.load_latent_upsampler(detail_upsampler) + _t_detail0 = _time.perf_counter() + sampled, _detail_note = _detailing.detail_refine_scene( + self, model, vae, sampler, scene_positive, scene_negative, sampled, + detail_targets, _detail_upsampler_model, scene_seed, cfg, + strength=detail_strength, debug=debug_log) + if _detail_note: + run_mechanisms.append(_detail_note) + _phase_sampling += _time.perf_counter() - _t_detail0 + except Exception as _detail_exc: + # A failed detail pass must never cost the scene itself. + print(f"[FunPackSceneChain] segmented_detailing failed, scene kept as-is: {_detail_exc}") if joyai_bank is not None: # Harvest from the clean, fully-cropped scene so injected memory tails never re-enter # the bank. Scene 0 seeds the pinned anchor (num_fix); later scenes roll in. The audio diff --git a/tests/test_detailing.py b/tests/test_detailing.py new file mode 100644 index 00000000..d168f04f --- /dev/null +++ b/tests/test_detailing.py @@ -0,0 +1,248 @@ +"""Unit tests for segmented detailing (detailing.py). + +Covers the pure pieces (target parsing, tube geometry + area cap, layout-cond +stripping, spatial downscale) and the detail_refine_scene contract: bit-identical +no-op when disabled / nothing detected, and the crop -> refine -> feathered +paste round trip with the model-touching stages stubbed out. +""" +import sys +import types +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +sys.modules.setdefault("comfy", types.ModuleType("comfy")) +_nested = types.ModuleType("comfy.nested_tensor") + + +class _FakeNested: + is_nested = True + + def __init__(self, tensors): + self._tensors = list(tensors) + + def unbind(self): + return list(self._tensors) + + +_nested.NestedTensor = _FakeNested +sys.modules.setdefault("comfy.nested_tensor", _nested) + + +@pytest.fixture(autouse=True) +def _nested_stub(monkeypatch): + """Re-install the usable NestedTensor stub around every test. + + Other test modules stub comfy.nested_tensor.NestedTensor = object on the SAME + shared sys.modules entry at collection time (whichever file is collected last + wins), and stubbed submodules never run real import machinery, so the parent + attribute must be bound by hand too.""" + monkeypatch.setattr(sys.modules["comfy.nested_tensor"], "NestedTensor", + _FakeNested, raising=False) + monkeypatch.setattr(sys.modules["comfy"], "nested_tensor", + sys.modules["comfy.nested_tensor"], raising=False) + +import detailing # noqa: E402 + + +# --------------------------------------------------------------------------- +# parse_targets +# --------------------------------------------------------------------------- + +def test_parse_targets_splits_and_strips(): + assert detailing.parse_targets(" hands , feet,,") == ["hands", "feet"] + assert detailing.parse_targets("") == [] + assert detailing.parse_targets(None) == [] + + +# --------------------------------------------------------------------------- +# find_tube +# --------------------------------------------------------------------------- + +def _heat_with_hot_square(h=64, w=64, y=(20, 30), x=(40, 50), value=0.9): + heat = torch.zeros(h, w) + heat[y[0]:y[1], x[0]:x[1]] = value + return heat + +def test_find_tube_boxes_hot_region_with_padding(): + heat = _heat_with_hot_square() + tube = detailing.find_tube(heat, 64, 64, threshold=0.35) + assert tube is not None + y0, y1, x0, x1, mask = tube + # Hot square plus TUBE_PAD on each side (interpolation is identity at same size). + assert y0 <= 20 - detailing.TUBE_PAD + 1 and y1 >= 30 + detailing.TUBE_PAD - 1 + assert x0 <= 40 - detailing.TUBE_PAD + 1 and x1 >= 50 + detailing.TUBE_PAD - 1 + assert mask.shape == (1, 1, 1, y1 - y0, x1 - x0) + assert float(mask.max()) <= 1.0 and float(mask.min()) >= 0.0 + # Mask peaks inside the hot region, feathers to ~0 at the padded border. + assert float(mask[..., (y1 - y0) // 2, (x1 - x0) // 2]) > 0.5 + assert float(mask[..., 0, 0]) < 0.2 + +def test_find_tube_none_below_threshold(): + heat = torch.full((64, 64), 0.1) + assert detailing.find_tube(heat, 64, 64, threshold=0.35) is None + assert detailing.find_tube(None, 64, 64) is None + +def test_find_tube_refuses_oversized_region(): + heat = torch.full((64, 64), 0.9) # whole frame hot -> a re-render, not a detail pass + assert detailing.find_tube(heat, 64, 64, threshold=0.35) is None + +def test_find_tube_enforces_minimum_edge(): + heat = _heat_with_hot_square(y=(31, 32), x=(31, 32)) # single hot cell + tube = detailing.find_tube(heat, 64, 64, threshold=0.35) + assert tube is not None + y0, y1, x0, x1, _ = tube + assert (y1 - y0) >= detailing.MIN_TUBE_EDGE + assert (x1 - x0) >= detailing.MIN_TUBE_EDGE + + +# --------------------------------------------------------------------------- +# _strip_layout_conds / _downscale_to +# --------------------------------------------------------------------------- + +def test_strip_layout_conds_removes_guide_keys_only(): + cond = [[torch.zeros(1, 4), {"keyframe_idxs": "kf", "guiding_latent": "g", + "pooled_output": "keep"}]] + out = detailing._strip_layout_conds(cond) + assert out[0][1] == {"pooled_output": "keep"} + # Original untouched (the scene's own conds must not lose their guides). + assert "keyframe_idxs" in cond[0][1] + +def test_downscale_to_shape_and_content(): + video = torch.arange(2 * 3 * 4 * 8 * 8, dtype=torch.float32).reshape(2, 3, 4, 8, 8) + down = detailing._downscale_to(video, 4, 4) + assert down.shape == (2, 3, 4, 4, 4) + # Area downscale of a constant tensor is exact. + const = torch.full((1, 2, 2, 6, 6), 3.5) + assert torch.allclose(detailing._downscale_to(const, 3, 3), torch.full((1, 2, 2, 3, 3), 3.5)) + + +# --------------------------------------------------------------------------- +# detail_refine_scene +# --------------------------------------------------------------------------- + +class _FakeChain: + """Just enough of FunPackLTXAVSceneChainSampler for detail_refine_scene.""" + + def __init__(self): + self.chunk_calls = [] + + def _latent_tensors(self, latent): + samples = latent["samples"] + if getattr(samples, "is_nested", False): + return list(samples.unbind()) + return [samples] + + def _sample_chunk(self, model, sampler, sigmas, seed, cfg, positive, negative, latent, + **kwargs): + self.chunk_calls.append({"sigmas": sigmas, "seed": seed, + "positive": positive, "latent": latent}) + return latent # identity "refine": returns the crop unchanged + + +def _scene_latent(nested=False): + video = torch.randn(1, 8, 5, 16, 24) + if not nested: + return {"samples": video}, video + audio = torch.randn(1, 4, 40) + return {"samples": _FakeNested([video, audio])}, video + +def test_refine_noop_without_targets_or_upsampler(): + chain = _FakeChain() + latent, _ = _scene_latent() + out, note = detailing.detail_refine_scene( + chain, None, None, None, [], [], latent, "", object(), 0, 1.0) + assert out is latent and note is None # same object: bit-identical guarantee + out, note = detailing.detail_refine_scene( + chain, None, None, None, [], [], latent, "hands", None, 0, 1.0) + assert out is latent and note is None + assert chain.chunk_calls == [] + +def test_refine_noop_when_nothing_detected(monkeypatch): + chain = _FakeChain() + latent, _ = _scene_latent() + monkeypatch.setattr(detailing, "_decode_detection_frames", lambda *a, **k: []) + monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: None) + out, note = detailing.detail_refine_scene( + chain, None, None, None, [], [], latent, "hands", object(), 0, 1.0) + assert out is latent and note is None + assert chain.chunk_calls == [] + +def _patch_detection_and_upsampler(monkeypatch): + """Hot square in pixel space; upsampler = exact 2x nearest (model-free).""" + heat = torch.zeros(160, 240) + heat[40:80, 120:180] = 0.9 + monkeypatch.setattr(detailing, "_decode_detection_frames", + lambda *a, **k: [(0, torch.zeros(160, 240, 3))]) + monkeypatch.setattr(detailing, "_clipseg_heat", lambda *a, **k: heat) + monkeypatch.setattr( + detailing, "_run_upsampler", + lambda upsampler, crop, vae, debug=False: torch.nn.functional.interpolate( + crop, scale_factor=(1, 2, 2), mode="nearest")) + return heat + +def test_refine_round_trip_touches_only_the_tube(monkeypatch): + chain = _FakeChain() + latent, video = _scene_latent() + _patch_detection_and_upsampler(monkeypatch) + pos = [[torch.zeros(1, 4), {"keyframe_idxs": "kf", "pooled_output": "keep"}]] + out, note = detailing.detail_refine_scene( + chain, "model", "vae", "sampler", pos, [], latent, "hands", object(), + seed=11, cfg=1.0) + assert note is not None and "segmented_detail" in note + assert len(chain.chunk_calls) == 1 + call = chain.chunk_calls[0] + # Official stage-2 tail schedule and a detail-offset seed. + assert torch.allclose(call["sigmas"], torch.tensor(detailing.STAGE2_SIGMAS)) + assert call["seed"] == 11 + 7777 + # Layout conds stripped for the crop pass. + assert "keyframe_idxs" not in call["positive"][0][1] + assert call["positive"][0][1]["pooled_output"] == "keep" + # The crop passed to the refine is 2x the tube size. + crop = call["latent"]["samples"] + out_video = out["samples"] + assert out_video.shape == video.shape + # With an identity refine, upsample(2x)->downscale(area) reproduces the crop, so + # the paste changes (nearly) nothing — but crucially everything OUTSIDE the tube + # must be exactly untouched regardless of what the refine did. + diff = (out_video - video).abs().amax(dim=(0, 1, 2)) # [H, W] + hot = diff > 1e-6 + assert not hot[:2].any() and not hot[:, :2].any() # far corners untouched + +def test_refine_preserves_audio_stream(monkeypatch): + chain = _FakeChain() + latent, video = _scene_latent(nested=True) + audio = chain._latent_tensors(latent)[1] + _patch_detection_and_upsampler(monkeypatch) + out, note = detailing.detail_refine_scene( + chain, "model", "vae", "sampler", [], [], latent, "hands", object(), + seed=0, cfg=1.0) + assert note is not None + out_tensors = chain._latent_tensors(out) + assert len(out_tensors) == 2 + # Audio comes back as the SAME tensor the scene finished with: protected by + # construction, not by a lossy round trip. + assert out_tensors[1] is audio + # The crop pass carried an audio stream so joint attention saw a legal AV latent. + crop_samples = chain.chunk_calls[0]["latent"]["samples"] + assert getattr(crop_samples, "is_nested", False) + +def test_refine_strength_scales_paste(monkeypatch): + chain = _FakeChain() + latent, video = _scene_latent() + _patch_detection_and_upsampler(monkeypatch) + # Non-identity refine: chunk returns a shifted crop so the paste has signal. + def _shifted_chunk(model, sampler, sigmas, seed, cfg, positive, negative, crop_latent, **kw): + return {"samples": crop_latent["samples"] + 10.0} + chain._sample_chunk = _shifted_chunk + out_full, _ = detailing.detail_refine_scene( + chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0, strength=1.0) + out_half, _ = detailing.detail_refine_scene( + chain, "m", "v", "s", [], [], latent, "hands", object(), 0, 1.0, strength=0.5) + d_full = (out_full["samples"] - video).abs().sum() + d_half = (out_half["samples"] - video).abs().sum() + assert d_full > 0 + assert abs(float(d_half / d_full) - 0.5) < 1e-3