591 lines
24 KiB
Python
591 lines
24 KiB
Python
"""
|
||
ComfyUI custom nodes for VOID — Video Object and Interaction Deletion.
|
||
|
||
Built on CogVideoX-Fun-V1.5-5b-InP fine-tuned for interaction-aware
|
||
video inpainting (netflix/void-model).
|
||
|
||
Workflow:
|
||
VOID Load VAE ─────────────────┐
|
||
VOID Load Text Encoder ─────────┤→ VOID Loader → VOID Inference
|
||
│
|
||
└────────────────────────────────
|
||
|
||
Model files — download once and place in the ComfyUI models folders:
|
||
|
||
VAE (~1.5 GB)
|
||
Source : https://huggingface.co/alibaba-pai/CogVideoX-Fun-V1.5-5b-InP
|
||
File : vae/diffusion_pytorch_model.safetensors
|
||
Save to: ComfyUI/models/vae/
|
||
|
||
T5-XXL text encoder (~5 GB, fp8)
|
||
Source : bundled with most FLUX setups as t5xxl_fp8_e4m3fn.safetensors
|
||
https://huggingface.co/Comfy-Org/flux1-dev
|
||
Save to: ComfyUI/models/clip/
|
||
|
||
VOID transformers (~10 GB each)
|
||
Source : https://huggingface.co/netflix/void-model
|
||
Files : void_pass1.safetensors, void_pass2.safetensors
|
||
Save to: ComfyUI/models/void-model/
|
||
|
||
Note on custom load nodes
|
||
--------------------------
|
||
ComfyUI's built-in "Load VAE" and "Load CLIP" nodes are incompatible here:
|
||
• CogVideoX uses a 3D temporal VAE (AutoencoderKLCogVideoX), not the SD VAE.
|
||
• The pipeline requires raw HuggingFace T5EncoderModel + T5Tokenizer objects,
|
||
not ComfyUI's CLIP abstraction.
|
||
"""
|
||
|
||
import os
|
||
import torch
|
||
import folder_paths
|
||
import comfy.model_management as mm
|
||
from comfy_api.latest import io
|
||
|
||
from .void_pipeline import (
|
||
load_transformer,
|
||
load_vae,
|
||
load_text_encoder,
|
||
assemble_pipeline,
|
||
preprocess_video,
|
||
preprocess_quadmask,
|
||
postprocess_video,
|
||
run_pass1,
|
||
run_pass2,
|
||
snap_to_valid_dims,
|
||
snap_to_valid_frames,
|
||
NEGATIVE_PROMPT_DEFAULT,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Custom data types (not compatible with core VAE / CLIP types)
|
||
# ---------------------------------------------------------------------------
|
||
VoidVAE = io.Custom("VOID_VAE")
|
||
VoidTextEncoder = io.Custom("VOID_TEXT_ENCODER")
|
||
VoidPipeline = io.Custom("VOID_PIPELINE")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Model file names expected in their respective models/ subfolders
|
||
# ---------------------------------------------------------------------------
|
||
_VOID_PASS1_VERSIONS = ["void_pass1.safetensors"]
|
||
_VOID_PASS2_VERSIONS = ["void_pass2.safetensors"]
|
||
|
||
_MEMORY_MODES = [
|
||
"fp8 + cpu offload (~12 GB)", # recommended default
|
||
"cpu offload (~16 GB)",
|
||
"full GPU (~40 GB)",
|
||
"sequential (~min VRAM)",
|
||
]
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Module-level model cache
|
||
# Keyed by absolute file path → loaded model object.
|
||
# Survives ComfyUI re-executions so loader nodes return immediately when the
|
||
# same file is selected again without re-reading from disk.
|
||
# ---------------------------------------------------------------------------
|
||
_CACHE: dict = {}
|
||
|
||
|
||
def _default_dir(subfolder: str) -> str:
|
||
return os.path.join(folder_paths.models_dir, subfolder)
|
||
|
||
|
||
def _cached_load(path: str, loader_fn, *args, **kwargs):
|
||
"""Return a cached model if already loaded, otherwise call loader_fn."""
|
||
key = os.path.normpath(os.path.abspath(path))
|
||
if key not in _CACHE:
|
||
_CACHE[key] = loader_fn(*args, **kwargs)
|
||
else:
|
||
print(f"[VOID] cache hit: {os.path.basename(path)}")
|
||
return _CACHE[key]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: VOID Load VAE
|
||
# ---------------------------------------------------------------------------
|
||
class VoidLoadVAE(io.ComfyNode):
|
||
"""
|
||
Load the CogVideoX-Fun 3D VAE.
|
||
Place diffusion_pytorch_model.safetensors (renamed as preferred) in
|
||
ComfyUI/models/vae/ and select it from the dropdown.
|
||
Model is cached in memory — subsequent runs return instantly.
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
options = folder_paths.get_filename_list("vae") or ["(empty — add files to models/vae/)"]
|
||
return io.Schema(
|
||
node_id="VoidLoadVAE",
|
||
display_name="VOID Load VAE",
|
||
category="VOID",
|
||
description=(
|
||
"Load the CogVideoX-Fun 3D VAE from models/vae/. "
|
||
"Source: huggingface.co/alibaba-pai/CogVideoX-Fun-V1.5-5b-InP "
|
||
"→ vae/diffusion_pytorch_model.safetensors"
|
||
),
|
||
inputs=[
|
||
io.Combo.Input("vae_name", options=options),
|
||
],
|
||
outputs=[VoidVAE.Output("VOID_VAE")],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, vae_name: str):
|
||
path = os.path.join(_default_dir("vae"), vae_name)
|
||
if not os.path.exists(path):
|
||
raise FileNotFoundError(
|
||
f"VAE not found: {path}\n"
|
||
"Source: https://huggingface.co/alibaba-pai/CogVideoX-Fun-V1.5-5b-InP\n"
|
||
"File : vae/diffusion_pytorch_model.safetensors"
|
||
)
|
||
vae = _cached_load(path, load_vae, path)
|
||
return io.NodeOutput({"vae": vae, "path": path})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: VOID Load Text Encoder
|
||
# ---------------------------------------------------------------------------
|
||
class VoidLoadTextEncoder(io.ComfyNode):
|
||
"""
|
||
Load the T5-v1_1-XXL text encoder (fp8 supported).
|
||
Select t5xxl_fp8_e4m3fn.safetensors from models/clip/ — this file ships
|
||
with most FLUX setups. Model is cached after first load.
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
all_clip = folder_paths.get_filename_list("clip")
|
||
t5_files = [f for f in all_clip if "t5" in f.lower()] or all_clip
|
||
options = t5_files or ["(empty — add t5xxl file to models/clip/)"]
|
||
return io.Schema(
|
||
node_id="VoidLoadTextEncoder",
|
||
display_name="VOID Load Text Encoder",
|
||
category="VOID",
|
||
description=(
|
||
"Load the T5-XXL encoder from models/clip/. "
|
||
"fp8 weights (t5xxl_fp8_e4m3fn.safetensors) are cast to bf16 on load. "
|
||
"Source: huggingface.co/Comfy-Org/flux1-dev"
|
||
),
|
||
inputs=[
|
||
io.Combo.Input("t5_name", options=options),
|
||
],
|
||
outputs=[VoidTextEncoder.Output("VOID_TEXT_ENCODER")],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, t5_name: str):
|
||
path = os.path.join(_default_dir("clip"), t5_name)
|
||
if not os.path.exists(path):
|
||
raise FileNotFoundError(
|
||
f"T5 encoder not found: {path}\n"
|
||
"Source: https://huggingface.co/Comfy-Org/flux1-dev\n"
|
||
"File : t5xxl_fp8_e4m3fn.safetensors"
|
||
)
|
||
result = _cached_load(path, load_text_encoder, path)
|
||
text_encoder, tokenizer = result
|
||
return io.NodeOutput({"text_encoder": text_encoder, "tokenizer": tokenizer, "path": path})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: VOID Loader
|
||
# ---------------------------------------------------------------------------
|
||
class VoidLoader(io.ComfyNode):
|
||
"""
|
||
Load both VOID transformer checkpoints and assemble the inference pipelines.
|
||
Place void_pass1.safetensors and void_pass2.safetensors in models/void-model/.
|
||
Transformers are cached after first load.
|
||
|
||
Memory modes
|
||
────────────
|
||
fp8 + cpu offload (~12 GB) Recommended. fp8-quantised transformer, CPU offload.
|
||
cpu offload (~16 GB) bf16 transformer, moves to CPU between steps.
|
||
full GPU (~40 GB) Everything on GPU — fastest, most VRAM.
|
||
sequential (~min VRAM) Layer-by-layer offload — slowest, least VRAM.
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="VoidLoader",
|
||
display_name="VOID Loader",
|
||
category="VOID",
|
||
description=(
|
||
"Assemble both VOID inference pipelines from models/void-model/. "
|
||
"Source: huggingface.co/netflix/void-model"
|
||
),
|
||
inputs=[
|
||
io.Combo.Input(
|
||
"pass1_model",
|
||
options=_VOID_PASS1_VERSIONS,
|
||
tooltip="Pass 1 checkpoint — base object removal.",
|
||
),
|
||
io.Combo.Input(
|
||
"pass2_model",
|
||
options=_VOID_PASS2_VERSIONS,
|
||
tooltip="Pass 2 checkpoint — warped-noise temporal refinement.",
|
||
),
|
||
VoidVAE.Input("void_vae"),
|
||
VoidTextEncoder.Input("void_text_encoder"),
|
||
io.Combo.Input(
|
||
"memory_mode",
|
||
options=_MEMORY_MODES,
|
||
default="fp8 + cpu offload (~12 GB)",
|
||
),
|
||
],
|
||
outputs=[VoidPipeline.Output("VOID_PIPELINE")],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
pass1_model: str,
|
||
pass2_model: str,
|
||
void_vae: dict,
|
||
void_text_encoder: dict,
|
||
memory_mode: str,
|
||
):
|
||
base = _default_dir("void-model")
|
||
p1_path = os.path.join(base, pass1_model)
|
||
p2_path = os.path.join(base, pass2_model)
|
||
|
||
for path, name in [(p1_path, pass1_model), (p2_path, pass2_model)]:
|
||
if not os.path.exists(path):
|
||
raise FileNotFoundError(
|
||
f"VOID checkpoint not found: {path}\n"
|
||
"Source: https://huggingface.co/netflix/void-model\n"
|
||
f"Save {name} to {base}"
|
||
)
|
||
|
||
vae = void_vae["vae"]
|
||
text_encoder = void_text_encoder["text_encoder"]
|
||
tokenizer = void_text_encoder["tokenizer"]
|
||
|
||
# Transformers and assembled pipelines are both cached.
|
||
# Pipelines must be cached (not just transformers) because enable_model_cpu_offload
|
||
# registers AlignDevicesHook on the transformer — re-calling assemble_pipeline on
|
||
# the same transformer object stacks hooks on every queue run, which corrupts
|
||
# device placement and leaks memory. Keying by (path, memory_mode) ensures a
|
||
# new pipeline is only built when the config actually changes.
|
||
transformer1 = _cached_load(p1_path, load_transformer, p1_path, use_vae_mask=True)
|
||
p1_pipeline_key = os.path.normpath(os.path.abspath(p1_path)) + ":pipeline:" + memory_mode
|
||
if p1_pipeline_key not in _CACHE:
|
||
_CACHE[p1_pipeline_key] = assemble_pipeline(transformer1, vae, text_encoder, tokenizer, memory_mode)
|
||
else:
|
||
print("[VOID] cache hit: Pass 1 pipeline")
|
||
pipeline1 = _CACHE[p1_pipeline_key]
|
||
|
||
transformer2 = _cached_load(p2_path, load_transformer, p2_path, use_vae_mask=True)
|
||
p2_pipeline_key = os.path.normpath(os.path.abspath(p2_path)) + ":pipeline:" + memory_mode
|
||
if p2_pipeline_key not in _CACHE:
|
||
_CACHE[p2_pipeline_key] = assemble_pipeline(transformer2, vae, text_encoder, tokenizer, memory_mode)
|
||
else:
|
||
print("[VOID] cache hit: Pass 2 pipeline")
|
||
pipeline2 = _CACHE[p2_pipeline_key]
|
||
|
||
return io.NodeOutput({
|
||
"pipeline1": pipeline1,
|
||
"pipeline2": pipeline2,
|
||
"vae": vae,
|
||
"memory_mode": memory_mode,
|
||
})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: VOID Inference
|
||
# ---------------------------------------------------------------------------
|
||
class VoidInference(io.ComfyNode):
|
||
"""
|
||
Run VOID — Pass 1 object removal followed by optional Pass 2 refinement.
|
||
|
||
Pass 1 Base inpainting using the quadmask. guidance_scale=1.0 (distillation
|
||
mode) is the setting from the paper and recommended for most clips.
|
||
|
||
Pass 2 Warped-noise refinement for better temporal consistency. Can run
|
||
with random init (enable_pass2=True, no warped_noise_path) or with
|
||
optical-flow-warped noise from make_warped_noise.py (void-model repo).
|
||
|
||
Quadmask pixel values
|
||
─────────────────────
|
||
0 primary object to remove
|
||
63 overlap of primary + affected region
|
||
127 affected region (interaction/shadow effects)
|
||
255 background — keep unchanged
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="VoidInference",
|
||
display_name="VOID Inference",
|
||
category="VOID",
|
||
description=(
|
||
"Run VOID Pass 1 inpainting + optional Pass 2 refinement. "
|
||
"Both outputs are always connected — Pass 2 mirrors Pass 1 when disabled."
|
||
),
|
||
inputs=[
|
||
VoidPipeline.Input("void_pipeline"),
|
||
io.Image.Input("video", tooltip=(
|
||
"Input video [T, H, W, C]. "
|
||
"Scaled down to fit within 480×720 if needed, "
|
||
"then snapped to the nearest multiple of 16. "
|
||
"Frame count is trimmed to (n−1) % 4 == 0, max 85."
|
||
)),
|
||
io.Image.Input("quadmask", tooltip=(
|
||
"Quadmask [T, H, W, C]. "
|
||
"0=remove 63=overlap 127=affected 255=background."
|
||
)),
|
||
# Prompts
|
||
io.String.Input("prompt", multiline=True, default="",
|
||
tooltip="Describe the scene after object removal (e.g. 'A table with a cup.')."),
|
||
io.String.Input("negative_prompt", multiline=True, default=NEGATIVE_PROMPT_DEFAULT,
|
||
tooltip="Used by both passes. Pass 2 ignores this (always runs with empty negative)."),
|
||
# Resolution / frame count
|
||
io.Int.Input("width", default=672, min=16, step=16,
|
||
tooltip="Output width in pixels. Snapped to nearest multiple of 16."),
|
||
io.Int.Input("height", default=384, min=16, step=16,
|
||
tooltip="Output height in pixels. Snapped to nearest multiple of 16."),
|
||
io.Int.Input("num_frames", default=85, min=5,
|
||
tooltip="Number of frames. Trimmed to nearest (n−1) % 4 == 0, max 85."),
|
||
# Shared / Pass 1 parameters
|
||
io.Int.Input("num_inference_steps", default=50, min=10, max=200,
|
||
tooltip="Denoising steps. Applies to both passes."),
|
||
io.Int.Input("seed", default=42, min=0, max=2**31 - 1,
|
||
tooltip="RNG seed. Applies to both passes."),
|
||
io.Float.Input("guidance_scale", default=1.0, min=1.0, max=20.0, step=0.5,
|
||
tooltip="Pass 1 CFG scale. 1.0 = distillation mode (paper default)."),
|
||
io.Float.Input("denoise_strength", default=1.0, min=0.1, max=1.0, step=0.05,
|
||
tooltip="Pass 1 denoise strength. 1.0 = full inpaint."),
|
||
# Pass 2 parameters
|
||
io.Boolean.Input("enable_pass2", default=True,
|
||
tooltip="Run Pass 2 refinement. When off, Pass 2 output mirrors Pass 1."),
|
||
io.Float.Input("pass2_guidance_scale", default=6.0, min=1.0, max=20.0, step=0.5,
|
||
tooltip="Pass 2 CFG scale. 6.0 recommended."),
|
||
io.String.Input("warped_noise_path", default="",
|
||
tooltip=(
|
||
"Optional path to noises.npy from make_warped_noise.py (void-model repo). "
|
||
"Seeds Pass 2 with optical-flow-warped noise for better temporal consistency. "
|
||
"Leave empty to use random init."
|
||
)),
|
||
],
|
||
outputs=[
|
||
io.Image.Output("Pass 1"),
|
||
io.Image.Output("Pass 2"),
|
||
],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
void_pipeline: dict,
|
||
video: torch.Tensor,
|
||
quadmask: torch.Tensor,
|
||
prompt: str,
|
||
negative_prompt: str,
|
||
width: int,
|
||
height: int,
|
||
num_frames: int,
|
||
num_inference_steps: int,
|
||
seed: int,
|
||
guidance_scale: float,
|
||
denoise_strength: float,
|
||
enable_pass2: bool,
|
||
pass2_guidance_scale: float,
|
||
warped_noise_path: str,
|
||
):
|
||
# Original input frame count — used to trim padded output back to this length.
|
||
T = video.shape[0]
|
||
|
||
# Cap to RoPE bounds (480×720) then snap to multiples of 16.
|
||
height, width = snap_to_valid_dims(height, width)
|
||
num_frames = snap_to_valid_frames(num_frames)
|
||
|
||
# Free VRAM held by other ComfyUI models before the pipeline runs.
|
||
mm.unload_all_models()
|
||
mm.soft_empty_cache()
|
||
|
||
pipeline1 = void_pipeline["pipeline1"]
|
||
device = _infer_device()
|
||
input_video = preprocess_video(video, height, width, num_frames).to(device)
|
||
input_mask = preprocess_quadmask(quadmask, height, width, num_frames).to(device)
|
||
|
||
interrupt_cb = _make_interrupt_callback()
|
||
pass1_raw = run_pass1(
|
||
pipeline=pipeline1,
|
||
input_video=input_video,
|
||
input_mask=input_mask,
|
||
prompt=prompt,
|
||
num_frames=num_frames,
|
||
height=height,
|
||
width=width,
|
||
guidance_scale=guidance_scale,
|
||
num_inference_steps=num_inference_steps,
|
||
seed=seed,
|
||
denoise_strength=denoise_strength,
|
||
negative_prompt=negative_prompt,
|
||
step_callback=interrupt_cb,
|
||
)
|
||
# Trim output to the original input frame count (temporal padding may have added frames).
|
||
pass1_frames = postprocess_video(pass1_raw)[:T]
|
||
|
||
if enable_pass2:
|
||
# Offload pass1 transformer to CPU — pass2 needs the full GPU budget.
|
||
if hasattr(pipeline1, "transformer"):
|
||
pipeline1.transformer.to("cpu")
|
||
torch.cuda.empty_cache()
|
||
|
||
pass2_raw = run_pass2(
|
||
pipeline=void_pipeline["pipeline2"],
|
||
input_video=input_video,
|
||
input_mask=input_mask,
|
||
prompt=prompt,
|
||
num_frames=num_frames,
|
||
height=height,
|
||
width=width,
|
||
guidance_scale=pass2_guidance_scale,
|
||
num_inference_steps=num_inference_steps,
|
||
seed=seed,
|
||
warped_noise_path=warped_noise_path.strip() or None,
|
||
negative_prompt="",
|
||
step_callback=interrupt_cb,
|
||
)
|
||
pass2_frames = postprocess_video(pass2_raw)[:T]
|
||
else:
|
||
pass2_frames = pass1_frames
|
||
|
||
return io.NodeOutput(pass1_frames, pass2_frames)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node: VOID Quadmask Builder
|
||
# ---------------------------------------------------------------------------
|
||
class VoidQuadmaskBuilder(io.ComfyNode):
|
||
"""
|
||
Compile a VOID quadmask from SAM2 / SAM3 segmentation masks.
|
||
|
||
Maps up to two binary masks onto the four VOID quadmask values:
|
||
|
||
0 — primary object (to remove)
|
||
63 — overlap: primary ∩ affected
|
||
127 — affected region only (shadows, contact, interaction effects)
|
||
255 — background (keep unchanged)
|
||
|
||
primary_mask (required) The object being removed.
|
||
Feed the first MASK output of SAM3VideoOutput,
|
||
SAM3Segmentation, or any ComfyUI MASK.
|
||
|
||
affected_mask (optional) The region physically affected by the object
|
||
(shadow, contact patch, wake, etc.).
|
||
Use a second SAM3 segmentation for this region.
|
||
If not connected the quadmask degrades gracefully
|
||
to a binary primary / background mask, which VOID
|
||
handles correctly.
|
||
|
||
Compatible with any node that outputs standard ComfyUI MASK type, including
|
||
ComfyUI-SAM3 (SAM3VideoOutput, SAM3Segmentation) and ComfyUI-SAM2.
|
||
|
||
Typical video workflow
|
||
──────────────────────
|
||
SAM3VideoSegmentation → SAM3Propagate → SAM3VideoOutput ─┐ primary_mask
|
||
SAM3VideoSegmentation → SAM3Propagate → SAM3VideoOutput ─┘ affected_mask (optional)
|
||
↓
|
||
VOID Quadmask Builder
|
||
↓
|
||
VOID Inference (quadmask)
|
||
"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="VoidQuadmaskBuilder",
|
||
display_name="VOID Quadmask Builder",
|
||
category="VOID",
|
||
description=(
|
||
"Compile a VOID quadmask from SAM2/SAM3 masks. "
|
||
"primary_mask = object to remove; "
|
||
"affected_mask = shadow/contact region (optional)."
|
||
),
|
||
inputs=[
|
||
io.Mask.Input(
|
||
"primary_mask",
|
||
tooltip=(
|
||
"Binary mask of the primary object to remove. "
|
||
"Accepts ComfyUI MASK [H,W] or [T,H,W] (video batch). "
|
||
"Maps to quadmask value 0."
|
||
),
|
||
),
|
||
io.Mask.Input(
|
||
"affected_mask",
|
||
optional=True,
|
||
tooltip=(
|
||
"Optional binary mask of the region physically affected by "
|
||
"the object (shadow, contact patch, wake, etc.). "
|
||
"Maps to quadmask value 127; overlap with primary → 63. "
|
||
"Leave disconnected for a simple binary quadmask."
|
||
),
|
||
),
|
||
],
|
||
outputs=[
|
||
io.Image.Output(
|
||
"QUADMASK",
|
||
tooltip=(
|
||
"Quadmask IMAGE [T, H, W, 3] with normalised values: "
|
||
"0=remove, 63/255=overlap, 127/255=affected, 1.0=background."
|
||
),
|
||
),
|
||
],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, primary_mask: torch.Tensor, affected_mask=None):
|
||
# Ensure [T, H, W] for both masks
|
||
if primary_mask.ndim == 2:
|
||
primary_mask = primary_mask.unsqueeze(0) # [H,W] → [1,H,W]
|
||
|
||
if affected_mask is not None and affected_mask.ndim == 2:
|
||
affected_mask = affected_mask.unsqueeze(0)
|
||
|
||
T, H, W = primary_mask.shape
|
||
p = primary_mask > 0.5
|
||
|
||
# Start: everything is background (1.0 = 255/255)
|
||
q = torch.ones(T, H, W, dtype=torch.float32, device=primary_mask.device)
|
||
|
||
if affected_mask is not None:
|
||
a = affected_mask > 0.5
|
||
# Affected only → 127/255
|
||
q = torch.where(a, torch.full_like(q, 127.0 / 255.0), q)
|
||
|
||
# Primary → 0.0
|
||
q = torch.where(p, torch.zeros_like(q), q)
|
||
|
||
if affected_mask is not None:
|
||
# Overlap (primary ∩ affected) → 63/255 (applied after primary so it
|
||
# correctly marks cells claimed by both masks)
|
||
q = torch.where(p & a, torch.full_like(q, 63.0 / 255.0), q)
|
||
|
||
# [T, H, W] → [T, H, W, 3] ComfyUI IMAGE
|
||
q = q.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous()
|
||
return io.NodeOutput(q)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _infer_device() -> torch.device:
|
||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
|
||
def _make_interrupt_callback():
|
||
"""
|
||
Return a diffusers callback_on_step_end that cancels inference when
|
||
the user clicks Cancel in ComfyUI.
|
||
|
||
ComfyUI sets an internal interrupt flag on cancel; mm.throw_exception_if_processing_interrupted
|
||
raises an exception if that flag is set. The exception propagates out of
|
||
the pipeline.__call__ loop, cleanly aborting the current denoising run.
|
||
"""
|
||
def _cb(pipe, step: int, timestep, callback_kwargs: dict) -> dict:
|
||
mm.throw_exception_if_processing_interrupted()
|
||
return {}
|
||
return _cb
|