Compare commits

..

6 Commits

Author SHA1 Message Date
85abace906 ComfyUI v0.22.2 2026-05-22 16:51:31 +00:00
f5d678d9ee [Partner Nodes] add new Rodin2.5 nodes (#14051)
* [Partner Nodes] add new Rodin2.5 nodes

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] fixed Quality Mesh Options

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] fix: remove non-supported "usdz"

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] fix: always pass seed to server

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] fix: set the default "material" value to "Shaded"

Signed-off-by: bigcat88 <bigcat88@icloud.com>

---------

Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-05-22 09:35:42 -07:00
59cafaf744 ComfyUI v0.22.1 2026-05-21 23:48:50 +00:00
13e2d133a6 [Partner Nodes] add widget for automatic upscaling for the ByteDance2Reference node (#14032)
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-05-21 15:39:13 -07:00
ef46f5de76 chore: update workflow templates to v0.9.82 (#14034) 2026-05-21 15:39:13 -07:00
7e02881b36 [Partner Nodes] add OpenRouter LLM node (#14007)
* [Partner Nodes] add reasoning widget to Anthropic node

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] add new OpenRouterLLM node

Signed-off-by: bigcat88 <bigcat88@icloud.com>

* [Partner Nodes] fix passing images to Grok LLM

Signed-off-by: bigcat88 <bigcat88@icloud.com>

---------

Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-05-21 15:39:12 -07:00
34 changed files with 1512 additions and 1834 deletions

View File

@ -1,5 +1,2 @@
# Admins
* @comfyanonymous @kosinkadink @guill @alexisrolland @rattus128 @kijai
/CODEOWNERS @comfyanonymous
/.ci/ @comfyanonymous
/.github/ @comfyanonymous

View File

@ -110,11 +110,13 @@ parser.add_argument("--preview-method", type=LatentPreviewMethod, default=Latent
parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
CACHE_RAM_AUTO_GB = -1.0
cache_group = parser.add_mutually_exclusive_group()
cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 25%% of system RAM (min 4GB, max 32GB), inactive 75%% of system RAM (min 12GB, max 96GB).")
cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
cache_group.add_argument("--cache-ram", nargs='?', const=CACHE_RAM_AUTO_GB, type=float, default=0, help="Use RAM pressure caching with the specified headroom threshold. If available RAM drops below the threshold the cache removes large items to free RAM. Default (when no value is provided): 25%% of system RAM (min 4GB, max 32GB).")
attn_group = parser.add_mutually_exclusive_group()
attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
@ -243,9 +245,6 @@ if comfy.options.args_parsing:
else:
args = parser.parse_args([])
if args.cache_ram is not None and len(args.cache_ram) > 2:
parser.error("--cache-ram accepts at most two values: active GB and inactive GB")
if args.windows_standalone_build:
args.auto_launch = True

View File

@ -484,23 +484,16 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori
return weight
def prefetch_prepared_value(value, counter, destination, stream, copy):
def prefetch_prepared_value(value, allocate_buffer, stream):
if isinstance(value, torch.Tensor):
size = comfy.memory_management.vram_aligned_size(value)
offset = counter[0]
counter[0] += size
if destination is None:
return value
dest = destination[offset:offset + size]
if copy:
comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
dest = allocate_buffer(comfy.memory_management.vram_aligned_size(value))
comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
return comfy.memory_management.interpret_gathered_like([value], dest)[0]
elif isinstance(value, weight_adapter.WeightAdapterBase):
return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, counter, destination, stream, copy))
return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, allocate_buffer, stream))
elif isinstance(value, tuple):
return tuple(prefetch_prepared_value(item, counter, destination, stream, copy) for item in value)
return tuple(prefetch_prepared_value(item, allocate_buffer, stream) for item in value)
elif isinstance(value, list):
return [prefetch_prepared_value(item, counter, destination, stream, copy) for item in value]
return [prefetch_prepared_value(item, allocate_buffer, stream) for item in value]
return value

View File

@ -15,7 +15,7 @@ class TensorFileSlice(NamedTuple):
size: int
def read_tensor_file_slice_into(tensor, destination, stream=None, destination2=None):
def read_tensor_file_slice_into(tensor, destination):
if isinstance(tensor, QuantizedTensor):
if not isinstance(destination, QuantizedTensor):
@ -23,17 +23,12 @@ def read_tensor_file_slice_into(tensor, destination, stream=None, destination2=N
if tensor._layout_cls != destination._layout_cls:
return False
if not read_tensor_file_slice_into(tensor._qdata, destination._qdata, stream=stream,
destination2=(destination2._qdata if destination2 is not None else None)):
if not read_tensor_file_slice_into(tensor._qdata, destination._qdata):
return False
dst_orig_dtype = destination._params.orig_dtype
destination._params.copy_from(tensor._params, non_blocking=False)
destination._params = dataclasses.replace(destination._params, orig_dtype=dst_orig_dtype)
if destination2 is not None:
dst_orig_dtype = destination2._params.orig_dtype
destination2._params.copy_from(destination._params, non_blocking=True)
destination2._params = dataclasses.replace(destination2._params, orig_dtype=dst_orig_dtype)
return True
info = getattr(tensor.untyped_storage(), "_comfy_tensor_file_slice", None)
@ -53,17 +48,6 @@ def read_tensor_file_slice_into(tensor, destination, stream=None, destination2=N
if info.size == 0:
return True
hostbuf = getattr(destination.untyped_storage(), "_comfy_hostbuf", None)
if hostbuf is not None:
stream_ptr = getattr(stream, "cuda_stream", 0) if stream is not None else 0
device_ptr = destination2.data_ptr() if destination2 is not None else 0
hostbuf.read_file_slice(file_obj, info.offset, info.size,
offset=destination.data_ptr() - hostbuf.get_raw_address(),
stream=stream_ptr,
device_ptr=device_ptr,
device=None if destination2 is None else destination2.device.index)
return True
buf_type = ctypes.c_ubyte * info.size
view = memoryview(buf_type.from_address(destination.data_ptr()))
@ -167,7 +151,7 @@ def set_ram_cache_release_state(callback, headroom):
extra_ram_release_callback = callback
RAM_CACHE_HEADROOM = max(0, int(headroom))
def extra_ram_release(target, free_active=False):
def extra_ram_release(target):
if extra_ram_release_callback is None:
return 0
return extra_ram_release_callback(target, free_active=free_active)
return extra_ram_release_callback(target)

View File

@ -31,7 +31,6 @@ from contextlib import nullcontext
import comfy.memory_management
import comfy.utils
import comfy.quant_ops
import comfy_aimdo.host_buffer
import comfy_aimdo.vram_buffer
class VRAMState(Enum):
@ -496,14 +495,6 @@ except:
current_loaded_models = []
DIRTY_MMAPS = set()
PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
#Freeing registerables on pressure does imply a GPU sync, so go big on
#the hysteresis so each expensive sync gives us back a good chunk.
REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
def module_size(module):
module_mem = 0
sd = module.state_dict()
@ -512,46 +503,27 @@ def module_size(module):
module_mem += t.nbytes
return module_mem
def mark_mmap_dirty(storage):
mmap_refs = getattr(storage, "_comfy_tensor_mmap_refs", None)
if mmap_refs is not None:
DIRTY_MMAPS.add(mmap_refs[0])
def free_pins(size, evict_active=False):
freed_total = 0
for loaded_model in reversed(current_loaded_models):
if size <= 0:
return freed_total
model = loaded_model.model
if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
freed = model.partially_unload_ram(size)
freed_total += freed
size -= freed
return freed_total
def ensure_pin_budget(size, evict_active=False):
shortfall = size + comfy.memory_management.RAM_CACHE_HEADROOM / 2 - psutil.virtual_memory().available
if shortfall <= 0:
return True
to_free = shortfall + PIN_PRESSURE_HYSTERESIS
return free_pins(to_free, evict_active=evict_active) >= shortfall
def ensure_pin_registerable(size, evict_active=False):
shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
if MAX_PINNED_MEMORY <= 0:
return False
if shortfall <= 0:
return True
shortfall += REGISTERABLE_PIN_HYSTERESIS
for loaded_model in reversed(current_loaded_models):
model = loaded_model.model
if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
shortfall -= model.unregister_inactive_pins(shortfall)
if shortfall <= 0:
return True
return shortfall <= REGISTERABLE_PIN_HYSTERESIS
def module_mmap_residency(module, free=False):
mmap_touched_mem = 0
module_mem = 0
bounced_mmaps = set()
sd = module.state_dict()
for k in sd:
t = sd[k]
module_mem += t.nbytes
storage = t._qdata.untyped_storage() if isinstance(t, comfy.quant_ops.QuantizedTensor) else t.untyped_storage()
if not getattr(storage, "_comfy_tensor_mmap_touched", False):
continue
mmap_touched_mem += t.nbytes
if not free:
continue
storage._comfy_tensor_mmap_touched = False
mmap_obj = storage._comfy_tensor_mmap_refs[0]
if mmap_obj in bounced_mmaps:
continue
mmap_obj.bounce()
bounced_mmaps.add(mmap_obj)
return mmap_touched_mem, module_mem
class LoadedModel:
def __init__(self, model):
@ -581,6 +553,9 @@ class LoadedModel:
def model_memory(self):
return self.model.model_size()
def model_mmap_residency(self, free=False):
return self.model.model_mmap_residency(free=free)
def model_loaded_memory(self):
return self.model.loaded_size()
@ -660,9 +635,15 @@ WINDOWS = any(platform.win32_ver())
EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
if WINDOWS:
import comfy.windows
EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue
if total_vram > (15 * 1024): # more extra reserved vram on 16GB+ cards
EXTRA_RESERVED_VRAM += 100 * 1024 * 1024
def get_free_ram():
return comfy.windows.get_free_ram()
else:
def get_free_ram():
return psutil.virtual_memory().available
if args.reserve_vram is not None:
EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024
@ -676,6 +657,7 @@ def minimum_inference_memory():
def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins_required=0, ram_required=0):
cleanup_models_gc()
comfy.memory_management.extra_ram_release(max(pins_required, ram_required))
unloaded_model = []
can_unload = []
unloaded_models = []
@ -691,9 +673,11 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
for x in can_unload_sorted:
i = x[-1]
memory_to_free = 1e32
if current_loaded_models[i].model.is_dynamic() and (not DISABLE_SMART_MEMORY or device is None):
pins_to_free = 1e32
if not DISABLE_SMART_MEMORY or device is None:
memory_to_free = 0 if device is None else memory_required - get_free_memory(device)
if for_dynamic:
pins_to_free = pins_required - get_free_ram()
if current_loaded_models[i].model.is_dynamic() and for_dynamic:
#don't actually unload dynamic models for the sake of other dynamic models
#as that works on-demand.
memory_required -= current_loaded_models[i].model.loaded_size()
@ -701,6 +685,18 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
if memory_to_free > 0 and current_loaded_models[i].model_unload(memory_to_free):
logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
unloaded_model.append(i)
if pins_to_free > 0:
logging.debug(f"PIN Unloading {current_loaded_models[i].model.model.__class__.__name__}")
current_loaded_models[i].model.partially_unload_ram(pins_to_free)
for x in can_unload_sorted:
i = x[-1]
ram_to_free = ram_required - psutil.virtual_memory().available
if ram_to_free <= 0 and i not in unloaded_model:
continue
resident_memory, _ = current_loaded_models[i].model_mmap_residency(free=True)
if resident_memory > 0:
logging.debug(f"RAM Unloading {current_loaded_models[i].model.model.__class__.__name__}")
for i in sorted(unloaded_model, reverse=True):
unloaded_models.append(current_loaded_models.pop(i))
@ -766,16 +762,29 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
model_to_unload.model.detach(unpatch_all=False)
model_to_unload.model_finalizer.detach()
total_memory_required = {}
total_pins_required = {}
total_ram_required = {}
for loaded_model in models_to_load:
device = loaded_model.device
total_memory_required[device] = total_memory_required.get(device, 0) + loaded_model.model_memory_required(device)
resident_memory, model_memory = loaded_model.model_mmap_residency()
pinned_memory = loaded_model.model.pinned_memory_size()
#FIXME: This can over-free the pins as it budgets to pin the entire model. We should
#make this JIT to keep as much pinned as possible.
pins_required = model_memory - pinned_memory
ram_required = model_memory - resident_memory
total_pins_required[device] = total_pins_required.get(device, 0) + pins_required
total_ram_required[device] = total_ram_required.get(device, 0) + ram_required
for device in total_memory_required:
if device != torch.device("cpu"):
free_memory(total_memory_required[device] * 1.1 + extra_mem,
device,
for_dynamic=free_for_dynamic)
for_dynamic=free_for_dynamic,
pins_required=total_pins_required[device],
ram_required=total_ram_required[device])
for device in total_memory_required:
if device != torch.device("cpu"):
@ -1171,7 +1180,6 @@ STREAM_CAST_BUFFERS = {}
LARGEST_CASTED_WEIGHT = (None, 0)
STREAM_AIMDO_CAST_BUFFERS = {}
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
STREAM_PIN_BUFFERS = {}
DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
@ -1212,66 +1220,21 @@ def get_aimdo_cast_buffer(offload_stream, device):
if cast_buffer is None:
cast_buffer = comfy_aimdo.vram_buffer.VRAMBuffer(DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE, device.index)
STREAM_AIMDO_CAST_BUFFERS[offload_stream] = cast_buffer
return cast_buffer
def get_pin_buffer(offload_stream):
pin_buffer = STREAM_PIN_BUFFERS.get(offload_stream, None)
if pin_buffer is None:
pin_buffer = comfy_aimdo.host_buffer.HostBuffer(0, 0, pinned_hostbuf_size(8 * 1024**3))
STREAM_PIN_BUFFERS[offload_stream] = pin_buffer
elif offload_stream is not None:
event = getattr(pin_buffer, "_comfy_event", None)
if event is not None:
event.synchronize()
delattr(pin_buffer, "_comfy_event")
return pin_buffer
def resize_pin_buffer(pin_buffer, size):
global TOTAL_PINNED_MEMORY
old_size = pin_buffer.size
if size <= old_size:
return True
growth = size - old_size
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
ensure_pin_budget(growth, evict_active=True)
ensure_pin_registerable(growth, evict_active=True)
try:
pin_buffer.extend(size=size, reallocate=True)
except RuntimeError:
return False
TOTAL_PINNED_MEMORY += pin_buffer.size - old_size
return True
def reset_cast_buffers():
global TOTAL_PINNED_MEMORY
global LARGEST_CASTED_WEIGHT
global LARGEST_AIMDO_CASTED_WEIGHT
LARGEST_CASTED_WEIGHT = (None, 0)
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS) | set(STREAM_PIN_BUFFERS):
for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS):
if offload_stream is not None:
offload_stream.synchronize()
synchronize()
for mmap_obj in DIRTY_MMAPS:
mmap_obj.bounce()
DIRTY_MMAPS.clear()
for pin_buffer in STREAM_PIN_BUFFERS.values():
TOTAL_PINNED_MEMORY -= pin_buffer.size
TOTAL_PINNED_MEMORY = max(0, TOTAL_PINNED_MEMORY)
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is not None and model.is_dynamic():
model.model.dynamic_pins[model.load_device]["active"] = False
model.partially_unload_ram(1e30, subsets=[ "patches" ])
model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0])
STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear()
STREAM_PIN_BUFFERS.clear()
soft_empty_cache()
def get_offload_stream(device):
@ -1317,7 +1280,7 @@ def sync_stream(device, stream):
current_stream(device).wait_stream(stream)
def cast_to_gathered(tensors, r, non_blocking=False, stream=None, r2=None):
def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
wf_context = nullcontext()
if stream is not None:
wf_context = stream
@ -1325,20 +1288,17 @@ def cast_to_gathered(tensors, r, non_blocking=False, stream=None, r2=None):
wf_context = wf_context.as_context(stream)
dest_views = comfy.memory_management.interpret_gathered_like(tensors, r)
dest2_views = comfy.memory_management.interpret_gathered_like(tensors, r2) if r2 is not None else None
with wf_context:
for tensor in tensors:
dest_view = dest_views.pop(0)
dest2_view = dest2_views.pop(0) if dest2_views is not None else None
if tensor is None:
continue
if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view, stream=stream, destination2=dest2_view):
if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view):
continue
storage = tensor._qdata.untyped_storage() if isinstance(tensor, comfy.quant_ops.QuantizedTensor) else tensor.untyped_storage()
mark_mmap_dirty(storage)
if hasattr(storage, "_comfy_tensor_mmap_touched"):
storage._comfy_tensor_mmap_touched = True
dest_view.copy_(tensor, non_blocking=non_blocking)
if dest2_view is not None:
dest2_view.copy_(dest_view, non_blocking=non_blocking)
def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None, r=None):
@ -1379,18 +1339,14 @@ TOTAL_PINNED_MEMORY = 0
MAX_PINNED_MEMORY = -1
if not args.disable_pinned_memory:
if is_nvidia() or is_amd():
ram = get_total_memory(torch.device("cpu"))
if WINDOWS:
MAX_PINNED_MEMORY = ram * 0.40 # Windows limit is apparently 50%
MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.40 # Windows limit is apparently 50%
else:
MAX_PINNED_MEMORY = ram * 0.90
MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.90
logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
def pinned_hostbuf_size(size):
return max(0, int(min(size, MAX_PINNED_MEMORY) * 2))
def discard_cuda_async_error():
try:
a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
@ -1422,8 +1378,8 @@ def pin_memory(tensor):
return False
size = tensor.nbytes
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
ensure_pin_registerable(size)
if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
return False
ptr = tensor.data_ptr()
if ptr == 0:
@ -1460,8 +1416,7 @@ def unpin_memory(tensor):
return False
if torch.cuda.cudart().cudaHostUnregister(ptr) == 0:
size = PINNED_MEMORY.pop(ptr)
TOTAL_PINNED_MEMORY -= size
TOTAL_PINNED_MEMORY -= PINNED_MEMORY.pop(ptr)
return True
else:
logging.warning("Unpin error.")

View File

@ -35,7 +35,6 @@ import comfy.model_management
import comfy.ops
import comfy.patcher_extension
import comfy.utils
import comfy_aimdo.host_buffer
from comfy.comfy_types import UnetWrapperFunction
from comfy.quant_ops import QuantizedTensor
from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
@ -118,8 +117,6 @@ def string_to_seed(data):
return comfy.utils.string_to_seed(data)
class LowVramPatch:
is_lowvram_patch = True
def __init__(self, key, patches, convert_func=None, set_func=None):
self.key = key
self.patches = patches
@ -127,21 +124,11 @@ class LowVramPatch:
self.set_func = set_func
self.prepared_patches = None
def memory_required(self):
counter = [0]
for patch in self.patches[self.key]:
comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False)
return counter[0]
def prepare(self, destination, stream, copy=True, commit=True):
counter = [0]
prepared_patches = [
(patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4])
def prepare(self, allocate_buffer, stream):
self.prepared_patches = [
(patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4])
for patch in self.patches[self.key]
]
if commit:
self.prepared_patches = prepared_patches
return prepared_patches
def clear_prepared(self):
self.prepared_patches = None
@ -354,6 +341,9 @@ class ModelPatcher:
self.size = comfy.model_management.module_size(self.model)
return self.size
def model_mmap_residency(self, free=False):
return comfy.model_management.module_mmap_residency(self.model, free=free)
def loaded_size(self):
return self.model.model_loaded_weight_memory
@ -1128,12 +1118,8 @@ class ModelPatcher:
# Pinned memory pressure tracking is only implemented for DynamicVram loading
return 0
def loaded_ram_size(self):
# Loaded RAM pressure tracking is only implemented for DynamicVram loading
return 0
def partially_unload_ram(self, ram_to_unload):
return 0
pass
def detach(self, unpatch_all=True):
self.eject_model()
@ -1564,16 +1550,6 @@ class ModelPatcherDynamic(ModelPatcher):
super().__init__(model, load_device, offload_device, size, weight_inplace_update)
if not hasattr(self.model, "dynamic_vbars"):
self.model.dynamic_vbars = {}
if not hasattr(self.model, "dynamic_pins"):
self.model.dynamic_pins = {}
if self.load_device not in self.model.dynamic_pins:
self.model.dynamic_pins[self.load_device] = {
"weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
"patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
"hostbufs_initialized": False,
"failed": False,
"active": False,
}
self.non_dynamic_delegate_model = None
assert load_device is not None
@ -1635,14 +1611,6 @@ class ModelPatcherDynamic(ModelPatcher):
self.unpatch_hooks()
vbar = self._vbar_get(create=True)
pin_state = self.model.dynamic_pins[self.load_device]
if not pin_state["hostbufs_initialized"]:
hostbuf_size = comfy.model_management.pinned_hostbuf_size(self.model_size())
pin_state["weights"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0])
pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0])
pin_state["hostbufs_initialized"] = True
pin_state["failed"] = False
pin_state["active"] = True
if vbar is not None:
vbar.prioritize()
@ -1668,9 +1636,7 @@ class ModelPatcherDynamic(ModelPatcher):
if key in self.patches:
if comfy.lora.calculate_shape(self.patches[key], weight, key) != weight.shape:
return (True, 0)
lowvram_patch = LowVramPatch(key, self.patches)
lowvram_patch._pin_state = pin_state
setattr(m, param_key + "_lowvram_function", lowvram_patch)
setattr(m, param_key + "_lowvram_function", LowVramPatch(key, self.patches))
num_patches += 1
else:
setattr(m, param_key + "_lowvram_function", None)
@ -1687,9 +1653,6 @@ class ModelPatcherDynamic(ModelPatcher):
def force_load_param(self, param_key, device_to):
key = key_param_name_to_key(n, param_key)
weight, _, _ = get_key_weight(self.model, key)
if weight is None:
return
if key in self.backup:
comfy.utils.set_attr_param(self.model, key, self.backup[key].weight)
self.patch_weight_to_device(key, device_to=device_to, force_cast=True)
@ -1699,23 +1662,17 @@ class ModelPatcherDynamic(ModelPatcher):
if hasattr(m, "comfy_cast_weights"):
m.comfy_cast_weights = True
m.pin_failed = False
m.seed_key = n
m._pin_state = pin_state
set_dirty(m, dirty)
#Models that mix tiny and giant weights can causing lopsided stream buffer
#rotations and stall. force the tinys over.
if module_mem > 16 * 1024:
force_load, v_weight_size = setup_param(self, m, n, "weight")
force_load_bias, v_weight_bias = setup_param(self, m, n, "bias")
force_load = force_load or force_load_bias
v_weight_size += v_weight_bias
if force_load:
logging.info(f"Module {n} has resizing Lora - force loading")
else:
force_load=True
force_load, v_weight_size = setup_param(self, m, n, "weight")
force_load_bias, v_weight_bias = setup_param(self, m, n, "bias")
force_load = force_load or force_load_bias
v_weight_size += v_weight_bias
if force_load:
logging.info(f"Module {n} has resizing Lora - force loading")
force_load_param(self, "weight", device_to)
force_load_param(self, "bias", device_to)
else:
@ -1783,58 +1740,23 @@ class ModelPatcherDynamic(ModelPatcher):
return freed
def loaded_ram_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][0].size +
self.model.dynamic_pins[self.load_device]["patches"][0].size)
def pinned_memory_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][3][0] +
self.model.dynamic_pins[self.load_device]["patches"][3][0])
total = 0
loading = self._load_list(for_dynamic=True)
for x in loading:
_, _, _, _, m, _ = x
pin = comfy.pinned_memory.get_pin(m)
if pin is not None:
total += pin.numel() * pin.element_size()
return total
def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
hostbuf, stack, stack_split, pinned_size = pin_state[subset]
split = stack_split[0]
while split >= 0:
module, offset = stack[split]
split -= 1
stack_split[0] = split
if not module._pin_registered:
continue
size = module._pin.numel() * module._pin.element_size()
if torch.cuda.cudart().cudaHostUnregister(module._pin.data_ptr()) != 0:
comfy.model_management.discard_cuda_async_error()
continue
module._pin_registered = False
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size
ram_to_unload -= size
if ram_to_unload <= 0:
return freed
return freed
def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
hostbuf, stack, stack_split, pinned_size = pin_state[subset]
while len(stack) > 0:
module, offset = stack.pop()
size = module._pin.numel() * module._pin.element_size()
del module._pin
hostbuf.truncate(offset, do_unregister=module._pin_registered)
stack_split[0] = min(stack_split[0], len(stack) - 1)
if module._pin_registered:
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size
ram_to_unload -= size
if ram_to_unload <= 0:
return freed
return freed
def partially_unload_ram(self, ram_to_unload):
loading = self._load_list(for_dynamic=True, default_device=self.offload_device)
for x in loading:
*_, m, _ = x
ram_to_unload -= comfy.pinned_memory.unpin_memory(m)
if ram_to_unload <= 0:
return
def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):
#This isn't used by the core at all and can only be to load a model out of

View File

@ -75,8 +75,6 @@ except:
cast_to = comfy.model_management.cast_to #TODO: remove once no more references
STREAM_PIN_BUFFER_HEADROOM = 8 * 1024 * 1024
def cast_to_input(weight, input, non_blocking=False, copy=True):
return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
@ -93,9 +91,6 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
offload_stream = None
cast_buffer = None
cast_buffer_offset = 0
stream_pin_hostbuf = None
stream_pin_offset = 0
stream_pin_queue = []
def ensure_offload_stream(module, required_size, check_largest):
nonlocal offload_stream
@ -129,22 +124,6 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
cast_buffer_offset += buffer_size
return buffer
def get_stream_pin_buffer_offset(buffer_size):
nonlocal stream_pin_hostbuf
nonlocal stream_pin_offset
if buffer_size == 0 or offload_stream is None:
return None
if stream_pin_hostbuf is None:
stream_pin_hostbuf = comfy.model_management.get_pin_buffer(offload_stream)
if stream_pin_hostbuf is None:
return None
offset = stream_pin_offset
stream_pin_offset += buffer_size
return offset
for s in comfy_modules:
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
@ -183,47 +162,23 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
if xfer_dest is None:
xfer_dest = get_cast_buffer(dest_size)
def cast_maybe_lowvram_patch(xfer_source, xfer_dest, stream):
if xfer_source is not None:
if getattr(xfer_source, "is_lowvram_patch", False):
xfer_source.prepare(xfer_dest, stream, copy=True, commit=False)
else:
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream)
if signature is None and pin is None:
comfy.pinned_memory.pin_memory(s)
pin = comfy.pinned_memory.get_pin(s)
else:
pin = None
def handle_pin(m, pin, source, dest, subset="weights", size=None):
if pin is not None:
cast_maybe_lowvram_patch([pin], dest, offload_stream)
return
if signature is None:
comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
pin = comfy.pinned_memory.get_pin(m, subset=subset)
if pin is not None:
if isinstance(source, list):
comfy.model_management.cast_to_gathered(source, pin, non_blocking=non_blocking, stream=offload_stream, r2=dest)
else:
cast_maybe_lowvram_patch(source, pin, None)
cast_maybe_lowvram_patch([ pin ], dest, offload_stream)
return
if pin is None:
pin_offset = get_stream_pin_buffer_offset(size)
if pin_offset is not None:
stream_pin_queue.append((source, pin_offset, size, dest))
return
cast_maybe_lowvram_patch(source, dest, offload_stream)
handle_pin(s, pin, xfer_source, xfer_dest, size=dest_size)
if pin is not None:
comfy.model_management.cast_to_gathered(xfer_source, pin)
xfer_source = [ pin ]
#send it over
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=offload_stream)
for param_key in ("weight", "bias"):
lowvram_source = getattr(s, param_key + "_lowvram_function", None)
if lowvram_source is not None:
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
if lowvram_fn is not None:
ensure_offload_stream(s, cast_buffer_offset, False)
lowvram_size = lowvram_source.memory_required()
lowvram_dest = get_cast_buffer(lowvram_size)
lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches")
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset="patches", size=lowvram_size)
lowvram_fn.prepare(lambda size: get_cast_buffer(size), offload_stream)
prefetch["xfer_dest"] = xfer_dest
prefetch["cast_dest"] = cast_dest
@ -231,23 +186,6 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
prefetch["needs_cast"] = needs_cast
s._prefetch = prefetch
if stream_pin_offset > 0:
if stream_pin_hostbuf.size < stream_pin_offset:
if not comfy.model_management.resize_pin_buffer(stream_pin_hostbuf, stream_pin_offset + STREAM_PIN_BUFFER_HEADROOM):
for xfer_source, _, _, xfer_dest in stream_pin_queue:
cast_maybe_lowvram_patch(xfer_source, xfer_dest, offload_stream)
return offload_stream
stream_pin_tensor = comfy_aimdo.torch.hostbuf_to_tensor(stream_pin_hostbuf)
stream_pin_tensor.untyped_storage()._comfy_hostbuf = stream_pin_hostbuf
for xfer_source, pin_offset, pin_size, xfer_dest in stream_pin_queue:
pin = stream_pin_tensor[pin_offset:pin_offset + pin_size]
if isinstance(xfer_source, list):
comfy.model_management.cast_to_gathered(xfer_source, pin, non_blocking=non_blocking, stream=offload_stream, r2=xfer_dest)
else:
cast_maybe_lowvram_patch(xfer_source, pin, None)
comfy.model_management.cast_to_gathered([ pin ], xfer_dest, non_blocking=non_blocking, stream=offload_stream)
stream_pin_hostbuf._comfy_event = offload_stream.record_event()
return offload_stream

View File

@ -2,62 +2,42 @@ import comfy.model_management
import comfy.memory_management
import comfy_aimdo.host_buffer
import comfy_aimdo.torch
import torch
from comfy.cli_args import args
def get_pin(module, subset="weights"):
pin = getattr(module, "_pin", None)
if pin is None or module._pin_registered or args.disable_pinned_memory:
return pin
def get_pin(module):
return getattr(module, "_pin", None)
_, _, stack_split, pinned_size = module._pin_state[subset]
size = pin.nbytes
comfy.model_management.ensure_pin_registerable(size)
if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
comfy.model_management.discard_cuda_async_error()
return pin
module._pin_registered = True
stack_split[0] = max(stack_split[0], module._pin_stack_index)
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
return pin
def pin_memory(module, subset="weights", size=None):
pin_state = module._pin_state
if args.disable_pinned_memory:
def pin_memory(module):
if module.pin_failed or args.disable_pinned_memory or get_pin(module) is not None:
return
pin = get_pin(module, subset)
if pin is not None or pin_state["failed"]:
return
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
hostbuf, stack, stack_split, pinned_size = pin_state[subset]
if size is None:
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
offset = hostbuf.size
registerable_size = size + max(0, hostbuf.size - pinned_size[0])
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
if (not comfy.model_management.ensure_pin_budget(size) or
not comfy.model_management.ensure_pin_registerable(registerable_size)):
pin_state["failed"] = True
if comfy.model_management.MAX_PINNED_MEMORY <= 0 or (comfy.model_management.TOTAL_PINNED_MEMORY + size) > comfy.model_management.MAX_PINNED_MEMORY:
module.pin_failed = True
return False
try:
hostbuf.extend(size=size)
hostbuf = comfy_aimdo.host_buffer.HostBuffer(size)
except RuntimeError:
pin_state["failed"] = True
module.pin_failed = True
return False
module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)[offset:offset + size]
module._pin.untyped_storage()._comfy_hostbuf = hostbuf
stack.append((module, offset))
module._pin_registered = True
module._pin_stack_index = len(stack) - 1
stack_split[0] = max(stack_split[0], module._pin_stack_index)
module._pin = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)
module._pin_hostbuf = hostbuf
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
return True
def unpin_memory(module):
if get_pin(module) is None:
return 0
size = module._pin.numel() * module._pin.element_size()
comfy.model_management.TOTAL_PINNED_MEMORY -= size
if comfy.model_management.TOTAL_PINNED_MEMORY < 0:
comfy.model_management.TOTAL_PINNED_MEMORY = 0
del module._pin
del module._pin_hostbuf
return size

View File

@ -113,6 +113,7 @@ def load_safetensors(ckpt):
"_comfy_tensor_file_slice",
comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start))
setattr(storage, "_comfy_tensor_mmap_refs", (model_mmap, mv))
setattr(storage, "_comfy_tensor_mmap_touched", False)
sd[name] = tensor
return sd, header.get("__metadata__", {}),
@ -1450,3 +1451,4 @@ def deepcopy_list_dict(obj, memo=None):
memo[obj_id] = res
return res

52
comfy/windows.py Normal file
View File

@ -0,0 +1,52 @@
import ctypes
import logging
import psutil
from ctypes import wintypes
import comfy_aimdo.control
psapi = ctypes.WinDLL("psapi")
kernel32 = ctypes.WinDLL("kernel32")
class PERFORMANCE_INFORMATION(ctypes.Structure):
_fields_ = [
("cb", wintypes.DWORD),
("CommitTotal", ctypes.c_size_t),
("CommitLimit", ctypes.c_size_t),
("CommitPeak", ctypes.c_size_t),
("PhysicalTotal", ctypes.c_size_t),
("PhysicalAvailable", ctypes.c_size_t),
("SystemCache", ctypes.c_size_t),
("KernelTotal", ctypes.c_size_t),
("KernelPaged", ctypes.c_size_t),
("KernelNonpaged", ctypes.c_size_t),
("PageSize", ctypes.c_size_t),
("HandleCount", wintypes.DWORD),
("ProcessCount", wintypes.DWORD),
("ThreadCount", wintypes.DWORD),
]
def get_free_ram():
#Windows is way too conservative and chalks recently used uncommitted model RAM
#as "in-use". So, calculate free RAM for the sake of general use as the greater of:
#
#1: What psutil says
#2: Total Memory - (Committed Memory - VRAM in use)
#
#We have to subtract VRAM in use from the comitted memory as WDDM creates a naked
#commit charge for all VRAM used just incase it wants to page it all out. This just
#isn't realistic so "overcommit" on our calculations by just subtracting it off.
pi = PERFORMANCE_INFORMATION()
pi.cb = ctypes.sizeof(pi)
if not psapi.GetPerformanceInfo(ctypes.byref(pi), pi.cb):
logging.warning("WARNING: Failed to query windows performance info. RAM usage may be sub optimal")
return psutil.virtual_memory().available
committed = pi.CommitTotal * pi.PageSize
total = pi.PhysicalTotal * pi.PageSize
return max(psutil.virtual_memory().available,
total - (committed - comfy_aimdo.control.get_total_vram_usage()))

View File

@ -21,10 +21,6 @@ class CacheProvider(ABC):
Exceptions from provider methods are caught by the caller and never break execution.
"""
async def on_set_prompt(self) -> None:
"""Called after prompt cache keys are prepared. Dispatched via asyncio.create_task."""
pass
@abstractmethod
async def on_lookup(self, context: CacheContext) -> Optional[CacheValue]:
"""Called on local cache miss. Return CacheValue if found, None otherwise."""

View File

@ -35,6 +35,19 @@ class AnthropicMessage(BaseModel):
content: list[AnthropicTextContent | AnthropicImageContent] = Field(...)
class AnthropicThinkingConfig(BaseModel):
type: Literal["enabled", "disabled", "adaptive"] = Field(...)
budget_tokens: int | None = Field(
None, ge=1024,
description="Reasoning budget in tokens. Used when type is 'enabled'. Must be less than max_tokens.",
)
class AnthropicOutputConfig(BaseModel):
"""Used with `thinking.type='adaptive'` on models like Opus 4.7."""
effort: Literal["low", "medium", "high"] | None = Field(None)
class AnthropicMessagesRequest(BaseModel):
model: str = Field(...)
messages: list[AnthropicMessage] = Field(...)
@ -44,6 +57,8 @@ class AnthropicMessagesRequest(BaseModel):
top_p: float | None = Field(None, ge=0.0, le=1.0)
top_k: int | None = Field(None, ge=0)
stop_sequences: list[str] | None = Field(None)
thinking: AnthropicThinkingConfig | None = Field(None)
output_config: AnthropicOutputConfig | None = Field(None)
class AnthropicResponseTextBlock(BaseModel):
@ -51,6 +66,14 @@ class AnthropicResponseTextBlock(BaseModel):
text: str = Field(...)
class AnthropicResponseThinkingBlock(BaseModel):
type: Literal["thinking"] = "thinking"
thinking: str = Field(...)
AnthropicResponseBlock = AnthropicResponseTextBlock | AnthropicResponseThinkingBlock
class AnthropicCacheCreationUsage(BaseModel):
ephemeral_5m_input_tokens: int | None = Field(None)
ephemeral_1h_input_tokens: int | None = Field(None)
@ -69,7 +92,7 @@ class AnthropicMessagesResponse(BaseModel):
type: str | None = Field(None)
role: str | None = Field(None)
model: str | None = Field(None)
content: list[AnthropicResponseTextBlock] | None = Field(None)
content: list[AnthropicResponseBlock] | None = Field(None)
stop_reason: str | None = Field(None)
stop_sequence: str | None = Field(None)
usage: AnthropicMessagesUsage | None = Field(None)

View File

@ -0,0 +1,93 @@
"""Pydantic models for the OpenRouter chat completions API.
See: https://openrouter.ai/docs/api/api-reference/chat/send-chat-completion-request
"""
from typing import Literal
from pydantic import BaseModel, Field
class OpenRouterTextContent(BaseModel):
type: Literal["text"] = "text"
text: str = Field(...)
class OpenRouterImageUrl(BaseModel):
url: str = Field(...)
class OpenRouterImageContent(BaseModel):
type: Literal["image_url"] = "image_url"
image_url: OpenRouterImageUrl = Field(...)
class OpenRouterVideoUrl(BaseModel):
url: str = Field(...)
class OpenRouterVideoContent(BaseModel):
type: Literal["video_url"] = "video_url"
video_url: OpenRouterVideoUrl = Field(...)
OpenRouterContentBlock = OpenRouterTextContent | OpenRouterImageContent | OpenRouterVideoContent
class OpenRouterMessage(BaseModel):
role: Literal["system", "user", "assistant"] = Field(...)
content: str | list[OpenRouterContentBlock] = Field(...)
class OpenRouterReasoningConfig(BaseModel):
effort: str | None = Field(None)
exclude: bool | None = Field(None, description="If true, model reasons but reasoning is excluded from response.")
class OpenRouterWebSearchOptions(BaseModel):
search_context_size: str | None = Field(None)
class OpenRouterChatRequest(BaseModel):
model: str = Field(...)
messages: list[OpenRouterMessage] = Field(...)
seed: int | None = Field(None)
reasoning: OpenRouterReasoningConfig | None = Field(None)
web_search_options: OpenRouterWebSearchOptions | None = Field(None)
stream: bool = Field(False)
class OpenRouterUsage(BaseModel):
prompt_tokens: int | None = Field(None)
completion_tokens: int | None = Field(None)
total_tokens: int | None = Field(None)
cost: float | None = Field(None, description="Server-side authoritative USD cost of the call.")
class OpenRouterResponseMessage(BaseModel):
role: str | None = Field(None)
content: str | None = Field(None)
reasoning: str | None = Field(None)
refusal: str | None = Field(None)
class OpenRouterChoice(BaseModel):
index: int | None = Field(None)
message: OpenRouterResponseMessage | None = Field(None)
finish_reason: str | None = Field(None)
class OpenRouterError(BaseModel):
code: int | str | None = Field(None)
message: str | None = Field(None)
metadata: dict | None = Field(None)
class OpenRouterChatResponse(BaseModel):
id: str | None = Field(None)
model: str | None = Field(None)
object: str | None = Field(None)
provider: str | None = Field(None)
choices: list[OpenRouterChoice] | None = Field(None)
usage: OpenRouterUsage | None = Field(None)
error: OpenRouterError | None = Field(None)

View File

@ -1,7 +1,5 @@
from __future__ import annotations
from enum import Enum
from typing import Optional, List
from pydantic import BaseModel, Field
@ -11,44 +9,76 @@ class Rodin3DGenerateRequest(BaseModel):
material: str = Field(..., description="The material type.")
quality_override: int = Field(..., description="The poly count of the mesh.")
mesh_mode: str = Field(..., description="It controls the type of faces of generated models.")
TAPose: Optional[bool] = Field(None, description="")
TAPose: bool | None = Field(None, description="")
class Rodin3DGen25Request(BaseModel):
tier: str = Field(..., description="Gen-2.5 tier (e.g. Gen-2.5-High).")
prompt: str | None = Field(None, description="Required for Text-to-3D; ignored otherwise.")
seed: int | None = Field(None, description="0-65535.")
material: str | None = Field(None, description="PBR | Shaded | All | None.")
geometry_file_format: str | None = Field(None, description="glb | usdz | fbx | obj | stl.")
texture_mode: str | None = Field(None, description="legacy | extreme-low | low | medium | high.")
mesh_mode: str | None = Field(None, description="Raw (triangular) | Quad.")
quality_override: int | None = Field(None, description="Mesh face count override.")
geometry_instruct_mode: str | None = Field(None, description="faithful | creative.")
bbox_condition: list[int] | None = Field(None, description="Bounding box [Width(Y), Height(Z), Length(X)] in cm.")
height: int | None = Field(None, description="Approximate model height in cm.")
TAPose: bool | None = Field(None, description="T/A pose for human-like models.")
hd_texture: bool | None = Field(None, description="Enhanced texture quality.")
texture_delight: bool | None = Field(None, description="Remove baked lighting from textures.")
is_micro: bool | None = Field(None, description="Micro detail (Extreme-High only).")
use_original_alpha: bool | None = Field(None, description="Preserve image transparency.")
preview_render: bool | None = Field(None, description="Generate high-quality preview render.")
addons: list[str] | None = Field(None, description='Optional addons, e.g. ["HighPack"].')
class GenerateJobsData(BaseModel):
uuids: List[str] = Field(..., description="str LIST")
uuids: list[str] = Field(..., description="str LIST")
subscription_key: str = Field(..., description="subscription key")
class Rodin3DGenerateResponse(BaseModel):
message: Optional[str] = Field(None, description="Return message.")
prompt: Optional[str] = Field(None, description="Generated Prompt from image.")
submit_time: Optional[str] = Field(None, description="Submit Time")
uuid: Optional[str] = Field(None, description="Task str")
jobs: Optional[GenerateJobsData] = Field(None, description="Details of jobs")
message: str | None = Field(None, description="Return message.")
prompt: str | None = Field(None, description="Generated Prompt from image.")
submit_time: str | None = Field(None, description="Submit Time")
uuid: str | None = Field(None, description="Task str")
jobs: GenerateJobsData | None = Field(None, description="Details of jobs")
class JobStatus(str, Enum):
"""
Status for jobs
"""
Done = "Done"
Failed = "Failed"
Generating = "Generating"
Waiting = "Waiting"
class Rodin3DCheckStatusRequest(BaseModel):
subscription_key: str = Field(..., description="subscription from generate endpoint")
class JobItem(BaseModel):
uuid: str = Field(..., description="uuid")
status: JobStatus = Field(...,description="Status Currently")
status: JobStatus = Field(..., description="Status Currently")
class Rodin3DCheckStatusResponse(BaseModel):
jobs: List[JobItem] = Field(..., description="Job status List")
jobs: list[JobItem] = Field(..., description="Job status List")
class Rodin3DDownloadRequest(BaseModel):
task_uuid: str = Field(..., description="Task str")
class RodinResourceItem(BaseModel):
url: str = Field(..., description="Download Url")
name: str = Field(..., description="File name with ext")
class Rodin3DDownloadResponse(BaseModel):
list: List[RodinResourceItem] = Field(..., description="Source List")
items: list[RodinResourceItem] = Field(..., alias="list", description="Source List")

View File

@ -9,8 +9,11 @@ from comfy_api_nodes.apis.anthropic import (
AnthropicMessage,
AnthropicMessagesRequest,
AnthropicMessagesResponse,
AnthropicOutputConfig,
AnthropicResponseTextBlock,
AnthropicRole,
AnthropicTextContent,
AnthropicThinkingConfig,
)
from comfy_api_nodes.util import (
ApiEndpoint,
@ -32,15 +35,29 @@ CLAUDE_MODELS: dict[str, str] = {
"Haiku 4.5": "claude-haiku-4-5-20251001",
}
_THINKING_UNSUPPORTED = {"Haiku 4.5"}
# Models that use the newer "adaptive" thinking mode (Opus 4.7 requires it; older models keep the explicit budget API).
# Anthropic decides the actual budget when adaptive is used, based on the `output_config.effort` hint.
_ADAPTIVE_THINKING_MODELS = {"Opus 4.7", "Opus 4.6", "Sonnet 4.6"}
def _claude_model_inputs():
return [
# Budget mode (Sonnet 4.5): effort -> reasoning budget in tokens. Must be < max_tokens.
# Sized so even the "high" budget fits comfortably under the default max_tokens=32768.
_REASONING_BUDGET: dict[str, int] = {
"low": 2048,
"medium": 8192,
"high": 16384,
}
_REASONING_EFFORTS = ["off", "low", "medium", "high"]
def _claude_model_inputs(model_label: str):
inputs: list = [
IO.Int.Input(
"max_tokens",
default=16000,
min=32,
max=32000,
tooltip="Maximum number of tokens to generate before stopping.",
default=32768,
min=4096,
max=64000,
tooltip="Maximum number of tokens to generate (includes reasoning tokens when enabled).",
advanced=True,
),
IO.Float.Input(
@ -49,10 +66,24 @@ def _claude_model_inputs():
min=0.0,
max=1.0,
step=0.01,
tooltip="Controls randomness. 0.0 is deterministic, 1.0 is most random. Ignored for Opus 4.7.",
tooltip=(
"Controls randomness. 0.0 is deterministic, 1.0 is most random. "
"Ignored for Opus 4.7 and any model when reasoning_effort is set."
),
advanced=True,
),
]
if model_label not in _THINKING_UNSUPPORTED:
inputs.append(
IO.Combo.Input(
"reasoning_effort",
options=_REASONING_EFFORTS,
default="off",
tooltip="Extended thinking effort. 'off' disables reasoning.",
advanced=True,
)
)
return inputs
def _model_price_per_million(model: str) -> tuple[float, float] | None:
@ -95,7 +126,11 @@ def calculate_tokens_price(response: AnthropicMessagesResponse) -> float | None:
def _get_text_from_response(response: AnthropicMessagesResponse) -> str:
if not response.content:
return ""
return "\n".join(block.text for block in response.content if block.text)
# Thinking blocks are silently dropped — we never want reasoning in the output.
return "\n".join(
block.text for block in response.content
if isinstance(block, AnthropicResponseTextBlock) and block.text
)
async def _build_image_content_blocks(
@ -133,7 +168,10 @@ class ClaudeNode(IO.ComfyNode):
),
IO.DynamicCombo.Input(
"model",
options=[IO.DynamicCombo.Option(label, _claude_model_inputs()) for label in CLAUDE_MODELS],
options=[
IO.DynamicCombo.Option(label, _claude_model_inputs(label))
for label in CLAUDE_MODELS
],
tooltip="The Claude model used to generate the response.",
),
IO.Int.Input(
@ -207,8 +245,29 @@ class ClaudeNode(IO.ComfyNode):
) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1)
model_label = model["model"]
max_tokens = model["max_tokens"]
temperature = None if model_label == "Opus 4.7" else model["temperature"]
max_tokens = model.get("max_tokens", 32768)
reasoning_effort = model.get("reasoning_effort", "off")
thinking_enabled = reasoning_effort not in ("off", None) and model_label not in _THINKING_UNSUPPORTED
# Anthropic requires temperature to be unset (defaults to 1.0) when thinking is enabled.
# Opus 4.7 also rejects user-supplied temperature.
if thinking_enabled or model_label == "Opus 4.7":
temperature = None
else:
temperature = model.get("temperature", 1.0)
thinking_cfg: AnthropicThinkingConfig | None = None
output_cfg: AnthropicOutputConfig | None = None
if thinking_enabled:
if model_label in _ADAPTIVE_THINKING_MODELS:
# Adaptive mode - Anthropic chooses the budget based on effort hint
thinking_cfg = AnthropicThinkingConfig(type="adaptive")
output_cfg = AnthropicOutputConfig(effort=reasoning_effort)
else:
# Budget mode (Sonnet 4.5). Leave at least 1024 tokens for the actual response
budget = _REASONING_BUDGET[reasoning_effort]
budget = min(budget, max(1024, max_tokens - 1024))
thinking_cfg = AnthropicThinkingConfig(type="enabled", budget_tokens=budget)
image_tensors: list[Input.Image] = [t for t in (images or {}).values() if t is not None]
if sum(get_number_of_images(t) for t in image_tensors) > CLAUDE_MAX_IMAGES:
@ -229,6 +288,8 @@ class ClaudeNode(IO.ComfyNode):
messages=[AnthropicMessage(role=AnthropicRole.user, content=content)],
system=system_prompt or None,
temperature=temperature,
thinking=thinking_cfg,
output_config=output_cfg,
),
price_extractor=calculate_tokens_price,
)

View File

@ -43,15 +43,16 @@ from comfy_api_nodes.util import (
ApiEndpoint,
download_url_to_image_tensor,
download_url_to_video_output,
downscale_video_to_max_pixels,
get_number_of_images,
image_tensor_pair_to_batch,
poll_op,
resize_video_to_pixel_budget,
sync_op,
upload_audio_to_comfyapi,
upload_image_to_comfyapi,
upload_images_to_comfyapi,
upload_video_to_comfyapi,
upscale_video_to_min_pixels,
validate_image_aspect_ratio,
validate_image_dimensions,
validate_string,
@ -110,12 +111,13 @@ def _validate_ref_video_pixels(video: Input.Video, model_id: str, resolution: st
max_px = limits.get("max")
if min_px and pixels < min_px:
raise ValueError(
f"Reference video {index} is too small: {w}x{h} = {pixels:,}px. " f"Minimum is {min_px:,}px for this model."
f"Reference video {index} is too small: {w}x{h} = {pixels:,} total pixels. "
f"Minimum for this model is {min_px:,} total pixels."
)
if max_px and pixels > max_px:
raise ValueError(
f"Reference video {index} is too large: {w}x{h} = {pixels:,}px. "
f"Maximum is {max_px:,}px for this model. Try downscaling the video."
f"Reference video {index} is too large: {w}x{h} = {pixels:,} total pixels. "
f"Maximum for this model is {max_px:,} total pixels. Try downscaling the video."
)
@ -1676,14 +1678,14 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
"first_frame_asset_id",
default="",
tooltip="Seedance asset_id to use as the first frame. "
"Mutually exclusive with the first_frame image input.",
"Mutually exclusive with the first_frame image input.",
optional=True,
),
IO.String.Input(
"last_frame_asset_id",
default="",
tooltip="Seedance asset_id to use as the last frame. "
"Mutually exclusive with the last_frame image input.",
"Mutually exclusive with the last_frame image input.",
optional=True,
),
IO.Int.Input(
@ -1865,11 +1867,20 @@ def _seedance2_reference_inputs(resolutions: list[str], default_ratio: str = "16
IO.Boolean.Input(
"auto_downscale",
default=False,
advanced=True,
optional=True,
tooltip="Automatically downscale reference videos that exceed the model's pixel budget "
"for the selected resolution. Aspect ratio is preserved; videos already within limits are untouched.",
),
IO.Boolean.Input(
"auto_upscale",
default=False,
advanced=True,
optional=True,
tooltip="Automatically upscale reference videos that are below the model's minimum pixel count "
"for the selected resolution. Aspect ratio is preserved; videos already meeting the minimum are "
"untouched. Note: upscaling a low-resolution source does not add real detail and may produce "
"lower-quality generations.",
),
IO.Autogrow.Input(
"reference_assets",
template=IO.Autogrow.TemplateNames(
@ -2030,7 +2041,13 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
max_px = SEEDANCE2_REF_VIDEO_PIXEL_LIMITS.get(model_id, {}).get(model["resolution"], {}).get("max")
if max_px:
for key in reference_videos:
reference_videos[key] = resize_video_to_pixel_budget(reference_videos[key], max_px)
reference_videos[key] = downscale_video_to_max_pixels(reference_videos[key], max_px)
if model.get("auto_upscale") and reference_videos:
min_px = SEEDANCE2_REF_VIDEO_PIXEL_LIMITS.get(model_id, {}).get(model["resolution"], {}).get("min")
if min_px:
for key in reference_videos:
reference_videos[key] = upscale_video_to_min_pixels(reference_videos[key], min_px)
total_video_duration = 0.0
for i, key in enumerate(reference_videos, 1):

View File

@ -0,0 +1,374 @@
"""API Nodes for OpenRouter LLM chat completions."""
from dataclasses import dataclass
from typing import Literal
from typing_extensions import override
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api_nodes.apis.openrouter import (
OpenRouterChatRequest,
OpenRouterChatResponse,
OpenRouterContentBlock,
OpenRouterImageContent,
OpenRouterImageUrl,
OpenRouterMessage,
OpenRouterReasoningConfig,
OpenRouterTextContent,
OpenRouterVideoContent,
OpenRouterVideoUrl,
OpenRouterWebSearchOptions,
)
from comfy_api_nodes.util import (
ApiEndpoint,
get_number_of_images,
sync_op,
upload_images_to_comfyapi,
upload_video_to_comfyapi,
validate_string,
)
OPENROUTER_CHAT_ENDPOINT = "/proxy/openrouter/api/v1/chat/completions"
Profile = Literal["standard", "reasoning", "frontier_reasoning", "perplexity", "perplexity_reasoning"]
@dataclass(frozen=True)
class _ModelSpec:
slug: str # exact OpenRouter model id
profile: Profile
price_in: float # USD per token (prompt)
price_out: float # USD per token (completion)
max_images: int = 0 # 0 = no image input; otherwise max URL-passed images supported
max_videos: int = 0 # 0 = no video input; otherwise max URL-passed videos supported
MODELS: list[_ModelSpec] = [
_ModelSpec("anthropic/claude-opus-4.7", "frontier_reasoning", 0.000005, 0.000025, max_images=20),
_ModelSpec("openai/gpt-5.5-pro", "frontier_reasoning", 0.00003, 0.00018, max_images=20),
_ModelSpec("openai/gpt-5.5", "frontier_reasoning", 0.000005, 0.00003, max_images=20),
_ModelSpec("google/gemini-3.5-flash", "reasoning", 0.0000015, 0.000009, max_images=20, max_videos=4),
_ModelSpec("x-ai/grok-4.20", "reasoning", 0.00000125, 0.0000025, max_images=20),
_ModelSpec("x-ai/grok-4.3", "reasoning", 0.00000125, 0.0000025, max_images=20),
_ModelSpec("deepseek/deepseek-v4-pro", "reasoning", 0.000000435, 0.00000087),
_ModelSpec("deepseek/deepseek-v4-flash", "reasoning", 0.000000112, 0.000000224),
_ModelSpec("deepseek/deepseek-v3.2", "reasoning", 0.000000252, 0.000000378),
_ModelSpec("qwen/qwen3.6-max-preview", "reasoning", 0.00000104, 0.00000624),
_ModelSpec("qwen/qwen3.6-plus", "reasoning", 0.000000325, 0.00000195, max_images=10, max_videos=4),
_ModelSpec("qwen/qwen3.6-flash", "reasoning", 0.0000001875, 0.000001125, max_images=10, max_videos=4),
_ModelSpec("mistralai/mistral-large-2512", "standard", 0.0000005, 0.0000015, max_images=8),
_ModelSpec("mistralai/mistral-medium-3-5", "reasoning", 0.0000015, 0.0000075, max_images=8),
_ModelSpec("z-ai/glm-4.6", "reasoning", 0.00000043, 0.00000174),
_ModelSpec("z-ai/glm-5", "reasoning", 0.0000006, 0.00000192),
_ModelSpec("moonshotai/kimi-k2.6", "reasoning", 0.00000073, 0.00000349, max_images=10),
_ModelSpec("moonshotai/kimi-k2-thinking", "reasoning", 0.0000006, 0.0000025),
_ModelSpec("perplexity/sonar-pro", "perplexity", 0.000003, 0.000015),
_ModelSpec("perplexity/sonar-reasoning-pro", "perplexity_reasoning", 0.000002, 0.000008),
_ModelSpec("perplexity/sonar-deep-research", "perplexity_reasoning", 0.000002, 0.000008),
]
_MODELS_BY_SLUG: dict[str, _ModelSpec] = {m.slug: m for m in MODELS}
_REASONING_EFFORTS = ["off", "low", "medium", "high"]
_SEARCH_CONTEXT_SIZES = ["low", "medium", "high"]
def _reasoning_extra_inputs() -> list:
return [
IO.Combo.Input(
"reasoning_effort",
options=_REASONING_EFFORTS,
default="off",
tooltip="Reasoning effort. 'off' disables reasoning entirely.",
advanced=True,
),
]
def _perplexity_extra_inputs() -> list:
return [
IO.Combo.Input(
"search_context_size",
options=_SEARCH_CONTEXT_SIZES,
default="medium",
tooltip="How much web search context to retrieve. Larger = more grounded but slower/pricier.",
advanced=True,
),
]
def _profile_inputs(profile: Profile) -> list:
if profile == "standard":
return []
if profile in ("reasoning", "frontier_reasoning"):
return _reasoning_extra_inputs()
if profile == "perplexity":
return _perplexity_extra_inputs()
if profile == "perplexity_reasoning":
return _perplexity_extra_inputs() + _reasoning_extra_inputs()
raise ValueError(f"Unknown profile: {profile}")
def _media_inputs(spec: _ModelSpec) -> list:
extras: list = []
if spec.max_images > 0:
extras.append(
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("image"),
names=[f"image_{i}" for i in range(1, spec.max_images + 1)],
min=0,
),
tooltip=f"Optional reference image(s) — up to {spec.max_images}. Sent as URLs.",
)
)
if spec.max_videos > 0:
extras.append(
IO.Autogrow.Input(
"videos",
template=IO.Autogrow.TemplateNames(
IO.Video.Input("video"),
names=[f"video_{i}" for i in range(1, spec.max_videos + 1)],
min=0,
),
tooltip=f"Optional reference video(s) — up to {spec.max_videos}. Sent as URLs.",
)
)
return extras
def _inputs_for_model(spec: _ModelSpec) -> list:
return _profile_inputs(spec.profile) + _media_inputs(spec)
def _build_model_options() -> list[IO.DynamicCombo.Option]:
return [IO.DynamicCombo.Option(spec.slug, _inputs_for_model(spec)) for spec in MODELS]
def _calculate_price(response: OpenRouterChatResponse) -> float | None:
if response.usage and response.usage.cost is not None:
return float(response.usage.cost)
return None
def _price_badge_jsonata() -> str:
rates_pairs = []
for spec in MODELS:
prompt_per_1k = spec.price_in * 1000
completion_per_1k = spec.price_out * 1000
rates_pairs.append(f' "{spec.slug}": [{prompt_per_1k:.8g}, {completion_per_1k:.8g}]')
rates_block = ",\n".join(rates_pairs)
return (
"(\n"
" $rates := {\n"
f"{rates_block}\n"
" };\n"
" $r := $lookup($rates, widgets.model);\n"
" $r ? {\n"
' "type": "list_usd",\n'
' "usd": $r,\n'
' "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }\n'
' } : {"type": "text", "text": "Token-based"}\n'
")"
)
async def _build_image_blocks(
cls: type[IO.ComfyNode], spec: _ModelSpec, images: list[Input.Image]
) -> list[OpenRouterImageContent]:
urls = await upload_images_to_comfyapi(
cls,
images,
max_images=spec.max_images,
total_pixels=2048 * 2048,
mime_type="image/png",
wait_label="Uploading reference images",
)
return [OpenRouterImageContent(image_url=OpenRouterImageUrl(url=url)) for url in urls]
async def _build_video_blocks(cls: type[IO.ComfyNode], videos: list[Input.Video]) -> list[OpenRouterVideoContent]:
blocks: list[OpenRouterVideoContent] = []
total = len(videos)
for idx, video in enumerate(videos):
label = "Uploading reference video"
if total > 1:
label = f"{label} ({idx + 1}/{total})"
url = await upload_video_to_comfyapi(cls, video, wait_label=label)
blocks.append(OpenRouterVideoContent(video_url=OpenRouterVideoUrl(url=url)))
return blocks
def _user_message(prompt: str, media_blocks: list[OpenRouterContentBlock]) -> OpenRouterMessage:
if not media_blocks:
return OpenRouterMessage(role="user", content=prompt)
blocks: list[OpenRouterContentBlock] = list(media_blocks)
blocks.append(OpenRouterTextContent(text=prompt))
return OpenRouterMessage(role="user", content=blocks)
def _build_messages(
system_prompt: str, prompt: str, media_blocks: list[OpenRouterContentBlock]
) -> list[OpenRouterMessage]:
messages: list[OpenRouterMessage] = []
if system_prompt:
messages.append(OpenRouterMessage(role="system", content=system_prompt))
messages.append(_user_message(prompt, media_blocks))
return messages
def _build_request(
slug: str,
system_prompt: str,
prompt: str,
media_blocks: list[OpenRouterContentBlock],
*,
seed: int,
reasoning_effort: str | None,
search_context_size: str | None,
) -> OpenRouterChatRequest:
reasoning_cfg: OpenRouterReasoningConfig | None = None
if reasoning_effort and reasoning_effort != "off":
# exclude=True asks providers to reason internally but not return the trace
reasoning_cfg = OpenRouterReasoningConfig(effort=reasoning_effort, exclude=True)
web_search_cfg: OpenRouterWebSearchOptions | None = None
if search_context_size:
web_search_cfg = OpenRouterWebSearchOptions(search_context_size=search_context_size)
return OpenRouterChatRequest(
model=slug,
messages=_build_messages(system_prompt, prompt, media_blocks),
seed=seed if seed > 0 else None,
reasoning=reasoning_cfg,
web_search_options=web_search_cfg,
)
def _extract_text(response: OpenRouterChatResponse) -> str:
if response.error:
code = response.error.code if response.error.code is not None else "unknown"
raise ValueError(f"OpenRouter error ({code}): {response.error.message or 'no message'}")
if not response.choices:
raise ValueError("Empty response from OpenRouter (no choices).")
message = response.choices[0].message
if not message:
raise ValueError("Empty response from OpenRouter (no message).")
if message.refusal:
raise ValueError(f"Model refused to respond: {message.refusal}")
return message.content or ""
class OpenRouterLLMNode(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="OpenRouterLLMNode",
display_name="OpenRouter LLM",
category="api node/text/OpenRouter",
essentials_category="Text Generation",
description=(
"Generate text responses through OpenRouter. Routes to a curated set of popular "
"models from xAI, DeepSeek, Qwen, Mistral, Z.AI (GLM), Moonshot (Kimi), and "
"Perplexity Sonar."
),
inputs=[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Text input to the model.",
),
IO.DynamicCombo.Input(
"model",
options=_build_model_options(),
tooltip="The OpenRouter model used to generate the response.",
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=2147483647,
control_after_generate=True,
tooltip="Seed for sampling. Set to 0 to omit. Most models treat this as a hint only.",
),
IO.String.Input(
"system_prompt",
multiline=True,
default="",
optional=True,
advanced=True,
tooltip="Foundational instructions that dictate the model's behavior.",
),
],
outputs=[IO.String.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(widgets=["model"]),
expr=_price_badge_jsonata(),
),
)
@classmethod
async def execute(
cls,
prompt: str,
model: dict,
seed: int,
system_prompt: str = "",
) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1)
slug: str = model["model"]
spec = _MODELS_BY_SLUG.get(slug)
if spec is None:
raise ValueError(f"Unknown OpenRouter model: {slug}")
reasoning_effort: str | None = model.get("reasoning_effort")
search_context_size: str | None = model.get("search_context_size")
image_tensors: list[Input.Image] = [t for t in (model.get("images") or {}).values() if t is not None]
if image_tensors and sum(get_number_of_images(t) for t in image_tensors) > spec.max_images:
raise ValueError(f"Up to {spec.max_images} images are supported for {slug}.")
video_inputs: list[Input.Video] = [v for v in (model.get("videos") or {}).values() if v is not None]
if video_inputs and len(video_inputs) > spec.max_videos:
raise ValueError(f"Up to {spec.max_videos} videos are supported for {slug}.")
media_blocks: list[OpenRouterContentBlock] = []
if image_tensors:
media_blocks.extend(await _build_image_blocks(cls, spec, image_tensors))
if video_inputs:
media_blocks.extend(await _build_video_blocks(cls, video_inputs))
request = _build_request(
slug,
system_prompt,
prompt,
media_blocks,
seed=seed,
reasoning_effort=reasoning_effort,
search_context_size=search_context_size,
)
response = await sync_op(
cls,
ApiEndpoint(path=OPENROUTER_CHAT_ENDPOINT, method="POST"),
response_model=OpenRouterChatResponse,
data=request,
price_extractor=_calculate_price,
)
return IO.NodeOutput(_extract_text(response))
class OpenRouterExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [OpenRouterLLMNode]
async def comfy_entrypoint() -> OpenRouterExtension:
return OpenRouterExtension()

View File

@ -5,32 +5,37 @@ Rodin API docs: https://developer.hyper3d.ai/
"""
from inspect import cleandoc
import folder_paths as comfy_paths
import os
import logging
import math
import os
from inspect import cleandoc
from io import BytesIO
from typing_extensions import override
from typing import Any
import aiohttp
from PIL import Image
from typing_extensions import override
import folder_paths as comfy_paths
from comfy_api.latest import IO, ComfyExtension, Types
from comfy_api_nodes.apis.rodin import (
Rodin3DGenerateRequest,
Rodin3DGenerateResponse,
JobStatus,
Rodin3DCheckStatusRequest,
Rodin3DCheckStatusResponse,
Rodin3DDownloadRequest,
Rodin3DDownloadResponse,
JobStatus,
Rodin3DGen25Request,
Rodin3DGenerateRequest,
Rodin3DGenerateResponse,
)
from comfy_api_nodes.util import (
sync_op,
poll_op,
ApiEndpoint,
download_url_to_bytesio,
download_url_to_file_3d,
poll_op,
sync_op,
validate_string,
)
from comfy_api.latest import ComfyExtension, IO, Types
COMMON_PARAMETERS = [
IO.Int.Input(
@ -51,40 +56,30 @@ COMMON_PARAMETERS = [
]
def get_quality_mode(poly_count):
polycount = poly_count.split("-")
poly = polycount[1]
count = polycount[0]
if poly == "Triangle":
mesh_mode = "Raw"
elif poly == "Quad":
mesh_mode = "Quad"
else:
mesh_mode = "Quad"
if count == "4K":
quality_override = 4000
elif count == "8K":
quality_override = 8000
elif count == "18K":
quality_override = 18000
elif count == "50K":
quality_override = 50000
elif count == "2K":
quality_override = 2000
elif count == "20K":
quality_override = 20000
elif count == "150K":
quality_override = 150000
elif count == "500K":
quality_override = 500000
else:
quality_override = 18000
return mesh_mode, quality_override
_QUALITY_MESH_OPTIONS: dict[str, tuple[str, int]] = {
"4K-Quad": ("Quad", 4000),
"8K-Quad": ("Quad", 8000),
"18K-Quad": ("Quad", 18000),
"50K-Quad": ("Quad", 50000),
"200K-Quad": ("Quad", 200000),
"2K-Triangle": ("Raw", 2000),
"20K-Triangle": ("Raw", 20000),
"150K-Triangle": ("Raw", 150000),
"200K-Triangle": ("Raw", 200000),
"500K-Triangle": ("Raw", 500000),
"1M-Triangle": ("Raw", 1000000),
}
def tensor_to_filelike(tensor, max_pixels: int = 2048*2048):
def get_quality_mode(poly_count: str) -> tuple[str, int]:
"""Map a polygon-count preset like '18K-Quad' to (mesh_mode, quality_override).
Falls back to ('Quad', 18000) for unknown labels; legacy parity.
"""
return _QUALITY_MESH_OPTIONS.get(poly_count, ("Quad", 18000))
def tensor_to_filelike(tensor, max_pixels: int = 2048 * 2048):
"""
Converts a PyTorch tensor to a file-like object.
@ -96,8 +91,8 @@ def tensor_to_filelike(tensor, max_pixels: int = 2048*2048):
- io.BytesIO: A file-like object containing the image data.
"""
array = tensor.cpu().numpy()
array = (array * 255).astype('uint8')
image = Image.fromarray(array, 'RGB')
array = (array * 255).astype("uint8")
image = Image.fromarray(array, "RGB")
original_width, original_height = image.size
original_pixels = original_width * original_height
@ -112,7 +107,7 @@ def tensor_to_filelike(tensor, max_pixels: int = 2048*2048):
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
img_byte_arr = BytesIO()
image.save(img_byte_arr, format='PNG') # PNG is used for lossless compression
image.save(img_byte_arr, format="PNG") # PNG is used for lossless compression
img_byte_arr.seek(0)
return img_byte_arr
@ -145,11 +140,9 @@ async def create_generate_task(
TAPose=ta_pose,
),
files=[
(
"images",
open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image)
)
for image in images if image is not None
("images", open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image))
for image in images
if image is not None
],
content_type="multipart/form-data",
)
@ -177,6 +170,7 @@ def check_rodin_status(response: Rodin3DCheckStatusResponse) -> str:
return "DONE"
return "Generating"
def extract_progress(response: Rodin3DCheckStatusResponse) -> int | None:
if not response.jobs:
return None
@ -214,7 +208,7 @@ async def download_files(url_list, task_uuid: str) -> tuple[str | None, Types.Fi
model_file_path = None
file_3d = None
for i in url_list.list:
for i in url_list.items:
file_path = os.path.join(save_path, i.name)
if i.name.lower().endswith(".glb"):
model_file_path = os.path.join(result_folder_name, i.name)
@ -489,7 +483,16 @@ class Rodin3D_Gen2(IO.ComfyNode):
IO.Combo.Input("Material_Type", options=["PBR", "Shaded"], default="PBR", optional=True),
IO.Combo.Input(
"Polygon_count",
options=["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "2K-Triangle", "20K-Triangle", "150K-Triangle", "500K-Triangle"],
options=[
"4K-Quad",
"8K-Quad",
"18K-Quad",
"50K-Quad",
"2K-Triangle",
"20K-Triangle",
"150K-Triangle",
"500K-Triangle",
],
default="500K-Triangle",
optional=True,
),
@ -542,6 +545,566 @@ class Rodin3D_Gen2(IO.ComfyNode):
return IO.NodeOutput(model_path, file_3d)
def _rodin_multipart_parser(data: dict[str, Any]) -> aiohttp.FormData:
"""Convert a Rodin request dict to an aiohttp form, fixing bool/list serialization.
Booleans --> "true"/"false". Lists --> one field per element.
"""
form = aiohttp.FormData(default_to_multipart=True)
for key, value in data.items():
if value is None:
continue
if isinstance(value, bool):
form.add_field(key, "true" if value else "false")
elif isinstance(value, list):
for item in value:
form.add_field(key, str(item))
elif isinstance(value, (bytes, bytearray)):
form.add_field(key, value)
else:
form.add_field(key, str(value))
return form
async def _create_gen25_task(
cls: type[IO.ComfyNode],
request: Rodin3DGen25Request,
images: list | None,
) -> tuple[str, str]:
"""Submit a Gen-2.5 generate job; returns (task_uuid, subscription_key)."""
if images is not None and len(images) > 5:
raise ValueError("Rodin Gen-2.5 supports at most 5 input images.")
files = None
if images:
files = [
(
"images",
open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image),
)
for image in images
if image is not None
]
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/rodin/api/v2/rodin", method="POST"),
response_model=Rodin3DGenerateResponse,
data=request,
files=files,
content_type="multipart/form-data",
multipart_parser=_rodin_multipart_parser,
)
if not response.uuid or not response.jobs or not response.jobs.subscription_key:
raise RuntimeError(f"Rodin Gen-2.5 submit failed: message={response.message!r}")
return response.uuid, response.jobs.subscription_key
_PREVIEWABLE_3D_EXTS = {".glb", ".obj", ".fbx", ".stl", ".gltf"}
async def _download_gen25_files(
download_list: Rodin3DDownloadResponse,
task_uuid: str,
geometry_file_format: str,
) -> Types.File3D | None:
"""Download every file in the list; return the File3D matching the chosen format."""
folder_name = f"Rodin3D_Gen25_{task_uuid}"
save_dir = os.path.join(comfy_paths.get_output_directory(), folder_name)
os.makedirs(save_dir, exist_ok=True)
target_ext = f".{geometry_file_format.lower().lstrip('.')}"
file_3d: Types.File3D | None = None
for item in download_list.items:
file_path = os.path.join(save_dir, item.name)
ext = os.path.splitext(item.name.lower())[1]
# Prefer the file matching the user's chosen format; fall back below.
if file_3d is None and ext == target_ext and ext in _PREVIEWABLE_3D_EXTS:
file_3d = await download_url_to_file_3d(item.url, target_ext.lstrip("."))
with open(file_path, "wb") as f:
f.write(file_3d.get_bytes())
continue
await download_url_to_bytesio(item.url, file_path)
# If the chosen format wasn't found, surface any model file we did get.
if file_3d is None:
for item in download_list.items:
ext = os.path.splitext(item.name.lower())[1]
if ext in _PREVIEWABLE_3D_EXTS:
file_3d = await download_url_to_file_3d(item.url, ext.lstrip("."))
break
return file_3d
_MODE_REGULAR = "Regular"
_MODE_FAST = "Fast"
_MODE_EXTREME_HIGH = "Extreme-High"
_REGULAR_POLY_OPTIONS = [
"Default",
"4K-Quad",
"8K-Quad",
"18K-Quad",
"50K-Quad",
"2K-Triangle",
"20K-Triangle",
"150K-Triangle",
"500K-Triangle",
"1M-Triangle",
]
_TEXTURE_MODE_OPTIONS = ["Default", "legacy", "extreme-low", "low", "medium", "high"]
_GEOMETRY_FORMAT_OPTIONS = ["glb", "fbx", "obj", "stl"]
_MATERIAL_OPTIONS = ["PBR", "Shaded", "All", "None"]
def _build_mode_input(name: str = "mode") -> IO.DynamicCombo.Input:
return IO.DynamicCombo.Input(
name,
options=[
IO.DynamicCombo.Option(
_MODE_REGULAR,
[
IO.Combo.Input(
"tier",
options=["Gen-2.5-Low", "Gen-2.5-Medium", "Gen-2.5-High"],
default="Gen-2.5-High",
tooltip="Quality tier. Higher tiers produce higher-fidelity geometry.",
),
IO.Combo.Input(
"polygon_count",
options=_REGULAR_POLY_OPTIONS,
default="Default",
tooltip="Preset face count. 'Default' uses the server's default for the selected tier.",
),
IO.Boolean.Input(
"creative",
default=False,
tooltip="Creative mode (Medium/High only). Enhances generative robustness.",
),
],
),
IO.DynamicCombo.Option(
_MODE_FAST,
[
IO.Combo.Input(
"tier",
options=[
"Gen-2.5-Extreme-Low",
"Gen-2.5-Low",
"Gen-2.5-Medium",
"Gen-2.5-High",
],
default="Gen-2.5-Low",
),
IO.Int.Input(
"mesh_faces",
default=20000,
min=1000,
max=20000,
display_mode=IO.NumberDisplay.number,
tooltip="Mesh face count (1K-20K in Fast mode).",
),
],
),
IO.DynamicCombo.Option(
_MODE_EXTREME_HIGH,
[
IO.Combo.Input("mesh_mode", options=["Raw", "Quad"], default="Raw"),
IO.Int.Input(
"mesh_faces",
default=1000000,
min=20000,
max=2000000,
display_mode=IO.NumberDisplay.number,
tooltip=(
"Mesh face count. Raw mode: 20K-2M. "
"Quad mode: keep under 200K (upstream may reject higher values)."
),
),
IO.Boolean.Input(
"is_micro",
default=False,
tooltip="Enable micro detail (Extreme-High only).",
),
IO.Boolean.Input(
"creative",
default=False,
tooltip="Creative mode. Enhances generative robustness.",
),
],
),
],
tooltip=(
"Generation mode. Regular = balanced. Fast = 1K-20K faces for rapid prototyping. "
"Extreme-High = 20K-2M faces with optional micro details."
),
)
def _build_common_inputs(*, include_image_only: bool) -> list:
inputs: list = [
IO.Combo.Input("material", options=_MATERIAL_OPTIONS, default="Shaded"),
IO.Combo.Input("geometry_file_format", options=_GEOMETRY_FORMAT_OPTIONS, default="glb"),
IO.Combo.Input(
"texture_mode",
options=_TEXTURE_MODE_OPTIONS,
default="Default",
optional=True,
tooltip="Texture quality preset. 'Default' uses the server's default for the selected tier.",
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=65535,
display_mode=IO.NumberDisplay.number,
control_after_generate=True,
optional=True,
),
IO.Boolean.Input(
"TAPose", default=False, optional=True, advanced=True, tooltip="T/A pose for human-like models."
),
IO.Boolean.Input(
"hd_texture", default=False, optional=True, advanced=True, tooltip="High-quality texture enhancement."
),
IO.Boolean.Input(
"texture_delight",
default=False,
optional=True,
advanced=True,
tooltip="Remove baked lighting from textures.",
),
]
if include_image_only:
inputs.append(
IO.Boolean.Input(
"use_original_alpha",
default=False,
optional=True,
advanced=True,
tooltip="Preserve image transparency.",
)
)
inputs.extend(
[
IO.Boolean.Input(
"addon_highpack",
default=False,
optional=True,
advanced=True,
tooltip="HighPack addon: 4K textures and ~16x faces in Quad mode.",
),
IO.Int.Input(
"bbox_width",
default=0,
min=0,
max=300,
display_mode=IO.NumberDisplay.number,
optional=True,
advanced=True,
tooltip="Bounding-box width (Y axis). Set to 0 with the others to skip bbox.",
),
IO.Int.Input(
"bbox_height",
default=0,
min=0,
max=300,
display_mode=IO.NumberDisplay.number,
optional=True,
advanced=True,
tooltip="Bounding-box height (Z axis).",
),
IO.Int.Input(
"bbox_length",
default=0,
min=0,
max=300,
display_mode=IO.NumberDisplay.number,
optional=True,
advanced=True,
tooltip="Bounding-box length (X axis).",
),
IO.Int.Input(
"height_cm",
default=0,
min=0,
max=10000,
display_mode=IO.NumberDisplay.number,
optional=True,
advanced=True,
tooltip="Approximate model height in centimeters (0 to skip).",
),
]
)
return inputs
_PRICE_EXPR = """
(
$baseCredits := widgets.mode = "extreme-high" ? 1.0 : 0.5;
$addonCredits := widgets.addon_highpack ? 1.0 : 0.0;
$total := ($baseCredits * 1.5) + ($addonCredits * 0.8);
{"type":"usd","usd": $total}
)
"""
def _resolve_mode_params(mode_input: dict) -> dict:
"""Translate the DynamicCombo `mode` payload into Gen-2.5 request fields.
Returns a dict with: tier, quality_override, mesh_mode, geometry_instruct_mode, is_micro.
Missing keys mean "do not send" (so we don't override server defaults).
"""
selected = mode_input["mode"]
out: dict = {}
if selected == _MODE_REGULAR:
out["tier"] = mode_input["tier"]
polygon = mode_input.get("polygon_count", "Default")
if polygon != "Default":
mesh_mode, faces = get_quality_mode(polygon)
out["mesh_mode"] = mesh_mode
out["quality_override"] = faces
if mode_input.get("creative"):
out["geometry_instruct_mode"] = "creative"
elif selected == _MODE_FAST:
out["tier"] = mode_input["tier"]
out["mesh_mode"] = "Raw"
out["quality_override"] = int(mode_input["mesh_faces"])
elif selected == _MODE_EXTREME_HIGH:
out["tier"] = "Gen-2.5-Extreme-High"
out["mesh_mode"] = mode_input["mesh_mode"]
out["quality_override"] = int(mode_input["mesh_faces"])
if mode_input.get("is_micro"):
out["is_micro"] = True
if mode_input.get("creative"):
out["geometry_instruct_mode"] = "creative"
return out
def _build_request(
*,
mode_input: dict,
material: str,
geometry_file_format: str,
texture_mode: str,
seed: int,
TAPose: bool,
hd_texture: bool,
texture_delight: bool,
addon_highpack: bool,
bbox_width: int,
bbox_height: int,
bbox_length: int,
height_cm: int,
prompt: str | None = None,
use_original_alpha: bool = False,
) -> Rodin3DGen25Request:
mode_params = _resolve_mode_params(mode_input)
bbox = None
if bbox_width and bbox_height and bbox_length:
bbox = [bbox_width, bbox_height, bbox_length]
return Rodin3DGen25Request(
tier=mode_params["tier"],
prompt=prompt or None,
seed=seed,
material=material,
geometry_file_format=geometry_file_format,
texture_mode=None if texture_mode == "Default" else texture_mode,
mesh_mode=mode_params.get("mesh_mode"),
quality_override=mode_params.get("quality_override"),
geometry_instruct_mode=mode_params.get("geometry_instruct_mode"),
bbox_condition=bbox,
height=height_cm or None,
TAPose=TAPose or None,
hd_texture=hd_texture or None,
texture_delight=texture_delight or None,
is_micro=mode_params.get("is_micro"),
use_original_alpha=use_original_alpha or None,
addons=["HighPack"] if addon_highpack else None,
)
class Rodin3D_Gen25_Image(IO.ComfyNode):
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="Rodin3D_Gen25_Image",
display_name="Rodin 3D Gen-2.5 - Image to 3D",
category="api node/3d/Rodin",
description=(
"Generate a 3D model from 1-5 reference images via Rodin Gen-2.5. "
"Pick a mode (Fast / Regular / Extreme-High) to tune quality vs. cost."
),
inputs=[
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplatePrefix(IO.Image.Input("image"), prefix="image", min=1, max=5),
tooltip="1-5 images. The first image is used for materials when multi-view.",
),
_build_mode_input(),
*_build_common_inputs(include_image_only=True),
],
outputs=[IO.File3DAny.Output(display_name="model_file")],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(widgets=["mode", "addon_highpack"]),
expr=_PRICE_EXPR,
),
)
@classmethod
async def execute(
cls,
images: IO.Autogrow.Type,
mode: dict,
material: str,
geometry_file_format: str,
texture_mode: str,
seed: int,
TAPose: bool,
hd_texture: bool,
texture_delight: bool,
use_original_alpha: bool,
addon_highpack: bool,
bbox_width: int,
bbox_height: int,
bbox_length: int,
height_cm: int,
) -> IO.NodeOutput:
image_tensors = [img for img in images.values() if img is not None]
if not image_tensors:
raise ValueError("Rodin Gen-2.5 Image-to-3D requires at least one image.")
# Flatten multi-image tensors into individual frames; the API accepts each as a separate part.
flat_images: list = []
for tensor in image_tensors:
if hasattr(tensor, "shape") and len(tensor.shape) == 4:
for i in range(tensor.shape[0]):
flat_images.append(tensor[i])
else:
flat_images.append(tensor)
if len(flat_images) > 5:
raise ValueError(f"Rodin Gen-2.5 accepts at most 5 images; received {len(flat_images)}.")
request = _build_request(
mode_input=mode,
material=material,
geometry_file_format=geometry_file_format,
texture_mode=texture_mode,
seed=seed,
TAPose=TAPose,
hd_texture=hd_texture,
texture_delight=texture_delight,
addon_highpack=addon_highpack,
bbox_width=bbox_width,
bbox_height=bbox_height,
bbox_length=bbox_length,
height_cm=height_cm,
prompt=None,
use_original_alpha=use_original_alpha,
)
task_uuid, subscription_key = await _create_gen25_task(cls, request, flat_images)
await poll_for_task_status(subscription_key, cls)
download_list = await get_rodin_download_list(task_uuid, cls)
file_3d = await _download_gen25_files(download_list, task_uuid, geometry_file_format)
return IO.NodeOutput(file_3d)
class Rodin3D_Gen25_Text(IO.ComfyNode):
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="Rodin3D_Gen25_Text",
display_name="Rodin 3D Gen-2.5 - Text to 3D",
category="api node/3d/Rodin",
description=(
"Generate a 3D model from a text prompt via Rodin Gen-2.5. "
"Pick a mode (Fast / Regular / Extreme-High) to tune quality vs. cost."
),
inputs=[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Text prompt for the 3D model.",
),
_build_mode_input(),
*_build_common_inputs(include_image_only=False),
],
outputs=[IO.File3DAny.Output(display_name="model_file")],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(widgets=["mode", "addon_highpack"]),
expr=_PRICE_EXPR,
),
)
@classmethod
async def execute(
cls,
prompt: str,
mode: dict,
material: str,
geometry_file_format: str,
texture_mode: str,
seed: int,
TAPose: bool,
hd_texture: bool,
texture_delight: bool,
addon_highpack: bool,
bbox_width: int,
bbox_height: int,
bbox_length: int,
height_cm: int,
) -> IO.NodeOutput:
validate_string(prompt, field_name="prompt", min_length=1, max_length=2500)
request = _build_request(
mode_input=mode,
material=material,
geometry_file_format=geometry_file_format,
texture_mode=texture_mode,
seed=seed,
TAPose=TAPose,
hd_texture=hd_texture,
texture_delight=texture_delight,
addon_highpack=addon_highpack,
bbox_width=bbox_width,
bbox_height=bbox_height,
bbox_length=bbox_length,
height_cm=height_cm,
prompt=prompt,
)
task_uuid, subscription_key = await _create_gen25_task(cls, request, images=None)
await poll_for_task_status(subscription_key, cls)
download_list = await get_rodin_download_list(task_uuid, cls)
file_3d = await _download_gen25_files(download_list, task_uuid, geometry_file_format)
return IO.NodeOutput(file_3d)
class Rodin3DExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@ -551,6 +1114,8 @@ class Rodin3DExtension(ComfyExtension):
Rodin3D_Smooth,
Rodin3D_Sketch,
Rodin3D_Gen2,
Rodin3D_Gen25_Image,
Rodin3D_Gen25_Text,
]

View File

@ -16,16 +16,17 @@ from .conversions import (
convert_mask_to_image,
downscale_image_tensor,
downscale_image_tensor_by_max_side,
downscale_video_to_max_pixels,
image_tensor_pair_to_batch,
pil_to_bytesio,
resize_mask_to_image,
resize_video_to_pixel_budget,
tensor_to_base64_string,
tensor_to_bytesio,
tensor_to_pil,
text_filepath_to_base64_string,
text_filepath_to_data_uri,
trim_video,
upscale_video_to_min_pixels,
video_to_base64_string,
)
from .download_helpers import (
@ -88,16 +89,17 @@ __all__ = [
"convert_mask_to_image",
"downscale_image_tensor",
"downscale_image_tensor_by_max_side",
"downscale_video_to_max_pixels",
"image_tensor_pair_to_batch",
"pil_to_bytesio",
"resize_mask_to_image",
"resize_video_to_pixel_budget",
"tensor_to_base64_string",
"tensor_to_bytesio",
"tensor_to_pil",
"text_filepath_to_base64_string",
"text_filepath_to_data_uri",
"trim_video",
"upscale_video_to_min_pixels",
"video_to_base64_string",
# Validation utilities
"get_image_dimensions",

View File

@ -415,14 +415,48 @@ def trim_video(video: Input.Video, duration_sec: float) -> Input.Video:
raise RuntimeError(f"Failed to trim video: {str(e)}") from e
def resize_video_to_pixel_budget(video: Input.Video, total_pixels: int) -> Input.Video:
"""Downscale a video to fit within ``total_pixels`` (w * h), preserving aspect ratio.
def downscale_video_to_max_pixels(video: Input.Video, max_pixels: int) -> Input.Video:
"""Downscale a video to fit within ``max_pixels`` (w * h), preserving aspect ratio.
Returns the original video object untouched when it already fits. Preserves frame rate, duration, and audio.
Aspect ratio is preserved up to a fraction of a percent (even-dim rounding).
"""
src_w, src_h = video.get_dimensions()
scale_dims = _compute_downscale_dims(src_w, src_h, total_pixels)
scale_dims = _compute_downscale_dims(src_w, src_h, max_pixels)
if scale_dims is None:
return video
return _apply_video_scale(video, scale_dims)
def _compute_upscale_dims(src_w: int, src_h: int, total_pixels: int) -> tuple[int, int] | None:
"""Return upscaled (w, h) with even dims meeting at least ``total_pixels``, or None if already large enough.
Source aspect ratio is preserved; output may drift by a fraction of a percent because both dimensions
are rounded up to even values (many codecs require divisible-by-2). The result is guaranteed to be at
least ``total_pixels``.
"""
pixels = src_w * src_h
if pixels >= total_pixels:
return None
scale = math.sqrt(total_pixels / pixels)
new_w = math.ceil(src_w * scale)
new_h = math.ceil(src_h * scale)
if new_w % 2:
new_w += 1
if new_h % 2:
new_h += 1
return new_w, new_h
def upscale_video_to_min_pixels(video: Input.Video, min_pixels: int) -> Input.Video:
"""Upscale a video to meet at least ``min_pixels`` (w * h), preserving aspect ratio.
Returns the original video object untouched when it already meets the minimum. Preserves frame rate,
duration, and audio. Aspect ratio is preserved up to a fraction of a percent (even-dim rounding).
Note: upscaling a low-resolution source does not add real detail; downstream model quality may suffer.
"""
src_w, src_h = video.get_dimensions()
scale_dims = _compute_upscale_dims(src_w, src_h, min_pixels)
if scale_dims is None:
return video
return _apply_video_scale(video, scale_dims)

View File

@ -164,7 +164,6 @@ class BasicCache:
await self.cache_key_set.add_keys(node_ids)
self.is_changed_cache = is_changed_cache
self.initialized = True
await self._notify_providers_set_prompt()
def all_node_ids(self):
assert self.initialized
@ -264,24 +263,6 @@ class BasicCache:
except Exception as e:
_logger.warning(f"Cache provider {provider.__class__.__name__} error on store: {e}")
async def _notify_providers_set_prompt(self):
from comfy_execution.cache_provider import (
_has_cache_providers, _get_cache_providers, _logger
)
if not self.enable_providers:
return
if not _has_cache_providers():
return
for provider in _get_cache_providers():
try:
task = asyncio.create_task(self._safe_provider_set_prompt(provider))
self._pending_store_tasks.add(task)
task.add_done_callback(self._pending_store_tasks.discard)
except Exception as e:
_logger.warning(f"Cache provider {provider.__class__.__name__} error on set_prompt: {e}")
@staticmethod
async def _safe_provider_store(provider, context, cache_value):
from comfy_execution.cache_provider import _logger
@ -290,14 +271,6 @@ class BasicCache:
except Exception as e:
_logger.warning(f"Cache provider {provider.__class__.__name__} async store error: {e}")
@staticmethod
async def _safe_provider_set_prompt(provider):
from comfy_execution.cache_provider import _logger
try:
await provider.on_set_prompt()
except Exception as e:
_logger.warning(f"Cache provider {provider.__class__.__name__} async set_prompt error: {e}")
async def _check_providers_lookup(self, node_id, cache_key):
from comfy_execution.cache_provider import (
_has_cache_providers, _get_cache_providers,

View File

@ -1,111 +0,0 @@
"""Pure-numpy port of MediaPipe's face_geometry (FACE_LANDMARK_PIPELINE mode)
+ weighted Procrustes solver. Computes the 4x4 facial transformation matrix.
"""
from __future__ import annotations
import math
import numpy as np
def _solve_weighted_orthogonal_problem(src: np.ndarray, tgt: np.ndarray, weights: np.ndarray) -> np.ndarray:
"""Weighted orthogonal Procrustes (similarity). Returns 4x4 M with
`target ≈ M @ homogeneous(source)` in the weighted LS sense. fp64 for
SVD stability. Port of procrustes_solver.cc."""
sqrt_w = np.sqrt(weights.astype(np.float64))
w_total = float((sqrt_w ** 2).sum())
ws = src.astype(np.float64) * sqrt_w
wt = tgt.astype(np.float64) * sqrt_w
c_w = (ws @ sqrt_w) / w_total
centered = ws - np.outer(c_w, sqrt_w)
U, _S, Vt = np.linalg.svd(wt @ centered.T, full_matrices=True)
# Disallow reflection: flip the least-significant axis when det(U)·det(V)<0.
post, pre = U.copy(), Vt.T.copy()
if np.linalg.det(post) * np.linalg.det(pre) < 0:
post[:, 2] *= -1.0
R = post @ pre.T
denom = float((centered * ws).sum())
if denom < 1e-12:
raise ValueError("Procrustes denominator collapsed (degenerate source).")
scale = float((R @ centered * wt).sum()) / denom
translation = ((wt - scale * (R @ ws)) @ sqrt_w) / w_total
M = np.eye(4, dtype=np.float64)
M[:3, :3] = scale * R
M[:3, 3] = translation
return M
def _estimate_scale(canonical: np.ndarray, runtime: np.ndarray, weights: np.ndarray) -> float:
"""scale = ‖first column of M[:3]‖ per geometry_pipeline.cc::EstimateScale."""
return float(np.linalg.norm(_solve_weighted_orthogonal_problem(canonical, runtime, weights)[:3, 0]))
def solve_facial_transformation_matrix(
landmarks_normalized: np.ndarray,
canonical_vertices: np.ndarray,
procrustes_indices: np.ndarray,
procrustes_weights: np.ndarray,
image_width: int,
image_height: int,
# face_geometry_calculator_options.pbtxt defaults
vertical_fov_degrees: float = 63.0,
near: float = 1.0,
) -> np.ndarray:
"""4x4 facial transformation matrix via two-pass scale recovery
`landmarks_normalized` is (N, 3) in MediaPipe normalized convention: x, y
in [0,1] with TOP-LEFT origin, z in width-scaled units.
"""
h_near = 2.0 * near * math.tan(0.5 * math.radians(vertical_fov_degrees))
w_near = image_width * h_near / image_height
sub = procrustes_indices.astype(np.int64)
screen = landmarks_normalized[sub].T.astype(np.float64).copy()
canon = canonical_vertices[sub].T.astype(np.float64).copy()
weights = procrustes_weights.astype(np.float64)
# ProjectXY (TOP_LEFT y-flip, then scale all 3 axes; z uses x-scale).
screen[1] = 1.0 - screen[1]
screen[0] = screen[0] * w_near - 0.5 * w_near
screen[1] = screen[1] * h_near - 0.5 * h_near
screen[2] = screen[2] * w_near
depth_offset = float(screen[2].mean())
def _unproject(s: np.ndarray, scale: float) -> np.ndarray:
s = s.copy()
s[2] = (s[2] - depth_offset + near) / scale
s[0] *= s[2] / near
s[1] *= s[2] / near
s[2] *= -1.0
return s
first = screen.copy()
first[2] *= -1.0
s1 = _estimate_scale(canon, first, weights) # 1st pass: Procrustes on projected XY
s2 = _estimate_scale(canon, _unproject(screen, s1), weights) # 2nd pass: rescale z by s1, un-project XY
return _solve_weighted_orthogonal_problem(canon, _unproject(screen, s1 * s2), weights).astype(np.float32)
def transformation_matrix_from_detection(face_dict: dict, image_width: int, image_height: int, canonical_data: dict) -> np.ndarray:
"""Adapt a FaceLandmarker face dict to MP's normalized convention and solve.
FaceMesh emits (x, y, z) in 192-canonical units; MP's geometry expects
z_norm = z_canonical * scale_x / image_width"""
lmks_xy, lmks_3d = face_dict["landmarks_xy"], face_dict["landmarks_3d"]
aug = np.concatenate([lmks_3d[:, :2].astype(np.float64), np.ones((lmks_xy.shape[0], 1))], axis=1)
M, *_ = np.linalg.lstsq(aug, lmks_xy.astype(np.float64), rcond=None)
scale_x = float(np.linalg.norm(M[0]))
z_scale = scale_x / image_width if scale_x > 1e-6 else 1.0 / image_width
normalized = np.empty((lmks_xy.shape[0], 3), dtype=np.float32)
normalized[:, 0] = lmks_xy[:, 0] / image_width
normalized[:, 1] = lmks_xy[:, 1] / image_height
normalized[:, 2] = lmks_3d[:, 2] * z_scale
return solve_facial_transformation_matrix(
normalized, canonical_data["canonical_vertices"],
canonical_data["procrustes_indices"], canonical_data["procrustes_weights"],
image_width=image_width, image_height=image_height,
)

View File

@ -1,682 +0,0 @@
"""Pure-PyTorch port of MediaPipe's face_landmarker_v2_with_blendshapes.task:
BlazeFace detector → FaceMesh v2 → ARKit-52 blendshapes."""
from __future__ import annotations
import math
from functools import lru_cache
from typing import List, Tuple
import numpy as np
import torch
import torch.nn.functional as F
from scipy.special import expit
from torch import Tensor, nn
# Values below must stay verbatim with the published face_landmarker_v2 graph
# face_blendshapes_graph.cc::kLandmarksSubsetIdxs
_BS_INPUT_INDICES: Tuple[int, ...] = (
0, 1, 4, 5, 6, 7, 8, 10, 13, 14, 17, 21, 33, 37, 39, 40, 46, 52, 53, 54,
55, 58, 61, 63, 65, 66, 67, 70, 78, 80, 81, 82, 84, 87, 88, 91, 93, 95,
103, 105, 107, 109, 127, 132, 133, 136, 144, 145, 146, 148, 149, 150, 152,
153, 154, 155, 157, 158, 159, 160, 161, 162, 163, 168, 172, 173, 176, 178,
181, 185, 191, 195, 197, 234, 246, 249, 251, 263, 267, 269, 270, 276, 282,
283, 284, 285, 288, 291, 293, 295, 296, 297, 300, 308, 310, 311, 312, 314,
317, 318, 321, 323, 324, 332, 334, 336, 338, 356, 361, 362, 365, 373, 374,
375, 377, 378, 379, 380, 381, 382, 384, 385, 386, 387, 388, 389, 390, 397,
398, 400, 402, 405, 409, 415, 454, 466, 468, 469, 470, 471, 472, 473, 474,
475, 476, 477,
)
# face_blendshapes_graph.cc::kCategoryNames
BLENDSHAPE_NAMES: Tuple[str, ...] = (
"_neutral", "browDownLeft", "browDownRight", "browInnerUp", "browOuterUpLeft",
"browOuterUpRight", "cheekPuff", "cheekSquintLeft", "cheekSquintRight",
"eyeBlinkLeft", "eyeBlinkRight", "eyeLookDownLeft", "eyeLookDownRight",
"eyeLookInLeft", "eyeLookInRight", "eyeLookOutLeft", "eyeLookOutRight",
"eyeLookUpLeft", "eyeLookUpRight", "eyeSquintLeft", "eyeSquintRight",
"eyeWideLeft", "eyeWideRight", "jawForward", "jawLeft", "jawOpen",
"jawRight", "mouthClose", "mouthDimpleLeft", "mouthDimpleRight",
"mouthFrownLeft", "mouthFrownRight", "mouthFunnel", "mouthLeft",
"mouthLowerDownLeft", "mouthLowerDownRight", "mouthPressLeft",
"mouthPressRight", "mouthPucker", "mouthRight", "mouthRollLower",
"mouthRollUpper", "mouthShrugLower", "mouthShrugUpper", "mouthSmileLeft",
"mouthSmileRight", "mouthStretchLeft", "mouthStretchRight",
"mouthUpperUpLeft", "mouthUpperUpRight", "noseSneerLeft", "noseSneerRight",
)
# face_detection.pbtxt — short-range BlazeFace.
_BF_NUM_LAYERS = 4
_BF_INPUT_SIZE = 128
_BF_STRIDES = (8, 16, 16, 16)
_BF_ANCHOR_OFFSET_X = 0.5
_BF_ANCHOR_OFFSET_Y = 0.5
_BF_ASPECT_RATIOS = (1.0,)
_BF_INTERP_SCALE_AR = 1.0
_BF_BOX_SCALE = 128.0
_BF_KP_OFFSET = 4
_BF_SCORE_CLIP = 100.0
_BF_MIN_SCORE = 0.5
# face_detection_full_range.pbtxt — 48x48 grid at stride 4, 1 anchor/cell.
_BF_FR_INPUT_SIZE = 192
_BF_FR_GRID = 48
_BF_FR_NUM_ANCHORS = _BF_FR_GRID * _BF_FR_GRID
_BF_FR_BOX_SCALE = 192.0
_BF_FR_SCORE_CLIP = 100.0
_FM_INPUT_SIZE = 192
# Face ROI: 1.5xbbox rect warped anisotropically into 192x192.
_FACE_LEFT_EYE_KP = 0
_FACE_RIGHT_EYE_KP = 1
_FACE_ROI_SCALE_X = 1.5
_FACE_ROI_SCALE_Y = 1.5
_FACE_ROI_TARGET_ANGLE = 0.0
def _tf_same_pad(x: Tensor, kernel: int, stride: int) -> Tensor:
"""TF SAME pad (asymmetric on stride-2; PyTorch's symmetric pad undershoots by 1 px)."""
H, W = x.shape[-2], x.shape[-1]
pad_h = max(((H + stride - 1) // stride - 1) * stride + kernel - H, 0)
pad_w = max(((W + stride - 1) // stride - 1) * stride + kernel - W, 0)
if pad_h == 0 and pad_w == 0:
return x
return F.pad(x, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2))
# BlazeFace short-range: stem 5x5/s2 → 16 BlazeBlocks → parallel heads at
# 16²x88 (2 anchors/cell) and 8²x96 (6/cell) = 896 anchors. (in, out, stride):
_BLAZEFACE_BLOCKS = [
(24, 24, 1), (24, 28, 1), (28, 32, 2), (32, 36, 1),
(36, 42, 1), (42, 48, 2), (48, 56, 1), (56, 64, 1),
(64, 72, 1), (72, 80, 1), (80, 88, 1), (88, 96, 2),
(96, 96, 1), (96, 96, 1), (96, 96, 1), (96, 96, 1),
]
class BlazeFaceBlock(nn.Module):
"""DW 3x3 + PW + residual. Residual max-pools on stride>1, channel-pads on out_ch>in_ch."""
def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride
self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, device=device, dtype=dtype)
self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, device=device, dtype=dtype)
def forward(self, x: Tensor) -> Tensor:
residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x
if self.out_ch > self.in_ch:
residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))
x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))
return F.relu(self.pointwise(self.depthwise(x)) + residual)
class BlazeFace(nn.Module):
"""Short-range BlazeFace: (B, 3, 128, 128) in [-1, 1] → 896 anchors x 17."""
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
self.stem = ops.Conv2d(3, 24, 5, stride=2, padding=0, bias=True, **kw)
self.blocks = nn.ModuleList(BlazeFaceBlock(i, o, s, device=device, dtype=dtype, operations=operations)
for (i, o, s) in _BLAZEFACE_BLOCKS)
# 16²x2 + 8²x6 = 512 + 384 = 896 anchors.
self.cls_16 = ops.Conv2d(88, 2, 1, padding=0, bias=True, **kw)
self.cls_8 = ops.Conv2d(96, 6, 1, padding=0, bias=True, **kw)
self.reg_16 = ops.Conv2d(88, 32, 1, padding=0, bias=True, **kw)
self.reg_8 = ops.Conv2d(96, 96, 1, padding=0, bias=True, **kw)
def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
x = F.relu(self.stem(_tf_same_pad(image_chw_normalized, 5, 2)))
# 16x16 tap is block-10 output (before the 88→96 stride-2 in block 11).
for i in range(11):
x = self.blocks[i](x)
feat_16 = x
for i in range(11, 16):
x = self.blocks[i](x)
feat_8 = x
def flat(t, a, k): # NHWC flatten → (B, H*W*A, K)
B, _, H, W = t.shape
return t.permute(0, 2, 3, 1).reshape(B, H * W * a, k)
cls = torch.cat([flat(self.cls_16(feat_16), 2, 1), flat(self.cls_8(feat_8), 6, 1)], dim=1)
reg = torch.cat([flat(self.reg_16(feat_16), 2, 16), flat(self.reg_8(feat_8), 6, 16)], dim=1)
return reg, cls
# BlazeFace full-range (face_detection_full_range_sparse.tflite): MobileNetV2-ish
# backbone + top-down FPN, 192² input → 2304 anchors at the 48x48 grid.
class FRBlock(nn.Module):
"""Double inverted residual: DW → PW(mid) → DW → PW(out) [+ residual].
Per source tflite: dw* have no fused activation, pw1 is always ReLU, pw2
is ReLU only when no residual (else ReLU fuses into the ADD).
"""
def __init__(self, in_ch: int, mid_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
self.has_residual = (in_ch == out_ch and stride == 1)
self.dw1 = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)
self.pw1 = ops.Conv2d(in_ch, mid_ch, 1, padding=0, bias=True, **kw)
self.dw2 = ops.Conv2d(mid_ch, mid_ch, 3, stride=1, padding=0, groups=mid_ch, bias=True, **kw)
self.pw2 = ops.Conv2d(mid_ch, out_ch, 1, padding=0, bias=True, **kw)
def forward(self, x: Tensor) -> Tensor:
residual = x if self.has_residual else None
x = F.relu(self.pw1(self.dw1(F.pad(x, (1, 1, 1, 1)))))
x = self.pw2(self.dw2(F.pad(x, (1, 1, 1, 1))))
return F.relu(x + residual) if residual is not None else F.relu(x)
# (in_ch, mid_ch, out_ch, stride). Stages downsample 96²x32 → 48²x64 → 24²x128
# → 12²x192 → 6²x384. Lateral taps at indices 4, 7, 10 (see _FR_LATERAL_*).
_FR_BACKBONE_BLOCKS = [
(32, 8, 32, 1), (32, 8, 32, 1), # 96²x32
(32, 16, 64, 2), (64, 16, 64, 1), (64, 16, 64, 1), # 48²x64 — tap[0]
(64, 32, 128, 2), (128, 32, 128, 1), (128, 32, 128, 1), # 24²x128 — tap[1]
(128, 48, 192, 2), (192, 48, 192, 1), (192, 48, 192, 1), # 12²x192 — tap[2]
(192, 96, 384, 2), (384, 96, 384, 1), (384, 96, 384, 1), (384, 96, 384, 1), # 6²x384
]
_FR_LATERAL_TAP_INDICES = (4, 7, 10)
_FR_LATERAL_CHANNELS = ((64, 48), (128, 64), (192, 96)) # (in, out) per side-conv
# Decoder blocks per FPN level (after upsample-and-merge with the lateral).
_FR_DECODER_BLOCKS = [
[(96, 48, 96, 1), (96, 48, 96, 1)], # 12²x96
[(64, 32, 64, 1), (64, 32, 64, 1)], # 24²x64
[(48, 24, 48, 1)], # 48²x48 — feeds the heads
]
def _dcr_depth_to_space(t: Tensor, r: int, c_out: int) -> Tensor:
"""TF DEPTH_TO_SPACE in DCR layout (input channels = (i, j, c_out)).
pixel_shuffle uses CRD which permutes output channels for c_out > 1."""
B_, _, H_, W_ = t.shape
t = t.reshape(B_, r, r, c_out, H_, W_)
t = t.permute(0, 3, 4, 1, 5, 2).contiguous()
return t.reshape(B_, c_out, H_ * r, W_ * r)
class BlazeFaceFullRange(nn.Module):
"""Full-range face detector: (B, 3, 192, 192) in [-1, 1] → 2304 anchors x 17 values."""
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
mk_block = lambda i, m, o, s: FRBlock(i, m, o, s, device=device, dtype=dtype, operations=operations)
self.stem = ops.Conv2d(3, 32, 3, stride=2, padding=0, bias=True, **kw)
self.backbone = nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in _FR_BACKBONE_BLOCKS)
self.lateral_convs = nn.ModuleList(ops.Conv2d(i, o, 1, padding=0, bias=True, **kw) for (i, o) in _FR_LATERAL_CHANNELS)
self.top_conv = ops.Conv2d(384, 96, 1, padding=0, bias=True, **kw)
self.decoder_levels = nn.ModuleList(
nn.ModuleList(mk_block(i, m, o, s) for (i, m, o, s) in lvl) for lvl in _FR_DECODER_BLOCKS
)
# 96→64 before 12→24, 64→48 before 24→48.
self.decoder_reduce_convs = nn.ModuleList([
ops.Conv2d(96, 64, 1, padding=0, bias=True, **kw),
ops.Conv2d(64, 48, 1, padding=0, bias=True, **kw),
])
# Heads mix 2x2-cell info via DW-stride-2 + depth_to_space block_size=2.
self.cls_conv = ops.Conv2d(48, 4, 1, padding=0, bias=True, **kw)
self.cls_dw = ops.Conv2d(4, 4, 3, stride=2, padding=0, groups=4, bias=True, **kw)
self.reg_conv = ops.Conv2d(48, 64, 1, padding=0, bias=True, **kw)
self.reg_dw = ops.Conv2d(64, 64, 3, stride=2, padding=0, groups=64, bias=True, **kw)
def forward(self, image_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
# Symmetric pad-1 throughout (full-range tflite uses explicit TF PAD, not SAME).
x = F.relu(self.stem(F.pad(image_chw_normalized, (1, 1, 1, 1))))
tap_set = set(_FR_LATERAL_TAP_INDICES)
laterals: list[Tensor] = []
for i, blk in enumerate(self.backbone):
x = blk(x)
if i in tap_set:
laterals.append(x)
# top_conv / lateral_convs / decoder_reduce_convs all have fused ReLU in the tflite.
p = F.relu(self.top_conv(x))
laterals_rev = list(reversed(laterals))
lateral_convs_rev = list(reversed(self.lateral_convs))
for level in range(len(self.decoder_levels)):
lateral = laterals_rev[level]
p = F.interpolate(p, size=lateral.shape[-2:], mode="bilinear", align_corners=False)
p = p + F.relu(lateral_convs_rev[level](lateral))
for blk in self.decoder_levels[level]:
p = blk(p)
if level < len(self.decoder_reduce_convs):
p = F.relu(self.decoder_reduce_convs[level](p))
c = self.cls_dw(F.pad(self.cls_conv(p), (1, 1, 1, 1)))
c = _dcr_depth_to_space(c, r=2, c_out=1)
r = self.reg_dw(F.pad(self.reg_conv(p), (1, 1, 1, 1)))
r = _dcr_depth_to_space(r, r=2, c_out=16)
B = c.shape[0]
cls_out = c.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 1)
reg_out = r.permute(0, 2, 3, 1).reshape(B, _BF_FR_NUM_ANCHORS, 16)
return reg_out, cls_out
@lru_cache(maxsize=1)
def _blazeface_full_range_anchors() -> np.ndarray:
"""2304 anchors over 48x48; anchor_w=anchor_h=1 (fixed_anchor_size)."""
feat = _BF_FR_GRID
yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")
cx, cy, ones = (xx + 0.5) / feat, (yy + 0.5) / feat, np.ones_like(xx)
return np.stack([cx, cy, ones, ones], axis=-1).reshape(_BF_FR_NUM_ANCHORS, 4)
def _decode_blazeface_full_range(regressors: np.ndarray, classificators: np.ndarray,
score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:
"""Same decode as short-range with 2304-anchor grid and box_scale=192."""
scores = expit(np.clip(classificators[:, 0], -_BF_FR_SCORE_CLIP, _BF_FR_SCORE_CLIP))
keep = scores >= score_thresh
if not keep.any():
return np.empty((0, 17), dtype=np.float32)
r = regressors[keep] / _BF_FR_BOX_SCALE
a = _blazeface_full_range_anchors()[keep]
cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys
w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs
out = np.empty((r.shape[0], 17), dtype=np.float32)
out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs
out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys
out[:, 16] = scores[keep]
return out
# FaceMesh (face_landmarks_detector.tflite): PReLU variant of BlazeBlock,
# 17 blocks, heads for 478x3 landmarks + presence.
_FACEMESH_BLOCKS = [ # (in_ch, out_ch, stride)
(16, 16, 1), (16, 16, 1), (16, 32, 2), (32, 32, 1), (32, 32, 1), (32, 64, 2),
(64, 64, 1), (64, 64, 1), (64, 128, 2), (128, 128, 1), (128, 128, 1), (128, 128, 2),
(128, 128, 1), (128, 128, 1), (128, 128, 2), (128, 128, 1), (128, 128, 1),
]
class FaceMeshBlock(nn.Module):
"""PReLU BlazeBlock: PReLU between DW and PW, and after the residual add."""
def __init__(self, in_ch: int, out_ch: int, stride: int, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
self.in_ch, self.out_ch, self.stride = in_ch, out_ch, stride
self.depthwise = ops.Conv2d(in_ch, in_ch, 3, stride=stride, padding=0, groups=in_ch, bias=True, **kw)
self.prelu_dwise = nn.PReLU(num_parameters=in_ch, **kw)
self.pointwise = ops.Conv2d(in_ch, out_ch, 1, padding=0, bias=True, **kw)
self.prelu_out = nn.PReLU(num_parameters=out_ch, **kw)
def forward(self, x: Tensor) -> Tensor:
residual = F.max_pool2d(x, 2, 2) if self.stride > 1 else x
if self.out_ch > self.in_ch:
residual = F.pad(residual, (0, 0, 0, 0, 0, self.out_ch - self.in_ch))
x = _tf_same_pad(x, 3, self.stride) if self.stride > 1 else F.pad(x, (1, 1, 1, 1))
return self.prelu_out(self.pointwise(self.prelu_dwise(self.depthwise(x))) + residual)
class FaceMesh(nn.Module):
NUM_LANDMARKS = 478
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
self.stem = ops.Conv2d(3, 16, 3, stride=2, padding=0, bias=True, **kw)
self.prelu_stem = nn.PReLU(num_parameters=16, **kw)
self.blocks = nn.ModuleList(FaceMeshBlock(i, o, s, device=device, dtype=dtype, operations=operations)
for (i, o, s) in _FACEMESH_BLOCKS)
self.head_reduce = ops.Conv2d(128, 8, 1, padding=0, bias=True, **kw)
self.prelu_head_reduce = nn.PReLU(num_parameters=8, **kw)
self.head_block = FaceMeshBlock(8, 8, 1, device=device, dtype=dtype, operations=operations)
self.head_presence = ops.Conv2d(8, 1, 3, padding=0, bias=True, **kw)
self.head_landmarks = ops.Conv2d(8, self.NUM_LANDMARKS * 3, 3, padding=0, bias=True, **kw)
def forward(self, face_chw_normalized: Tensor) -> tuple[Tensor, Tensor]:
"""(B, 3, 192, 192) in [0, 1] → ((B, 478, 3) landmarks in 192-canonical, (B,) presence)."""
x = self.prelu_stem(self.stem(_tf_same_pad(face_chw_normalized, 3, 2)))
for blk in self.blocks:
x = blk(x)
x = self.prelu_head_reduce(self.head_reduce(x))
x = self.head_block(x)
B = x.shape[0]
presence = self.head_presence(x).reshape(B)
lmks = self.head_landmarks(x).reshape(B, self.NUM_LANDMARKS, 3)
return lmks, presence
# FaceBlendshapes (MLP-Mixer "GhumMarkerPoserMlpMixerGeneral"):
# 146x2 → token-reduce 146→96 → embed 2→64 → +cls token → 4x mixer → cls→52.
_BS_NUM_INPUT_LANDMARKS = 146
_BS_NUM_TOKENS_REDUCED = 96
_BS_NUM_TOKENS = 97 # +1 cls
_BS_TOKEN_DIM = 64
_BS_TOKEN_MIX_HIDDEN = 384
_BS_CHANNEL_MIX_HIDDEN = 256
_BS_NUM_BLENDSHAPES = 52
_BS_LN_EPS = 1e-6
class MlpMixerBlock(nn.Module):
"""MLP-Mixer block: token-mixing MLP (over tokens) → channel-mixing MLP (over dim).
Both pre-LN, both residual. LN has no beta (bias=False) to match MP."""
def __init__(self, num_tokens: int, token_dim: int, token_hidden: int, channel_hidden: int,
device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
# bias=False → no LN beta (matches MP).
self.ln1 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)
self.ln2 = ops.LayerNorm(token_dim, eps=_BS_LN_EPS, bias=False, **kw)
self.token_mlp1 = ops.Linear(num_tokens, token_hidden, bias=True, **kw)
self.token_mlp2 = ops.Linear(token_hidden, num_tokens, bias=True, **kw)
self.channel_mlp1 = ops.Linear(token_dim, channel_hidden, bias=True, **kw)
self.channel_mlp2 = ops.Linear(channel_hidden, token_dim, bias=True, **kw)
def forward(self, x: Tensor) -> Tensor:
y = self.ln1(x).transpose(1, 2)
x = x + self.token_mlp2(F.relu(self.token_mlp1(y))).transpose(1, 2)
return x + self.channel_mlp2(F.relu(self.channel_mlp1(self.ln2(x))))
class FaceBlendshapes(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
ops = operations if operations is not None else nn
kw = dict(device=device, dtype=dtype)
self.token_reduce = ops.Linear(_BS_NUM_INPUT_LANDMARKS, _BS_NUM_TOKENS_REDUCED, bias=True, **kw)
self.token_embed = ops.Linear(2, _BS_TOKEN_DIM, bias=True, **kw)
self.cls_token = nn.Parameter(torch.zeros(1, 1, _BS_TOKEN_DIM, **kw))
self.blocks = nn.ModuleList(
MlpMixerBlock(_BS_NUM_TOKENS, _BS_TOKEN_DIM, _BS_TOKEN_MIX_HIDDEN, _BS_CHANNEL_MIX_HIDDEN,
device=device, dtype=dtype, operations=operations) for _ in range(4)
)
self.head = ops.Linear(_BS_TOKEN_DIM, _BS_NUM_BLENDSHAPES, bias=True, **kw)
@staticmethod
def _input_normalize(landmarks_2d: Tensor) -> Tensor:
# Centroid-subtract → L2 scale → x0.5. The 0.5 is baked into training.
centroid = landmarks_2d.mean(dim=1, keepdim=True)
x = landmarks_2d - centroid
mag = torch.sqrt((x * x).sum(dim=-1, keepdim=True))
scale = mag.mean(dim=1, keepdim=True)
return (x / scale.clamp(min=1e-12)) * 0.5
def forward(self, landmarks_2d: Tensor) -> Tensor:
"""(B, 146, 2) → (B, 52) in [0, 1]. Input units don't matter (centroid + L2 normalize)."""
x = self._input_normalize(landmarks_2d)
x = self.token_reduce(x.transpose(1, 2)).transpose(1, 2)
x = self.token_embed(x)
cls = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat([cls, x], dim=1)
for blk in self.blocks:
x = blk(x)
return torch.sigmoid(self.head(x[:, 0]))
@lru_cache(maxsize=1)
def _blazeface_anchors() -> np.ndarray:
"""896 anchors per SsdAnchorsCalculator (fixed_anchor_size → anchor_w=anchor_h=1)."""
per_ar = len(_BF_ASPECT_RATIOS) + (1 if _BF_INTERP_SCALE_AR > 0 else 0)
layer_anchors: List[np.ndarray] = []
layer = 0
while layer < _BF_NUM_LAYERS:
stride = _BF_STRIDES[layer]
last = layer
while last < _BF_NUM_LAYERS and _BF_STRIDES[last] == stride:
last += 1
per_cell = per_ar * (last - layer)
feat = (_BF_INPUT_SIZE + stride - 1) // stride
yy, xx = np.meshgrid(np.arange(feat, dtype=np.float32), np.arange(feat, dtype=np.float32), indexing="ij")
cx, cy, ones = (xx + _BF_ANCHOR_OFFSET_X) / feat, (yy + _BF_ANCHOR_OFFSET_Y) / feat, np.ones_like(xx)
cell = np.stack([cx, cy, ones, ones], axis=-1).reshape(-1, 4)
layer_anchors.append(np.repeat(cell, per_cell, axis=0))
layer = last
out = np.concatenate(layer_anchors, axis=0)
assert out.shape == (896, 4), out.shape
return out
def _decode_blazeface(regressors: np.ndarray, classificators: np.ndarray,
score_thresh: float = _BF_MIN_SCORE) -> np.ndarray:
"""Decode (regs (896,16), cls (896,1)) → (N, 17) = [xyxy, kp0x..kp5y, score] in [0, 1]."""
scores = expit(np.clip(classificators[:, 0], -_BF_SCORE_CLIP, _BF_SCORE_CLIP))
keep = scores >= score_thresh
if not keep.any():
return np.empty((0, 17), dtype=np.float32)
r = regressors[keep] / _BF_BOX_SCALE
a = _blazeface_anchors()[keep] # (N, 4) cx, cy, 1, 1
cxs, cys, aws, ahs = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
xc, yc = r[:, 0:1] * aws + cxs, r[:, 1:2] * ahs + cys
w, h = r[:, 2:3] * aws, r[:, 3:4] * ahs
out = np.empty((r.shape[0], 17), dtype=np.float32)
out[:, 0:1], out[:, 1:2], out[:, 2:3], out[:, 3:4] = xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2
out[:, 4:16:2] = r[:, _BF_KP_OFFSET::2] * aws + cxs
out[:, 5:16:2] = r[:, _BF_KP_OFFSET + 1::2] * ahs + cys
out[:, 16] = scores[keep]
return out
def _weighted_nms(detections: np.ndarray, iou_thresh: float = 0.5) -> np.ndarray:
"""MP weighted NMS — kept boxes are score-weighted averages of overlapping detections."""
if detections.shape[0] == 0:
return detections
dets = detections[np.argsort(-detections[:, 16])]
N = dets.shape[0]
areas = np.clip(dets[:, 2] - dets[:, 0], 0, None) * np.clip(dets[:, 3] - dets[:, 1], 0, None)
kept: List[np.ndarray] = []
used = np.zeros(N, dtype=bool)
for i in range(N):
if used[i]:
continue
ax1, ay1, ax2, ay2 = dets[i, 0:4]
merge_idx = [i]
for j in range(i + 1, N):
if used[j]:
continue
bx1, by1, bx2, by2 = dets[j, 0:4]
iw = max(0.0, min(ax2, bx2) - max(ax1, bx1))
ih = max(0.0, min(ay2, by2) - max(ay1, by1))
inter = iw * ih
union = areas[i] + areas[j] - inter
if union > 0 and inter / union > iou_thresh: # strict > matches MP
merge_idx.append(j)
used[j] = True
used[i] = True
cluster = dets[merge_idx]
ws = cluster[:, 16:17]
ws_sum = ws.sum()
merged = np.copy(cluster[0])
if ws_sum > 0:
merged[:16] = (cluster[:, :16] * ws).sum(axis=0) / ws_sum
kept.append(merged)
return np.stack(kept, axis=0) if kept else np.empty((0, 17), dtype=np.float32)
def _detection_to_face_rect(detection: np.ndarray, image_w: int, image_h: int) -> Tuple[float, float, float, float, float]:
"""Detection (normalized) → rotated 1.5xbbox ROI in image pixels (anisotropic)."""
xmin, ymin, xmax, ymax = detection[0:4]
lx = detection[4 + _FACE_LEFT_EYE_KP * 2 + 0] * image_w
ly = detection[4 + _FACE_LEFT_EYE_KP * 2 + 1] * image_h
rx = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 0] * image_w
ry = detection[4 + _FACE_RIGHT_EYE_KP * 2 + 1] * image_h
# Image-y-down convention: angle = target - atan2(-dy, dx).
angle = _FACE_ROI_TARGET_ANGLE - math.atan2(ly - ry, rx - lx)
return (float((xmin + xmax) * 0.5 * image_w),
float((ymin + ymax) * 0.5 * image_h),
float((xmax - xmin) * image_w * _FACE_ROI_SCALE_X),
float((ymax - ymin) * image_h * _FACE_ROI_SCALE_Y),
float(angle))
def _sample_warp(image_chw: Tensor, src_x: Tensor, src_y: Tensor, padding_mode: str) -> Tensor:
"""Bilinear-sample image_chw at corner-aligned (src_x, src_y)."""
H, W = int(image_chw.shape[-2]), int(image_chw.shape[-1])
grid = torch.stack([(2.0 * src_x + 1.0) / W - 1.0,
(2.0 * src_y + 1.0) / H - 1.0], dim=-1).unsqueeze(0)
return F.grid_sample(image_chw.unsqueeze(0), grid, mode="bilinear",
align_corners=False, padding_mode=padding_mode).squeeze(0)
def _warp_face_crop(image_chw: Tensor, cx: float, cy: float, width: float, height: float,
angle: float, output_size: int = _FM_INPUT_SIZE) -> Tensor:
"""Rotated rect → output_size² with BORDER_REPLICATE. image_chw must be in [0, 1]."""
s_x, s_y = width / output_size, height / output_size
cos_a, sin_a = math.cos(angle), math.sin(angle)
arange = torch.arange(output_size, dtype=image_chw.dtype, device=image_chw.device) - output_size * 0.5
v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")
src_x = cx + u_grid * s_x * cos_a - v_grid * s_y * sin_a
src_y = cy + u_grid * s_x * sin_a + v_grid * s_y * cos_a
return _sample_warp(image_chw, src_x, src_y, "border")
def _blazeface_input_warp(image_chw_raw: Tensor, target: int = _BF_INPUT_SIZE) -> Tuple[Tensor, float, float, float]:
"""Centered max(W,H) square → target² with BORDER_ZERO + [-1, 1] norm.
Sub-pixel grid_sample matters; integer-pad-then-resize drifts the bbox ~5%.
Returns (warped, sub_rect_cx, sub_rect_cy, sub_rect_size) — the triplet maps
tensor-normalized [0,1] detections back to image pixels.
"""
H, W = int(image_chw_raw.shape[1]), int(image_chw_raw.shape[2])
sub_rect_size = float(max(W, H))
sub_rect_cx, sub_rect_cy = W * 0.5, H * 0.5
s = sub_rect_size / target
arange = torch.arange(target, dtype=image_chw_raw.dtype, device=image_chw_raw.device) - target * 0.5
v_grid, u_grid = torch.meshgrid(arange, arange, indexing="ij")
out = _sample_warp(image_chw_raw, sub_rect_cx + u_grid * s, sub_rect_cy + v_grid * s, "zeros")
return (out / 127.5) - 1.0, sub_rect_cx, sub_rect_cy, sub_rect_size
class FaceLandmarker(nn.Module):
"""BlazeFace → FaceMesh v2 → blendshapes. `detector_variant` selects 'short'
(128², ≤2m) or 'full' (192² FPN, ≤5m). State dict uses inner-module prefixes
`detector.*` / `mesh.*` / `blendshapes.*`; the outer FaceLandmarkerModel
wrapper rewrites `detector_{variant}.*` keys to `detector.*` before loading.
"""
def __init__(self, device=None, dtype=None, operations=None, detector_variant: str = "short"):
super().__init__()
det_cls = {"short": BlazeFace, "full": BlazeFaceFullRange}.get(detector_variant)
self.detector_variant = detector_variant
self.detector = det_cls(device=device, dtype=dtype, operations=operations)
self.mesh = FaceMesh(device=device, dtype=dtype, operations=operations)
self.blendshapes = FaceBlendshapes(device=device, dtype=dtype, operations=operations)
self.register_buffer("_bs_idx", torch.tensor(_BS_INPUT_INDICES, dtype=torch.long), persistent=False)
def run_detector_batch(self, images_rgb_uint8: List[np.ndarray],
score_thresh: float = _BF_MIN_SCORE,
iou_thresh: float = 0.5):
"""Batched detector pass. Returns (img_raws, sub_rects, sizes, per_frame_decoded)
where per_frame_decoded[b] is (N, 17) in tensor-normalized [0,1] coords."""
if not images_rgb_uint8:
return [], [], [], []
device, dtype = self.detector.stem.weight.device, self.detector.stem.weight.dtype
det_input_size, decode_fn = ((_BF_FR_INPUT_SIZE, _decode_blazeface_full_range)
if self.detector_variant == "full"
else (_BF_INPUT_SIZE, _decode_blazeface))
# Same-size frames: stack once and transfer once. Variable size falls back
# to per-image (only triggers for SAM3DBody's head crops).
sizes = [tuple(img.shape[:2]) for img in images_rgb_uint8]
if len(set(sizes)) == 1:
batch_chw = torch.from_numpy(np.stack(images_rgb_uint8, axis=0)).to(device, dtype).movedim(-1, -3).contiguous()
img_raws = [batch_chw[bi] for bi in range(batch_chw.shape[0])]
else:
img_raws = [torch.from_numpy(img).to(device, dtype).movedim(-1, -3).contiguous() for img in images_rgb_uint8]
warps = [_blazeface_input_warp(img_raw, det_input_size) for img_raw in img_raws]
det_crops = [w[0] for w in warps]
sub_rects = [(w[1], w[2], w[3]) for w in warps]
regs_b, cls_b = self.detector(torch.stack(det_crops, dim=0))
regs_np, cls_np = regs_b.float().cpu().numpy(), cls_b.float().cpu().numpy()
per_frame = []
for b in range(len(images_rgb_uint8)):
decoded = decode_fn(regs_np[b], cls_np[b], score_thresh=score_thresh)
per_frame.append(_weighted_nms(decoded, iou_thresh=iou_thresh) if decoded.shape[0] > 0 else decoded)
return img_raws, sub_rects, sizes, per_frame
def detect_batch(self, images_rgb_uint8: List[np.ndarray], num_faces: int = 1,
score_thresh: float = _BF_MIN_SCORE) -> List[List[dict]]:
"""Full pipeline batched across `images_rgb_uint8`. Returns one face-dict
list per image (empty if nothing detected). Face dict:
bbox_xyxy (4,) image pixels, blendshapes {52} ∈ [0,1],
landmarks_xy (478, 2) image pixels, landmarks_3d (478, 3) in
192-canonical (pre-transformation) units, presence float (raw logit).
"""
img_raws, sub_rects, sizes, per_frame_dets = self.run_detector_batch(
images_rgb_uint8, score_thresh=score_thresh,
)
# tensor-normalized → image-normalized [0,1] for _detection_to_face_rect.
for b, decoded in enumerate(per_frame_dets):
if decoded.shape[0] == 0:
continue
cx, cy, size = sub_rects[b]
H, W = sizes[b]
sx0, sy0 = cx - size * 0.5, cy - size * 0.5
decoded[:, 0:16:2] = (sx0 + size * decoded[:, 0:16:2]) / W
decoded[:, 1:16:2] = (sy0 + size * decoded[:, 1:16:2]) / H
if num_faces > 0:
per_frame_dets[b] = decoded[: int(num_faces)]
# Collect every detected face across all frames into one mesh input.
face_params: List[Tuple[int, float, float, float, float, float, float]] = []
mesh_crops: List[Tensor] = []
for b, dets in enumerate(per_frame_dets):
if dets.shape[0] == 0:
continue
H, W = sizes[b]
img_for_mesh = img_raws[b] / 255.0
for det in dets:
cx, cy, w, h, angle = _detection_to_face_rect(det, W, H)
mesh_crops.append(_warp_face_crop(img_for_mesh, cx, cy, w, h, angle, _FM_INPUT_SIZE))
face_params.append((b, float(det[16]), cx, cy, w, h, angle))
results: List[List[dict]] = [[] for _ in range(len(images_rgb_uint8))]
if not mesh_crops:
return results
lmks_canon_b, presence_b = self.mesh(torch.stack(mesh_crops, dim=0))
bs_out_b = self.blendshapes(lmks_canon_b[:, self._bs_idx, :2])
# Batched canonical→image affine
params_t = torch.tensor(
[(cx, cy, w, h, math.cos(a), math.sin(a)) for (_b, _s, cx, cy, w, h, a) in face_params],
device=lmks_canon_b.device, dtype=lmks_canon_b.dtype,
)
cxs, cys, ws, hs, cos_a, sin_a = params_t.unbind(dim=1)
inv = 1.0 / _FM_INPUT_SIZE
u = lmks_canon_b[..., 0] - _FM_INPUT_SIZE * 0.5
v = lmks_canon_b[..., 1] - _FM_INPUT_SIZE * 0.5
lmks_xy_t = torch.stack([
cxs[:, None] + u * (ws * inv * cos_a)[:, None] - v * (hs * inv * sin_a)[:, None],
cys[:, None] + u * (ws * inv * sin_a)[:, None] + v * (hs * inv * cos_a)[:, None],
], dim=-1)
lmks_xy_np = lmks_xy_t.float().cpu().numpy()
lmks_canon_np = lmks_canon_b.float().cpu().numpy()
presence_np = presence_b.float().cpu().numpy()
bs_np = bs_out_b.float().cpu().numpy()
for i, (b, score, *_) in enumerate(face_params):
lmks_xy = lmks_xy_np[i]
mn, mx = lmks_xy.min(0), lmks_xy.max(0)
results[b].append({
"bbox_xyxy": np.array([mn[0], mn[1], mx[0], mx[1]], dtype=np.float32),
"blendshapes": dict(zip(BLENDSHAPE_NAMES, bs_np[i].tolist())),
"landmarks_xy": lmks_xy,
"landmarks_3d": lmks_canon_np[i],
"presence": float(presence_np[i]),
"score": score,
})
return results

View File

@ -1,502 +0,0 @@
"""ComfyUI nodes for the pure-PyTorch MediaPipe Face Landmarker port.
Custom IO types:
FACE_LANDMARKER — FaceLandmarkerModel wrapper (ModelPatcher inside)
FACE_LANDMARKS — {"frames": List[List[face_dict]], "image_size": (H, W),
"connection_sets": dict[str, frozenset[(int, int)]]}
face_dict: bbox_xyxy, blendshapes, landmarks_xy,
landmarks_3d, presence, score, transformation_matrix
MediaPipeFaceLandmarker also emits the core BOUNDING_BOX type — pair with DrawBBoxes.
"""
from __future__ import annotations
import numpy as np
import torch
from PIL import Image, ImageColor, ImageDraw
from tqdm.auto import tqdm
from typing_extensions import override
import comfy.model_management
import comfy.model_patcher
import comfy.utils
import folder_paths
from comfy_api.latest import ComfyExtension, io
from comfy_extras.mediapipe.face_landmarker import FaceLandmarker
from comfy_extras.mediapipe.face_geometry import transformation_matrix_from_detection
FaceLandmarkerType = io.Custom("FACE_LANDMARKER")
FaceLandmarksType = io.Custom("FACE_LANDMARKS")
_CANONICAL_KEYS = ("canonical_vertices", "procrustes_indices", "procrustes_weights")
_CONTOUR_PARTS = ("face_oval", "left_eye", "right_eye", "left_eyebrow", "right_eyebrow", "lips")
class FaceLandmarkerModel:
"""Loaded FaceLandmarker variants + ModelPatcher per variant.
Safetensors layout: `detector_short.*` / `detector_full.*` plus shared
`mesh.*`, `blendshapes.*`, `canonical_*`, and `topology.*`.
PReLU forces plain-nn / fp32 (manual_cast strands buffers across devices).
"""
def __init__(self, state_dict: dict):
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = torch.float32
# FACEMESH_* connection sets, embedded as int32 (N, 2) under topology.*.
base: dict[str, frozenset] = {}
for k in [k for k in state_dict if k.startswith("topology.")]:
base[k[len("topology."):]] = frozenset(map(tuple, state_dict.pop(k).tolist()))
base["contours"] = frozenset().union(*(base[p] for p in _CONTOUR_PARTS))
base["all"] = base["contours"] | base["irises"] | base["nose"]
self.connection_sets: dict[str, frozenset] = base
self.canonical_data: dict[str, np.ndarray] = {k: state_dict.pop(k).numpy() for k in _CANONICAL_KEYS}
shared = {k: v for k, v in state_dict.items() if k.startswith(("mesh.", "blendshapes."))}
self.models: dict[str, FaceLandmarker] = {}
self.patchers: dict[str, comfy.model_patcher.ModelPatcher] = {}
for variant in ("short", "full"):
prefix = f"detector_{variant}."
sub = dict(shared)
sub.update({f"detector.{k[len(prefix):]}": v for k, v in state_dict.items() if k.startswith(prefix)})
fl = FaceLandmarker(device=offload_device, dtype=self.dtype, operations=None, detector_variant=variant).eval()
fl.load_state_dict(sub, strict=False)
self.models[variant] = fl
self.patchers[variant] = comfy.model_patcher.CoreModelPatcher(
fl, load_device=self.load_device, offload_device=offload_device,
size=comfy.model_management.module_size(fl),
)
def detect_batch(self, images, num_faces: int, score_thresh: float, variant: str):
comfy.model_management.load_model_gpu(self.patchers[variant])
return self.models[variant].detect_batch(images, num_faces=num_faces, score_thresh=score_thresh)
def _image_to_uint8(image: torch.Tensor) -> np.ndarray:
return image[..., :3].mul(255.0).add_(0.5).clamp_(0, 255).to(torch.uint8).cpu().numpy()
def _parse_color(color: str) -> tuple[int, int, int]:
try:
return ImageColor.getrgb(color)[:3]
except ValueError:
return (0, 255, 0)
def _copy_face(face: dict) -> dict:
"""Shallow copy of a face_dict with array-fields cloned so callers can mutate."""
return {
"bbox_xyxy": face["bbox_xyxy"].copy(),
"blendshapes": dict(face["blendshapes"]),
"landmarks_xy": face["landmarks_xy"].copy(),
"landmarks_3d": face["landmarks_3d"].copy(),
"presence": face["presence"],
"score": face["score"],
}
def _lerp_face(a: dict, b: dict, t: float) -> dict:
return {
"bbox_xyxy": (1 - t) * a["bbox_xyxy"] + t * b["bbox_xyxy"],
"blendshapes": {k: (1 - t) * a["blendshapes"][k] + t * b["blendshapes"][k] for k in a["blendshapes"]},
"landmarks_xy": (1 - t) * a["landmarks_xy"] + t * b["landmarks_xy"],
"landmarks_3d": (1 - t) * a["landmarks_3d"] + t * b["landmarks_3d"],
"presence": (1 - t) * a["presence"] + t * b["presence"],
"score": (1 - t) * a["score"] + t * b["score"],
}
def _match_faces(a: list[dict], b: list[dict]) -> list[tuple[int, int]]:
"""Greedy nearest-neighbour pairing of faces between two frames by bbox
centre distance. Unmatched (when counts differ) are dropped."""
if not a or not b:
return []
centers_a = np.array([(0.5 * (f["bbox_xyxy"][0] + f["bbox_xyxy"][2]),
0.5 * (f["bbox_xyxy"][1] + f["bbox_xyxy"][3])) for f in a])
centers_b = np.array([(0.5 * (f["bbox_xyxy"][0] + f["bbox_xyxy"][2]),
0.5 * (f["bbox_xyxy"][1] + f["bbox_xyxy"][3])) for f in b])
dists = np.linalg.norm(centers_a[:, None] - centers_b[None], axis=-1)
pairs: list[tuple[int, int]] = []
used_a: set[int] = set()
used_b: set[int] = set()
candidates = sorted((dists[ia, ib], ia, ib) for ia in range(len(a)) for ib in range(len(b)))
for _, ia, ib in candidates:
if ia in used_a or ib in used_b:
continue
pairs.append((ia, ib))
used_a.add(ia)
used_b.add(ib)
return pairs
def _fill_missing_frames(frames: list[list[dict]], mode: str) -> None:
"""In-place fill empty frame slots from neighbouring detections. Multi-face
aware: pairs faces across bracketing frames by greedy bbox-centre NN.
When counts differ, unmatched faces are dropped from the synthesised frame."""
if mode == "empty":
return
valid = [i for i, fr in enumerate(frames) if fr]
if not valid:
return # nothing to fill from
if mode == "previous":
last: list[dict] = []
for i, fr in enumerate(frames):
if fr:
last = fr
elif last:
frames[i] = [_copy_face(f) for f in last]
return
# interpolate: lerp between bracketing valid frames; clamp at ends.
for i in range(len(frames)):
if frames[i]:
continue
prev_i = max((v for v in valid if v < i), default=None)
next_i = min((v for v in valid if v > i), default=None)
if prev_i is None:
frames[i] = [_copy_face(f) for f in frames[next_i]]
elif next_i is None:
frames[i] = [_copy_face(f) for f in frames[prev_i]]
else:
t = (i - prev_i) / (next_i - prev_i)
pairs = _match_faces(frames[prev_i], frames[next_i])
frames[i] = [_lerp_face(frames[prev_i][a], frames[next_i][b], t) for a, b in pairs]
def _ordered_rings(edges: frozenset[tuple[int, int]]) -> list[list[int]]:
"""Walk an unordered edge set into one or more closed-loop vertex rings
(handles multi-loop sets like FACEMESH_LIPS: outer + inner)."""
adj: dict[int, set[int]] = {}
for a, b in edges:
adj.setdefault(a, set()).add(b)
adj.setdefault(b, set()).add(a)
visited: set[int] = set()
rings: list[list[int]] = []
for start in adj:
if start in visited:
continue
ring = [start]
visited.add(start)
prev, cur = -1, start
while True:
nxt = next((v for v in adj[cur] if v != prev), None)
if nxt is None or nxt == start:
break
ring.append(nxt)
visited.add(nxt)
prev, cur = cur, nxt
rings.append(ring)
return rings
class LoadMediaPipeFaceLandmarker(io.ComfyNode):
"""Load MediaPipe Face Landmarker v2 weights. Contains both detector variants
(short / full), shared mesh, blendshapes, and canonical geometry."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="LoadMediaPipeFaceLandmarker",
display_name="Load MediaPipe Face Landmarker",
category="loaders",
inputs=[
io.Combo.Input("model_name", options=folder_paths.get_filename_list("mediapipe"),
tooltip="Face Landmarker safetensors from models/mediapipe/."),
],
outputs=[FaceLandmarkerType.Output()],
)
@classmethod
def execute(cls, model_name) -> io.NodeOutput:
sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("mediapipe", model_name), safe_load=True)
wrapper = FaceLandmarkerModel(sd)
return io.NodeOutput(wrapper)
# Per-frame fallback modes for detection failures in a batch.
_FALLBACK_MODES = ("empty", "previous", "interpolate")
class MediaPipeFaceLandmarker(io.ComfyNode):
"""BlazeFace → FaceMesh v2 → ARKit-52 blendshapes, batched across the
input. Also emits a BOUNDING_BOX list (landmark-extent bbox per face) —
pair with DrawBBoxes for detector-only viz or MediaPipeFaceMeshVisualize
for the mesh overlay."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MediaPipeFaceLandmarker",
display_name="MediaPipe Face Landmarker",
category="image/detection",
inputs=[
FaceLandmarkerType.Input("face_landmarker"),
io.Image.Input("image"),
io.Combo.Input("detector_variant", options=["short", "full", "both"], default="short",
tooltip="Face detector range. 'short' is tuned for close-up faces "
"(within ~2 m of the camera); 'full' covers farther / smaller "
"faces (up to ~5 m) but is slower. 'both' runs both detectors and "
"keeps whichever found more faces per frame (~2× detection cost)."),
io.Int.Input("num_faces", default=1, min=0, max=16, step=1,
tooltip="Maximum faces to return per frame. 0 = no cap (return all detected)."),
io.Float.Input("min_confidence", default=0.5, min=0.0, max=1.0, step=0.01, advanced=True,
tooltip="BlazeFace score threshold. Lower to catch small/occluded faces."),
io.Combo.Input("missing_frame_fallback", options=list(_FALLBACK_MODES), default="empty", advanced=True,
tooltip="Per-frame behaviour when detection fails in a batch. "
"'empty' leaves the frame faceless. 'previous' copies the most recent successful "
"detection. 'interpolate' lerps landmarks/bbox/blendshapes between bracketing "
"successful frames. Multi-face: pairs faces across frames by greedy bbox-centre NN."),
],
outputs=[
FaceLandmarksType.Output(display_name="face_landmarks"),
io.BoundingBox.Output("bboxes"),
],
)
@classmethod
def execute(cls, face_landmarker, image, detector_variant, num_faces, min_confidence,
missing_frame_fallback) -> io.NodeOutput:
canonical = face_landmarker.canonical_data
img_np = _image_to_uint8(image)
B, H, W = img_np.shape[:3]
chunk = 16
is_both = detector_variant == "both"
total_work = 2 * B if is_both else B
pbar = comfy.utils.ProgressBar(total_work)
def _run(variant: str) -> list[list[dict]]:
res: list[list[dict]] = []
with tqdm(total=B, desc=f"MediaPipe Face Landmarker ({variant})") as tq:
for i in range(0, B, chunk):
end = min(i + chunk, B)
res.extend(face_landmarker.detect_batch(
[img_np[bi] for bi in range(i, end)],
num_faces=int(num_faces),
score_thresh=float(min_confidence),
variant=variant,
))
pbar.update_absolute(min(pbar.current + (end - i), total_work))
tq.update(end - i)
return res
if is_both:
short_res = _run("short")
full_res = _run("full")
# Per-frame keep whichever found more faces (tie → short).
frames: list[list[dict]] = [
short_res[bi] if len(short_res[bi]) >= len(full_res[bi]) else full_res[bi]
for bi in range(B)
]
else:
frames = _run(detector_variant)
_fill_missing_frames(frames, missing_frame_fallback)
bboxes = []
for per_frame in frames:
per_bb = []
for f in per_frame:
f["transformation_matrix"] = transformation_matrix_from_detection(f, W, H, canonical)
x1, y1, x2, y2 = (float(v) for v in f["bbox_xyxy"])
per_bb.append({"x": x1, "y": y1, "width": x2 - x1, "height": y2 - y1, "label": "face", "score": float(f["score"])})
bboxes.append(per_bb)
return io.NodeOutput({"frames": frames, "image_size": (H, W),
"connection_sets": face_landmarker.connection_sets}, bboxes)
# Topology keys unioned by the 'all' connections preset (contour parts + irises + nose).
_ALL_CONNECTION_PARTS: tuple[str, ...] = (*_CONTOUR_PARTS, "irises", "nose")
_CUSTOM_FEATURES: tuple[tuple[str, bool], ...] = (
("face_oval", True),
("lips", True),
("left_eye", True),
("right_eye", True),
("left_eyebrow", True),
("right_eyebrow", True),
("irises", True),
("nose", True),
("tesselation", False),
)
class MediaPipeFaceMeshVisualize(io.ComfyNode):
"""Draw a FACEMESH_* subset over an image. Topology travels with the
FACE_LANDMARKS payload (set at detection time)."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MediaPipeFaceMeshVisualize",
display_name="MediaPipe Face Mesh Visualize",
category="image/detection",
inputs=[
FaceLandmarksType.Input("face_landmarks"),
io.Image.Input("image", optional=True, tooltip="If not connected, a black canvas will be used."),
io.DynamicCombo.Input(
"connections",
tooltip="'all' = oval+eyes+brows+lips+irises+nose. 'fill' = solid face_oval polygon (silhouette mask). 'custom' = toggle each feature individually (including 'tesselation', the full 2547-edge wireframe).",
options=[
io.DynamicCombo.Option("all", []),
io.DynamicCombo.Option("fill", []),
io.DynamicCombo.Option("custom", [
io.Boolean.Input(feat, default=default,
tooltip=f"Draw the '{feat}' connection set.")
for feat, default in _CUSTOM_FEATURES
]),
],
),
io.Color.Input("color", default="#00ff00"),
io.Int.Input("thickness", default=1, min=0, max=8, step=1,
tooltip="Edge line thickness in pixels. 0 disables edge drawing."),
io.Int.Input("point_size", default=2, min=0, max=16, step=1,
tooltip="Landmark dot radius in pixels. 0 disables point drawing."),
],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, face_landmarks, connections, color, thickness, point_size, image=None) -> io.NodeOutput:
sets = face_landmarks["connection_sets"]
sel = connections["connections"]
fill_rings: list[list[int]] | None = None
if sel == "fill":
fill_rings = _ordered_rings(sets["face_oval"])
edges = frozenset()
elif sel == "custom":
parts = [feat for feat, _ in _CUSTOM_FEATURES if connections.get(feat, False)]
edges = frozenset().union(*(sets[p] for p in parts))
else: # "all"
edges = frozenset().union(*(sets[p] for p in _ALL_CONNECTION_PARTS))
rgb, thick, psize = _parse_color(color), int(thickness), int(point_size)
frames = face_landmarks["frames"]
if image is None:
H, W = face_landmarks["image_size"]
img_np = np.zeros((len(frames), H, W, 3), dtype=np.uint8)
else:
img_np = _image_to_uint8(image)
B = img_np.shape[0]
n_frames = len(frames)
pbar = comfy.utils.ProgressBar(B)
out = np.empty_like(img_np)
for bi in range(B):
faces = frames[bi] if bi < n_frames else []
out[bi] = _draw_mesh(img_np[bi], faces, edges, rgb, thick, psize, fill_rings)
pbar.update_absolute(bi + 1)
return io.NodeOutput(torch.from_numpy(out).to(
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
).div_(255.0))
def _draw_mesh(image_rgb: np.ndarray, faces: list, edges,
rgb: tuple[int, int, int], thickness: int,
point_size: int, fill_rings: list[list[int]] | None = None) -> np.ndarray:
draw_edges = thickness > 0 and edges
if not faces or (fill_rings is None and not draw_edges and point_size <= 0):
return image_rgb.copy()
pil = Image.fromarray(image_rgb)
draw = ImageDraw.Draw(pil)
r = point_size * 0.5
if fill_rings is not None:
for f in faces:
lmks = f["landmarks_xy"]
for ring in fill_rings:
draw.polygon([(float(lmks[i, 0]), float(lmks[i, 1])) for i in ring], fill=rgb)
return np.asarray(pil)
for f in faces:
lmks = f["landmarks_xy"]
n = lmks.shape[0]
if draw_edges:
for a, b in edges:
if a < n and b < n:
draw.line([(float(lmks[a, 0]), float(lmks[a, 1])),
(float(lmks[b, 0]), float(lmks[b, 1]))], fill=rgb, width=thickness)
if point_size == 1:
draw.point(lmks.flatten().tolist(), fill=rgb)
elif point_size > 1:
for x, y in lmks:
draw.ellipse((float(x) - r, float(y) - r, float(x) + r, float(y) + r), fill=rgb)
return np.asarray(pil)
# Mask region presets — closed-loop topologies only.
_MASK_REGIONS: tuple[str, ...] = ("face_oval", "lips", "left_eye", "right_eye", "irises")
_MASK_CUSTOM_FEATURES: tuple[tuple[str, bool], ...] = (
("face_oval", True),
("lips", False),
("left_eye", False),
("right_eye", False),
("irises", False),
)
class MediaPipeFaceMask(io.ComfyNode):
"""Binary mask from face landmarks, filled polygon per face. One mask per
frame in the batch; faces in the same frame composite (union)."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="MediaPipeFaceMask",
display_name="MediaPipe Face Mask",
category="image/detection",
inputs=[
FaceLandmarksType.Input("face_landmarks"),
io.DynamicCombo.Input(
"regions",
tooltip="'all' = union of face_oval+lips+eyes+irises (which collapses to face_oval since it encloses the rest). 'custom' = toggle each region individually for combos like lips+eyes.",
options=[
io.DynamicCombo.Option("all", []),
io.DynamicCombo.Option("custom", [
io.Boolean.Input(reg, default=default,
tooltip=f"Include the '{reg}' region in the mask.")
for reg, default in _MASK_CUSTOM_FEATURES
]),
],
),
],
outputs=[io.Mask.Output()],
)
@classmethod
def execute(cls, face_landmarks, regions) -> io.NodeOutput:
sets = face_landmarks["connection_sets"]
sel = regions["regions"]
if sel == "custom":
picked = [reg for reg, _ in _MASK_CUSTOM_FEATURES if regions.get(reg, False)]
else:
picked = list(_MASK_REGIONS)
rings = [r for reg in picked for r in _ordered_rings(sets[reg])]
frames = face_landmarks["frames"]
H, W = face_landmarks["image_size"]
masks = np.zeros((len(frames), H, W), dtype=np.uint8)
pbar = comfy.utils.ProgressBar(len(frames))
for bi, per_frame in enumerate(frames):
if per_frame:
pil = Image.new("L", (W, H), 0)
draw = ImageDraw.Draw(pil)
for f in per_frame:
lmks = f["landmarks_xy"]
for ring in rings:
draw.polygon([(float(lmks[i, 0]), float(lmks[i, 1])) for i in ring], fill=255)
masks[bi] = np.asarray(pil)
pbar.update_absolute(bi + 1)
return io.NodeOutput(torch.from_numpy(masks).to(
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
).div_(255.0))
class MediaPipeFaceExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [LoadMediaPipeFaceLandmarker, MediaPipeFaceLandmarker, MediaPipeFaceMeshVisualize, MediaPipeFaceMask]
async def comfy_entrypoint() -> MediaPipeFaceExtension:
return MediaPipeFaceExtension()

View File

@ -1,3 +1,3 @@
# This file is automatically generated by the build process when version is
# updated in pyproject.toml.
__version__ = "0.22.0"
__version__ = "0.22.2"

View File

@ -2,7 +2,6 @@ import copy
import heapq
import inspect
import logging
import psutil
import sys
import threading
import time
@ -728,7 +727,6 @@ class PromptExecutor:
self._notify_prompt_lifecycle("start", prompt_id)
ram_headroom = int(self.cache_args["ram"] * (1024 ** 3))
ram_inactive_headroom = int(self.cache_args["ram_inactive"] * (1024 ** 3))
ram_release_callback = self.caches.outputs.ram_release if self.cache_type == CacheType.RAM_PRESSURE else None
comfy.memory_management.set_ram_cache_release_state(ram_release_callback, ram_headroom)
@ -782,14 +780,8 @@ class PromptExecutor:
execution_list.complete_node_execution()
if self.cache_type == CacheType.RAM_PRESSURE:
ram_release_callback(ram_inactive_headroom)
ram_shortfall = ram_headroom - psutil.virtual_memory().available
freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
if freed < ram_shortfall:
if freed > 64 * (1024 ** 2):
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
time.sleep(0.05)
ram_release_callback(ram_headroom, free_active=True)
comfy.model_management.free_memory(0, None, pins_required=ram_headroom, ram_required=ram_headroom)
ram_release_callback(ram_headroom, free_active=True)
else:
# Only execute when the while-loop ends without break
# Send cached UI for intermediate output nodes that weren't executed

View File

@ -60,8 +60,6 @@ folder_names_and_paths["geometry_estimation"] = ([os.path.join(models_dir, "geom
folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flow")], supported_pt_extensions)
folder_names_and_paths["mediapipe"] = ([os.path.join(models_dir, "mediapipe")], supported_pt_extensions)
output_directory = os.path.join(base_path, "output")
temp_directory = os.path.join(base_path, "temp")
input_directory = os.path.join(base_path, "input")

20
main.py
View File

@ -283,25 +283,19 @@ def _collect_output_absolute_paths(history_result: dict) -> list[str]:
def prompt_worker(q, server_instance):
current_time: float = 0.0
cache_ram = 0
cache_ram_inactive = 0
if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
cache_ram = args.cache_ram
if cache_ram < 0:
cache_ram = min(32.0, max(4.0, comfy.model_management.total_ram * 0.25 / 1024.0))
cache_ram_inactive = min(96.0, max(12.0, comfy.model_management.total_ram * 0.75 / 1024.0))
if len(args.cache_ram) > 0:
cache_ram = args.cache_ram[0]
if len(args.cache_ram) > 1:
cache_ram_inactive = args.cache_ram[1]
cache_type = execution.CacheType.RAM_PRESSURE
if args.cache_classic:
cache_type = execution.CacheType.CLASSIC
elif args.cache_lru > 0:
cache_type = execution.CacheType.CLASSIC
if args.cache_lru > 0:
cache_type = execution.CacheType.LRU
elif cache_ram > 0:
cache_type = execution.CacheType.RAM_PRESSURE
elif args.cache_none:
cache_type = execution.CacheType.NONE
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive } )
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram } )
last_gc_collect = 0
need_gc = False
gc_collect_interval = 10.0

View File

@ -2444,7 +2444,6 @@ async def init_builtin_extra_nodes():
"nodes_hidream_o1.py",
"nodes_save_3d.py",
"nodes_moge.py",
"nodes_mediapipe.py",
]
import_failed = []

View File

@ -1,6 +1,6 @@
[project]
name = "ComfyUI"
version = "0.22.0"
version = "0.22.2"
readme = "README.md"
license = { file = "LICENSE" }
requires-python = ">=3.10"

View File

@ -1,5 +1,5 @@
comfyui-frontend-package==1.43.18
comfyui-workflow-templates==0.9.79
comfyui-workflow-templates==0.9.82
comfyui-embedded-docs==0.5.0
torch
torchsde
@ -23,7 +23,7 @@ SQLAlchemy>=2.0.0
filelock
av>=14.2.0
comfy-kitchen>=0.2.8
comfy-aimdo==0.4.3
comfy-aimdo==0.3.0
requests
simpleeval>=1.0.0
blake3

View File

@ -14,6 +14,7 @@ from tests.execution.test_execution import ComfyClient, run_warmup
class TestAsyncNodes:
@fixture(scope="class", autouse=True, params=[
(False, 0),
(True, 0),
(True, 100),
])
def _server(self, args_pytest, request):
@ -28,8 +29,6 @@ class TestAsyncNodes:
use_lru, lru_size = request.param
if use_lru:
pargs += ['--cache-lru', str(lru_size)]
else:
pargs += ['--cache-classic']
# Running server with args: pargs
p = subprocess.Popen(pargs)
yield

View File

@ -183,7 +183,8 @@ class TestExecution:
# Initialize server and client
#
@fixture(scope="class", autouse=True, params=[
{ "extra_args" : ["--cache-classic"], "should_cache_results" : True },
{ "extra_args" : [], "should_cache_results" : True },
{ "extra_args" : ["--cache-lru", 0], "should_cache_results" : True },
{ "extra_args" : ["--cache-lru", 100], "should_cache_results" : True },
{ "extra_args" : ["--cache-none"], "should_cache_results" : False },
])