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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Fable 5
Claude Fable 5
parent
302028a491
commit
a6e4a934f1
+349
@@ -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
|
||||
+50
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user